{-# LANGUAGE BangPatterns, MagicHash, OverloadedStrings #-}
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 ()
{-# 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')
{-# INLINE shiftr_w32 #-}
shiftr_w32 :: Word32 -> Int -> Word32
shiftr_w32 (W32# w) (I# i) = W32# (w `uncheckedShiftRL#` i)
{-# 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)
{-# INLINE word32HexLength #-}
word32HexLength :: Word32 -> Int
word32HexLength = max 1 . iterationsUntilZero (`shiftr_w32` 4)
chunkedTransferEncoding :: Builder -> Builder
chunkedTransferEncoding innerBuilder =
B.builder transferEncodingStep
where
transferEncodingStep k =
go (B.runBuilder innerBuilder)
where
go innerStep (BufferRange op ope)
| outRemaining < minimalBufferSize =
pure $ B.bufferFull minimalBufferSize op (go innerStep)
| otherwise = do
let !brInner@(BufferRange opInner _) = BufferRange
(op `F.plusPtr` (maxChunkSizeLength + 2))
(ope `F.plusPtr` (-maxAfterBufferOverhead))
{-# 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
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 =
wrapChunk opInner' $ \op' ->
pure $! B.insertChunk op' S.empty (go nextInnerStep)
| otherwise =
wrapChunk opInner' $ \op' -> do
let chunkSize = fromIntegral $ S.length bs
hexLength = word32HexLength chunkSize
!op'' <- writeWord32Hex hexLength chunkSize op'
!op''' <- writeCRLF op''
pure $! B.insertChunk
op''' bs
(B.runBuilderWith crlfBuilder $ go nextInnerStep)
B.fillWithBuildStep innerStep doneH fullH insertChunkH brInner
where
minimalChunkSize = 1
maxBeforeBufferOverhead = F.sizeOf (undefined :: Int) + 2
maxAfterBufferOverhead = 2 +
F.sizeOf (undefined :: Int) + 2
maxEncodingOverhead = maxBeforeBufferOverhead + maxAfterBufferOverhead
minimalBufferSize = minimalChunkSize + maxEncodingOverhead
outRemaining = ope `F.minusPtr` op
maxChunkSizeLength = word32HexLength $ fromIntegral outRemaining
chunkedTransferTerminator :: Builder
chunkedTransferTerminator = B.byteStringCopy "0\r\n\r\n"