on master help sign in
Using more enumerator functions, added test suite
8d1ba29Michael Snoyman
Network/Wai/Handler/Warp.hs+48−73
@@ -20,6 +20,11 @@ module Network.Wai.Handler.Warp
20 ( run
20 ( run
21 , sendResponse
21 , sendResponse
22 , parseRequest
22 , parseRequest
23+#if TEST
24+ , takeLineMax
25+ , takeUntilBlank
26+ , InvalidRequest (..)
27+#endif
23 ) where
28 ) where
24
29
25import Prelude hiding (catch)
30import Prelude hiding (catch)
@@ -47,6 +52,7 @@ import Data.Typeable (Typeable)
47import Data.Enumerator (($$), (>>==))
52import Data.Enumerator (($$), (>>==))
48import qualified Data.Enumerator as E
53import qualified Data.Enumerator as E
49import qualified Data.Enumerator.List as EL
54import qualified Data.Enumerator.List as EL
55+import qualified Data.Enumerator.Binary as EB
50import Blaze.ByteString.Builder.Enumerator (builderToByteString)
56import Blaze.ByteString.Builder.Enumerator (builderToByteString)
51import Blaze.ByteString.Builder.HTTP
57import Blaze.ByteString.Builder.HTTP
52 (chunkedTransferEncoding, chunkedTransferTerminator)
58 (chunkedTransferEncoding, chunkedTransferTerminator)
@@ -55,6 +61,7 @@ import Blaze.ByteString.Builder
55import Blaze.ByteString.Builder.Char8 (fromChar, fromShow)
61import Blaze.ByteString.Builder.Char8 (fromChar, fromShow)
56import Data.Monoid (mconcat, mappend)
62import Data.Monoid (mconcat, mappend)
57import Network.Socket.SendFile (sendFile)
63import Network.Socket.SendFile (sendFile)
64+import Network.Socket.Enumerator (iterSocket)
58
65
59import Control.Monad.IO.Class (liftIO)
66import Control.Monad.IO.Class (liftIO)
60import System.Timeout (timeout)
67import System.Timeout (timeout)
@@ -103,38 +110,6 @@ maxHeaderLength = 1024
103bytesPerRead = 4096
110bytesPerRead = 4096
104readTimeout = 3000000
111readTimeout = 3000000
105
112
106−takeUntilBlank :: Int
107− -> ([ByteString] -> [ByteString])
108− -> E.Iteratee S.ByteString IO [ByteString]
109−takeUntilBlank count _
110− | count > maxHeaders = E.throwError TooManyHeaders
111−takeUntilBlank count front = do
112− l <- takeLine 0 id
113− if B.null l
114− then return $ front []
115− else takeUntilBlank (count + 1) $ front . (:) l
116−
117−takeLine :: Int
118− -> ([ByteString] -> [ByteString])
119− -> E.Iteratee ByteString IO ByteString
120−takeLine len front = do
121− mbs <- EL.head
122− case mbs of
123− Nothing -> E.throwError IncompleteHeaders
124− Just bs -> do
125− let (x, y) = S.breakByte 10 bs
126− x' = if S.length x > 0 && S.last x == 13
127− then S.init x
128− else x
129− let len' = len + B.length x
130− case () of
131− ()
132− | len' > maxHeaderLength -> E.throwError OverLargeHeader
133− | B.null y -> takeLine len' $ front . (:) x
134− | otherwise -> do
135− E.yield () $ E.Chunks [B.drop 1 y]
136− return $ B.concat $ front [x']
137−
138data InvalidRequest =
113data InvalidRequest =
139 NotEnoughLines [String]
114 NotEnoughLines [String]
140 | BadFirstLine String
115 | BadFirstLine String
@@ -143,7 +118,7 @@ data InvalidRequest =
143 | IncompleteHeaders
118 | IncompleteHeaders
144 | OverLargeHeader
119 | OverLargeHeader
145 | SocketTimeout
120 | SocketTimeout
146− deriving (Show, Typeable)
121+ deriving (Show, Typeable, Eq)
147instance Exception InvalidRequest
122instance Exception InvalidRequest
148
123
149-- | Parse a set of header lines and body into a 'Request'.
124-- | Parse a set of header lines and body into a 'Request'.
@@ -170,7 +145,9 @@ parseRequest' port (firstLine:otherLines) remoteHost' = do
170 (x, _):_ -> x
145 (x, _):_ -> x
171 [] -> 0
146 [] -> 0
172 let serverName' = takeUntil 58 host -- ':'
147 let serverName' = takeUntil 58 host -- ':'
148+ -- FIXME isolate takes an Integer instead of Int or Int64. If this is a
149+ -- performance penalty, we may need our own version.
173− return (requestBodyHandle len, Request
150+ return (EB.isolate len, Request
174 { requestMethod = method
151 { requestMethod = method
175 , httpVersion = httpversion
152 , httpVersion = httpversion
176 , pathInfo = rpath
153 , pathInfo = rpath
@@ -297,45 +274,8 @@ parseHeaderNoAttr s =
297 else rest
274 else rest
298 in (mkCIByteString k, rest')
275 in (mkCIByteString k, rest')
299
276
300−requestBodyHandle :: Int
277+-- FIXME when we switch to Jeremy's timeout code, we can probably start using
301− -> E.Enumeratee ByteString ByteString IO a
278+-- network-enumerator's enumSocket
302−requestBodyHandle initLen =
303− go initLen
304− where
305− go 0 step = return step
306− go len (E.Continue k) = do
307− x <- E.head
308− case x of
309− Nothing -> return $ E.Continue k
310− Just bs -> do
311− (bs', newlen) <- yieldExtra len bs
312− k (E.Chunks [bs']) >>== go newlen
313− go len step = do
314− drain len
315− return step
316− drain 0 = return ()
317− drain len = do
318− mbs <- E.head
319− case mbs of
320− Nothing -> return ()
321− Just bs -> do
322− (_, newlen) <- yieldExtra len bs
323− drain newlen
324− yieldExtra len bs
325− | B.length bs == len = return (bs, 0)
326− | B.length bs < len = return (bs, len - B.length bs)
327− | otherwise = do
328− let (x, y) = B.splitAt len bs
329− E.yield () $ E.Chunks [y]
330− return (x, 0)
331−
332−iterSocket :: Socket -> E.Iteratee ByteString IO ()
333−iterSocket socket =
334− E.continue go
335− where
336− go E.EOF = E.yield () E.EOF
337− go (E.Chunks cs) = liftIO (Sock.sendMany socket cs) >> E.continue go
338−
339enumSocket :: Int -> Socket -> E.Enumerator ByteString IO a
279enumSocket :: Int -> Socket -> E.Enumerator ByteString IO a
340enumSocket len socket (E.Continue k) = do
280enumSocket len socket (E.Continue k) = do
341#if NO_TIMEOUT_PROTECTION
281#if NO_TIMEOUT_PROTECTION
@@ -352,3 +292,38 @@ enumSocket len socket (E.Continue k) = do
352 | S.length bs == 0 = E.continue k
292 | S.length bs == 0 = E.continue k
353 | otherwise = k (E.Chunks [bs]) >>== enumSocket len socket
293 | otherwise = k (E.Chunks [bs]) >>== enumSocket len socket
354enumSocket _ _ step = E.returnI step
294enumSocket _ _ step = E.returnI step
295+
296+------ The functions below are not warp-specific and could be split out into a
297+--separate package.
298+
299+takeUntilBlank :: Int
300+ -> ([ByteString] -> [ByteString])
301+ -> E.Iteratee S.ByteString IO [ByteString]
302+takeUntilBlank count _
303+ | count > maxHeaders = E.throwError TooManyHeaders
304+takeUntilBlank count front = do
305+ l <- takeLineMax 0 id
306+ if B.null l
307+ then return $ front []
308+ else takeUntilBlank (count + 1) $ front . (:) l
309+
310+takeLineMax :: Int
311+ -> ([ByteString] -> [ByteString])
312+ -> E.Iteratee ByteString IO ByteString
313+takeLineMax len front = do
314+ mbs <- EL.head
315+ case mbs of
316+ Nothing -> E.throwError IncompleteHeaders
317+ Just bs -> do
318+ let (x, y) = S.breakByte 10 bs
319+ x' = if S.length x > 0 && S.last x == 13
320+ then S.init x
321+ else x
322+ let len' = len + B.length x
323+ case () of
324+ ()
325+ | len' > maxHeaderLength -> E.throwError OverLargeHeader
326+ | B.null y -> takeLineMax len' $ front . (:) x
327+ | otherwise -> do
328+ E.yield () $ E.Chunks [B.drop 1 y]
329+ return $ B.concat $ front [x']
runtests.hsnew file+71
@@ -0,0 +1,71 @@
1+{-# LANGUAGE QuasiQuotes #-}
2+{-# LANGUAGE TemplateHaskell #-}
3+{-# LANGUAGE OverloadedStrings #-}
4+import Test.Framework (defaultMain, testGroup, Test)
5+import Test.Framework.Providers.HUnit
6+import Test.HUnit hiding (Test)
7+
8+import Network.Wai.Handler.Warp (takeLineMax, takeUntilBlank, InvalidRequest (..))
9+import Data.Enumerator (run_, ($$), enumList, run)
10+import Control.Exception (fromException)
11+
12+main :: IO ()
13+main = defaultMain [testSuite]
14+
15+testSuite :: Test
16+testSuite = testGroup "Text.Hamlet"
17+ [ testCase "takeLineMax safe" caseTakeLineMaxSafe
18+ , testCase "takeUntilBlank safe" caseTakeUntilBlankSafe
19+ , testCase "takeLineMax unsafe" caseTakeLineMaxUnsafe
20+ , testCase "takeLineMax incomplete" caseTakeLineMaxIncomplete
21+ , testCase "takeUntilBlank too many lines" caseTakeUntilBlankTooMany
22+ , testCase "takeUntilBlank too large" caseTakeUntilBlankTooLarge
23+ ]
24+
25+caseTakeLineMaxSafe = do
26+ x <- run_ $ (enumList 1 ["f", "oo\n\n", "bar\n\nbaz\n"]) $$ do
27+ a <- takeLineMax 0 id
28+ b <- takeLineMax 0 id
29+ c <- takeLineMax 0 id
30+ d <- takeLineMax 0 id
31+ e <- takeLineMax 0 id
32+ return (a, b, c, d, e)
33+ x @?= ("foo", "", "bar", "", "baz")
34+
35+caseTakeUntilBlankSafe = do
36+ x <- run_ $ (enumList 1 ["f", "oo\n", "bar\nbaz\n\n"]) $$ takeUntilBlank 0 id
37+ x @?= ["foo", "bar", "baz"]
38+
39+caseTakeLineMaxUnsafe = do
40+ x <- run $ (enumList 1 $ repeat "abc") $$ do
41+ a <- takeLineMax 0 id
42+ b <- takeLineMax 0 id
43+ c <- takeLineMax 0 id
44+ d <- takeLineMax 0 id
45+ e <- takeLineMax 0 id
46+ return (a, b, c, d, e)
47+ assertException OverLargeHeader x
48+
49+assertException x (Left se) =
50+ case fromException se of
51+ Just e -> e @?= x
52+ Nothing -> assertFailure "Not an exception"
53+assertException _ _ = assertFailure "Not an exception"
54+
55+caseTakeLineMaxIncomplete = do
56+ x <- run $ (enumList 1 ["f", "oo\n\n", "bar\n\nbaz"]) $$ do
57+ a <- takeLineMax 0 id
58+ b <- takeLineMax 0 id
59+ c <- takeLineMax 0 id
60+ d <- takeLineMax 0 id
61+ e <- takeLineMax 0 id
62+ return (a, b, c, d, e)
63+ assertException IncompleteHeaders x
64+
65+caseTakeUntilBlankTooMany = do
66+ x <- run $ (enumList 1 $ repeat "f\n") $$ takeUntilBlank 0 id
67+ assertException TooManyHeaders x
68+
69+caseTakeUntilBlankTooLarge = do
70+ x <- run $ (enumList 1 $ repeat "f") $$ takeUntilBlank 0 id
71+ assertException OverLargeHeader x
warp.cabal+1
@@ -25,6 +25,7 @@ Library
25 , enumerator >= 0.4.5 && < 0.5
25 , enumerator >= 0.4.5 && < 0.5
26 , blaze-builder >= 0.2.1.4 && < 0.3
26 , blaze-builder >= 0.2.1.4 && < 0.3
27 , sendfile >= 0.7.2 && < 0.8
27 , sendfile >= 0.7.2 && < 0.8
28+ , network-enumerator >= 0.1.1 && < 0.2
28 if flag(network-bytestring)
29 if flag(network-bytestring)
29 build-depends: network >= 2.2.1 && < 2.2.3
30 build-depends: network >= 2.2.1 && < 2.2.3
30 , network-bytestring >= 0.1.3 && < 0.1.4
31 , network-bytestring >= 0.1.3 && < 0.1.4