HTTP Header builder

gistfile1.hs

{-# LANGUAGE BangPatterns, OverloadedStrings, FlexibleInstances, CPP #-}

module Main where

import Data.ByteString (ByteString)
import qualified Data.ByteString as S
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Lazy as L

import Blaze.ByteString.Builder (copyByteString, Builder, toLazyByteString)
import Blaze.ByteString.Builder.Char8 (fromChar, fromShow)
import Data.Monoid (mappend)

import Data.Word (Word8)
import Data.List (foldl')
import Control.Monad
import qualified Network.HTTP.Types as H
import qualified Data.CaseInsensitive as CI

import Data.ByteString.Internal (ByteString(..), unsafeCreate, memcpy)
import Foreign.ForeignPtr
import Foreign.Ptr
import GHC.Storable

import Criterion.Main

----------------------------------------------------------------

testHdrs :: H.ResponseHeaders
testHdrs = [
    ("Date", "Mon, 13 Aug 2012 04:22:55 GMT")
  , ("Content-Lenght", "151")
  , ("Server", "Mighttpd/2.5.8")
  , ("Last-Modified", "Fri, 22 Jun 2012 01:18:08 GMT")
  , ("Content-Type", "text/html")
  ]
           
main :: IO ()
main = defaultMain $ [
    bench "blaze builder" $ whnf blazeBuidler testHdrs
  , bench "my builder"    $ whnf myBuidler testHdrs
  , bench "lookup"        $ whnf (lookup "Foo") testHdrs
  ]

----------------------------------------------------------------

myBuidler :: H.ResponseHeaders -> ByteString
myBuidler hdr = headerBuilder H.http11 H.ok200 hdr

headerBuilder :: H.HttpVersion -> H.Status -> H.ResponseHeaders -> ByteString
headerBuilder !httpversion !status !responseHeaders = unsafeCreate len $ \ptr -> do
    ptr1 <- copyStatus ptr httpversion status
    ptr2 <- copyHeaders ptr1 responseHeaders
    void $ copyCRLF ptr2
  where
#ifdef BYPASS
    len = 172
#else
    !len = 17 + slen + foldl' fieldLength 0 responseHeaders
    fieldLength !l !(k,v) = l + S.length (CI.original k) + S.length v + 4
    !slen = S.length $ H.statusMessage status
#endif

{-# INLINE copy #-}
copy :: Ptr Word8 -> ByteString -> IO (Ptr Word8)
copy !ptr !(PS fp o l) = withForeignPtr fp $ \p -> do
    memcpy ptr (p `plusPtr` o) (fromIntegral l)
    return $! ptr `plusPtr` l

httpVer11 :: ByteString
httpVer11 = "HTTP/1.1 "

httpVer10 :: ByteString
httpVer10 = "HTTP/1.0 "

{-# INLINE copyStatus #-}
copyStatus :: Ptr Word8 -> H.HttpVersion -> H.Status -> IO (Ptr Word8)
copyStatus !ptr !httpversion !status = do
    ptr1 <- copy ptr httpVer
#ifdef BYPASS
    copy ptr1 "200"
#else
    writeWord8OffPtr ptr1 0 (zero + fromIntegral r2)
    writeWord8OffPtr ptr1 1 (zero + fromIntegral r1)
    writeWord8OffPtr ptr1 2 (zero + fromIntegral r0)
#endif
    writeWord8OffPtr ptr1 3 spc
    ptr2 <- copy (ptr1 `plusPtr` 4) (H.statusMessage status)
    copyCRLF ptr2
  where
    httpVer
      | httpversion == H.HttpVersion 1 1 = httpVer11
      | otherwise                        = httpVer10
#ifndef BYPASS
    (q0,r0) = H.statusCode status `divMod` 10
    (q1,r1) = q0 `divMod` 10
    r2      = q1 `mod` 10
#endif

{-# INLINE copyHeaders #-}
copyHeaders :: Ptr Word8 -> [H.Header] -> IO (Ptr Word8)
copyHeaders !ptr [] = return ptr
copyHeaders !ptr (h:hs) = do
    ptr1 <- copyHeader ptr h
    copyHeaders ptr1 hs

{-# INLINE copyHeader #-}
copyHeader :: Ptr Word8 -> H.Header -> IO (Ptr Word8)
copyHeader !ptr (k,v) = do
    ptr1 <- copy ptr (CI.original k)
    writeWord8OffPtr ptr1 0 colon
    writeWord8OffPtr ptr1 1 spc
    ptr2 <- copy (ptr1 `plusPtr` 2) v
    copyCRLF ptr2

{-# INLINE copyCRLF #-}
copyCRLF :: Ptr Word8 -> IO (Ptr Word8)
copyCRLF !ptr = do
    writeWord8OffPtr ptr 0 cr
    writeWord8OffPtr ptr 1 lf
    return $! ptr `plusPtr` 2

zero :: Word8
zero = 48
spc :: Word8
spc = 32
colon :: Word8
colon = 58
cr :: Word8
cr = 13
lf :: Word8
lf = 10

----------------------------------------------------------------

blazeBuidler :: H.ResponseHeaders -> ByteString
blazeBuidler hdr = S.concat $ L.toChunks $ toLazyByteString $ headers H.http11 H.ok200 hdr False

httpBuilder, spaceBuilder, newlineBuilder, transferEncodingBuilder
           , colonSpaceBuilder :: Builder
httpBuilder = copyByteString "HTTP/"
spaceBuilder = fromChar ' '
newlineBuilder = copyByteString "\r\n"
transferEncodingBuilder = copyByteString "Transfer-Encoding: chunked\r\n\r\n"
colonSpaceBuilder = copyByteString ": "

headers :: H.HttpVersion -> H.Status -> H.ResponseHeaders -> Bool -> Builder
headers !httpversion !status !responseHeaders !isChunked' = {-# SCC "headers" #-}
    let !start = httpBuilder
                `mappend` copyByteString
                            (case httpversion of
                                H.HttpVersion 1 1 -> "1.1"
                                _ -> "1.0")
                `mappend` spaceBuilder
                `mappend` fromShow (H.statusCode status)
                `mappend` spaceBuilder
                `mappend` copyByteString (H.statusMessage status)
                `mappend` newlineBuilder
        !start' = foldl' responseHeaderToBuilder start (serverHeader responseHeaders)
        !end = if isChunked'
                 then transferEncodingBuilder
                 else newlineBuilder
    in start' `mappend` end

responseHeaderToBuilder :: Builder -> H.Header -> Builder
responseHeaderToBuilder b (x, y) = b
  `mappend` copyByteString (CI.original x)
  `mappend` colonSpaceBuilder
  `mappend` copyByteString y
  `mappend` newlineBuilder

serverHeader :: H.RequestHeaders -> H.RequestHeaders
serverHeader hdrs = case lookup key hdrs of
    Nothing -> server : hdrs
    Just _  -> hdrs
  where
    key = "Server"
    server = (key, severVersion)
    
warpVersion :: String
warpVersion = "1.3"    

severVersion :: ByteString
severVersion = B.pack $ "Warp/" ++ warpVersion