Frame.hs

{-# LANGUAGE BinaryLiterals #-}
{-# LANGUAGE OverloadedStrings #-}

module Network.HTTP3.Frame (
    H3Frame (..),
    H3FrameType (..),
    fromH3FrameType,
    toH3FrameType,
    encodeH3Frame,
    encodeH3Frames,
    decodeH3Frame,
    IFrame (..),
    parseH3Frame,
    frameTypeOf,
    QInt (..),
    parseQInt,
    recvQInt,
    permittedInControlStream,
    permittedInRequestStream,
    permittedInPushStream,
) where

import qualified Data.ByteString as BS
import Network.ByteOrder
import Network.QUIC.Internal

import Imports

data H3Frame = H3Frame H3FrameType ByteString

data H3FrameType
    = H3FrameData
    | H3FrameHeaders
    | H3FrameCancelPush
    | H3FrameSettings
    | H3FramePushPromise
    | H3FrameGoaway
    | H3FrameMaxPushId
    | H3FrameUnknown Int64
    deriving (Eq, Show)

{- FOURMOLU_DISABLE -}
fromH3FrameType :: H3FrameType -> Int64
fromH3FrameType H3FrameData        = 0x00
fromH3FrameType H3FrameHeaders     = 0x01
fromH3FrameType H3FrameCancelPush  = 0x03
fromH3FrameType H3FrameSettings    = 0x04
fromH3FrameType H3FramePushPromise = 0x05
fromH3FrameType H3FrameGoaway      = 0x07
fromH3FrameType H3FrameMaxPushId   = 0x0d
fromH3FrameType (H3FrameUnknown i) = i

toH3FrameType :: Int64 -> H3FrameType
toH3FrameType 0x00 = H3FrameData
toH3FrameType 0x01 = H3FrameHeaders
toH3FrameType 0x03 = H3FrameCancelPush
toH3FrameType 0x04 = H3FrameSettings
toH3FrameType 0x05 = H3FramePushPromise
toH3FrameType 0x07 = H3FrameGoaway
toH3FrameType 0x0d = H3FrameMaxPushId
toH3FrameType i   = H3FrameUnknown i

permittedInControlStream :: H3FrameType -> Bool
permittedInControlStream H3FrameData        = False
permittedInControlStream H3FrameHeaders     = False
permittedInControlStream H3FrameCancelPush  = True
permittedInControlStream H3FrameSettings    = True
permittedInControlStream H3FramePushPromise = False
permittedInControlStream H3FrameGoaway      = True
permittedInControlStream H3FrameMaxPushId   = True
permittedInControlStream (H3FrameUnknown i)
    | i <= 0x09 = False
    | otherwise = True

permittedInRequestStream :: H3FrameType -> Bool
permittedInRequestStream H3FrameData        = True
permittedInRequestStream H3FrameHeaders     = True
permittedInRequestStream H3FrameCancelPush  = False
permittedInRequestStream H3FrameSettings    = False
permittedInRequestStream H3FramePushPromise = True
permittedInRequestStream H3FrameGoaway      = False
permittedInRequestStream H3FrameMaxPushId   = False
permittedInRequestStream (H3FrameUnknown i)
    | i <= 0x09 = False
    | otherwise = True

permittedInPushStream :: H3FrameType -> Bool
permittedInPushStream H3FrameData        = True
permittedInPushStream H3FrameHeaders     = True
permittedInPushStream H3FrameCancelPush  = False
permittedInPushStream H3FrameSettings    = False
permittedInPushStream H3FramePushPromise = False
permittedInPushStream H3FrameGoaway      = False
permittedInPushStream H3FrameMaxPushId   = False
permittedInPushStream (H3FrameUnknown i)
    | i <= 0x09 = False
    | otherwise = True
{- FOURMOLU_ENABLE -}

encodeH3Frame :: H3Frame -> IO ByteString
encodeH3Frame (H3Frame typ bs) = do
    tl <- withWriteBuffer 16 $ \wbuf -> do
        encodeInt' wbuf $ fromIntegral $ fromH3FrameType typ
        encodeInt' wbuf $ fromIntegral $ BS.length bs
    return $ tl `BS.append` bs

encodeH3Frames :: [H3Frame] -> [ByteString]
encodeH3Frames fs0 = loop fs0 id
  where
    loop [] build = build []
    loop (H3Frame ty val : fs) build = loop fs (build . (typ :) . (len :) . (val :))
      where
        typ = encodeInt $ fromIntegral $ fromH3FrameType ty
        len = encodeInt $ fromIntegral $ BS.length val

decodeH3Frame :: ByteString -> IO H3Frame
decodeH3Frame hf = withReadBuffer hf $ \rbuf -> do
    typ <- toH3FrameType . fromIntegral <$> decodeInt' rbuf
    len <- fromIntegral <$> decodeInt' rbuf
    bs <- extractByteString rbuf len
    return $ H3Frame typ bs

data QInt
    = QInit
    | QMore
        Word8 -- Masked first byte
        Int -- Bytes required
        Int -- Bytes received so far. (sum . map length)
        [ByteString] -- Reverse order
    | QDone
        Int64 -- Result
        ByteString -- leftover
    deriving (Eq, Show)

parseQInt :: QInt -> ByteString -> QInt
parseQInt st "" = st
parseQInt QInit bs0
    | len1 < reqLen = QMore ft reqLen len1 [bs1]
    | otherwise =
        let (bs2, bs3) = BS.splitAt reqLen bs1
         in QDone (toLen ft bs2) bs3
  where
    hd = BS.head bs0
    reqLen = requiredLen (hd .&. 0b11000000)
    ft = hd .&. 0b00111111
    bs1 = BS.tail bs0
    len1 = BS.length bs1
parseQInt (QMore ft reqLen len0 bss0) bs0
    | len1 < reqLen = QMore ft reqLen len1 (bs0 : bss0)
    | otherwise =
        let (bs2, bs3) = BS.splitAt reqLen $ compose bs0 bss0
         in QDone (toLen ft bs2) bs3
  where
    len1 = len0 + BS.length bs0
parseQInt (QDone _ _) _ = error "parseQInt"

requiredLen :: Word8 -> Int
requiredLen 0b00000000 = 0
requiredLen 0b01000000 = 1
requiredLen 0b10000000 = 3
requiredLen _ = 7

toLen :: Word8 -> ByteString -> Int64
toLen w0 bs = BS.foldl (\n w -> n * 256 + fromIntegral w) (fromIntegral w0) bs

-- | Read one variable-length integer from a byte source.
--
-- The source is asked for a byte at a time and answers with an empty string at
-- end of input, which is 'recvStream'\'s contract.  'Nothing' means the input
-- ended before a whole integer arrived.
--
-- A unidirectional stream announces its type this way (RFC 9114, section 6.2),
-- which is one, two, four or eight octets -- not the single one it is tempting
-- to read.
recvQInt :: (Int -> IO ByteString) -> IO (Maybe Int64)
recvQInt recv = loop QInit
  where
    loop st = do
        bs <- recv 1
        if BS.null bs
            then return Nothing
            else case parseQInt st bs of
                QDone i _ -> return $ Just i
                st' -> loop st'

data IFrame
    = -- | Parsing is about to start
      IInit
    | -- | Parsing type
      IType QInt
    | -- | Parsing length
      ILen H3FrameType QInt
    | -- | Parsing payload
      IPay
        H3FrameType
        Int -- Bytes required
        Int -- Bytes received so far.  (sum . map length)
        [ByteString] -- Reverse order
    | -- | Parsing done
      IDone
        H3FrameType
        ByteString -- Payload (entire or sentinel)
        ByteString -- Leftover
    | -- | The frame says it is longer than we are willing to hold
      ITooLong
        H3FrameType
        Int -- The length it claimed
    deriving (Eq, Show)

-- | The frame type, once the parse has got far enough to know it.
frameTypeOf :: IFrame -> Maybe H3FrameType
frameTypeOf IInit = Nothing
frameTypeOf (IType _) = Nothing
frameTypeOf (ILen typ _) = Just typ
frameTypeOf (IPay typ _ _ _) = Just typ
frameTypeOf (IDone typ _ _) = Just typ
frameTypeOf (ITooLong typ _) = Just typ

-- | Feed bytes to a frame parse.
--
-- The first argument caps the payload of any frame that has to be held whole
-- before it can be used -- everything but DATA, whose payload is handed to the
-- caller as it arrives.  A length is a variable-length integer, so without a
-- cap a peer can announce up to 2^62-1 octets and have us buffer whatever it
-- then sends towards that.  DATA is exempt because a large body is a perfectly
-- ordinary thing to send; a caller that does /not/ drain DATA must refuse it
-- on sight instead.
parseH3Frame :: Int -> IFrame -> ByteString -> IFrame
parseH3Frame _ st "" = st
parseH3Frame lim IInit bs = case parseQInt QInit bs of
    QDone i bs' ->
        let typ = toH3FrameType i
         in parseH3Frame lim (ILen typ QInit) bs'
    ist -> IType ist
parseH3Frame lim (IType ist) bs = case parseQInt ist bs of
    QDone i bs' ->
        let typ = toH3FrameType i
         in parseH3Frame lim (ILen typ QInit) bs'
    ist' -> IType ist'
parseH3Frame lim (ILen typ ist) bs = case parseQInt ist bs of
    QDone i bs'
        | reqLen == 0 -> IDone typ "" bs'
        | typ /= H3FrameData && reqLen > lim -> ITooLong typ reqLen
        | otherwise -> parseH3Frame lim (IPay typ reqLen 0 []) bs'
      where
        reqLen = fromIntegral i
    ist' -> ILen typ ist'
parseH3Frame _ (IPay typ reqLen len0 bss0) bs0 = case len1 `compare` reqLen of
    LT -> IPay typ reqLen len1 (bs0 : bss0)
    EQ -> IDone typ (compose bs0 bss0) ""
    GT ->
        let (bs2, leftover) = BS.splitAt (reqLen - len0) bs0
         in IDone typ (compose bs2 bss0) leftover
  where
    len1 = len0 + BS.length bs0
parseH3Frame _ st _ = st

compose :: ByteString -> [ByteString] -> ByteString
compose bs bss = BS.concat $ reverse (bs : bss)

{-
test :: Int64 -> QInt
tset i = loop QInit bss0
  where
    loop st [] = st
    loop st (bs:bss) = case parseQInt st bs of
        st1@(QDone _ _) -> st1
        st1             -> loop st1 bss
    bs0 = encodeInt i
    bss0 = map BS.singleton $ BS.unpack bs0
-}