From 64a9e779d79f9d126121544ee9f4d56e8fe1cd6d Mon Sep 17 00:00:00 2001 From: Gregory Collins Date: Thu, 5 Jul 2012 20:46:41 +0200 Subject: [PATCH] Move response fixup code from snap-server to snap-core. 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. --- src/Snap/Internal/Http/Server.hs | 128 ++++++------------ test/suite/Snap/Internal/Http/Server/Tests.hs | 48 +++---- 2 files changed, 58 insertions(+), 118 deletions(-) diff --git a/src/Snap/Internal/Http/Server.hs b/src/Snap/Internal/Http/Server.hs index e46c6dd4..08a8f1ed 100644 --- a/src/Snap/Internal/Http/Server.hs +++ b/src/Snap/Internal/Http/Server.hs @@ -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 @@ -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 @@ -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 " ++ @@ -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' @@ -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 @@ -807,7 +833,7 @@ sendResponse req rsp' buffer writeEnd' onSendFile = do -------------------------------------------------------------------------- - (major,minor) = rspHttpVersion rsp' + (major,minor) = rspHttpVersion rsp0 -------------------------------------------------------------------------- @@ -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 @@ -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 @@ -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" #-} @@ -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 diff --git a/test/suite/Snap/Internal/Http/Server/Tests.hs b/test/suite/Snap/Internal/Http/Server/Tests.hs index cc43a113..49764199 100644 --- a/test/suite/Snap/Internal/Http/Server/Tests.hs +++ b/test/suite/Snap/Internal/Http/Server/Tests.hs @@ -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 @@ -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) @@ -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")) $ @@ -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" @@ -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")) $ @@ -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 @@ -498,13 +488,13 @@ 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) }) @@ -512,10 +502,8 @@ 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 [ @@ -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 @@ -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 @@ -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 @@ -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" @@ -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