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