1{-# LANGUAGE RecordWildCards #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Network.HTTP.Client.Body
4    ( makeChunkedReader
5    , makeLengthReader
6    , makeGzipReader
7    , makeUnlimitedReader
8    , brConsume
9    , brEmpty
10    , constBodyReader
11    , brReadSome
12    , brRead
13    ) where
14
15import Network.HTTP.Client.Connection
16import Network.HTTP.Client.Types
17import Control.Exception (assert)
18import Data.ByteString (empty, uncons)
19import Data.IORef
20import qualified Data.ByteString as S
21import qualified Data.ByteString.Lazy as L
22import Control.Monad (unless, when)
23import qualified Data.Streaming.Zlib as Z
24
25-- | Get a single chunk of data from the response body, or an empty
26-- bytestring if no more data is available.
27--
28-- Note that in order to consume the entire request body, you will need to
29-- repeatedly call this function until you receive an empty @ByteString@ as a
30-- result.
31--
32-- Since 0.1.0
33brRead :: BodyReader -> IO S.ByteString
34brRead = id
35
36-- | Continuously call 'brRead', building up a lazy ByteString until a chunk is
37-- constructed that is at least as many bytes as requested.
38--
39-- Since 0.4.20
40brReadSome :: BodyReader -> Int -> IO L.ByteString
41brReadSome brRead' =
42    loop id
43  where
44    loop front rem'
45        | rem' <= 0 = return $ L.fromChunks $ front []
46        | otherwise = do
47            bs <- brRead'
48            if S.null bs
49                then return $ L.fromChunks $ front []
50                else loop (front . (bs:)) (rem' - S.length bs)
51
52brEmpty :: BodyReader
53brEmpty = return S.empty
54
55constBodyReader :: [S.ByteString] -> IO BodyReader
56constBodyReader input = do
57  iinput <- newIORef input
58  return $ atomicModifyIORef iinput $ \input' ->
59        case input' of
60            [] -> ([], S.empty)
61            x:xs -> (xs, x)
62
63-- | Strictly consume all remaining chunks of data from the stream.
64--
65-- Since 0.1.0
66brConsume :: BodyReader -> IO [S.ByteString]
67brConsume brRead' =
68    go id
69  where
70    go front = do
71        x <- brRead'
72        if S.null x
73            then return $ front []
74            else go (front . (x:))
75
76makeGzipReader :: BodyReader -> IO BodyReader
77makeGzipReader brRead' = do
78    inf <- Z.initInflate $ Z.WindowBits 31
79    istate <- newIORef Nothing
80    let goPopper popper = do
81            res <- popper
82            case res of
83                Z.PRNext bs -> do
84                    writeIORef istate $ Just popper
85                    return bs
86                Z.PRDone -> do
87                    bs <- Z.flushInflate inf
88                    if S.null bs
89                        then start
90                        else do
91                            writeIORef istate Nothing
92                            return bs
93                Z.PRError e -> throwHttp $ HttpZlibException e
94        start = do
95            bs <- brRead'
96            if S.null bs
97                then return S.empty
98                else do
99                    popper <- Z.feedInflate inf bs
100                    goPopper popper
101    return $ do
102        state <- readIORef istate
103        case state of
104            Nothing -> start
105            Just popper -> goPopper popper
106
107makeUnlimitedReader
108  :: IO () -- ^ cleanup
109  -> Connection
110  -> IO BodyReader
111makeUnlimitedReader cleanup Connection {..} = do
112    icomplete <- newIORef False
113    return $ do
114        bs <- connectionRead
115        when (S.null bs) $ do
116          writeIORef icomplete True
117          cleanup
118        return bs
119
120makeLengthReader
121  :: IO () -- ^ cleanup
122  -> Int
123  -> Connection
124  -> IO BodyReader
125makeLengthReader cleanup count0 Connection {..} = do
126    icount <- newIORef count0
127    return $ do
128        count <- readIORef icount
129        if count <= 0
130            then return empty
131            else do
132                bs <- connectionRead
133                when (S.null bs) $ throwHttp $ ResponseBodyTooShort (fromIntegral count0) (fromIntegral $ count0 - count)
134                case compare count $ S.length bs of
135                    LT -> do
136                        let (x, y) = S.splitAt count bs
137                        connectionUnread y
138                        writeIORef icount (-1)
139                        cleanup
140                        return x
141                    EQ -> do
142                        writeIORef icount (-1)
143                        cleanup
144                        return bs
145                    GT -> do
146                        writeIORef icount (count - S.length bs)
147                        return bs
148
149makeChunkedReader
150  :: IO () -- ^ cleanup
151  -> Bool -- ^ raw
152  -> Connection
153  -> IO BodyReader
154makeChunkedReader cleanup raw conn@Connection {..} = do
155    icount <- newIORef 0
156    return $ do
157      bs <- go icount
158      when (S.null bs) cleanup
159      pure bs
160  where
161    go icount = do
162        count0 <- readIORef icount
163        (rawCount, count) <-
164            if count0 == 0
165                then readHeader
166                else return (empty, count0)
167        if count <= 0
168            then do
169                writeIORef icount (-1)
170                return $ if count /= (-1) && raw then rawCount else empty
171            else do
172                (bs, count') <- readChunk count
173                writeIORef icount count'
174                return $ appendHeader rawCount bs
175
176    appendHeader
177      | raw = S.append
178      | otherwise = flip const
179
180    readChunk 0 = return (empty, 0)
181    readChunk remainder = do
182        bs <- connectionRead
183        when (S.null bs) $ throwHttp InvalidChunkHeaders
184        case compare remainder $ S.length bs of
185            LT -> do
186                let (x, y) = S.splitAt remainder bs
187                assert (not $ S.null y) $ connectionUnread y
188                requireNewline
189                done x
190            EQ -> do
191                requireNewline
192                done bs
193            GT -> return (bs, remainder - S.length bs)
194      where
195        done x
196          | raw = return (x `S.append` "\r\n", 0)
197          | otherwise = return (x, 0)
198
199    requireNewline = do
200        bs <- connectionReadLine conn
201        unless (S.null bs) $ throwHttp InvalidChunkHeaders
202
203    readHeader = do
204        bs <- connectionReadLine conn
205        case parseHex bs of
206            Nothing -> throwHttp InvalidChunkHeaders
207            Just hex -> return (bs `S.append` "\r\n", hex)
208
209    parseHex bs0 =
210        case uncons bs0 of
211            Just (w0, bs')
212                | Just i0 <- toI w0 -> Just $ parseHex' i0 bs'
213            _ -> Nothing
214    parseHex' i bs =
215        case uncons bs of
216            Just (w, bs')
217                | Just i' <- toI w -> parseHex' (i * 16 + i') bs'
218            _ -> i
219
220    toI w
221        | 48 <= w && w <= 57  = Just $ fromIntegral w - 48
222        | 65 <= w && w <= 70  = Just $ fromIntegral w - 55
223        | 97 <= w && w <= 102 = Just $ fromIntegral w - 87
224        | otherwise = Nothing
225