{-# LANGUAGE BangPatterns, CPP, MagicHash, OverloadedStrings #-}
------------------------------------------------------------------------------
-- |
-- Module:      Blaze.ByteString.Builder.HTTP
-- Copyright:   (c) 2013 Simon Meier
-- License:     BSD3
-- Maintainer:  https://github.com/blaze-builder
-- Stability:   stable
--
-- Support for HTTP response encoding.
--
------------------------------------------------------------------------------

module Blaze.ByteString.Builder.HTTP (
  -- * Chunked HTTP transfer encoding
    chunkedTransferEncoding
  , chunkedTransferTerminator
  ) where

#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
#include "MachDeps.h"
#endif

#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
import GHC.Base
import GHC.Word (Word32(..))
#else
import Data.Word
#endif

import Foreign

import qualified Data.ByteString       as S
import Data.ByteString.Char8 ()

import Blaze.ByteString.Builder.Internal.Write
import Data.ByteString.Builder
import Data.ByteString.Builder.Internal
import Blaze.ByteString.Builder.ByteString (copyByteString)

import qualified Blaze.ByteString.Builder.Char8 as Char8

#if !MIN_VERSION_base(4,8,0)
import Data.Monoid
#endif


{-# INLINE shiftr_w32 #-}
shiftr_w32 :: Word32 -> Int -> Word32

#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
#if MIN_VERSION_ghc_prim(0,8,0)
shiftr_w32 :: Word32 -> Int -> Word32
shiftr_w32 (W32# Word32#
w) (I# Int#
i) = Word32# -> Word32
W32# (Word# -> Word32#
wordToWord32# ((Word32# -> Word#
word32ToWord# Word32#
w) Word# -> Int# -> Word#
`uncheckedShiftRL#` Int#
i))
#else
shiftr_w32 (W32# w) (I# i) = W32# (w `uncheckedShiftRL#`   i)
#endif
#else
shiftr_w32 = shiftR
#endif


-- | Write a CRLF sequence.
writeCRLF :: Write
writeCRLF :: Write
writeCRLF = Char -> Write
Char8.writeChar Char
'\r' forall a. Monoid a => a -> a -> a
`mappend` Char -> Write
Char8.writeChar Char
'\n'
{-# INLINE writeCRLF #-}

-- | Execute a write
{-# INLINE execWrite #-}
execWrite :: Write -> Ptr Word8 -> IO ()
execWrite :: Write -> Ptr Word8 -> IO ()
execWrite Write
w Ptr Word8
op = do
    Ptr Word8
_ <- Poke -> Ptr Word8 -> IO (Ptr Word8)
runPoke (Write -> Poke
getPoke Write
w) Ptr Word8
op
    forall (m :: * -> *) a. Monad m => a -> m a
return ()


------------------------------------------------------------------------------
-- Hex Encoding Infrastructure
------------------------------------------------------------------------------

pokeWord32HexN :: Int -> Word32 -> Ptr Word8 -> IO ()
pokeWord32HexN :: Int -> Word32 -> Ptr Word8 -> IO ()
pokeWord32HexN Int
n0 Word32
w0 Ptr Word8
op0 =
    Word32 -> Ptr Word8 -> IO ()
go Word32
w0 (Ptr Word8
op0 forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
n0 forall a. Num a => a -> a -> a
- Int
1))
  where
    go :: Word32 -> Ptr Word8 -> IO ()
go !Word32
w !Ptr Word8
op
      | Ptr Word8
op forall a. Ord a => a -> a -> Bool
< Ptr Word8
op0  = forall (m :: * -> *) a. Monad m => a -> m a
return ()
      | Bool
otherwise = do
          let nibble :: Word8
              nibble :: Word8
nibble = forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
w forall a. Bits a => a -> a -> a
.&. Word8
0xF
              hex :: Word8
hex | Word8
nibble forall a. Ord a => a -> a -> Bool
< Word8
10 = Word8
48 forall a. Num a => a -> a -> a
+ Word8
nibble
                  | Bool
otherwise   = Word8
55 forall a. Num a => a -> a -> a
+ Word8
nibble
          forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr Word8
op Word8
hex
          Word32 -> Ptr Word8 -> IO ()
go (Word32
w Word32 -> Int -> Word32
`shiftr_w32` Int
4) (Ptr Word8
op forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (-Int
1))
{-# INLINE pokeWord32HexN #-}

iterationsUntilZero :: Integral a => (a -> a) -> a -> Int
iterationsUntilZero :: forall a. Integral a => (a -> a) -> a -> Int
iterationsUntilZero a -> a
f = forall {t}. Num t => t -> a -> t
go Int
0
  where
    go :: t -> a -> t
go !t
count a
0  = t
count
    go !t
count !a
x = t -> a -> t
go (t
countforall a. Num a => a -> a -> a
+t
1) (a -> a
f a
x)
{-# INLINE iterationsUntilZero #-}

-- | Length of the hex-string required to encode the given 'Word32'.
word32HexLength :: Word32 -> Int
word32HexLength :: Word32 -> Int
word32HexLength = forall a. Ord a => a -> a -> a
max Int
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Integral a => (a -> a) -> a -> Int
iterationsUntilZero (Word32 -> Int -> Word32
`shiftr_w32` Int
4)
{-# INLINE word32HexLength #-}

writeWord32Hex :: Word32 -> Write
writeWord32Hex :: Word32 -> Write
writeWord32Hex Word32
w =
    Int -> Poke -> Write
boundedWrite (Int
2 forall a. Num a => a -> a -> a
* forall a. Storable a => a -> Int
sizeOf Word32
w) (Int -> (Ptr Word8 -> IO ()) -> Poke
pokeN Int
len forall a b. (a -> b) -> a -> b
$ Int -> Word32 -> Ptr Word8 -> IO ()
pokeWord32HexN Int
len Word32
w)
  where
    len :: Int
len = Word32 -> Int
word32HexLength Word32
w
{-# INLINE writeWord32Hex #-}


------------------------------------------------------------------------------
-- Chunked transfer encoding
------------------------------------------------------------------------------

-- | Transform a builder such that it uses chunked HTTP transfer encoding.
chunkedTransferEncoding :: Builder -> Builder
chunkedTransferEncoding :: Builder -> Builder
chunkedTransferEncoding Builder
innerBuilder =
    (forall r. BuildStep r -> BuildStep r) -> Builder
builder forall r. BuildStep r -> BuildStep r
transferEncodingStep
  where
    transferEncodingStep :: (BufferRange -> IO (BuildSignal a))
-> BufferRange -> IO (BuildSignal a)
transferEncodingStep BufferRange -> IO (BuildSignal a)
k =
        forall {a}.
(BufferRange -> IO (BuildSignal a))
-> BufferRange -> IO (BuildSignal a)
go (Builder -> BuildStep ()
runBuilder Builder
innerBuilder)
      where
        go :: (BufferRange -> IO (BuildSignal a))
-> BufferRange -> IO (BuildSignal a)
go BufferRange -> IO (BuildSignal a)
innerStep !(BufferRange Ptr Word8
op Ptr Word8
ope)
          -- FIXME: Assert that outRemaining < maxBound :: Word32
          | Int
outRemaining forall a. Ord a => a -> a -> Bool
< Int
minimalBufferSize =
              forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
bufferFull Int
minimalBufferSize Ptr Word8
op ((BufferRange -> IO (BuildSignal a))
-> BufferRange -> IO (BuildSignal a)
go BufferRange -> IO (BuildSignal a)
innerStep)
          | Bool
otherwise = do
              let !brInner :: BufferRange
brInner@(BufferRange Ptr Word8
opInner Ptr Word8
_) = Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange
                     (Ptr Word8
op  forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
chunkSizeLength forall a. Num a => a -> a -> a
+ Int
2))     -- leave space for chunk header
                     (Ptr Word8
ope forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (-Int
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 :: forall a.
Ptr Word8
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
wrapChunk !Ptr Word8
opInner' Ptr Word8 -> IO (BuildSignal a)
mkSignal
                    | Ptr Word8
opInner' forall a. Eq a => a -> a -> Bool
== Ptr Word8
opInner = Ptr Word8 -> IO (BuildSignal a)
mkSignal Ptr Word8
op
                    | Bool
otherwise           = do
                        Int -> Word32 -> Ptr Word8 -> IO ()
pokeWord32HexN Int
chunkSizeLength
                            (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Ptr Word8
opInner' forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
opInner)
                            Ptr Word8
op
                        Write -> Ptr Word8 -> IO ()
execWrite Write
writeCRLF (Ptr Word8
opInner forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (-Int
2))
                        Write -> Ptr Word8 -> IO ()
execWrite Write
writeCRLF Ptr Word8
opInner'
                        Ptr Word8 -> IO (BuildSignal a)
mkSignal (Ptr Word8
opInner' forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
2)

                  -- prepare handlers
                  doneH :: Ptr Word8 -> p -> IO (BuildSignal a)
doneH Ptr Word8
opInner' p
_ = forall a.
Ptr Word8
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
wrapChunk Ptr Word8
opInner' forall a b. (a -> b) -> a -> b
$ \Ptr Word8
op' -> do
                                         let !br' :: BufferRange
br' = Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
op' Ptr Word8
ope
                                         BufferRange -> IO (BuildSignal a)
k BufferRange
br'

                  fullH :: Ptr Word8
-> Int -> (BufferRange -> IO (BuildSignal a)) -> IO (BuildSignal a)
fullH Ptr Word8
opInner' Int
minRequiredSize BufferRange -> IO (BuildSignal a)
nextInnerStep =
                      forall a.
Ptr Word8
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
wrapChunk Ptr Word8
opInner' forall a b. (a -> b) -> a -> b
$ \Ptr Word8
op' ->
                        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
bufferFull
                          (Int
minRequiredSize forall a. Num a => a -> a -> a
+ Int
maxEncodingOverhead)
                          Ptr Word8
op'
                          ((BufferRange -> IO (BuildSignal a))
-> BufferRange -> IO (BuildSignal a)
go BufferRange -> IO (BuildSignal a)
nextInnerStep)

                  insertChunkH :: Ptr Word8
-> ByteString
-> (BufferRange -> IO (BuildSignal a))
-> IO (BuildSignal a)
insertChunkH Ptr Word8
opInner' ByteString
bs BufferRange -> IO (BuildSignal a)
nextInnerStep
                    | ByteString -> Bool
S.null ByteString
bs =                         -- flush
                        forall a.
Ptr Word8
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
wrapChunk Ptr Word8
opInner' forall a b. (a -> b) -> a -> b
$ \Ptr Word8
op' ->
                          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. Ptr Word8 -> ByteString -> BuildStep a -> BuildSignal a
insertChunk Ptr Word8
op' ByteString
S.empty ((BufferRange -> IO (BuildSignal a))
-> BufferRange -> IO (BuildSignal a)
go BufferRange -> IO (BuildSignal a)
nextInnerStep)

                    | Bool
otherwise =                         -- insert non-empty bytestring
                        forall a.
Ptr Word8
-> (Ptr Word8 -> IO (BuildSignal a)) -> IO (BuildSignal a)
wrapChunk Ptr Word8
opInner' forall a b. (a -> b) -> a -> b
$ \Ptr Word8
op' -> do
                          -- add header for inserted bytestring
                          -- FIXME: assert(S.length bs < maxBound :: Word32)
                          !Ptr Word8
op'' <- (Poke -> Ptr Word8 -> IO (Ptr Word8)
`runPoke` Ptr Word8
op') forall a b. (a -> b) -> a -> b
$ Write -> Poke
getPoke forall a b. (a -> b) -> a -> b
$
                              Word32 -> Write
writeWord32Hex (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ ByteString -> Int
S.length ByteString
bs)
                              forall a. Monoid a => a -> a -> a
`mappend` Write
writeCRLF

                          -- insert bytestring and write CRLF in next buildstep
                          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. Ptr Word8 -> ByteString -> BuildStep a -> BuildSignal a
insertChunk
                            Ptr Word8
op'' ByteString
bs
                            (forall a. Builder -> BuildStep a -> BuildStep a
runBuilderWith (Write -> Builder
fromWrite Write
writeCRLF) forall a b. (a -> b) -> a -> b
$ (BufferRange -> IO (BuildSignal a))
-> BufferRange -> IO (BuildSignal a)
go BufferRange -> IO (BuildSignal a)
nextInnerStep)

              -- execute inner builder with reduced boundaries
              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
fillWithBuildStep BufferRange -> IO (BuildSignal a)
innerStep forall {p}. Ptr Word8 -> p -> IO (BuildSignal a)
doneH Ptr Word8
-> Int -> (BufferRange -> IO (BuildSignal a)) -> IO (BuildSignal a)
fullH Ptr Word8
-> ByteString
-> (BufferRange -> IO (BuildSignal a))
-> IO (BuildSignal a)
insertChunkH BufferRange
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 :: Int
minimalChunkSize  = Int
1

            -- overhead computation
            maxBeforeBufferOverhead :: Int
maxBeforeBufferOverhead = forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: Int) forall a. Num a => a -> a -> a
+ Int
2 -- max chunk size and CRLF after header
            maxAfterBufferOverhead :: Int
maxAfterBufferOverhead  = Int
2 forall a. Num a => a -> a -> a
+                           -- CRLF after data
                                      forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: Int) forall a. Num a => a -> a -> a
+ Int
2 -- max bytestring size, CRLF after header

            maxEncodingOverhead :: Int
maxEncodingOverhead = Int
maxBeforeBufferOverhead forall a. Num a => a -> a -> a
+ Int
maxAfterBufferOverhead

            minimalBufferSize :: Int
minimalBufferSize = Int
minimalChunkSize forall a. Num a => a -> a -> a
+ Int
maxEncodingOverhead

            -- remaining and required space computation
            outRemaining :: Int
            outRemaining :: Int
outRemaining    = Ptr Word8
ope forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
op
            chunkSizeLength :: Int
chunkSizeLength = Word32 -> Int
word32HexLength forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
outRemaining


-- | The zero-length chunk '0\r\n\r\n' signaling the termination of the data transfer.
chunkedTransferTerminator :: Builder
chunkedTransferTerminator :: Builder
chunkedTransferTerminator = ByteString -> Builder
copyByteString ByteString
"0\r\n\r\n"