{-# 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
brRead :: BodyReader -> IO S.ByteString
brRead :: BodyReader -> BodyReader
brRead = BodyReader -> BodyReader
forall a. a -> a
id
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)
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 ()
-> 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 ()
-> 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 ()
-> Bool
-> 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
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