{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Network.HTTP.Semantics.FillBuf (
    -- * Filling a buffer
    Next (..),
    DynaNext,
    BytesFilled,
    StreamingChunk (..),
    IsEndOfStream (..),
    CleanupStream,
    fillBuilderBodyGetNext,
    fillFileBodyGetNext,
    fillStreamBodyGetNext,
) where

import Control.Exception (SomeException)
import Control.Monad
import qualified Data.ByteString as BS
import Data.ByteString.Builder (Builder)
import qualified Data.ByteString.Builder.Extra as B
import Data.ByteString.Internal
import Data.Int (Int64)
import Data.Maybe
import Foreign.Ptr (plusPtr)
import Network.ByteOrder
import Network.HTTP.Semantics.Client

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

-- | Write part of a streaming response to the write buffer
--
-- In @http2@ this will be used to construct a single HTTP2 @DATA@ frame
-- (see discussion of the maximum number of bytes, below).
type DynaNext =
       Buffer
    -- ^ Write buffer
    -> Int
    -- ^ Maximum number of bytes we are allowed to write
    --
    -- In @http2@, this maximum will be set to the space left in the write
    -- buffer. Implicitly this also means that this maximum cannot exceed the
    -- maximum size of a HTTP2 frame, since in @http2@ the size of the write
    -- buffer is also used to set @SETTINGS_MAX_FRAME_SIZE@ (see
    -- @confBufferSize@).
    -> IO Next
    -- ^ Information on the data written, and on how to continue if not all data
    -- was written

type BytesFilled = Int

data Next
    = Next
        BytesFilled -- payload length
        Bool -- require flushing
        (Maybe DynaNext)
    | CancelNext (Maybe SomeException)

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

data StreamingChunk
    = -- | Indicate that the stream is finished
      StreamingFinished (Maybe CleanupStream)
    | -- | Indicate that the stream is cancelled
      StreamingCancelled (Maybe SomeException)
    | -- | Flush the stream
      --
      -- This will cause the write buffer to be written to the network socket,
      -- without waiting for more data.
      StreamingFlush
    | -- | Construct a DATA frame, optionally terminating the stream
      StreamingBuilder Builder IsEndOfStream

-- | Action to run prior to terminating the stream
type CleanupStream = IO ()

data IsEndOfStream =
    -- | The stream is not yet terminated
    NotEndOfStream

    -- | The stream is terminated
    --
    -- In addition to indicating that the stream is terminated, we can also
    -- specify an optional `Cleanup` handler to be run.
  | EndOfStream (Maybe CleanupStream)

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

fillBuilderBodyGetNext :: Builder -> DynaNext
fillBuilderBodyGetNext :: Builder -> DynaNext
fillBuilderBodyGetNext Builder
bb Buffer
buf Int
room = do
    (Int
len, Next
signal) <- Builder -> BufferWriter
B.runBuilder Builder
bb Buffer
buf Int
room
    Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Next -> Next
nextForBuilder Int
len Next
signal

fillFileBodyGetNext
    :: PositionRead -> FileOffset -> ByteCount -> IO () -> DynaNext
fillFileBodyGetNext :: PositionRead -> Int64 -> Int64 -> IO () -> DynaNext
fillFileBodyGetNext PositionRead
pread Int64
start Int64
bytecount IO ()
refresh Buffer
buf Int
room = do
    Int64
len <- PositionRead
pread Int64
start (Int -> Int64 -> Int64
mini Int
room Int64
bytecount) Buffer
buf
    let len' :: Int
len' = Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
len
    Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> PositionRead -> Int64 -> Int64 -> IO () -> Next
nextForFile Int
len' PositionRead
pread (Int64
start Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
len) (Int64
bytecount Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
len) IO ()
refresh

fillStreamBodyGetNext :: IO (Maybe StreamingChunk) -> DynaNext
fillStreamBodyGetNext :: IO (Maybe StreamingChunk) -> DynaNext
fillStreamBodyGetNext IO (Maybe StreamingChunk)
takeQ = NextWithTotal
loop Int
0
  where
    loop :: NextWithTotal
    loop :: NextWithTotal
loop Int
total Buffer
buf Int
room = do
        Maybe StreamingChunk
mChunk <- IO (Maybe StreamingChunk)
takeQ
        case Maybe StreamingChunk
mChunk of
            Just StreamingChunk
chunk -> StreamingChunk -> NextWithTotal -> NextWithTotal
runStreamingChunk StreamingChunk
chunk NextWithTotal
loop Int
total Buffer
buf Int
room
            Maybe StreamingChunk
Nothing -> Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Bool -> Maybe DynaNext -> Next
Next Int
total Bool
False (DynaNext -> Maybe DynaNext
forall a. a -> Maybe a
Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
forall a b. (a -> b) -> a -> b
$ NextWithTotal
loop Int
0)

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

fillBufBuilderOne :: Int -> B.BufferWriter -> DynaNext
fillBufBuilderOne :: Int -> BufferWriter -> DynaNext
fillBufBuilderOne Int
minReq BufferWriter
writer Buffer
buf0 Int
room = do
    if Int
room Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
minReq
        then do
          (Int
len, Next
signal) <- BufferWriter
writer Buffer
buf0 Int
room
          Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Next -> Next
nextForBuilder Int
len Next
signal
        else do
          Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Bool -> Maybe DynaNext -> Next
Next Int
0 Bool
True (DynaNext -> Maybe DynaNext
forall a. a -> Maybe a
Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
forall a b. (a -> b) -> a -> b
$ Int -> BufferWriter -> DynaNext
fillBufBuilderOne Int
minReq BufferWriter
writer)

fillBufBuilderTwo :: ByteString -> B.BufferWriter -> DynaNext
fillBufBuilderTwo :: ByteString -> BufferWriter -> DynaNext
fillBufBuilderTwo ByteString
bs BufferWriter
writer Buffer
buf0 Int
room
    | ByteString -> Int
BS.length ByteString
bs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
room = do
        Buffer
buf1 <- Buffer -> ByteString -> IO Buffer
copy Buffer
buf0 ByteString
bs
        let len1 :: Int
len1 = ByteString -> Int
BS.length ByteString
bs
        (Int
len2, Next
signal) <- BufferWriter
writer Buffer
buf1 (Int
room Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len1)
        Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Next -> Next
nextForBuilder (Int
len1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len2) Next
signal
    | Bool
otherwise = do
        let (ByteString
bs1, ByteString
bs2) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
room ByteString
bs
        IO Buffer -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Buffer -> IO ()) -> IO Buffer -> IO ()
forall a b. (a -> b) -> a -> b
$ Buffer -> ByteString -> IO Buffer
copy Buffer
buf0 ByteString
bs1
        Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Next -> Next
nextForBuilder Int
room (ByteString -> BufferWriter -> Next
B.Chunk ByteString
bs2 BufferWriter
writer)

nextForBuilder :: BytesFilled -> B.Next -> Next
nextForBuilder :: Int -> Next -> Next
nextForBuilder Int
len Next
B.Done =
    Int -> Bool -> Maybe DynaNext -> Next
Next Int
len Bool
True Maybe DynaNext
forall a. Maybe a
Nothing -- let's flush
nextForBuilder Int
len (B.More Int
minReq BufferWriter
writer) =
    Int -> Bool -> Maybe DynaNext -> Next
Next Int
len Bool
False (Maybe DynaNext -> Next) -> Maybe DynaNext -> Next
forall a b. (a -> b) -> a -> b
$ DynaNext -> Maybe DynaNext
forall a. a -> Maybe a
Just (Int -> BufferWriter -> DynaNext
fillBufBuilderOne Int
minReq BufferWriter
writer)
nextForBuilder Int
len (B.Chunk ByteString
bs BufferWriter
writer) =
    Int -> Bool -> Maybe DynaNext -> Next
Next Int
len Bool
False (Maybe DynaNext -> Next) -> Maybe DynaNext -> Next
forall a b. (a -> b) -> a -> b
$ DynaNext -> Maybe DynaNext
forall a. a -> Maybe a
Just (ByteString -> BufferWriter -> DynaNext
fillBufBuilderTwo ByteString
bs BufferWriter
writer)

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

-- | Like 'DynaNext', but with additional argument indicating total bytes written
--
-- Since @http2@ uses @DynaNext@ to construct a /single/ @DATA@ frame, the
-- \"total number of bytes written\" refers to the current size of the payload
-- of that @DATA@ frame.
type NextWithTotal = Int -> DynaNext

-- | Run the chunk, then continue as specified, unless streaming is finished
runStreamingChunk :: StreamingChunk -> NextWithTotal -> NextWithTotal
runStreamingChunk :: StreamingChunk -> NextWithTotal -> NextWithTotal
runStreamingChunk StreamingChunk
chunk NextWithTotal
next =
    case StreamingChunk
chunk of
        StreamingFinished Maybe (IO ())
mdec -> Maybe (IO ()) -> NextWithTotal
finished Maybe (IO ())
mdec
        StreamingCancelled Maybe SomeException
mErr -> Maybe SomeException -> NextWithTotal
cancel Maybe SomeException
mErr
        StreamingChunk
StreamingFlush -> NextWithTotal
flush
        StreamingBuilder Builder
builder IsEndOfStream
NotEndOfStream -> Builder -> NextWithTotal -> NextWithTotal
runStreamingBuilder Builder
builder NextWithTotal
next
        StreamingBuilder Builder
builder (EndOfStream Maybe (IO ())
mdec) -> Builder -> NextWithTotal -> NextWithTotal
runStreamingBuilder Builder
builder (Maybe (IO ()) -> NextWithTotal
finished Maybe (IO ())
mdec)
  where
    finished :: Maybe CleanupStream -> NextWithTotal
    finished :: Maybe (IO ()) -> NextWithTotal
finished Maybe (IO ())
mdec = \Int
total Buffer
_buf Int
_room -> do
        IO () -> Maybe (IO ()) -> IO ()
forall a. a -> Maybe a -> a
fromMaybe (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()) Maybe (IO ())
mdec
        Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Bool -> Maybe DynaNext -> Next
Next Int
total Bool
True Maybe DynaNext
forall a. Maybe a
Nothing

    flush :: NextWithTotal
    flush :: NextWithTotal
flush = \Int
total Buffer
_buf Int
_room -> do
        Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Bool -> Maybe DynaNext -> Next
Next Int
total Bool
True (DynaNext -> Maybe DynaNext
forall a. a -> Maybe a
Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
forall a b. (a -> b) -> a -> b
$ NextWithTotal
next Int
0)

    -- Cancel streaming
    --
    -- The @_total@ number of bytes written refers to the @DATA@ frame currently
    -- under construction, but not yet sent (see discussion at 'DynaNext' and
    -- 'NextWithTotal'). Moreover, the documentation of 'outBodyCancel'
    -- explicitly states that such a partially constructed frame, if one exists,
    -- will be discarded on cancellation. We can therefore simply ignore
    -- @_total@ here.
    cancel :: Maybe SomeException -> NextWithTotal
    cancel :: Maybe SomeException -> NextWithTotal
cancel Maybe SomeException
mErr = \Int
_total Buffer
_buf Int
_room -> Next -> IO Next
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Maybe SomeException -> Next
CancelNext Maybe SomeException
mErr

-- | Run 'Builder' until completion, then continue as specified
runStreamingBuilder :: Builder -> NextWithTotal -> NextWithTotal
runStreamingBuilder :: Builder -> NextWithTotal -> NextWithTotal
runStreamingBuilder Builder
builder NextWithTotal
next = \Int
total Buffer
buf Int
room -> do
    (Int, Next)
writeResult <- Builder -> BufferWriter
B.runBuilder Builder
builder Buffer
buf Int
room
    (Int, Next) -> NextWithTotal
ranWriter (Int, Next)
writeResult Int
total Buffer
buf Int
room
  where
    ranWriter :: (Int, B.Next) -> NextWithTotal
    ranWriter :: (Int, Next) -> NextWithTotal
ranWriter (Int
len, Next
signal) = \Int
total Buffer
buf Int
room -> do
        let total' :: Int
total' = Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len
        case Next
signal of
            Next
B.Done ->
                NextWithTotal
next Int
total' (Buffer
buf Buffer -> Int -> Buffer
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len) (Int
room Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len)
            B.More Int
minReq BufferWriter
writer ->
                Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Bool -> Maybe DynaNext -> Next
Next Int
total' Bool
False (DynaNext -> Maybe DynaNext
forall a. a -> Maybe a
Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
forall a b. (a -> b) -> a -> b
$ Maybe Int -> BufferWriter -> NextWithTotal
goMore (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
minReq) BufferWriter
writer Int
0)
            B.Chunk ByteString
bs BufferWriter
writer ->
                Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Bool -> Maybe DynaNext -> Next
Next Int
total' Bool
False (DynaNext -> Maybe DynaNext
forall a. a -> Maybe a
Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
forall a b. (a -> b) -> a -> b
$ ByteString -> BufferWriter -> NextWithTotal
goChunk ByteString
bs BufferWriter
writer Int
0)

    goMore :: Maybe Int -> B.BufferWriter -> NextWithTotal
    goMore :: Maybe Int -> BufferWriter -> NextWithTotal
goMore Maybe Int
mMinReq BufferWriter
writer = \Int
total Buffer
buf Int
room -> do
        let enoughRoom :: Bool
enoughRoom = Bool -> (Int -> Bool) -> Maybe Int -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (Int
room Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>=) Maybe Int
mMinReq
        if Bool
enoughRoom
            then do
              (Int, Next)
writeResult <- BufferWriter
writer Buffer
buf Int
room
              (Int, Next) -> NextWithTotal
ranWriter (Int, Next)
writeResult Int
total Buffer
buf Int
room
            else do
              Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Bool -> Maybe DynaNext -> Next
Next Int
total Bool
True (DynaNext -> Maybe DynaNext
forall a. a -> Maybe a
Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
forall a b. (a -> b) -> a -> b
$ Maybe Int -> BufferWriter -> NextWithTotal
goMore Maybe Int
mMinReq BufferWriter
writer Int
0)

    goChunk :: ByteString -> B.BufferWriter -> NextWithTotal
    goChunk :: ByteString -> BufferWriter -> NextWithTotal
goChunk ByteString
bs BufferWriter
writer = \Int
total Buffer
buf Int
room ->
        if ByteString -> Int
BS.length ByteString
bs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
room
            then do
                Buffer
buf' <- Buffer -> ByteString -> IO Buffer
copy Buffer
buf ByteString
bs
                let len :: Int
len = ByteString -> Int
BS.length ByteString
bs
                Maybe Int -> BufferWriter -> NextWithTotal
goMore Maybe Int
forall a. Maybe a
Nothing BufferWriter
writer (Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) Buffer
buf' (Int
room Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len)
            else do
                let (ByteString
bs1, ByteString
bs2) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
room ByteString
bs
                IO Buffer -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Buffer -> IO ()) -> IO Buffer -> IO ()
forall a b. (a -> b) -> a -> b
$ Buffer -> ByteString -> IO Buffer
copy Buffer
buf ByteString
bs1
                Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> Bool -> Maybe DynaNext -> Next
Next (Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
room) Bool
False (DynaNext -> Maybe DynaNext
forall a. a -> Maybe a
Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
forall a b. (a -> b) -> a -> b
$ ByteString -> BufferWriter -> NextWithTotal
goChunk ByteString
bs2 BufferWriter
writer Int
0)

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

fillBufFile :: PositionRead -> FileOffset -> ByteCount -> IO () -> DynaNext
fillBufFile :: PositionRead -> Int64 -> Int64 -> IO () -> DynaNext
fillBufFile PositionRead
pread Int64
start Int64
bytes IO ()
refresh Buffer
buf Int
room = do
    Int64
len <- PositionRead
pread Int64
start (Int -> Int64 -> Int64
mini Int
room Int64
bytes) Buffer
buf
    IO ()
refresh
    let len' :: Int
len' = Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
len
    Next -> IO Next
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Next -> IO Next) -> Next -> IO Next
forall a b. (a -> b) -> a -> b
$ Int -> PositionRead -> Int64 -> Int64 -> IO () -> Next
nextForFile Int
len' PositionRead
pread (Int64
start Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
len) (Int64
bytes Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
len) IO ()
refresh

nextForFile
    :: BytesFilled -> PositionRead -> FileOffset -> ByteCount -> IO () -> Next
nextForFile :: Int -> PositionRead -> Int64 -> Int64 -> IO () -> Next
nextForFile Int
0 PositionRead
_ Int64
_ Int64
_ IO ()
_ = Int -> Bool -> Maybe DynaNext -> Next
Next Int
0 Bool
True Maybe DynaNext
forall a. Maybe a
Nothing -- let's flush
nextForFile Int
len PositionRead
_ Int64
_ Int64
0 IO ()
_ = Int -> Bool -> Maybe DynaNext -> Next
Next Int
len Bool
False Maybe DynaNext
forall a. Maybe a
Nothing
nextForFile Int
len PositionRead
pread Int64
start Int64
bytes IO ()
refresh =
    Int -> Bool -> Maybe DynaNext -> Next
Next Int
len Bool
False (Maybe DynaNext -> Next) -> Maybe DynaNext -> Next
forall a b. (a -> b) -> a -> b
$ DynaNext -> Maybe DynaNext
forall a. a -> Maybe a
Just (DynaNext -> Maybe DynaNext) -> DynaNext -> Maybe DynaNext
forall a b. (a -> b) -> a -> b
$ PositionRead -> Int64 -> Int64 -> IO () -> DynaNext
fillBufFile PositionRead
pread Int64
start Int64
bytes IO ()
refresh

{-# INLINE mini #-}
mini :: Int -> Int64 -> Int64
mini :: Int -> Int64 -> Int64
mini Int
i Int64
n
    | Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
< Int64
n = Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
    | Bool
otherwise = Int64
n