http3-0.0.0: util/qif.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import Conduit hiding (yield)
import Control.Concurrent
import Control.Concurrent.STM
import Control.Monad
import Data.Attoparsec.ByteString (Parser)
import qualified Data.Attoparsec.ByteString as P
import Data.ByteString (ByteString, ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BS8
import Data.Conduit.Attoparsec
import System.Environment
import System.IO
import qualified UnliftIO.Exception as E

import Network.QPACK

data Block = Block Int ByteString deriving Show

----------------------------------------------------------------

main :: IO ()
main = do
    args <- getArgs
    case args of
      [size,efile]       -> dump (read size) efile
      [size,efile,qfile] -> test (read size) efile qfile
      _ -> putStrLn "qif size <encode-file> [<qif-file>]"

----------------------------------------------------------------

dump :: Int -> FilePath -> IO ()
dump size efile = do
    (dec, insthdr) <- newQDecoderS defaultQDecoderConfig { dcDynamicTableSize = size } True
    runConduitRes (sourceFile efile .| conduitParser block .| mapM_C (liftIO . dumpSwitch dec insthdr))

dumpSwitch :: (ByteString -> IO HeaderList)
       -> EncoderInstructionHandlerS
       -> (a, Block)
       -> IO ()
dumpSwitch dec insthdr (_, Block n bs)
  | n == 0    = do
        putStrLn "---- Stream 0"
        insthdr bs
  | otherwise = do
        putStrLn $ "---- Stream " ++ show n
        _ <- dec bs
        return ()

----------------------------------------------------------------

test :: Int -> FilePath -> FilePath -> IO ()
test size efile qfile = do
    (dec, insthdr') <- newQDecoderS defaultQDecoderConfig { dcDynamicTableSize = size } False
    q <- newTQueueIO
    let recv   = atomically $ readTQueue q
        send x = atomically $ writeTQueue q x
        insthdr bs = do
            emp <- atomically $ isEmptyTQueue q
            unless emp yield
            insthdr' bs
            yield
    mvar <- newEmptyMVar
    withFile qfile ReadMode $ \h -> do
        tid <- forkIO $ decode dec h recv mvar
        runConduitRes (sourceFile efile .| conduitParser block .| mapM_C (liftIO . testSwitch send insthdr))
        takeMVar mvar
        killThread tid

testSwitch :: (Block -> IO ())
           -> EncoderInstructionHandlerS
           -> (a, Block)
           -> IO ()
testSwitch send insthdr (_, blk@(Block n bs))
  | n == 0    = insthdr bs
  | otherwise = send blk

decode :: QDecoderS -> Handle -> IO Block -> MVar () -> IO ()
decode dec h recv mvar = loop
  where
    loop = do
        hdr' <- fromCaseSensitive <$> headerlist h
        if hdr' == [] then
            putMVar mvar ()
          else do
            Block n bs <- recv
            hdr <- dec bs
            if hdr == hdr' then
                loop
              else do
                putStrLn $ "---- Stream " ++ show n
                mapM_ print hdr
                putStrLn "----"
                mapM_ print hdr'
                putStrLn "----"
                putMVar mvar ()

fromCaseSensitive :: HeaderList -> HeaderList
fromCaseSensitive = map (\(k,v) -> (foldedCase $ mk k,v))

----------------------------------------------------------------

block :: Parser Block
block = do
    num <- toInt <$> P.take 8
    len <- toInt <$> P.take 4
    dat <- P.take len
    return $ Block num dat

toInt :: ByteString -> Int
toInt bs = BS.foldl' f 0 bs
  where
    f n w8 = n * 256 + fromIntegral w8

----------------------------------------------------------------

headerlist :: Handle -> IO HeaderList
headerlist h = loop id
  where
    loop b = do
        ml <- line h
        case ml of
          Nothing -> return $ b []
          Just l
            | l == ""   -> return $ b []
            | otherwise -> do
                  let (k,v0) = BS8.break (== '\t') l
                      v = BS8.drop 1 v0
                  loop (b . ((k,v) :))

line :: Handle -> IO (Maybe ByteString)
line h = do
    el <- E.try $ BS8.hGetLine h
    case el of
      Left (_ :: E.IOException) -> return Nothing
      Right l
        | BS8.take 1 l == "#" -> line h
        | otherwise           -> return $ Just l