1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
|
{-# LANGUAGE BangPatterns, MagicHash, OverloadedStrings, ScopedTypeVariables #-}
-- | 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.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 () -- 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` crlfLength
{-# INLINE crlfBuilder #-}
crlfBuilder :: Builder
crlfBuilder = P.primFixed (P.char8 P.>*< P.char8) ('\r', '\n')
------------------------------------------------------------------------------
-- Hex Encoding Infrastructure
------------------------------------------------------------------------------
-- | Pad the chunk size with leading zeros?
data Padding
= NoPadding
| PadTo !Int
{-# INLINE writeWord32Hex #-}
writeWord32Hex :: Padding -> Word32 -> Ptr Word8 -> IO (Ptr Word8)
writeWord32Hex NoPadding w op = writeWord32Hex' (word32HexLength w) w op
writeWord32Hex (PadTo len) w op = writeWord32Hex' len w op
-- | @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 =
when (op >= op0) $ do
let nibble :: Word8
nibble = fromIntegral w .&. 0xF
hex | nibble < 10 = 48 + nibble
| otherwise = 55 + nibble
F.poke op hex
go (w `F.unsafeShiftR` 4) (op `F.plusPtr` (-1))
-- | Length of the hex-string required to encode the given 'Word32'.
{-# INLINE word32HexLength #-}
word32HexLength :: Word32 -> Int
word32HexLength w = maxW32HexLength - (F.countLeadingZeros w `F.unsafeShiftR` 2)
------------------------------------------------------------------------------
-- Constants
------------------------------------------------------------------------------
crlfLength, maxW32HexLength, minimalChunkSize, maxBeforeBufferOverhead,
maxAfterBufferOverhead, maxEncodingOverhead, minimalBufferSize :: Int
crlfLength = 2
maxW32HexLength = 8 -- 4 bytes, 2 hex digits per byte
minimalChunkSize = 1
maxBeforeBufferOverhead = maxW32HexLength + crlfLength
maxAfterBufferOverhead = crlfLength + maxW32HexLength + crlfLength
maxEncodingOverhead = maxBeforeBufferOverhead + maxAfterBufferOverhead
minimalBufferSize = minimalChunkSize + maxEncodingOverhead
------------------------------------------------------------------------------
-- 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 :: forall a. BuildStep a -> BuildStep a
transferEncodingStep k =
go (B.runBuilder innerBuilder)
where
go :: (BufferRange -> IO (BuildSignal _x)) -> BuildStep a
go innerStep (BufferRange op ope)
-- FIXME: Assert that outRemaining < maxBound :: Word32
| outRemaining < minimalBufferSize =
pure $ B.bufferFull minimalBufferSize op (go innerStep)
| otherwise =
-- execute inner builder with reduced boundaries
B.fillWithBuildStep innerStep doneH fullH insertChunkH brInner
where
outRemaining = ope `F.minusPtr` op
maxChunkSizeLength = word32HexLength $ fromIntegral outRemaining
!brInner@(BufferRange opInner _) = BufferRange
(op `F.plusPtr` (maxChunkSizeLength + crlfLength)) -- leave space for chunk header
(ope `F.plusPtr` (-maxAfterBufferOverhead)) -- leave space at end of data
doneH :: Ptr Word8 -> _x
-> IO (BuildSignal a)
doneH opInner' _ =
wrapChunk opInner' $ \op' ->
k $! BufferRange op' ope
fullH :: Ptr Word8 -> Int -> BuildStep _x
-> IO (BuildSignal a)
fullH opInner' minRequiredSize nextInnerStep =
wrapChunk opInner' $ \op' ->
pure $! B.bufferFull
(minRequiredSize + maxEncodingOverhead)
op'
(go nextInnerStep)
insertChunkH :: Ptr Word8 -> ByteString -> BuildStep _x
-> IO (BuildSignal a)
insertChunkH opInner' bs nextInnerStep =
wrapChunk opInner' $ \op' ->
if S.null bs -- flush
then pure $! B.insertChunk op' S.empty (go nextInnerStep)
else do -- insert non-empty bytestring
-- add header for inserted bytestring
-- FIXME: assert(S.length bs < maxBound :: Word32)
let chunkSize = fromIntegral $ S.length bs
!op'' <- writeWord32Hex NoPadding chunkSize op'
!op''' <- writeCRLF op''
-- insert bytestring and write CRLF in next buildstep
pure $! B.insertChunk
op''' bs
(B.runBuilderWith crlfBuilder $ go nextInnerStep)
-- 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
void $ writeWord32Hex (PadTo maxChunkSizeLength) chunkSize op
void $ writeCRLF (opInner `F.plusPtr` (-crlfLength))
void $ writeCRLF chunkDataEnd
mkSignal (chunkDataEnd `F.plusPtr` crlfLength)
-- | The zero-length chunk @0\\r\\n\\r\\n@ signalling the termination of the data transfer.
chunkedTransferTerminator :: Builder
chunkedTransferTerminator = B.byteStringCopy "0\r\n\r\n"
|
