module Network.HTTP.Toolkit.Body (
BodyReader
, BodyType(..)
, bodyTypeFromHeaders
, makeBodyReader
, consumeBody
, sendBody
, fromByteString
, maxChunkSize
, makeChunkedReader
, readChunkSize
, makeLengthReader
, makeUnlimitedReader
) where
import Control.Applicative
import Control.Monad
import Control.Exception
import Text.Read (readMaybe)
import Data.Maybe
import Data.Char
import Data.Bits
import Data.IORef
import Numeric
import Data.ByteString (ByteString, breakByte)
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as B8
import Network.HTTP.Types
import Network.HTTP.Toolkit.Util
import Network.HTTP.Toolkit.Error
import Network.HTTP.Toolkit.Connection
data BodyType =
None
| Chunked
| Length Int
| Unlimited
deriving (Eq, Show)
bodyTypeFromHeaders :: [Header] -> Maybe BodyType
bodyTypeFromHeaders headers = chunked <|> length_
where
chunked = lookup "Transfer-Encoding" headers >>= guard . (/= "identity") >> Just Chunked
length_ = Length <$> (lookup "Content-Length" headers >>= readMaybe . B8.unpack)
makeBodyReader :: Connection -> BodyType -> IO BodyReader
makeBodyReader c bodyType = case bodyType of
Chunked -> makeChunkedReader c
Length n -> makeLengthReader n c
Unlimited -> makeUnlimitedReader c
None -> return (pure "")
maxChunkSize :: Int
maxChunkSize = pred $ 2 ^ (maxChunkSizeDigits * 4)
maxChunkSizeDigits :: Int
maxChunkSizeDigits = pred (bitSize (undefined :: Int) `div` 4)
type BodyReader = IO ByteString
consumeBody :: BodyReader -> IO ByteString
consumeBody bodyReader = B.concat <$> go
where
go :: IO [ByteString]
go = do
bs <- bodyReader
case bs of
"" -> return []
_ -> (bs:) <$> go
sendBody :: (ByteString -> IO ()) -> BodyReader -> IO ()
sendBody send body = while (not . B.null) body send
fromByteString :: ByteString -> IO BodyReader
fromByteString input = do
ref <- newIORef (Just input)
return $ atomicModifyIORef ref $ ((,) Nothing) . fromMaybe ""
makeUnlimitedReader :: Connection -> IO BodyReader
makeUnlimitedReader c = do
ref <- newIORef False
return $ do
done <- readIORef ref
if done
then return ""
else do
xs <- connectionRead c `catchOnly` UnexpectedEndOfInput $ do
writeIORef ref True
return ""
return xs
makeLengthReader :: Int -> Connection -> IO BodyReader
makeLengthReader total c = do
ref <- newIORef total
return $ do
n <- readIORef ref
if n == 0
then return ""
else do
bs <- connectionRead c
case B.splitAt n bs of
(xs, ys) -> do
writeIORef ref (n B.length xs)
connectionUnread c ys
return xs
data Where = Data | Extension
data State = More Int Where | Trailer | Done
makeChunkedReader :: Connection -> IO BodyReader
makeChunkedReader conn = do
ref <- newIORef (More 0 Data)
return $ go ref `catchOnly` UnexpectedEndOfInput $ do
writeIORef ref Done
return ""
where
go ref = do
c <- readIORef ref
case c of
More 0 Data -> do
(n, xs) <- readChunkSize conn
writeIORef ref (More n Extension)
return xs
More n Extension -> do
bs <- connectionRead conn
case breakOnNewline bs of
("", _) ->
if n > 0
then do
handleChunkData ref (n + 3) bs
else do
writeIORef ref Trailer
connectionUnread conn bs
readTrailer ref
(xs, ys) -> do
connectionUnread conn ys
return xs
More n Data -> do
connectionRead conn >>= handleChunkData ref n
Trailer -> readTrailer ref
Done -> return ""
handleChunkData :: IORef State -> Int -> ByteString -> IO ByteString
handleChunkData ref n bs = do
let (xs, ys) = B.splitAt n bs
connectionUnread conn ys
writeIORef ref (More (n B.length xs) Data)
return xs
readTrailer :: IORef State -> IO ByteString
readTrailer ref = do
xs <- connectionReadAtLeast conn 3
if "\n\r\n" `B.isPrefixOf` xs
then do
writeIORef ref Done
let (ys, zs) = B.splitAt 3 xs
connectionUnread conn zs
return ys
else do
let Just (y, ys) = B.uncons xs
case breakOnNewline ys of
(zs, rest) -> do
connectionUnread conn rest
return (y `B.cons` zs)
breakOnNewline :: ByteString -> (ByteString, ByteString)
breakOnNewline = breakByte 10
readChunkSize :: Connection -> IO (Int, ByteString)
readChunkSize conn = do
xs <- go 0
case readHex (B8.unpack xs) of
[(n, "")] -> return (n, xs)
_ -> throwIO InvalidChunk
where
go :: Int -> IO ByteString
go n = do
bs <- connectionRead conn
case B8.span isHexDigit bs of
(xs, ys) -> do
let m = (n + B.length xs)
when (m > maxChunkSizeDigits) $
throwIO ChunkTooLarge
case ys of
"" -> do
zs <- go m
return (xs `B.append` zs)
_ -> do
_unread conn ys
return xs