{-# LANGUAGE BangPatterns, MagicHash, OverloadedStrings, ScopedTypeVariables #-}
module Data.ByteString.Builder.HTTP.Chunked (
chunkedTransferEncoding
, chunkedTransferTerminator
) where
import Control.Monad (void, when)
import Foreign (Ptr, Word8, Word32, (.&.))
import qualified Foreign as F
import Data.ByteString (ByteString)
import qualified Data.ByteString as S
import Data.ByteString.Builder (Builder)
import Data.ByteString.Builder.Internal (BufferRange(..), BuildSignal, BuildStep)
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 :: Ptr Word8 -> IO (Ptr Word8)
writeCRLF Ptr Word8
op = do
FixedPrim (Char, Char) -> (Char, Char) -> Ptr Word8 -> IO ()
forall a. FixedPrim a -> a -> Ptr Word8 -> IO ()
P.runF (FixedPrim Char
P.char8 FixedPrim Char -> FixedPrim Char -> FixedPrim (Char, Char)
forall (f :: * -> *) a b. Monoidal f => f a -> f b -> f (a, b)
P.>*< FixedPrim Char
P.char8) (Char
'\r', Char
'\n') Ptr Word8
op
Ptr Word8 -> IO (Ptr Word8)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ptr Word8 -> IO (Ptr Word8)) -> Ptr Word8 -> IO (Ptr Word8)
forall a b. (a -> b) -> a -> b
$ Ptr Word8
op Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`F.plusPtr` Int
crlfLength
{-# INLINE crlfBuilder #-}
crlfBuilder :: Builder
crlfBuilder :: Builder
crlfBuilder = FixedPrim (Char, Char) -> (Char, Char) -> Builder
forall a. FixedPrim a -> a -> Builder
P.primFixed (FixedPrim Char
P.char8 FixedPrim Char -> FixedPrim Char -> FixedPrim (Char, Char)
forall (f :: * -> *) a b. Monoidal f => f a -> f b -> f (a, b)
P.>*< FixedPrim Char
P.char8) (Char
'\r', Char
'\n')
data Padding
= NoPadding
| PadTo !Int
{-# INLINE writeWord32Hex #-}
writeWord32Hex :: Padding -> Word32 -> Ptr Word8 -> IO (Ptr Word8)
writeWord32Hex :: Padding -> Word32 -> Ptr Word8 -> IO (Ptr Word8)
writeWord32Hex Padding
NoPadding Word32
w Ptr Word8
op = Int -> Word32 -> Ptr Word8 -> IO (Ptr Word8)
writeWord32Hex' (Word32 -> Int
word32HexLength Word32
w) Word32
w Ptr Word8
op
writeWord32Hex (PadTo Int
len) Word32
w Ptr Word8
op = Int -> Word32 -> Ptr Word8 -> IO (Ptr Word8)
writeWord32Hex' Int
len Word32
w Ptr Word8
op
{-# INLINE writeWord32Hex' #-}
writeWord32Hex' :: Int -> Word32 -> Ptr Word8 -> IO (Ptr Word8)
writeWord32Hex' :: Int -> Word32 -> Ptr Word8 -> IO (Ptr Word8)
writeWord32Hex' Int
len Word32
w0 Ptr Word8
op0 = do
Word32 -> Ptr Word8 -> IO ()
forall {t}. (Integral t, Bits t) => t -> Ptr Word8 -> IO ()
go Word32
w0 (Ptr Word8
op0 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`F.plusPtr` (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
Ptr Word8 -> IO (Ptr Word8)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ptr Word8 -> IO (Ptr Word8)) -> Ptr Word8 -> IO (Ptr Word8)
forall a b. (a -> b) -> a -> b
$ Ptr Word8
op0 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`F.plusPtr` Int
len
where
go :: t -> Ptr Word8 -> IO ()
go !t
w !Ptr Word8
op =
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Ptr Word8
op Ptr Word8 -> Ptr Word8 -> Bool
forall a. Ord a => a -> a -> Bool
>= Ptr Word8
op0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
let nibble :: Word8
nibble :: Word8
nibble = t -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral t
w Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0xF
hex :: Word8
hex | Word8
nibble Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
< Word8
10 = Word8
48 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
nibble
| Bool
otherwise = Word8
55 Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
nibble
Ptr Word8 -> Word8 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
F.poke Ptr Word8
op Word8
hex
t -> Ptr Word8 -> IO ()
go (t
w t -> Int -> t
forall a. Bits a => a -> Int -> a
`F.unsafeShiftR` Int
4) (Ptr Word8
op Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`F.plusPtr` (-Int
1))
{-# INLINE word32HexLength #-}
word32HexLength :: Word32 -> Int
word32HexLength :: Word32 -> Int
word32HexLength Word32
w = Int
maxW32HexLength Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Word32 -> Int
forall b. FiniteBits b => b -> Int
F.countLeadingZeros Word32
w Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`F.unsafeShiftR` Int
2)
crlfLength, maxW32HexLength, minimalChunkSize, maxBeforeBufferOverhead,
maxAfterBufferOverhead, maxEncodingOverhead, minimalBufferSize :: Int
crlfLength :: Int
crlfLength = Int
2
maxW32HexLength :: Int
maxW32HexLength = Int
8
minimalChunkSize :: Int
minimalChunkSize = Int
1
maxBeforeBufferOverhead :: Int
maxBeforeBufferOverhead = Int
maxW32HexLength Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
crlfLength
maxAfterBufferOverhead :: Int
maxAfterBufferOverhead = Int
crlfLength Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
maxW32HexLength Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
crlfLength
maxEncodingOverhead :: Int
maxEncodingOverhead = Int
maxBeforeBufferOverhead Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
maxAfterBufferOverhead
minimalBufferSize :: Int
minimalBufferSize = Int
minimalChunkSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
maxEncodingOverhead
chunkedTransferEncoding :: Builder -> Builder
chunkedTransferEncoding :: Builder -> Builder
chunkedTransferEncoding Builder
innerBuilder =
(forall r. BuildStep r -> BuildStep r) -> Builder
B.builder BuildStep r -> BuildStep r
forall r. BuildStep r -> BuildStep r
transferEncodingStep
where
transferEncodingStep :: forall a. BuildStep a -> BuildStep a
transferEncodingStep :: forall r. BuildStep r -> BuildStep r
transferEncodingStep BuildStep a
k =
(BufferRange -> IO (BuildSignal ())) -> BuildStep a
forall _x. (BufferRange -> IO (BuildSignal _x)) -> BuildStep a
go (Builder -> BufferRange -> IO (BuildSignal ())
B.runBuilder Builder
innerBuilder)
where
go :: (BufferRange -> IO (BuildSignal _x)) -> BuildStep a
go :: forall _x. (BufferRange -> IO (BuildSignal _x)) -> BuildStep a
go BufferRange -> IO (BuildSignal _x)
innerStep (BufferRange Ptr Word8
op Ptr Word8
ope)
| Int
outRemaining Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
minimalBufferSize =
BuildSignal a -> IO (BuildSignal a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BuildSignal a -> IO (BuildSignal a))
-> BuildSignal a -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$ Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
B.bufferFull Int
minimalBufferSize Ptr Word8
op ((BufferRange -> IO (BuildSignal _x)) -> BuildStep a
forall _x. (BufferRange -> IO (BuildSignal _x)) -> BuildStep a
go BufferRange -> IO (BuildSignal _x)
innerStep)
| Bool
otherwise =
(BufferRange -> IO (BuildSignal _x))
-> (Ptr Word8 -> _x -> IO (BuildSignal a))
-> (Ptr Word8
-> Int
-> (BufferRange -> IO (BuildSignal _x))
-> IO (BuildSignal a))
-> (Ptr Word8
-> ByteString
-> (BufferRange -> IO (BuildSignal _x))
-> IO (BuildSignal a))
-> BuildStep a
forall a b.
BuildStep a
-> (Ptr Word8 -> a -> IO b)
-> (Ptr Word8 -> Int -> BuildStep a -> IO b)
-> (Ptr Word8 -> ByteString -> BuildStep a -> IO b)
-> BufferRange
-> IO b
B.fillWithBuildStep BufferRange -> IO (BuildSignal _x)
innerStep Ptr Word8 -> _x -> IO (BuildSignal a)
forall _x. Ptr Word8 -> _x -> IO (BuildSignal a)
doneH Ptr Word8
-> Int
-> (BufferRange -> IO (BuildSignal _x))
-> IO (BuildSignal a)
forall _x. Ptr Word8 -> Int -> BuildStep _x -> IO (BuildSignal a)
fullH Ptr Word8
-> ByteString
-> (BufferRange -> IO (BuildSignal _x))
-> IO (BuildSignal a)
forall _x.
Ptr Word8 -> ByteString -> BuildStep _x -> IO (BuildSignal a)
insertChunkH BufferRange
brInner
where
outRemaining :: Int
outRemaining = Ptr Word8
ope Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`F.minusPtr` Ptr Word8
op
maxChunkSizeLength :: Int
maxChunkSizeLength = Word32 -> Int
word32HexLength (Word32 -> Int) -> Word32 -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
outRemaining
!brInner :: BufferRange
brInner@(BufferRange Ptr Word8
opInner Ptr Word8
_) = Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange
(Ptr Word8
op Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`F.plusPtr` (Int
maxChunkSizeLength Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
crlfLength))
(Ptr Word8
ope Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`F.plusPtr` (-Int
maxAfterBufferOverhead))
doneH :: Ptr Word8 -> _x
-> IO (BuildSignal a)
doneH :: forall _x. Ptr Word8 -> _x -> IO (BuildSignal a)
doneH Ptr Word8
opInner' _x
_ =
Ptr Word8
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
wrapChunk Ptr Word8
opInner' ((Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a))
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
op' ->
BuildStep a
k BuildStep a -> BuildStep a
forall a b. (a -> b) -> a -> b
$! Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
op' Ptr Word8
ope
fullH :: Ptr Word8 -> Int -> BuildStep _x
-> IO (BuildSignal a)
fullH :: forall _x. Ptr Word8 -> Int -> BuildStep _x -> IO (BuildSignal a)
fullH Ptr Word8
opInner' Int
minRequiredSize BuildStep _x
nextInnerStep =
Ptr Word8
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
wrapChunk Ptr Word8
opInner' ((Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a))
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
op' ->
BuildSignal a -> IO (BuildSignal a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BuildSignal a -> IO (BuildSignal a))
-> BuildSignal a -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$! Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
B.bufferFull
(Int
minRequiredSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
maxEncodingOverhead)
Ptr Word8
op'
(BuildStep _x -> BuildStep a
forall _x. (BufferRange -> IO (BuildSignal _x)) -> BuildStep a
go BuildStep _x
nextInnerStep)
insertChunkH :: Ptr Word8 -> ByteString -> BuildStep _x
-> IO (BuildSignal a)
insertChunkH :: forall _x.
Ptr Word8 -> ByteString -> BuildStep _x -> IO (BuildSignal a)
insertChunkH Ptr Word8
opInner' ByteString
bs BuildStep _x
nextInnerStep =
Ptr Word8
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
wrapChunk Ptr Word8
opInner' ((Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a))
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
op' ->
if ByteString -> Bool
S.null ByteString
bs
then BuildSignal a -> IO (BuildSignal a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BuildSignal a -> IO (BuildSignal a))
-> BuildSignal a -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$! Ptr Word8 -> ByteString -> BuildStep a -> BuildSignal a
forall a. Ptr Word8 -> ByteString -> BuildStep a -> BuildSignal a
B.insertChunk Ptr Word8
op' ByteString
S.empty (BuildStep _x -> BuildStep a
forall _x. (BufferRange -> IO (BuildSignal _x)) -> BuildStep a
go BuildStep _x
nextInnerStep)
else do
let chunkSize :: Word32
chunkSize = Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word32) -> Int -> Word32
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
S.length ByteString
bs
!Ptr Word8
op'' <- Padding -> Word32 -> Ptr Word8 -> IO (Ptr Word8)
writeWord32Hex Padding
NoPadding Word32
chunkSize Ptr Word8
op'
!Ptr Word8
op''' <- Ptr Word8 -> IO (Ptr Word8)
writeCRLF Ptr Word8
op''
BuildSignal a -> IO (BuildSignal a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BuildSignal a -> IO (BuildSignal a))
-> BuildSignal a -> IO (BuildSignal a)
forall a b. (a -> b) -> a -> b
$! Ptr Word8 -> ByteString -> BuildStep a -> BuildSignal a
forall a. Ptr Word8 -> ByteString -> BuildStep a -> BuildSignal a
B.insertChunk
Ptr Word8
op''' ByteString
bs
(Builder -> BuildStep a -> BuildStep a
forall a. Builder -> BuildStep a -> BuildStep a
B.runBuilderWith Builder
crlfBuilder (BuildStep a -> BuildStep a) -> BuildStep a -> BuildStep a
forall a b. (a -> b) -> a -> b
$ BuildStep _x -> BuildStep a
forall _x. (BufferRange -> IO (BuildSignal _x)) -> BuildStep a
go BuildStep _x
nextInnerStep)
{-# INLINE wrapChunk #-}
wrapChunk :: Ptr Word8 -> (Ptr Word8 -> IO (BuildSignal a))
-> IO (BuildSignal a)
wrapChunk :: Ptr Word8
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
wrapChunk !Ptr Word8
chunkDataEnd Ptr Word8 -> IO (BuildSignal a)
mkSignal
| Ptr Word8
chunkDataEnd Ptr Word8 -> Ptr Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr Word8
opInner = Ptr Word8 -> IO (BuildSignal a)
mkSignal Ptr Word8
op
| Bool
otherwise = do
let chunkSize :: Word32
chunkSize = Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word32) -> Int -> Word32
forall a b. (a -> b) -> a -> b
$ Ptr Word8
chunkDataEnd Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`F.minusPtr` Ptr Word8
opInner
IO (Ptr Word8) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Ptr Word8) -> IO ()) -> IO (Ptr Word8) -> IO ()
forall a b. (a -> b) -> a -> b
$ Padding -> Word32 -> Ptr Word8 -> IO (Ptr Word8)
writeWord32Hex (Int -> Padding
PadTo Int
maxChunkSizeLength) Word32
chunkSize Ptr Word8
op
IO (Ptr Word8) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Ptr Word8) -> IO ()) -> IO (Ptr Word8) -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> IO (Ptr Word8)
writeCRLF (Ptr Word8
opInner Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`F.plusPtr` (-Int
crlfLength))
IO (Ptr Word8) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Ptr Word8) -> IO ()) -> IO (Ptr Word8) -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> IO (Ptr Word8)
writeCRLF Ptr Word8
chunkDataEnd
Ptr Word8 -> IO (BuildSignal a)
mkSignal (Ptr Word8
chunkDataEnd Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`F.plusPtr` Int
crlfLength)
chunkedTransferTerminator :: Builder
chunkedTransferTerminator :: Builder
chunkedTransferTerminator = ByteString -> Builder
B.byteStringCopy ByteString
"0\r\n\r\n"