http3-0.1.0: Network/QPACK.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

-- | Thread-safe QPACK encoder/decoder.
module Network.QPACK (
    -- * Encoder
    QEncoderConfig (..),
    defaultQEncoderConfig,
    QEncoder,
    newQEncoder,
    TableOperation (..),

    -- ** Encoder for debugging
    QEncoderS,
    newQEncoderS,

    -- * Decoder
    QDecoderConfig (..),
    defaultQDecoderConfig,
    QDecoder,
    newQDecoder,

    -- ** Decoder for debugging
    QDecoderS,
    newQDecoderS,

    -- * Types
    EncodedEncoderInstruction,
    EncoderInstructionHandler,
    EncoderInstructionHandlerS,
    EncodedDecoderInstruction,
    DecoderInstructionHandler,
    InstructionHandler,
    Size,

    -- * Re-exports
    TokenHeaderTable,
    TokenHeaderList,
    ValueTable,
    Header,
    getFieldValue,
    toTokenHeaderTable,
    original,
    foldedCase,
    mk,
) where

import Control.Concurrent
import Control.Concurrent.STM
import qualified Control.Exception as E
import qualified Data.ByteString as BS
import Data.CaseInsensitive hiding (map)
import qualified Data.CaseInsensitive as CI
import Network.ByteOrder
import Network.HPACK.Internal (
    GCBuffer,
    Size,
    entryToken,
    toEntryToken,
    toTokenHeaderTable,
 )
import Network.HTTP.Types
import Network.QUIC.Internal (StreamId)

import Imports
import Network.QPACK.Error
import Network.QPACK.HeaderBlock
import Network.QPACK.Instruction
import Network.QPACK.Table
import Network.QPACK.Types

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

-- | QPACK encoder.
type QEncoder = StreamId -> TokenHeaderList -> IO EncodedFieldSection

-- | QPACK simple encoder.
type QEncoderS = StreamId -> [Header] -> IO EncodedFieldSection

-- | QPACK decoder.
type QDecoder = StreamId -> EncodedFieldSection -> IO TokenHeaderTable

-- | QPACK simple decoder.
type QDecoderS = StreamId -> EncodedFieldSection -> IO (Maybe [Header])

-- | Encoder instruction handler.
type EncoderInstructionHandler = (Int -> IO EncodedEncoderInstruction) -> IO ()

-- | Simple encoder instruction handler.
--   Leftover is returned.
type EncoderInstructionHandlerS =
    EncodedEncoderInstruction -> IO EncodedEncoderInstruction

-- | Encoded decoder instruction.
type EncodedDecoderInstruction = ByteString

-- | Decoder instruction handler.
type DecoderInstructionHandler = (Int -> IO EncodedDecoderInstruction) -> IO ()

-- | A type to integrating handlers.
type InstructionHandler = (Int -> IO ByteString) -> IO ()

data TableOperation = TableOperation
    { setCapacity :: Int -> IO ()
    , setBlockedStreams :: Int -> IO ()
    , setHeaderSize :: Int -> IO ()
    }

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

-- | Configuration for QPACK encoder.
data QEncoderConfig = QEncoderConfig
    { ecMaxTableCapacity :: Size
    , ecHeaderBlockBufferSize :: Size
    , ecInstructionBufferSize :: Size
    }
    deriving (Show)

-- | Default configuration for QPACK encoder.
--
-- >>> defaultQEncoderConfig
-- QEncoderConfig {ecMaxTableCapacity = 4096, ecHeaderBlockBufferSize = 4096, ecInstructionBufferSize = 4096}
defaultQEncoderConfig :: QEncoderConfig
defaultQEncoderConfig =
    QEncoderConfig
        { ecMaxTableCapacity = 4096
        , ecHeaderBlockBufferSize = 4096
        , ecInstructionBufferSize = 4096
        }

-- | Creating a new QPACK encoder.
newQEncoder
    :: QEncoderConfig
    -> (EncodedEncoderInstruction -> IO ())
    -> IO (QEncoder, DecoderInstructionHandler, TableOperation)
newQEncoder QEncoderConfig{..} sendEI = do
    let bufsiz1 = ecHeaderBlockBufferSize
        bufsiz2 = ecInstructionBufferSize
    gcbuf1 <- mallocPlainForeignPtrBytes bufsiz1
    gcbuf2 <- mallocPlainForeignPtrBytes bufsiz2
    dyntbl <- newDynamicTableForEncoding sendEI
    lock <- newMVar ()
    let enc =
            qpackEncoder
                gcbuf1
                bufsiz1
                gcbuf2
                bufsiz2
                dyntbl
                lock
        handler = decoderInstructionHandler dyntbl
        ctl =
            TableOperation
                { setCapacity = \n -> do
                    -- "n" is decoder-proposed size via settings.
                    let tableSize = min ecMaxTableCapacity n
                    setTableCapacity dyntbl tableSize
                    ins <- encodeEncoderInstructions [SetDynamicTableCapacity tableSize] False
                    sendIns dyntbl ins
                , setBlockedStreams = setMaxBlockedStreams dyntbl
                , setHeaderSize = setMaxHeaderSize dyntbl
                }
    return (enc, handler, ctl)

tokenHeaderSize :: TokenHeader -> Int
tokenHeaderSize (t, v) = BS.length (CI.original (tokenKey t)) + BS.length v + 8 -- adhoc overhead

split :: Int -> TokenHeaderList -> (TokenHeaderList, TokenHeaderList)
split lim ts = split' 0 ts
  where
    split' _ [] = ([], [])
    split' s xxs@(x : xs)
        | siz > lim = E.throw BufferOverrun
        | s' < lim =
            let (ys, zs) = split' s' xs
             in (x : ys, zs)
        | otherwise = ([], xxs)
      where
        siz = tokenHeaderSize x
        s' = s + siz

splitThrough :: Int -> TokenHeaderList -> [TokenHeaderList]
splitThrough lim ts0 = loop ts0 id
  where
    loop [] builder = builder []
    loop ts builder = loop ts2 (builder . (ts1 :))
      where
        (ts1, ts2) = split lim ts

qpackEncoder
    :: GCBuffer
    -> Int
    -> GCBuffer
    -> Int
    -> DynamicTable
    -> MVar ()
    -> QEncoder
qpackEncoder gcbuf1 bufsiz1 gcbuf2 bufsiz2 dyntbl lock sid ts =
    withMVar lock $ \_ ->
        withForeignPtr gcbuf1 $ \buf1 ->
            withForeignPtr gcbuf2 $ \buf2 -> do
                siz <- getTableCapacity dyntbl
                qpackDebug dyntbl $
                    putStrLn $
                        "---- Stream " ++ show sid ++ " " ++ "tblsiz: " ++ show siz
                setBasePointToInsersionPoint dyntbl
                clearRequiredInsertCount dyntbl
                let tss = splitThrough bufsiz1 ts
                his <- mapM (qpackEncodeHeader buf1 bufsiz1 buf2 bufsiz2 dyntbl) tss
                let (hbs, daiss) = unzip his
                prefix <- qpackEncodePrefix buf1 bufsiz1 dyntbl
                let section = BS.concat (prefix : hbs)
                reqInsCnt <- getRequiredInsertCount dyntbl
                -- To count only blocked sections,
                -- dont' register this section if reqInsCnt == 0.
                when (reqInsCnt /= 0) $
                    insertSection dyntbl sid $
                        Section reqInsCnt $
                            concat daiss
                return section

qpackEncoderS
    :: GCBuffer
    -> Int
    -> GCBuffer
    -> Int
    -> DynamicTable
    -> MVar ()
    -> QEncoderS
qpackEncoderS gcbuf1 bufsiz1 gcbuf2 bufsiz2 dyntbl lock sid hs =
    withMVar lock $ \_ ->
        withForeignPtr gcbuf1 $ \buf1 ->
            withForeignPtr gcbuf2 $ \buf2 -> do
                siz <- getTableCapacity dyntbl
                qpackDebug dyntbl $
                    putStrLn $
                        "---- Stream " ++ show sid ++ " " ++ "tblsiz: " ++ show siz
                setBasePointToInsersionPoint dyntbl
                clearRequiredInsertCount dyntbl
                let tss = splitThrough bufsiz1 ts
                his <- mapM (qpackEncodeHeader buf1 bufsiz1 buf2 bufsiz2 dyntbl) tss
                let (hbs, daiss) = unzip his
                prefix <- qpackEncodePrefix buf1 bufsiz1 dyntbl
                let section = BS.concat (prefix : hbs)
                reqInsCnt <- getRequiredInsertCount dyntbl
                -- To count only blocked sections,
                -- dont' register this section if reqInsCnt == 0.
                immAck <- getImmediateAck dyntbl
                when (reqInsCnt /= 0) $ do
                    blocked <- wouldSectionBeBlocked dyntbl reqInsCnt
                    when blocked $ insertBlockedStreamE dyntbl sid
                    let dais = concat daiss
                    insertSection dyntbl sid $ Section reqInsCnt dais
                    when immAck $ do
                        -- The same logic of SectionAcknowledgement.
                        updateKnownReceivedCount dyntbl reqInsCnt
                        mapM_ (decreaseReference dyntbl) dais
                        deleteBlockedStreamE dyntbl sid
                -- Need to emulate InsertCountIncrement since
                -- SectionAcknowledgement is not returned if
                -- RequiredInsertCount is 0.
                when immAck $ setInsersionPointToKnownReceivedCount dyntbl
                return section
  where
    mk' (k, v) = (t, v)
      where
        t = toToken $ foldedCase k
    ts = map mk' hs

qpackEncodeHeader
    :: Buffer
    -> BufferSize
    -> Buffer
    -> BufferSize
    -> DynamicTable
    -> TokenHeaderList
    -> IO (ByteString, [AbsoluteIndex])
qpackEncodeHeader buf1 bufsiz1 buf2 bufsiz2 dyntbl ts = do
    wbuf1 <- newWriteBuffer buf1 bufsiz1
    wbuf2 <- newWriteBuffer buf2 bufsiz2
    dais <- encodeTokenHeader wbuf1 wbuf2 dyntbl ts
    hb <- toByteString wbuf1
    ins <- toByteString wbuf2
    when (ins /= "") $ sendIns dyntbl ins
    return (hb, dais)

qpackEncodePrefix :: Buffer -> BufferSize -> DynamicTable -> IO ByteString
qpackEncodePrefix buf1 bufsiz1 dyntbl = do
    wbuf1 <- newWriteBuffer buf1 bufsiz1
    encodePrefix wbuf1 dyntbl
    toByteString wbuf1

-- Note: dyntbl for encoder
decoderInstructionHandler :: DynamicTable -> DecoderInstructionHandler
decoderInstructionHandler dyntbl recv = loop ""
  where
    loop bs0 = do
        bs1 <- recv 1024
        let bs
                | bs0 == "" = bs1
                | otherwise = bs0 <> bs1
        when (bs /= "") $ do
            (ins, leftover) <- decodeDecoderInstructions bs
            qpackDebug dyntbl $ mapM_ print ins
            mapM_ handle ins
            loop leftover
    handle (SectionAcknowledgement sid) = do
        msec <- getAndDelSection dyntbl sid
        case msec of
            Nothing -> E.throwIO DecoderInstructionError
            Just (Section reqInsCnt ais) -> do
                updateKnownReceivedCount dyntbl reqInsCnt
                mapM_ (decreaseReference dyntbl) ais
                deleteBlockedStreamE dyntbl sid
    handle (StreamCancellation _n) = return () -- fixme
    handle (InsertCountIncrement n)
        | n == 0 = E.throwIO DecoderInstructionError
        | otherwise = incrementKnownReceivedCount dyntbl n

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

newQEncoderS
    :: QEncoderConfig -- capacity
    -> (EncodedEncoderInstruction -> IO ())
    -> Int -- blocked stream
    -> Bool -- immediate Acks
    -> Bool -- debug
    -> IO QEncoderS
newQEncoderS QEncoderConfig{..} saveEI blocked immediateAck debug = do
    let bufsiz1 = ecHeaderBlockBufferSize
        bufsiz2 = ecInstructionBufferSize
    gcbuf1 <- mallocPlainForeignPtrBytes bufsiz1
    gcbuf2 <- mallocPlainForeignPtrBytes bufsiz2
    dyntbl <- newDynamicTableForEncoding saveEI
    setTableCapacity dyntbl ecMaxTableCapacity
    setMaxBlockedStreams dyntbl blocked
    setImmediateAck dyntbl immediateAck
    setDebugQPACK dyntbl debug
    lock <- newMVar ()
    let enc =
            qpackEncoderS
                gcbuf1
                bufsiz1
                gcbuf2
                bufsiz2
                dyntbl
                lock
    return enc

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

-- | Configuration for QPACK decoder.
data QDecoderConfig = QDecoderConfig
    { dcMaxTableCapacity :: Size
    , dcHuffmanBufferSize :: Size -- for encoder insteruction handler
    , dcBlockedSterams :: Int
    , dcMaxFieldSectionSize :: Int
    }
    deriving (Show)

-- | Default configuration for QPACK decoder.
--
-- >>> defaultQDecoderConfig
-- QDecoderConfig {dcMaxTableCapacity = 4096, dcHuffmanBufferSize = 2048, dcBlockedSterams = 100, dcMaxFieldSectionSize = 32768}
defaultQDecoderConfig :: QDecoderConfig
defaultQDecoderConfig =
    QDecoderConfig
        { dcMaxTableCapacity = 4096
        , dcHuffmanBufferSize = 2048 -- no global locking
        , dcBlockedSterams = 100
        , dcMaxFieldSectionSize = 32768
        }

-- | Creating a new QPACK decoder.
newQDecoder
    :: QDecoderConfig
    -> (EncodedDecoderInstruction -> IO ())
    -> IO (QDecoder, EncoderInstructionHandler)
newQDecoder QDecoderConfig{..} sendDI = do
    dyntbl <-
        newDynamicTableForDecoding dcHuffmanBufferSize sendDI
    setMaxBlockedStreams dyntbl dcBlockedSterams
    let dec = qpackDecoder dyntbl
        handler = encoderInstructionHandler dcMaxTableCapacity dyntbl
    return (dec, handler)

-- | Creating a new simple QPACK decoder.
newQDecoderS
    :: QDecoderConfig
    -> (EncodedDecoderInstruction -> IO ())
    -> Bool
    -> IO (QDecoderS, EncoderInstructionHandlerS)
newQDecoderS QDecoderConfig{..} sendDI debug = do
    dyntbl <-
        newDynamicTableForDecoding dcHuffmanBufferSize sendDI
    setMaxBlockedStreams dyntbl dcBlockedSterams
    setDebugQPACK dyntbl debug
    let dec = qpackDecoderS dyntbl
        handler = encoderInstructionHandlerS dcMaxTableCapacity dyntbl
    return (dec, handler)

qpackDecoder
    :: DynamicTable -> StreamId -> EncodedFieldSection -> IO TokenHeaderTable
qpackDecoder dyntbl sid bs = do
    (tbl, needAck) <- withReadBuffer bs $ \rbuf -> decodeTokenHeader dyntbl rbuf
    when needAck $
        encodeDecoderInstructions [SectionAcknowledgement sid] >>= sendIns dyntbl
    return tbl

qpackDecoderS
    :: DynamicTable -> StreamId -> EncodedFieldSection -> IO (Maybe [Header])
qpackDecoderS dyntbl sid bs = do
    qpackDebug dyntbl $ putStrLn $ "---- Stream " ++ show sid
    mhs <- withReadBuffer bs $ \rbuf -> decodeTokenHeaderS dyntbl rbuf
    case mhs of
        Nothing -> return Nothing
        Just (hs, needAck) -> do
            when needAck $
                encodeDecoderInstructions [SectionAcknowledgement sid] >>= sendIns dyntbl
            return $ Just hs

-- Note: dyntbl for decoder
encoderInstructionHandler :: Int -> DynamicTable -> EncoderInstructionHandler
encoderInstructionHandler decCapLim dyntbl recv = loop ""
  where
    loop bs0 = do
        bs1 <- recv 1024
        let bs
                | bs0 == "" = bs1
                | otherwise = bs0 <> bs1
        when (bs /= "") $ do
            leftover <- encoderInstructionHandlerS decCapLim dyntbl bs
            loop leftover

-- Note: dyntbl for decoder
encoderInstructionHandlerS :: Int -> DynamicTable -> EncoderInstructionHandlerS
encoderInstructionHandlerS _ _dyntbl "" = return ""
encoderInstructionHandlerS decCapLim dyntbl bs = do
    (ins, leftover) <- decodeEncoderInstructions hufdec bs
    cnt <- sum <$> mapM handle ins
    when (cnt /= 0) $
        encodeDecoderInstructions [InsertCountIncrement cnt] >>= sendIns dyntbl
    return leftover
  where
    hufdec = getHuffmanDecoder dyntbl -- only for encoder instruction handler
    handle ins@(SetDynamicTableCapacity n)
        | n > decCapLim = E.throwIO EncoderInstructionError
        | otherwise = do
            setTableCapacity dyntbl n
            qpackDebug dyntbl $ print ins
            return 0
    handle ins@(InsertWithNameReference ii val) = do
        ready <- isTableReady dyntbl
        unless ready $ E.throwIO EncoderInstructionError
        dai <- atomically $ do
            idx <- case ii of
                Left ai -> return $ SIndex ai
                Right ri -> do
                    ip <- getInsertionPointSTM dyntbl
                    return $ DIndex $ fromInsRelativeIndex ri ip
            ent0 <- toIndexedEntry dyntbl idx
            let ent = toEntryToken (entryToken ent0) val
            _ <- insertEntryToDecoder ent dyntbl
            return idx
        qpackDebug dyntbl $ putStrLn $ show ins ++ ": " ++ show dai
        return 1
    handle ins@(InsertWithLiteralName t val) = do
        ready <- isTableReady dyntbl
        unless ready $ E.throwIO EncoderInstructionError
        dai <- atomically $ do
            let ent = toEntryToken t val
            insertEntryToDecoder ent dyntbl
        qpackDebug dyntbl $ putStrLn $ show ins ++ ": " ++ show dai
        return 1
    handle ins@(Duplicate ri) = do
        ready <- isTableReady dyntbl
        unless ready $ E.throwIO EncoderInstructionError
        (dai, dai') <- atomically $ do
            ip <- getInsertionPointSTM dyntbl
            let ai = fromInsRelativeIndex ri ip
                idx = DIndex ai
            ent <- toIndexedEntry dyntbl idx
            ai' <- insertEntryToDecoder ent dyntbl
            return (ai, ai')
        qpackDebug dyntbl $
            putStrLn $
                show ins ++ ": " ++ show dai ++ " -> " ++ show dai'
        return 1