{-# 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 :: 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')

------------------------------------------------------------------------------
-- 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 :: 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

-- | @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' :: 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))

-- | Length of the hex-string required to encode the given 'Word32'.
{-# 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)

------------------------------------------------------------------------------
-- Constants
------------------------------------------------------------------------------

crlfLength, maxW32HexLength, minimalChunkSize, maxBeforeBufferOverhead,
  maxAfterBufferOverhead, maxEncodingOverhead, minimalBufferSize :: Int
crlfLength :: Int
crlfLength = Int
2
maxW32HexLength :: Int
maxW32HexLength = Int
8 -- 4 bytes, 2 hex digits per byte
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

------------------------------------------------------------------------------
-- 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 :: 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)
          -- FIXME: Assert that outRemaining < maxBound :: Word32
          | 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 =
              -- execute inner builder with reduced boundaries
              (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)) -- leave space for chunk header
                (Ptr Word8
ope Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`F.plusPtr` (-Int
maxAfterBufferOverhead))         -- leave space at end of data

            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                      -- flush
                  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                           -- insert non-empty bytestring
                    -- add header for inserted bytestring
                    -- FIXME: assert(S.length bs < maxBound :: Word32)
                    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''
                    -- insert bytestring and write CRLF in next buildstep
                    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)

            -- 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 :: 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)


-- | The zero-length chunk @0\\r\\n\\r\\n@ signalling the termination of the data transfer.
chunkedTransferTerminator :: Builder
chunkedTransferTerminator :: Builder
chunkedTransferTerminator = ByteString -> Builder
B.byteStringCopy ByteString
"0\r\n\r\n"