    {-# 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