{-# 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 = 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 forall a. a -> a
id
  where
    loop :: ([ByteString] -> [ByteString]) -> Int -> IO ByteString
loop [ByteString] -> [ByteString]
front Int
rem'
        | Int
rem' forall a. Ord a => a -> a -> Bool
<= Int
0 = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
L.fromChunks forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
front []
        | Bool
otherwise = do
            ByteString
bs <- BodyReader
brRead'
            if ByteString -> Bool
S.null ByteString
bs
                then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
L.fromChunks forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
front []
                else ([ByteString] -> [ByteString]) -> Int -> IO ByteString
loop ([ByteString] -> [ByteString]
front forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString
bsforall a. a -> [a] -> [a]
:)) (Int
rem' forall a. Num a => a -> a -> a
- ByteString -> Int
S.length ByteString
bs)

brEmpty :: BodyReader
brEmpty :: BodyReader
brEmpty = 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
  IORef [ByteString]
iinput <- forall a. a -> IO (IORef a)
newIORef [ByteString]
input
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef [ByteString]
iinput forall a b. (a -> b) -> a -> b
$ \[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' =
    forall {c}. ([ByteString] -> c) -> IO c
go forall a. a -> a
id
  where
    go :: ([ByteString] -> c) -> IO c
go [ByteString] -> c
front = do
        ByteString
x <- BodyReader
brRead'
        if ByteString -> Bool
S.null ByteString
x
            then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [ByteString] -> c
front []
            else ([ByteString] -> c) -> IO c
go ([ByteString] -> c
front forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString
xforall a. a -> [a] -> [a]
:))

makeGzipReader :: BodyReader -> IO BodyReader
makeGzipReader :: BodyReader -> IO BodyReader
makeGzipReader BodyReader
brRead' = do
    Inflate
inf <- WindowBits -> IO Inflate
Z.initInflate forall a b. (a -> b) -> a -> b
$ Int -> WindowBits
Z.WindowBits Int
31
    IORef (Maybe (IO PopperRes))
istate <- forall a. a -> IO (IORef a)
newIORef forall a. Maybe a
Nothing
    let goPopper :: IO PopperRes -> BodyReader
goPopper IO PopperRes
popper = do
            PopperRes
res <- IO PopperRes
popper
            case PopperRes
res of
                Z.PRNext ByteString
bs -> do
                    forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (IO PopperRes))
istate forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just IO PopperRes
popper
                    forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs
                PopperRes
Z.PRDone -> do
                    ByteString
bs <- Inflate -> BodyReader
Z.flushInflate Inflate
inf
                    if ByteString -> Bool
S.null ByteString
bs
                        then BodyReader
start
                        else do
                            forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (IO PopperRes))
istate forall a. Maybe a
Nothing
                            forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs
                Z.PRError ZlibException
e -> forall a. HttpExceptionContent -> IO a
throwHttp forall a b. (a -> b) -> a -> b
$ ZlibException -> HttpExceptionContent
HttpZlibException ZlibException
e
        start :: BodyReader
start = do
            ByteString
bs <- BodyReader
brRead'
            if ByteString -> Bool
S.null ByteString
bs
                then forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
S.empty
                else do
                    IO PopperRes
popper <- Inflate -> ByteString -> IO (IO PopperRes)
Z.feedInflate Inflate
inf ByteString
bs
                    IO PopperRes -> BodyReader
goPopper IO PopperRes
popper
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ do
        Maybe (IO PopperRes)
state <- forall a. IORef a -> IO a
readIORef IORef (Maybe (IO PopperRes))
istate
        case Maybe (IO PopperRes)
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 ()
connectionClose :: Connection -> IO ()
connectionWrite :: Connection -> ByteString -> IO ()
connectionUnread :: Connection -> ByteString -> IO ()
connectionRead :: Connection -> BodyReader
connectionClose :: IO ()
connectionWrite :: ByteString -> IO ()
connectionUnread :: ByteString -> IO ()
connectionRead :: BodyReader
..} = do
    IORef Bool
icomplete <- forall a. a -> IO (IORef a)
newIORef Bool
False
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ do
        ByteString
bs <- BodyReader
connectionRead
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ByteString -> Bool
S.null ByteString
bs) forall a b. (a -> b) -> a -> b
$ do
          forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
icomplete Bool
True
          IO ()
cleanup
        forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
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 :: IO ()
connectionWrite :: ByteString -> IO ()
connectionUnread :: ByteString -> IO ()
connectionRead :: BodyReader
connectionClose :: Connection -> IO ()
connectionWrite :: Connection -> ByteString -> IO ()
connectionUnread :: Connection -> ByteString -> IO ()
connectionRead :: Connection -> BodyReader
..} = do
    IORef Int
icount <- forall a. a -> IO (IORef a)
newIORef Int
count0
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ do
        Int
count <- forall a. IORef a -> IO a
readIORef IORef Int
icount
        if Int
count forall a. Ord a => a -> a -> Bool
<= Int
0
            then forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
empty
            else do
                ByteString
bs <- BodyReader
connectionRead
                forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ByteString -> Bool
S.null ByteString
bs) forall a b. (a -> b) -> a -> b
$ forall a. HttpExceptionContent -> IO a
throwHttp forall a b. (a -> b) -> a -> b
$ Word64 -> Word64 -> HttpExceptionContent
ResponseBodyTooShort (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
count0) (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Int
count0 forall a. Num a => a -> a -> a
- Int
count)
                case forall a. Ord a => a -> a -> Ordering
compare Int
count forall a b. (a -> b) -> a -> b
$ ByteString -> Int
S.length ByteString
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
                        forall a. IORef a -> a -> IO ()
writeIORef IORef Int
icount (-Int
1)
                        IO ()
cleanup
                        forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
x
                    Ordering
EQ -> do
                        forall a. IORef a -> a -> IO ()
writeIORef IORef Int
icount (-Int
1)
                        IO ()
cleanup
                        forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs
                    Ordering
GT -> do
                        forall a. IORef a -> a -> IO ()
writeIORef IORef Int
icount (Int
count forall a. Num a => a -> a -> a
- ByteString -> Int
S.length ByteString
bs)
                        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 :: IO ()
connectionWrite :: ByteString -> IO ()
connectionUnread :: ByteString -> IO ()
connectionRead :: BodyReader
connectionClose :: Connection -> IO ()
connectionWrite :: Connection -> ByteString -> IO ()
connectionUnread :: Connection -> ByteString -> IO ()
connectionRead :: Connection -> BodyReader
..} = do
    IORef Int
icount <- forall a. a -> IO (IORef a)
newIORef Int
0
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ do
      ByteString
bs <- IORef Int -> BodyReader
go IORef Int
icount
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ByteString -> Bool
S.null ByteString
bs) IO ()
cleanup
      forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
bs
  where
    go :: IORef Int -> BodyReader
go IORef Int
icount = do
        Int
count0 <- forall a. IORef a -> IO a
readIORef IORef Int
icount
        (ByteString
rawCount, Int
count) <-
            if Int
count0 forall a. Eq a => a -> a -> Bool
== Int
0
                then IO (ByteString, Int)
readHeader
                else forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
empty, Int
count0)
        if Int
count forall a. Ord a => a -> a -> Bool
<= Int
0
            then do
                -- count == -1 indicates that all chunks have been consumed
                forall a. IORef a -> a -> IO ()
writeIORef IORef Int
icount (-Int
1)
                if | Int
count forall a. Eq a => a -> a -> Bool
/= -Int
1 Bool -> Bool -> Bool
&& Bool
raw -> ByteString -> ByteString -> ByteString
S.append ByteString
rawCount forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BodyReader
readTrailersRaw
                   | Int
count forall a. Eq a => a -> a -> Bool
/= -Int
1        -> IO ()
consumeTrailers forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
empty
                   | Bool
otherwise          -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
empty
            else do
                (ByteString
bs, Int
count') <- Int -> IO (ByteString, Int)
readChunk Int
count
                forall a. IORef a -> a -> IO ()
writeIORef IORef Int
icount Int
count'
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString -> ByteString
appendHeader ByteString
rawCount ByteString
bs

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

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

    requireNewline :: IO ()
requireNewline = do
        ByteString
bs <- Maybe MaxHeaderLength -> Connection -> BodyReader
connectionReadLine Maybe MaxHeaderLength
mhl Connection
conn
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
S.null ByteString
bs) forall a b. (a -> b) -> a -> b
$ forall a. HttpExceptionContent -> IO a
throwHttp HttpExceptionContent
InvalidChunkHeaders

    readHeader :: IO (ByteString, Int)
readHeader = do
        ByteString
bs <- Maybe MaxHeaderLength -> Connection -> BodyReader
connectionReadLine Maybe MaxHeaderLength
mhl Connection
conn
        case forall {a}. Num a => ByteString -> Maybe a
parseHex ByteString
bs of
            Maybe Int
Nothing -> forall a. HttpExceptionContent -> IO a
throwHttp HttpExceptionContent
InvalidChunkHeaders
            Just Int
hex -> 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 <- forall {a} {a}. (Num a, Integral a) => a -> Maybe a
toI Word8
w0 -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t}. Num t => t -> ByteString -> t
parseHex' a
i0 ByteString
bs'
            Maybe (Word8, ByteString)
_ -> 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' <- forall {a} {a}. (Num a, Integral a) => a -> Maybe a
toI Word8
w -> t -> ByteString -> t
parseHex' (t
i forall a. Num a => a -> a -> a
* t
16 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 forall a. Ord a => a -> a -> Bool
<= a
w Bool -> Bool -> Bool
&& a
w forall a. Ord a => a -> a -> Bool
<= a
57  = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
w forall a. Num a => a -> a -> a
- a
48
        | a
65 forall a. Ord a => a -> a -> Bool
<= a
w Bool -> Bool -> Bool
&& a
w forall a. Ord a => a -> a -> Bool
<= a
70  = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
w forall a. Num a => a -> a -> a
- a
55
        | a
97 forall a. Ord a => a -> a -> Bool
<= a
w Bool -> Bool -> Bool
&& a
w forall a. Ord a => a -> a -> Bool
<= a
102 = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
w forall a. Num a => a -> a -> a
- a
87
        | Bool
otherwise = forall a. Maybe a
Nothing

    readTrailersRaw :: BodyReader
readTrailersRaw = do
        ByteString
bs <- Maybe MaxHeaderLength -> Connection -> BodyReader
connectionReadLine Maybe MaxHeaderLength
mhl Connection
conn
        if ByteString -> Bool
S.null ByteString
bs
        then forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
"\r\n"
        else (ByteString
bs ByteString -> ByteString -> ByteString
`S.append` ByteString
"\r\n" ByteString -> ByteString -> ByteString
`S.append`) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BodyReader
readTrailersRaw

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