{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE MultiWayIf #-}
module Network.HTTP.Client.Body
    ( makeChunkedReader
    , makeLengthReader
    , makeGzipReader
    , makeUnlimitedReader
    , brConsume
    , brEmpty
    , constBodyReader
    , brReadSome
    , brRead
    ) where

import Network.HTTP.Client.Connection
import Network.HTTP.Client.Types
import Control.Exception (assert)
import Data.ByteString (empty, uncons)
import Data.IORef
import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import Control.Monad (unless, when)
import qualified Data.Streaming.Zlib as Z

-- | Get a single chunk of data from the response body, or an empty
-- bytestring if no more data is available.
--
-- Note that in order to consume the entire request body, you will need to
-- repeatedly call this function until you receive an empty @ByteString@ as a
-- result.
--
-- Since 0.1.0
brRead :: BodyReader -> IO S.ByteString
brRead :: BodyReader -> BodyReader
brRead = BodyReader -> BodyReader
forall a. a -> a
id

-- | Continuously call 'brRead', building up a lazy ByteString until a chunk is
-- constructed that is at least as many bytes as requested.
--
-- Since 0.4.20
brReadSome :: BodyReader -> Int -> IO L.ByteString
brReadSome :: BodyReader -> Int -> IO ByteString
brReadSome BodyReader
brRead' =
    ([ByteString] -> [ByteString]) -> Int -> IO ByteString
loop [ByteString] -> [ByteString]
forall a. a -> a
id
  where
    loop :: ([ByteString] -> [ByteString]) -> Int -> IO ByteString
loop [ByteString] -> [ByteString]
front Int
rem'
        | Int
rem' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
L.fromChunks ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
front []
        | Bool
otherwise = do
            bs <- BodyReader
brRead'
            if S.null bs
                then return $ L.fromChunks $ front []
                else loop (front . (bs:)) (rem' - S.length bs)

brEmpty :: BodyReader
brEmpty :: BodyReader
brEmpty = ByteString -> BodyReader
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
S.empty

constBodyReader :: [S.ByteString] -> IO BodyReader
constBodyReader :: [ByteString] -> IO BodyReader
constBodyReader [ByteString]
input = do
  iinput <- [ByteString] -> IO (IORef [ByteString])
forall a. a -> IO (IORef a)
newIORef [ByteString]
input
  return $ atomicModifyIORef iinput $ \[ByteString]
input' ->
        case [ByteString]
input' of
            [] -> ([], ByteString
S.empty)
            ByteString
x:[ByteString]
xs -> ([ByteString]
xs, ByteString
x)

-- | Strictly consume all remaining chunks of data from the stream.
--
-- Since 0.1.0
brConsume :: BodyReader -> IO [S.ByteString]
brConsume :: BodyReader -> IO [ByteString]
brConsume BodyReader
brRead' =
    ([ByteString] -> [ByteString]) -> IO [ByteString]
forall {c}. ([ByteString] -> c) -> IO c
go [ByteString] -> [ByteString]
forall a. a -> a
id
  where
    go :: ([ByteString] -> c) -> IO c
go [ByteString] -> c
front = do
        x <- BodyReader
brRead'
        if S.null x
            then return $ front []
            else go (front . (x:))

makeGzipReader :: BodyReader -> IO BodyReader
makeGzipReader :: BodyReader -> IO BodyReader
makeGzipReader BodyReader
brRead' = do
    inf <- WindowBits -> IO Inflate
Z.initInflate (WindowBits -> IO Inflate) -> WindowBits -> IO Inflate
forall a b. (a -> b) -> a -> b
$ Int -> WindowBits
Z.WindowBits Int
31
    istate <- newIORef Nothing
    let goPopper IO PopperRes
popper = do
            res <- IO PopperRes
popper
            case res of
                Z.PRNext ByteString
bs -> do
                    IORef (Maybe (IO PopperRes)) -> Maybe (IO PopperRes) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (IO PopperRes))
istate (Maybe (IO PopperRes) -> IO ()) -> Maybe (IO PopperRes) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO PopperRes -> Maybe (IO PopperRes)
forall a. a -> Maybe a
Just IO PopperRes
popper
                    ByteString -> BodyReader
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs
                PopperRes
Z.PRDone -> do
                    bs <- Inflate -> BodyReader
Z.flushInflate Inflate
inf
                    if S.null bs
                        then start
                        else do
                            writeIORef istate Nothing
                            return bs
                Z.PRError ZlibException
e -> HttpExceptionContent -> BodyReader
forall a. HttpExceptionContent -> IO a
throwHttp (HttpExceptionContent -> BodyReader)
-> HttpExceptionContent -> BodyReader
forall a b. (a -> b) -> a -> b
$ ZlibException -> HttpExceptionContent
HttpZlibException ZlibException
e
        start = do
            bs <- BodyReader
brRead'
            if S.null bs
                then return S.empty
                else do
                    popper <- Z.feedInflate inf bs
                    goPopper popper
    return $ do
        state <- readIORef istate
        case state of
            Maybe (IO PopperRes)
Nothing -> BodyReader
start
            Just IO PopperRes
popper -> IO PopperRes -> BodyReader
goPopper IO PopperRes
popper

makeUnlimitedReader
  :: IO () -- ^ cleanup
  -> Connection
  -> IO BodyReader
makeUnlimitedReader :: IO () -> Connection -> IO BodyReader
makeUnlimitedReader IO ()
cleanup Connection {IO ()
BodyReader
ByteString -> IO ()
connectionRead :: BodyReader
connectionUnread :: ByteString -> IO ()
connectionWrite :: ByteString -> IO ()
connectionClose :: IO ()
connectionClose :: Connection -> IO ()
connectionWrite :: Connection -> ByteString -> IO ()
connectionUnread :: Connection -> ByteString -> IO ()
connectionRead :: Connection -> BodyReader
..} = do
    icomplete <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
    return $ do
        bs <- connectionRead
        when (S.null bs) $ do
          writeIORef icomplete True
          cleanup
        return bs

makeLengthReader
  :: IO () -- ^ cleanup
  -> Int
  -> Connection
  -> IO BodyReader
makeLengthReader :: IO () -> Int -> Connection -> IO BodyReader
makeLengthReader IO ()
cleanup Int
count0 Connection {IO ()
BodyReader
ByteString -> IO ()
connectionClose :: Connection -> IO ()
connectionWrite :: Connection -> ByteString -> IO ()
connectionUnread :: Connection -> ByteString -> IO ()
connectionRead :: Connection -> BodyReader
connectionRead :: BodyReader
connectionUnread :: ByteString -> IO ()
connectionWrite :: ByteString -> IO ()
connectionClose :: IO ()
..} = do
    icount <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
count0
    return $ do
        count <- readIORef icount
        if count <= 0
            then return empty
            else do
                bs <- connectionRead
                when (S.null bs) $ throwHttp $ ResponseBodyTooShort (fromIntegral count0) (fromIntegral $ count0 - count)
                case compare count $ S.length bs of
                    Ordering
LT -> do
                        let (ByteString
x, ByteString
y) = Int -> ByteString -> (ByteString, ByteString)
S.splitAt Int
count ByteString
bs
                        ByteString -> IO ()
connectionUnread ByteString
y
                        IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int
icount (-Int
1)
                        IO ()
cleanup
                        ByteString -> BodyReader
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
x
                    Ordering
EQ -> do
                        IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int
icount (-Int
1)
                        IO ()
cleanup
                        ByteString -> BodyReader
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs
                    Ordering
GT -> do
                        IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int
icount (Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
S.length ByteString
bs)
                        ByteString -> BodyReader
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs

makeChunkedReader
  :: Maybe MaxHeaderLength
  -> IO () -- ^ cleanup
  -> Bool -- ^ raw
  -> Connection
  -> IO BodyReader
makeChunkedReader :: Maybe MaxHeaderLength
-> IO () -> Bool -> Connection -> IO BodyReader
makeChunkedReader Maybe MaxHeaderLength
mhl IO ()
cleanup Bool
raw conn :: Connection
conn@Connection {IO ()
BodyReader
ByteString -> IO ()
connectionClose :: Connection -> IO ()
connectionWrite :: Connection -> ByteString -> IO ()
connectionUnread :: Connection -> ByteString -> IO ()
connectionRead :: Connection -> BodyReader
connectionRead :: BodyReader
connectionUnread :: ByteString -> IO ()
connectionWrite :: ByteString -> IO ()
connectionClose :: IO ()
..} = do
    icount <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0
    return $ do
      bs <- go icount
      when (S.null bs) cleanup
      pure bs
  where
    go :: IORef Int -> BodyReader
go IORef Int
icount = do
        count0 <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
icount
        (rawCount, count) <-
            if count0 == 0
                then readHeader
                else return (empty, count0)
        if count <= 0
            then do
                -- count == -1 indicates that all chunks have been consumed
                writeIORef icount (-1)
                if | count /= -1 && raw -> S.append rawCount <$> readTrailersRaw
                   | count /= -1        -> consumeTrailers *> pure empty
                   | otherwise          -> pure empty
            else do
                (bs, count') <- readChunk count
                writeIORef icount count'
                return $ appendHeader rawCount bs

    appendHeader :: ByteString -> ByteString -> ByteString
appendHeader
      | Bool
raw = ByteString -> ByteString -> ByteString
S.append
      | Bool
otherwise = (ByteString -> ByteString -> ByteString)
-> ByteString -> ByteString -> ByteString
forall a b c. (a -> b -> c) -> b -> a -> c
flip ByteString -> ByteString -> ByteString
forall a b. a -> b -> a
const

    readChunk :: Int -> IO (ByteString, Int)
readChunk Int
0 = (ByteString, Int) -> IO (ByteString, Int)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
empty, Int
0)
    readChunk Int
remainder = do
        bs <- BodyReader
connectionRead
        when (S.null bs) $ throwHttp InvalidChunkHeaders
        case compare remainder $ S.length bs of
            Ordering
LT -> do
                let (ByteString
x, ByteString
y) = Int -> ByteString -> (ByteString, ByteString)
S.splitAt Int
remainder ByteString
bs
                Bool -> IO () -> IO ()
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ByteString -> Bool
S.null ByteString
y) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> IO ()
connectionUnread ByteString
y
                IO ()
requireNewline
                ByteString -> IO (ByteString, Int)
forall {m :: * -> *} {b}.
(Monad m, Num b) =>
ByteString -> m (ByteString, b)
done ByteString
x
            Ordering
EQ -> do
                IO ()
requireNewline
                ByteString -> IO (ByteString, Int)
forall {m :: * -> *} {b}.
(Monad m, Num b) =>
ByteString -> m (ByteString, b)
done ByteString
bs
            Ordering
GT -> (ByteString, Int) -> IO (ByteString, Int)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
bs, Int
remainder Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
S.length ByteString
bs)
      where
        done :: ByteString -> m (ByteString, b)
done ByteString
x
          | Bool
raw = (ByteString, b) -> m (ByteString, b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
x ByteString -> ByteString -> ByteString
`S.append` ByteString
"\r\n", b
0)
          | Bool
otherwise = (ByteString, b) -> m (ByteString, b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
x, b
0)

    requireNewline :: IO ()
requireNewline = do
        bs <- Maybe MaxHeaderLength -> Connection -> BodyReader
connectionReadLine Maybe MaxHeaderLength
mhl Connection
conn
        unless (S.null bs) $ throwHttp InvalidChunkHeaders

    readHeader :: IO (ByteString, Int)
readHeader = do
        bs <- Maybe MaxHeaderLength -> Connection -> BodyReader
connectionReadLine Maybe MaxHeaderLength
mhl Connection
conn
        case parseHex bs of
            Maybe Int
Nothing -> HttpExceptionContent -> IO (ByteString, Int)
forall a. HttpExceptionContent -> IO a
throwHttp HttpExceptionContent
InvalidChunkHeaders
            Just Int
hex -> (ByteString, Int) -> IO (ByteString, Int)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
bs ByteString -> ByteString -> ByteString
`S.append` ByteString
"\r\n", Int
hex)

    parseHex :: ByteString -> Maybe a
parseHex ByteString
bs0 =
        case ByteString -> Maybe (Word8, ByteString)
uncons ByteString
bs0 of
            Just (Word8
w0, ByteString
bs')
                | Just a
i0 <- Word8 -> Maybe a
forall {a} {a}. (Integral a, Num a) => a -> Maybe a
toI Word8
w0 -> a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ a -> ByteString -> a
forall {t}. Num t => t -> ByteString -> t
parseHex' a
i0 ByteString
bs'
            Maybe (Word8, ByteString)
_ -> Maybe a
forall a. Maybe a
Nothing
    parseHex' :: t -> ByteString -> t
parseHex' t
i ByteString
bs =
        case ByteString -> Maybe (Word8, ByteString)
uncons ByteString
bs of
            Just (Word8
w, ByteString
bs')
                | Just t
i' <- Word8 -> Maybe t
forall {a} {a}. (Integral a, Num a) => a -> Maybe a
toI Word8
w -> t -> ByteString -> t
parseHex' (t
i t -> t -> t
forall a. Num a => a -> a -> a
* t
16 t -> t -> t
forall a. Num a => a -> a -> a
+ t
i') ByteString
bs'
            Maybe (Word8, ByteString)
_ -> t
i

    toI :: a -> Maybe a
toI a
w
        | a
48 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
w Bool -> Bool -> Bool
&& a
w a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
57  = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
w a -> a -> a
forall a. Num a => a -> a -> a
- a
48
        | a
65 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
w Bool -> Bool -> Bool
&& a
w a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
70  = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
w a -> a -> a
forall a. Num a => a -> a -> a
- a
55
        | a
97 a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
w Bool -> Bool -> Bool
&& a
w a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
102 = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
w a -> a -> a
forall a. Num a => a -> a -> a
- a
87
        | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing

    readTrailersRaw :: BodyReader
readTrailersRaw = do
        bs <- Maybe MaxHeaderLength -> Connection -> BodyReader
connectionReadLine Maybe MaxHeaderLength
mhl Connection
conn
        if S.null bs
        then pure "\r\n"
        else (bs `S.append` "\r\n" `S.append`) <$> readTrailersRaw

    consumeTrailers :: IO ()
consumeTrailers = Maybe MaxHeaderLength -> Connection -> IO ()
connectionDropTillBlankLine Maybe MaxHeaderLength
mhl Connection
conn