bsb-http-chunked-0.0.0.3: Data/ByteString/Builder/HTTP/Chunked.hs
{-# LANGUAGE BangPatterns, MagicHash, OverloadedStrings #-}
-- | HTTP/1.1 chunked transfer encoding as defined
-- in [RFC 7230 Section 4.1](https://tools.ietf.org/html/rfc7230#section-4.1)
module Data.ByteString.Builder.HTTP.Chunked (
chunkedTransferEncoding
, chunkedTransferTerminator
) where
import Control.Applicative (pure)
import Control.Monad (void)
import Foreign (Ptr, Word8, (.&.))
import qualified Foreign as F
import GHC.Base (Int(..), uncheckedShiftRL#)
import GHC.Word (Word32(..))
import qualified Data.ByteString as S
import Data.ByteString.Builder (Builder)
import Data.ByteString.Builder.Internal (BufferRange(..), BuildSignal)
import qualified Data.ByteString.Builder.Internal as B
import qualified Data.ByteString.Builder.Prim as P
import qualified Data.ByteString.Builder.Prim.Internal as P
import Data.ByteString.Char8 () -- For the IsString instance
------------------------------------------------------------------------------
-- CRLF utils
------------------------------------------------------------------------------
{-# INLINE writeCRLF #-}
writeCRLF :: Ptr Word8 -> IO (Ptr Word8)
writeCRLF op = do
P.runF (P.char8 P.>*< P.char8) ('\r', '\n') op
pure $! op `F.plusPtr` 2
{-# INLINE crlfBuilder #-}
crlfBuilder :: Builder
crlfBuilder = P.primFixed (P.char8 P.>*< P.char8) ('\r', '\n')
------------------------------------------------------------------------------
-- Hex Encoding Infrastructure
------------------------------------------------------------------------------
{-# INLINE shiftr_w32 #-}
shiftr_w32 :: Word32 -> Int -> Word32
shiftr_w32 (W32# w) (I# i) = W32# (w `uncheckedShiftRL#` i)
-- | @writeWord32Hex len w op@ writes the hex encoding of @w@ to @op@ and
-- returns @op `'F.plusPtr'` len@.
--
-- If writing @w@ doesn't consume all @len@ bytes, leading zeros are added.
{-# INLINE writeWord32Hex #-}
writeWord32Hex :: Int -> Word32 -> Ptr Word8 -> IO (Ptr Word8)
writeWord32Hex len w0 op0 = do
go w0 (op0 `F.plusPtr` (len - 1))
pure $! op0 `F.plusPtr` len
where
go !w !op
| op < op0 = pure ()
| otherwise = do
let nibble :: Word8
nibble = fromIntegral w .&. 0xF
hex | nibble < 10 = 48 + nibble
| otherwise = 55 + nibble
F.poke op hex
go (w `shiftr_w32` 4) (op `F.plusPtr` (-1))
{-# INLINE iterationsUntilZero #-}
iterationsUntilZero :: Integral a => (a -> a) -> a -> Int
iterationsUntilZero f = go 0
where
go !count 0 = count
go !count !x = go (count+1) (f x)
-- | Length of the hex-string required to encode the given 'Word32'.
{-# INLINE word32HexLength #-}
word32HexLength :: Word32 -> Int
word32HexLength = max 1 . iterationsUntilZero (`shiftr_w32` 4)
------------------------------------------------------------------------------
-- Chunked transfer encoding
------------------------------------------------------------------------------
-- | Transform a builder such that it uses chunked HTTP transfer encoding.
--
-- >>> :set -XOverloadedStrings
-- >>> import Data.ByteString.Builder as B
-- >>> let f = B.toLazyByteString . chunkedTransferEncoding . B.lazyByteString
-- >>> f "data"
-- "004\r\ndata\r\n"
--
-- >>> f ""
-- ""
--
-- /Note/: While for many inputs, the bytestring chunks that can be obtained from the output
-- via @'Data.ByteString.Lazy.toChunks' . 'Data.ByteString.Builder.toLazyByteString'@
-- each form a chunk in the sense
-- of [RFC 7230 Section 4.1](https://tools.ietf.org/html/rfc7230#section-4.1),
-- this correspondence doesn't hold in general.
chunkedTransferEncoding :: Builder -> Builder
chunkedTransferEncoding innerBuilder =
B.builder transferEncodingStep
where
transferEncodingStep k =
go (B.runBuilder innerBuilder)
where
go innerStep (BufferRange op ope)
-- FIXME: Assert that outRemaining < maxBound :: Word32
| outRemaining < minimalBufferSize =
pure $ B.bufferFull minimalBufferSize op (go innerStep)
| otherwise = do
let !brInner@(BufferRange opInner _) = BufferRange
(op `F.plusPtr` (maxChunkSizeLength + 2)) -- leave space for chunk header
(ope `F.plusPtr` (-maxAfterBufferOverhead)) -- leave space at end of data
-- wraps the chunk, if it is non-empty, and returns the
-- signal constructed with the correct end-of-data pointer
{-# INLINE wrapChunk #-}
wrapChunk :: Ptr Word8 -> (Ptr Word8 -> IO (BuildSignal a))
-> IO (BuildSignal a)
wrapChunk !chunkDataEnd mkSignal
| chunkDataEnd == opInner = mkSignal op
| otherwise = do
let chunkSize = fromIntegral $ chunkDataEnd `F.minusPtr` opInner
-- If the hex of chunkSize requires less space than
-- maxChunkSizeLength, we get leading zeros.
void $ writeWord32Hex maxChunkSizeLength chunkSize op
void $ writeCRLF (opInner `F.plusPtr` (-2))
void $ writeCRLF chunkDataEnd
mkSignal (chunkDataEnd `F.plusPtr` 2)
doneH opInner' _ = wrapChunk opInner' $ \op' -> do
let !br' = BufferRange op' ope
k br'
fullH opInner' minRequiredSize nextInnerStep =
wrapChunk opInner' $ \op' ->
pure $! B.bufferFull
(minRequiredSize + maxEncodingOverhead)
op'
(go nextInnerStep)
insertChunkH opInner' bs nextInnerStep
| S.null bs = -- flush
wrapChunk opInner' $ \op' ->
pure $! B.insertChunk op' S.empty (go nextInnerStep)
| otherwise = -- insert non-empty bytestring
wrapChunk opInner' $ \op' -> do
-- add header for inserted bytestring
-- FIXME: assert(S.length bs < maxBound :: Word32)
let chunkSize = fromIntegral $ S.length bs
hexLength = word32HexLength chunkSize
!op'' <- writeWord32Hex hexLength chunkSize op'
!op''' <- writeCRLF op''
-- insert bytestring and write CRLF in next buildstep
pure $! B.insertChunk
op''' bs
(B.runBuilderWith crlfBuilder $ go nextInnerStep)
-- execute inner builder with reduced boundaries
B.fillWithBuildStep innerStep doneH fullH insertChunkH brInner
where
-- minimal size guaranteed for actual data no need to require more
-- than 1 byte to guarantee progress the larger sizes will be
-- hopefully provided by the driver or requested by the wrapped
-- builders.
minimalChunkSize = 1
-- overhead computation
maxBeforeBufferOverhead = F.sizeOf (undefined :: Int) + 2 -- max chunk size and CRLF after header
maxAfterBufferOverhead = 2 + -- CRLF after data
F.sizeOf (undefined :: Int) + 2 -- max bytestring size, CRLF after header
maxEncodingOverhead = maxBeforeBufferOverhead + maxAfterBufferOverhead
minimalBufferSize = minimalChunkSize + maxEncodingOverhead
-- remaining and required space computation
outRemaining = ope `F.minusPtr` op
maxChunkSizeLength = word32HexLength $ fromIntegral outRemaining
-- | The zero-length chunk @0\\r\\n\\r\\n@ signaling the termination of the data transfer.
chunkedTransferTerminator :: Builder
chunkedTransferTerminator = B.byteStringCopy "0\r\n\r\n"