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

module Network.HTTP3.Recv (
    Source,
    newSource,
    readSource,
    readSource',
    recvHeader,
    recvBody,
) where

import qualified Data.ByteString as BS
import Data.IORef
import Network.QUIC

import Imports
import Network.HTTP3.Context
import Network.HTTP3.Error
import Network.HTTP3.Frame

data Source = Source
    { Source -> IO ByteString
sourceRead :: IO ByteString
    , Source -> IORef (Maybe ByteString)
sourcePending :: IORef (Maybe ByteString)
    }

newSource :: Stream -> IO Source
newSource :: Stream -> IO Source
newSource Stream
strm = IO ByteString -> IORef (Maybe ByteString) -> Source
Source (Stream -> Int -> IO ByteString
recvStream Stream
strm Int
1024) (IORef (Maybe ByteString) -> Source)
-> IO (IORef (Maybe ByteString)) -> IO Source
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe ByteString -> IO (IORef (Maybe ByteString))
forall a. a -> IO (IORef a)
newIORef Maybe ByteString
forall a. Maybe a
Nothing

readSource :: Source -> IO ByteString
readSource :: Source -> IO ByteString
readSource Source{IO ByteString
IORef (Maybe ByteString)
sourceRead :: Source -> IO ByteString
sourcePending :: Source -> IORef (Maybe ByteString)
sourceRead :: IO ByteString
sourcePending :: IORef (Maybe ByteString)
..} = do
    Maybe ByteString
mx <- IORef (Maybe ByteString) -> IO (Maybe ByteString)
forall a. IORef a -> IO a
readIORef IORef (Maybe ByteString)
sourcePending
    case Maybe ByteString
mx of
        Maybe ByteString
Nothing -> IO ByteString
sourceRead
        Just ByteString
x -> do
            IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe ByteString)
sourcePending Maybe ByteString
forall a. Maybe a
Nothing
            ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
x

readSource' :: Source -> IO (ByteString, Bool)
readSource' :: Source -> IO (ByteString, Bool)
readSource' Source
src = do
    ByteString
x <- Source -> IO ByteString
readSource Source
src
    (ByteString, Bool) -> IO (ByteString, Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((ByteString, Bool) -> IO (ByteString, Bool))
-> (ByteString, Bool) -> IO (ByteString, Bool)
forall a b. (a -> b) -> a -> b
$ if ByteString
x ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
"" then (ByteString
x, Bool
True) else (ByteString
x, Bool
False)

pushbackSource :: Source -> ByteString -> IO ()
pushbackSource :: Source -> ByteString -> IO ()
pushbackSource Source
_ ByteString
"" = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
pushbackSource Source{IO ByteString
IORef (Maybe ByteString)
sourceRead :: Source -> IO ByteString
sourcePending :: Source -> IORef (Maybe ByteString)
sourceRead :: IO ByteString
sourcePending :: IORef (Maybe ByteString)
..} ByteString
bs = IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe ByteString)
sourcePending (Maybe ByteString -> IO ()) -> Maybe ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
bs

recvHeader :: Context -> Source -> IO (Maybe TokenHeaderTable)
recvHeader :: Context -> Source -> IO (Maybe TokenHeaderTable)
recvHeader Context
ctx Source
src = IFrame -> IO (Maybe TokenHeaderTable)
loop IFrame
IInit
  where
    loop :: IFrame -> IO (Maybe TokenHeaderTable)
loop IFrame
st = do
        ByteString
bs <- Source -> IO ByteString
readSource Source
src
        if ByteString
bs ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
""
            then Maybe TokenHeaderTable -> IO (Maybe TokenHeaderTable)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe TokenHeaderTable
forall a. Maybe a
Nothing
            else case IFrame -> ByteString -> IFrame
parseH3Frame IFrame
st ByteString
bs of
                IDone H3FrameType
typ ByteString
payload ByteString
leftover
                    | H3FrameType
typ H3FrameType -> H3FrameType -> Bool
forall a. Eq a => a -> a -> Bool
== H3FrameType
H3FrameHeaders -> do
                        Source -> ByteString -> IO ()
pushbackSource Source
src ByteString
leftover
                        TokenHeaderTable -> Maybe TokenHeaderTable
forall a. a -> Maybe a
Just (TokenHeaderTable -> Maybe TokenHeaderTable)
-> IO TokenHeaderTable -> IO (Maybe TokenHeaderTable)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Context -> QDecoder
qpackDecode Context
ctx ByteString
payload
                    | H3FrameType
typ H3FrameType -> H3FrameType -> Bool
forall a. Eq a => a -> a -> Bool
== H3FrameType
H3FrameData -> do
                        Context -> ApplicationProtocolError -> IO ()
abort Context
ctx ApplicationProtocolError
H3FrameUnexpected
                        IFrame -> IO (Maybe TokenHeaderTable)
loop IFrame
IInit -- dummy
                    | H3FrameType -> Bool
permittedInRequestStream H3FrameType
typ -> do
                        Source -> ByteString -> IO ()
pushbackSource Source
src ByteString
leftover
                        IFrame -> IO (Maybe TokenHeaderTable)
loop IFrame
IInit
                    | Bool
otherwise -> do
                        Context -> ApplicationProtocolError -> IO ()
abort Context
ctx ApplicationProtocolError
H3FrameUnexpected
                        IFrame -> IO (Maybe TokenHeaderTable)
loop IFrame
IInit -- dummy
                IFrame
st' -> IFrame -> IO (Maybe TokenHeaderTable)
loop IFrame
st'

recvBody
    :: Context
    -> Source
    -> IORef IFrame
    -> IORef (Maybe TokenHeaderTable)
    -> IO (ByteString, Bool)
recvBody :: Context
-> Source
-> IORef IFrame
-> IORef (Maybe TokenHeaderTable)
-> IO (ByteString, Bool)
recvBody Context
ctx Source
src IORef IFrame
refI IORef (Maybe TokenHeaderTable)
refH = do
    IFrame
st <- IORef IFrame -> IO IFrame
forall a. IORef a -> IO a
readIORef IORef IFrame
refI
    IFrame -> IO (ByteString, Bool)
loop IFrame
st
  where
    loop :: IFrame -> IO (ByteString, Bool)
loop IFrame
st = do
        ByteString
bs <- Source -> IO ByteString
readSource Source
src
        if ByteString
bs ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
""
            then (ByteString, Bool) -> IO (ByteString, Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
"", Bool
True)
            else case IFrame -> ByteString -> IFrame
parseH3Frame IFrame
st ByteString
bs of
                IPay H3FrameType
H3FrameData Int
siz Int
received [ByteString]
bss -> do
                    let st' :: IFrame
st' = H3FrameType -> Int -> Int -> [ByteString] -> IFrame
IPay H3FrameType
H3FrameData Int
siz Int
received []
                    if [ByteString] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [ByteString]
bss
                        then IFrame -> IO (ByteString, Bool)
loop IFrame
st'
                        else do
                            IORef IFrame -> IFrame -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef IFrame
refI IFrame
st'
                            let ret :: ByteString
ret = [ByteString] -> ByteString
BS.concat ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
forall a. [a] -> [a]
reverse [ByteString]
bss
                            (ByteString, Bool) -> IO (ByteString, Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
ret, Bool
False)
                IDone H3FrameType
typ ByteString
payload ByteString
leftover
                    | H3FrameType
typ H3FrameType -> H3FrameType -> Bool
forall a. Eq a => a -> a -> Bool
== H3FrameType
H3FrameHeaders -> do
                        IORef IFrame -> IFrame -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef IFrame
refI IFrame
IInit
                        -- pushbackSource src leftover -- fixme
                        TokenHeaderTable
hdr <- Context -> QDecoder
qpackDecode Context
ctx ByteString
payload
                        IORef (Maybe TokenHeaderTable) -> Maybe TokenHeaderTable -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe TokenHeaderTable)
refH (Maybe TokenHeaderTable -> IO ())
-> Maybe TokenHeaderTable -> IO ()
forall a b. (a -> b) -> a -> b
$ TokenHeaderTable -> Maybe TokenHeaderTable
forall a. a -> Maybe a
Just TokenHeaderTable
hdr
                        (ByteString, Bool) -> IO (ByteString, Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
"", Bool
True)
                    | H3FrameType
typ H3FrameType -> H3FrameType -> Bool
forall a. Eq a => a -> a -> Bool
== H3FrameType
H3FrameData -> do
                        IORef IFrame -> IFrame -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef IFrame
refI IFrame
IInit
                        Source -> ByteString -> IO ()
pushbackSource Source
src ByteString
leftover
                        (ByteString, Bool) -> IO (ByteString, Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
payload, Bool
False)
                    | H3FrameType -> Bool
permittedInRequestStream H3FrameType
typ -> do
                        Source -> ByteString -> IO ()
pushbackSource Source
src ByteString
leftover
                        IFrame -> IO (ByteString, Bool)
loop IFrame
IInit
                    | Bool
otherwise -> do
                        Context -> ApplicationProtocolError -> IO ()
abort Context
ctx ApplicationProtocolError
H3FrameUnexpected
                        (ByteString, Bool) -> IO (ByteString, Bool)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
payload, Bool
False) -- dummy
                IFrame
st' -> IFrame -> IO (ByteString, Bool)
loop IFrame
st'