Skip to content

Commit

Permalink
Move response fixup code from snap-server to snap-core.
Browse files Browse the repository at this point in the history 
This ensures that as far as HTTP headers go, the response generated by
Snap.Test.runHandler is as close as possible to what will be sent out over the
socket (excepting chunked transfer-encoding).

Fixes #145.
  • Loading branch information
@gregorycollins
gregorycollins committed Jul 5, 2012
1 parent 4473e1c commit 64a9e77
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 118 deletions.
128 changes: 38 additions & 90 deletions src/Snap/Internal/Http/Server.hs
View file
Expand Up @@ -34,7 +34,8 @@ import Data.Int
import Data.IORef
import Data.List (foldl')
import qualified Data.Map as Map
import Data.Maybe (catMaybes, fromJust, fromMaybe, isJust)
import Data.Maybe ( catMaybes, fromJust, fromMaybe, isJust
, isNothing )
import Data.Monoid
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
Expand All @@ -46,8 +47,6 @@ import Network.Socket (withSocketsDo, Socket)
import Prelude hiding (catch)
import System.IO
import System.Locale
import System.PosixCompat.Files hiding (setFileSize)
import System.Posix.Types (FileOffset)
------------------------------------------------------------------------------
import System.FastLogger (timestampedLogEntry, combinedLogEntry)
import Snap.Internal.Http.Types
Expand Down Expand Up @@ -430,7 +429,7 @@ httpSession defaultTimeout writeEnd' buffer onSendFile tickle handler = do
then id
else insHeader
let rsp' = updateHeaders ins rsp
(bytesSent,_) <- sendResponse req rsp' buffer writeEnd onSendFile
(bytesSent,_) <- sendResponse rsp' buffer writeEnd onSendFile
`catch` errCatch "sending response" req

debug $ "Server.httpSession: sent " ++
Expand Down Expand Up @@ -715,19 +714,20 @@ receiveRequest writeEnd = do

------------------------------------------------------------------------------
-- Response must be well-formed here
sendResponse :: forall a . Request
-> Response
sendResponse :: forall a . Response
-> Buffer
-> Iteratee ByteString IO a -- ^ iteratee write end
-> (FilePath -> Int64 -> Int64 -> IO a) -- ^ function to call on
-- sendfile
-> ServerMonad (Int64, a)
sendResponse req rsp' buffer writeEnd' onSendFile = do
let rsp'' = renderCookies rsp'
(rsp, shouldClose) <- liftIO $ fixupResponse rsp''
sendResponse rsp0 buffer writeEnd' onSendFile = do
let rsp1 = renderCookies rsp0

let (rsp, shouldClose) = if isNothing $ rspContentLength rsp1
then noCL rsp1
else (rsp1, False)

when shouldClose $
modify $! \s -> s { _forceConnectionClose = True }
when shouldClose $ modify $! \s -> s { _forceConnectionClose = True }

let (!headerString,!hlen) = mkHeaderBuilder rsp
let writeEnd = fixCLIteratee hlen rsp writeEnd'
Expand All @@ -746,6 +746,32 @@ sendResponse req rsp' buffer writeEnd' onSendFile = do
return $! (bs,x)

where
--------------------------------------------------------------------------
noCL :: Response -> (Response, Bool)
noCL r =
if rspHttpVersion r >= (1,1)
then
let r' = setHeader "Transfer-Encoding" "chunked" r
origE = rspBodyToEnum $ rspBody r
e = \i -> joinI $ origE $$ chunkIt i
in (r' { rspBody = Enum e }, False)
else
-- HTTP/1.0 and no content-length? We'll have to close the
-- socket.
(setHeader "Connection" "close" r, True)
{-# INLINE noCL #-}


--------------------------------------------------------------------------
chunkIt :: forall x . Enumeratee Builder Builder IO x
chunkIt = checkDone $ continue . step
where
step k EOF = k (Chunks [chunkedTransferTerminator]) >>== return
step k (Chunks []) = continue $ step k
step k (Chunks xs) = k (Chunks [chunkedTransferEncoding $ mconcat xs])
>>== chunkIt


--------------------------------------------------------------------------
whenEnum :: Iteratee ByteString IO a
-> Builder
Expand Down Expand Up @@ -807,7 +833,7 @@ sendResponse req rsp' buffer writeEnd' onSendFile = do


--------------------------------------------------------------------------
(major,minor) = rspHttpVersion rsp'
(major,minor) = rspHttpVersion rsp0


--------------------------------------------------------------------------
Expand Down Expand Up @@ -835,31 +861,6 @@ sendResponse req rsp' buffer writeEnd' onSendFile = do
klen = S.length k' + 2


--------------------------------------------------------------------------
noCL :: Response -> (Response, Bool)
noCL r =
-- are we in HTTP/1.1?
let sendChunked = (rspHttpVersion r) == (1,1)
in if sendChunked
then
let r' = setHeader "Transfer-Encoding" "chunked" r
origE = rspBodyToEnum $ rspBody r
e = \i -> joinI $ origE $$ chunkIt i
in (r' { rspBody = Enum e }, False)
else
-- HTTP/1.0 and no content-length? We'll have to close the
-- socket.
(setHeader "Connection" "close" r, True)

--------------------------------------------------------------------------
chunkIt :: forall x . Enumeratee Builder Builder IO x
chunkIt = checkDone $ continue . step
where
step k EOF = k (Chunks [chunkedTransferTerminator]) >>== return
step k (Chunks []) = continue $ step k
step k (Chunks xs) = k (Chunks [chunkedTransferEncoding $ mconcat xs])
>>== chunkIt

--------------------------------------------------------------------------
fixCLIteratee :: Int -- ^ header length
-> Response -- ^ response
Expand All @@ -874,24 +875,6 @@ sendResponse req rsp' buffer writeEnd' onSendFile = do

mbCL = rspContentLength resp

--------------------------------------------------------------------------
hasCL :: Int64 -> Response -> Response
hasCL cl r = setHeader "Content-Length" (toByteString $ fromShow cl) r


--------------------------------------------------------------------------
setFileSize :: FilePath -> Response -> IO Response
setFileSize fp r = {-# SCC "setFileSize" #-} do
fs <- liftM fromIntegral $ getFileSize fp
return $! r { rspContentLength = Just fs }


--------------------------------------------------------------------------
handle304 :: Response -> Response
handle304 r = setResponseBody (enumBuilder mempty) $
updateHeaders (H.delete "Transfer-Encoding") $
setContentLength 0 r


--------------------------------------------------------------------------
renderCookies :: Response -> Response
Expand All @@ -903,36 +886,6 @@ sendResponse req rsp' buffer writeEnd' onSendFile = do
cookies = fmap cookieToBS . Map.elems $ rspCookies r


--------------------------------------------------------------------------
fixupResponse :: Response
-> IO (Response, Bool) -- ^ if the Bool is true then we
-- should close the connection.
fixupResponse r = {-# SCC "fixupResponse" #-} do
let r' = deleteHeader "Content-Length" r
let code = rspStatus r'
let r'' = if code == 204 || code == 304
then handle304 r'
else r'

(r''', shouldClose) <- do
z <- case rspBody r'' of
(Enum _) -> return r''
(SendFile f Nothing) -> setFileSize f r''
(SendFile _ (Just (s,e))) -> return $!
setContentLength (e-s) r''

return $! case rspContentLength z of
Nothing -> noCL z
(Just sz) -> (hasCL sz z, False)

-- HEAD requests cannot have bodies per RFC 2616 sec. 9.4
r'''' <- if rqMethod req == HEAD
then return $! deleteHeader "Transfer-Encoding" $
r''' { rspBody = Enum $ enumBuilder mempty }
else return $! r'''
return $! (r'''', shouldClose)


--------------------------------------------------------------------------
mkHeaderBuilder :: Response -> (Builder,Int)
mkHeaderBuilder r = {-# SCC "mkHeaderBuilder" #-}
Expand Down Expand Up @@ -1002,11 +955,6 @@ cookieToBS (Cookie k v mbExpTime mbDomain mbPath isSec isHOnly) = cookie
"%a, %d-%b-%Y %H:%M:%S GMT"


------------------------------------------------------------------------------
getFileSize :: FilePath -> IO FileOffset
getFileSize fp = liftM fileSize $ getFileStatus fp


------------------------------------------------------------------------------
l2s :: L.ByteString -> S.ByteString
l2s = S.concat . L.toChunks
Expand Down
48 changes: 20 additions & 28 deletions test/suite/Snap/Internal/Http/Server/Tests.hs
View file
Expand Up @@ -160,13 +160,6 @@ dummyIter :: Iteratee ByteString IO ()
dummyIter = consume >> return ()


mkRequest :: ByteString -> IO Request
mkRequest s = do
step <- runIteratee $ liftM fromJust $ rsm $ receiveRequest dummyIter
let iter = enumBS s step
run_ iter


testReceiveRequest :: Iteratee ByteString IO (Request,L.ByteString)
testReceiveRequest = do
r <- liftM fromJust $ rsm $ receiveRequest dummyIter
Expand Down Expand Up @@ -383,11 +376,10 @@ rsm = runServerMonad "localhost" (SessionInfo "127.0.0.1" 80 "127.0.0.1" 58382 F

testHttpResponse1 :: Test
testHttpResponse1 = testCase "server/HttpResponse1" $ do
req <- mkRequest sampleRequest
buf <- allocBuffer 16384

b <- run_ $ rsm $
sendResponse req rsp1 buf copyingStream2Stream testOnSendFile >>=
sendResponse rsp1 buf copyingStream2Stream testOnSendFile >>=
return . snd

assertBool "http response" (b == text1 || b == text2)
Expand All @@ -406,7 +398,7 @@ testHttpResponse1 = testCase "server/HttpResponse1" $ do
]

rsp1 = updateHeaders (H.insert "Foo" "Bar") $
setContentLength 10 $
setContentLength' 10 $
setResponseStatus 600 "Test" $
modifyResponseBody (>==> (enumBuilder $
fromByteString "0123456789")) $
Expand Down Expand Up @@ -435,10 +427,9 @@ testOnSendFile f st sz = do

testHttpResponse2 :: Test
testHttpResponse2 = testCase "server/HttpResponse2" $ do
req <- mkRequest sampleRequest
buf <- allocBuffer 16384
b2 <- liftM (S.concat . L.toChunks) $ run_ $ rsm $
sendResponse req rsp2 buf copyingStream2Stream testOnSendFile >>=
sendResponse rsp2 buf copyingStream2Stream testOnSendFile >>=
return . snd

assertBool "http prefix"
Expand All @@ -455,7 +446,7 @@ testHttpResponse2 = testCase "server/HttpResponse2" $ do

where
rsp1 = updateHeaders (H.insert "Foo" "Bar") $
setContentLength 10 $
setContentLength' 10 $
setResponseStatus 600 "Test" $
modifyResponseBody (>==> (enumBuilder $
fromByteString "0123456789")) $
Expand All @@ -466,11 +457,10 @@ testHttpResponse2 = testCase "server/HttpResponse2" $ do

testHttpResponse3 :: Test
testHttpResponse3 = testCase "server/HttpResponse3" $ do
req <- mkRequest sampleRequest
buf <- allocBuffer 16384

b3 <- run_ $ rsm $
sendResponse req rsp3 buf copyingStream2Stream testOnSendFile >>=
sendResponse rsp3 buf copyingStream2Stream testOnSendFile >>=
return . snd

let lns = LC.lines b3
Expand Down Expand Up @@ -498,24 +488,22 @@ testHttpResponse3 = testCase "server/HttpResponse3" $ do
hdrs = strToHeaders s

rsp1 = updateHeaders (H.insert "Foo" "Bar") $
setContentLength 10 $
setContentLength' 10 $
setResponseStatus 600 "Test" $
modifyResponseBody (>==> (enumBuilder $
fromByteString "0123456789")) $
setResponseBody returnI $
emptyResponse { rspHttpVersion = (1,0) }
rsp2 = rsp1 { rspContentLength = Nothing }
rsp2 = deleteHeader "Content-Length" $ rsp1 { rspContentLength = Nothing }
rsp3 = setContentType "text/plain" $ (rsp2 { rspHttpVersion = (1,1) })


testHttpResponse4 :: Test
testHttpResponse4 = testCase "server/HttpResponse4" $ do
buf <- allocBuffer 16384

req <- mkRequest sampleRequest

b <- run_ $ rsm $
sendResponse req rsp1 buf copyingStream2Stream testOnSendFile >>=
sendResponse rsp1 buf copyingStream2Stream testOnSendFile >>=
return . snd

assertEqual "http response" (L.concat [
Expand All @@ -525,16 +513,15 @@ testHttpResponse4 = testCase "server/HttpResponse4" $ do

where
rsp1 = setResponseStatus 304 "Test" $
setContentLength' 0 $
emptyResponse { rspHttpVersion = (1,0) }


testHttpResponseCookies :: Test
testHttpResponseCookies = testCase "server/HttpResponseCookies" $ do
buf <- allocBuffer 16384
req <- mkRequest sampleRequest

b <- run_ $ rsm $
sendResponse req rsp2 buf copyingStream2Stream testOnSendFile >>=
sendResponse rsp2 buf copyingStream2Stream testOnSendFile >>=
return . snd

let lns = LC.lines b
Expand All @@ -550,7 +537,7 @@ testHttpResponseCookies = testCase "server/HttpResponseCookies" $ do
assertBool "http response" ok

where
check s = (H.lookup "Content-Length" hdrs == Just ["0\r"]) &&
check s = (H.lookup "Connection" hdrs == Just ["close\r"]) &&
(ch $ H.lookup "Set-Cookie" hdrs)
where
hdrs = strToHeaders s
Expand Down Expand Up @@ -589,9 +576,9 @@ echoServer _ _ req = do
liftIO $ writeIORef (rqBody req) (SomeEnumerator $ joinI . I.take 0)
return (req, rsp b cl)
where
rsp s cl = emptyResponse { rspBody = Enum $
enumBuilder (fromLazyByteString s)
, rspContentLength = Just $ fromIntegral cl }
rsp s cl = setContentLength' cl $
emptyResponse { rspBody = Enum $
enumBuilder (fromLazyByteString s) }


echoServer2 :: ServerHandler
Expand Down Expand Up @@ -928,7 +915,7 @@ testExpectGarbage = testCase "server/expectGarbage" $ do
pongServer :: Snap ()
pongServer = modifyResponse $ setResponseBody enum .
setContentType "text/plain" .
setContentLength 4
setContentLength' 4
where
enum = enumBuilder $ fromByteString "PONG"

Expand Down Expand Up @@ -1069,3 +1056,8 @@ copyingStream2Stream = go []
maybe (return $ L.fromChunks $ reverse l)
(\x -> let !z = S.copy x in go (z:l))
mbx


setContentLength' :: Int64 -> Response -> Response
setContentLength' cl = setHeader "Content-Length" (S.pack $ show cl) .
setContentLength cl

0 comments on commit 64a9e77

Please sign in to comment.