http3-0.1.6: Network/QPACK.hs

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

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

    -- ** Encoder for debugging
    QEncoderS,
    newQEncoderS,

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

    -- ** 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 ()
    , getHeaderSize :: IO Int
    -- ^ The peer's SETTINGS_MAX_FIELD_SECTION_SIZE, or 'maxBound' until it
    -- is known
    }

-- | The size of a field section as SETTINGS_MAX_FIELD_SECTION_SIZE counts
--   it: the lengths of each name and value plus 32 for every field (RFC
--   9114, section 4.2.2).
fieldSectionSize :: TokenHeaderList -> Int
fieldSectionSize = foldl' (\n (t, v) -> n + keyLength t + BS.length v + 32) 0
  where
    keyLength = BS.length . CI.original . tokenKey

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

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

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

-- | 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
                ecUseHuffman
                dyntbl
                lock
        handler = decoderInstructionHandler dyntbl lock
        ctl =
            TableOperation
                { setCapacity = \n -> withMVar lock $ \_ -> do
                    -- "n" is decoder-proposed size via settings.
                    -- It, not the capacity chosen below, determines
                    -- MaxEntries.
                    setMaxEntries dyntbl n
                    let tableSize = min ecMaxTableCapacity n
                    setTableCapacity dyntbl tableSize
                    ins <- encodeEncoderInstructions [SetDynamicTableCapacity tableSize] False
                    sendIns dyntbl ins
                , setBlockedStreams = setMaxBlockedStreams dyntbl
                , setHeaderSize = setMaxHeaderSize dyntbl
                , getHeaderSize = getMaxHeaderSize dyntbl
                }
    return (enc, handler, ctl)

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

-- | Taking as many fields as fit in @lim@, but always at least one.
--
-- A field that does not fit on its own makes a chunk by itself, which
-- 'encodeChunk' gives buffers of its own.  It used to be refused with
-- 'BufferOverrun' -- so a value of 4K or so could not be sent at all -- and
-- one that came to exactly @lim@ produced an empty chunk and left the rest
-- as it was, so 'splitThrough' went round forever.
split :: Int -> TokenHeaderList -> (TokenHeaderList, TokenHeaderList)
split lim ts = split' 0 ts
  where
    split' _ [] = ([], [])
    split' s xxs@(x : xs)
        | s /= 0 && s' > lim = ([], xxs)
        | otherwise =
            let (ys, zs) = split' s' xs
             in (x : ys, zs)
      where
        s' = s + tokenHeaderSize x

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

-- | Encoding a chunk from 'splitThrough', in the shared buffers when it fits.
--
-- One that does not is a single large field, and gets buffers of its own.
-- Their size allows for Huffman coding being tried in place before it is known
-- to be the shorter: a code is at most 30 bits, so under four octets per octet.
encodeChunk
    :: Buffer
    -> BufferSize
    -> Buffer
    -> BufferSize
    -> Bool
    -> DynamicTable
    -> TokenHeaderList
    -> IO (ByteString, [AbsoluteIndex])
encodeChunk buf1 bufsiz1 buf2 bufsiz2 huff dyntbl ts
    | siz <= min bufsiz1 bufsiz2 =
        qpackEncodeHeader buf1 bufsiz1 buf2 bufsiz2 huff dyntbl ts
    | otherwise = do
        let bufsiz = siz * 4 + 64
        gcbuf1 <- mallocPlainForeignPtrBytes bufsiz
        gcbuf2 <- mallocPlainForeignPtrBytes bufsiz
        withForeignPtr gcbuf1 $ \b1 ->
            withForeignPtr gcbuf2 $ \b2 ->
                qpackEncodeHeader b1 bufsiz b2 bufsiz huff dyntbl ts
  where
    siz = sum $ map tokenHeaderSize ts

qpackEncoder
    :: GCBuffer
    -> Int
    -> GCBuffer
    -> Int
    -> Bool
    -> DynamicTable
    -> MVar ()
    -> QEncoder
qpackEncoder gcbuf1 bufsiz1 gcbuf2 bufsiz2 huff dyntbl lock sid ts = do
    -- Before anything is inserted or counted.
    lim <- getMaxHeaderSize dyntbl
    let siz0 = fieldSectionSize ts
    when (siz0 > lim) $ E.throwIO $ FieldSectionTooLargeForPeer siz0 lim
    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 (min bufsiz1 bufsiz2) ts
                his <- mapM (encodeChunk buf1 bufsiz1 buf2 bufsiz2 huff 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) $ do
                    -- Counted against the decoder's
                    -- SETTINGS_QPACK_BLOCKED_STREAMS, which
                    -- 'checkBlockedStreams' consults while encoding.
                    blocked <- wouldSectionBeBlocked dyntbl reqInsCnt
                    when blocked $ insertBlockedStreamE dyntbl sid
                    insertSection dyntbl sid $
                        Section reqInsCnt $
                            concat daiss
                return section

qpackEncoderS
    :: GCBuffer
    -> Int
    -> GCBuffer
    -> Int
    -> Bool
    -> DynamicTable
    -> MVar ()
    -> QEncoderS
qpackEncoderS gcbuf1 bufsiz1 gcbuf2 bufsiz2 huff 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 (min bufsiz1 bufsiz2) ts
                his <- mapM (encodeChunk buf1 bufsiz1 buf2 bufsiz2 huff 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
                        _ <- getAndDelSection dyntbl sid
                        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
    -> Bool
    -> DynamicTable
    -> TokenHeaderList
    -> IO (ByteString, [AbsoluteIndex])
qpackEncodeHeader buf1 bufsiz1 buf2 bufsiz2 huff dyntbl ts = do
    wbuf1 <- newWriteBuffer buf1 bufsiz1
    wbuf2 <- newWriteBuffer buf2 bufsiz2
    dais <- encodeTokenHeader wbuf1 wbuf2 huff 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
--
-- The lock is the encoder's.  Acknowledgements change the same reference
-- counts, sections and blocked streams that encoding reads and writes, with
-- plain reads and writes rather than atomic ones; run beside an encoding in
-- progress, an update could be lost and an entry either kept forever or
-- evicted while a field section still referred to it.
decoderInstructionHandler
    :: DynamicTable -> MVar () -> DecoderInstructionHandler
decoderInstructionHandler dyntbl lock recv = loop ""
  where
    -- Returns when the stream ends, with an instruction cut short or not.
    -- Going on with what was left over used to spin: the same incomplete
    -- instruction, decoded again after every empty read.
    loop bs0 = do
        bs1 <- recv 1024
        let bs
                | bs0 == "" = bs1
                | otherwise = bs0 <> bs1
        when (bs1 /= "") $ do
            (ins, leftover) <- decodeDecoderInstructions bs
            qpackDebug dyntbl $ mapM_ print ins
            withMVar lock $ \_ -> 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
                -- Not simply deleted: a later section on the same stream
                -- may still be blocked.
                unblockStreamsE dyntbl
    -- The decoder will not process the rest of the stream (RFC 9204, section
    -- 4.4.2), so its outstanding field sections will never be acknowledged.
    -- Their references are released as an acknowledgement would release
    -- them, but the Known Received Count stays where it is: nothing about
    -- which insertions arrived is learnt from this.  Ignoring it used to keep
    -- the entries those sections referred to from ever being evicted, and the
    -- stream counted as blocked for good.
    handle (StreamCancellation sid) = do
        secs <- getAndDelSections dyntbl sid
        forM_ secs $ \(Section _ ais) -> mapM_ (decreaseReference dyntbl) ais
        unblockStreamsE dyntbl
    handle (InsertCountIncrement n)
        | n == 0 = E.throwIO DecoderInstructionError
        | otherwise = do
            incrementKnownReceivedCount dyntbl n
            unblockStreamsE dyntbl

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

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
    setMaxEntries dyntbl ecMaxTableCapacity
    setTableCapacity dyntbl ecMaxTableCapacity
    setMaxBlockedStreams dyntbl blocked
    setImmediateAck dyntbl immediateAck
    setDebugQPACK dyntbl debug
    lock <- newMVar ()
    let enc =
            qpackEncoderS
                gcbuf1
                bufsiz1
                gcbuf2
                bufsiz2
                ecUseHuffman
                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
    setMaxEntries dyntbl dcMaxTableCapacity
    setMaxBlockedStreams dyntbl dcBlockedSterams
    -- What we announce as SETTINGS_MAX_FIELD_SECTION_SIZE.
    setMaxHeaderSize dyntbl dcMaxFieldSectionSize
    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
    setMaxEntries dyntbl dcMaxTableCapacity
    setMaxBlockedStreams dyntbl dcBlockedSterams
    -- What we announce as SETTINGS_MAX_FIELD_SECTION_SIZE.
    setMaxHeaderSize dyntbl dcMaxFieldSectionSize
    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
    -- Returns when the stream ends; see 'decoderInstructionHandler'.
    loop bs0 = do
        bs1 <- recv 1024
        let bs
                | bs0 == "" = bs1
                | otherwise = bs0 <> bs1
        when (bs1 /= "") $ 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
            setDecoderTableCapacity 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