@@ -18,20 +18,19 @@ module Network.Wai.Handler.Warp
1818 ( run
1919 , sendResponse
2020 , parseRequest
21− , drainRequestBody
2221 ) where
2322
2423import Network.Wai
2524import qualified System.IO
2625
2726import Data.ByteString (ByteString)
28−import qualified Data.ByteString.Char8 as B
2927import qualified Data.ByteString as S
28+import qualified Data.ByteString.Char8 as B
3029import qualified Data.ByteString.Lazy as L
3130import Network
3231 ( listenOn, accept, sClose, PortID(PortNumber), Socket
3332 , withSocketsDo)
34−import Control.Exception (bracket, finally, Exception, throwIO)
33+import Control.Exception (bracket, finally, Exception)
3534import System.IO (Handle, hClose, hFlush)
3635import System.IO.Error (isEOFError, ioeGetHandle)
3736import Control.Concurrent (forkIO)
@@ -44,7 +43,7 @@ import Control.Arrow (first)
4443
4544import Data.Enumerator (($$), (>>==))
4645import qualified Data.Enumerator as E
47−import Data.Enumerator.IO (iterHandle)
46+import Data.Enumerator.IO (iterHandle, enumHandle)
4847import Blaze.ByteString.Builder.Enumerator (builderToByteString)
4948import Blaze.ByteString.Builder.HTTP
5049 (chunkedTransferEncoding, chunkedTransferTerminator)
@@ -54,9 +53,6 @@ import Data.Monoid (mconcat)
5453import Network.Socket.SendFile (unsafeSendFile)
5554
5655import Control.Monad.IO.Class (liftIO)
57−import Data.ByteString.Lazy.Internal (defaultChunkSize)
58−
59−import Data.IORef (IORef, writeIORef, readIORef, newIORef)
6056
6157run :: Port -> Application -> IO ()
6258run port = withSocketsDo .
@@ -69,74 +65,96 @@ type Port = Int
6965serveConnections :: Port -> Application -> Socket -> IO ()
7066serveConnections port app socket = do
7167 (conn, remoteHost', _) <- accept socket -- FIXME use sockets directly instead of Handlers?
68+ -- FIXME disable buffering
7269 _ <- forkIO $ serveConnection port app conn remoteHost'
7370 serveConnections port app socket
7471
7572serveConnection :: Port -> Application -> Handle -> String -> IO ()
7673serveConnection port app conn remoteHost' = do
77− catch
78− (finally
79− serveConnection'
80− (hClose conn))
81− catchEOFError
82− where serveConnection' = do
83− (ilen, env) <- parseRequest port conn remoteHost'
84− res <- app env
85− drainRequestBody conn ilen
86− keepAlive <- sendResponse env (httpVersion env) conn res
87− hFlush conn
88− when keepAlive serveConnection'
74+ catch
75+ (finally
76+ (E.run_ $ fromClient $$ serveConnection')
77+ (hClose conn))
78+ catchEOFError
79+ where
80+ fromClient = enumHandle 4096 conn
81+ serveConnection' = do
82+ (enumeratee, env) <- parseRequest port remoteHost'
83+ res <- E.joinI $ enumeratee $$ app env
84+ keepAlive <- liftIO $ sendResponse env (httpVersion env) conn res
85+ liftIO $ hFlush conn
86+ when keepAlive serveConnection'
8987
90− catchEOFError :: IOError -> IO ()
91− catchEOFError e | isEOFError e = case ioeGetHandle e of
92− Just h -> unless (h == conn) (ioError e)
93− Nothing -> ioError e
94− | otherwise = ioError e
88+ catchEOFError :: IOError -> IO ()
89+ catchEOFError e
90+ | isEOFError e =
91+ case ioeGetHandle e of
92+ Just h -> unless (h == conn) (ioError e)
93+ Nothing -> ioError e
94+ | otherwise = ioError e
9595
96−drainRequestBody :: Handle -> IORef Int -> IO ()
97−drainRequestBody conn ilen = do
98− remaining <- readIORef ilen
99− _ <- B.hGet conn remaining -- FIXME just skip, don't read
100− return ()
96+parseRequest :: Port -> String -> E.Iteratee S.ByteString IO (E.Enumeratee ByteString ByteString IO a, Request)
97+parseRequest port remoteHost' = do
98+ headers' <- takeUntilBlank 0 id
99+ parseRequest' port headers' remoteHost'
101100
102−parseRequest :: Port -> Handle -> String -> IO (IORef Int, Request)
103−parseRequest port conn remoteHost' = do
104− headers' <- takeUntilBlank conn id
105− parseRequest' port headers' conn remoteHost'
101+-- FIXME come up with good values here
102+maxHeaders, maxHeaderLength :: Int
103+maxHeaders = 30
104+maxHeaderLength = 1024
106105
107−takeUntilBlank :: Handle
106+takeUntilBlank :: Int
108107 -> ([ByteString] -> [ByteString])
109− -> IO [ByteString]
110−takeUntilBlank h front = do
111− l <- stripCR `fmap` B.hGetLine h
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
112113 if B.null l
113114 then return $ front []
114− else takeUntilBlank h $ front . (:) l
115+ else takeUntilBlank (count + 1) $ front . (:) l
115116
116−stripCR :: ByteString -> ByteString
117−stripCR bs
118− | B.null bs = bs
119− | B.last bs == '\r' = B.init bs
120− | otherwise = bs
117+takeLine :: Int
118+ -> ([ByteString] -> [ByteString])
119+ -> E.Iteratee ByteString IO ByteString
120+takeLine len front = do
121+ mbs <- E.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']
121137
122138data InvalidRequest =
123139 NotEnoughLines [String]
124140 | BadFirstLine String
125141 | NonHttp
142+ | TooManyHeaders
143+ | IncompleteHeaders
144+ | OverLargeHeader
126145 deriving (Show, Typeable)
127146instance Exception InvalidRequest
128147
129148-- | Parse a set of header lines and body into a 'Request'.
130149parseRequest' :: Port
131150 -> [ByteString]
132− -> Handle
133151 -> String
134− -> IO (IORef Int, Request)
135−parseRequest' port lines' handle remoteHost' = do
152+ -> E.Iteratee S.ByteString IO (E.Enumeratee S.ByteString S.ByteString IO a, Request)
153+parseRequest' port lines' remoteHost' = do
136154 (firstLine, otherLines) <-
137155 case lines' of
138156 x:xs -> return (x, xs)
139− [] -> throwIO $ NotEnoughLines $ map B.unpack lines'
157+ [] -> E.throwError $ NotEnoughLines $ map B.unpack lines'
140158 (method, rpath', gets, httpversion) <- parseFirst firstLine
141159 let rpath = '/' : case B.unpack rpath' of
142160 ('/':x) -> x
@@ -150,8 +168,7 @@ parseRequest' port lines' handle remoteHost' = do
150168 (x, _):_ -> Just x
151169 _ -> Nothing
152170 let (serverName', _) = B.break (== ':') host
153− ilen <- newIORef len
154− return (ilen, Request
171+ return (requestBodyHandle len, Request
155172 { requestMethod = method
156173 , httpVersion = httpversion
157174 , pathInfo = B.pack rpath
@@ -160,21 +177,20 @@ parseRequest' port lines' handle remoteHost' = do
160177 , serverPort = port
161178 , requestHeaders = heads
162179 , isSecure = False
163− , requestBody = requestBodyHandle handle ilen len
164180 , errorHandler = System.IO.hPutStr System.IO.stderr
165181 , remoteHost = B.pack remoteHost'
166182 })
167183
168184parseFirst :: ByteString
169− -> IO (ByteString, ByteString, ByteString, HttpVersion)
185+ -> E.Iteratee S.ByteString IO (ByteString, ByteString, ByteString, HttpVersion)
170186parseFirst s = do
171187 let pieces = B.words s
172188 (method, query, http') <-
173189 case pieces of
174190 [x, y, z] -> return (x, y, z)
175− _ -> throwIO $ BadFirstLine $ B.unpack s
191+ _ -> E.throwError $ BadFirstLine $ B.unpack s
176192 let (hfirst, hsecond) = B.splitAt 5 http'
177− unless (hfirst == B.pack "HTTP/") $ throwIO NonHttp
193+ unless (hfirst == "HTTP/") $ E.throwError NonHttp
178194 let (rpath, qstring) = B.break (== '?') query
179195 return (method, rpath, qstring, hsecond)
180196
@@ -256,17 +272,29 @@ parseHeaderNoAttr s =
256272 else rest
257273 in (k, rest')
258274
259−requestBodyHandle :: Handle
260− -> IORef Int
261− -> Int
262− -> E.Enumerator B.ByteString IO a
263−requestBodyHandle handle ilen initLen =
275+requestBodyHandle :: Int
276+ -> E.Enumeratee ByteString ByteString IO a
277+requestBodyHandle initLen =
264278 go initLen
265279 where
266− go 0 step = E.returnI step
280+ go 0 step = return step
267281 go len (E.Continue k) = do
268− bs <- liftIO $ B.hGet handle (min len defaultChunkSize) -- FIXME play with defaultChunkSize
269− let newLen = len - B.length bs
270− liftIO $ writeIORef ilen newLen
271− k (E.Chunks [bs]) >>== go newLen
272− go _ step = E.returnI step
282+ x <- E.head
283+ case x of
284+ Nothing -> return $ E.Continue k
285+ Just bs -> do
286+ let newlen = max 0 $ len - B.length bs
287+ k (E.Chunks [bs]) >>== go newlen
288+ go len step = do
289+ drain len
290+ return step
291+ drain 0 = return ()
292+ drain len = do
293+ mbs <- E.head
294+ case mbs of
295+ Nothing -> return ()
296+ Just bs -> do
297+ let newlen = len - B.length bs
298+ if newlen <= 0
299+ then return ()
300+ else drain newlen