@@ -95,7 +95,7 @@ import Network.Sendfile
9595
9696import qualified System.PosixCompat.Files as P
9797
98−import Control.Monad.IO.Class (liftIO)
98+import Control.Monad.IO.Class (MonadIO, liftIO)
9999import qualified Timeout as T
100100import Timeout (Manager, registerKillThread, pause, resume)
101101import Data.Word (Word8)
@@ -210,6 +210,65 @@ runSettingsConnection set getConn app = do
210210 T.cancel th
211211 return ()
212212
213+-- | Contains a @Source@ and a byte count that is still to be read in.
214+newtype IsolatedBSSource = IsolatedBSSource (I.IORef (Int, C.Source (ResourceT IO) ByteString))
215+
216+-- | Given an @IsolatedBSSource@ provide a @Source@ that only allows up to the
217+-- specified number of bytes to be passed downstream. All leftovers should be
218+-- retained within the @Source@. If there are not enough bytes available,
219+-- throws a @ConnectionClosedByPeer@ exception.
220+ibsIsolate :: IsolatedBSSource -> C.Source (ResourceT IO) ByteString
221+ibsIsolate ibs@(IsolatedBSSource ref) =
222+ C.PipeM pull (return ())
223+ where
224+ pull = do
225+ (count, src) <- liftIO $ I.readIORef ref
226+ if count == 0
227+ -- No more bytes wanted downstream, so we're done.
228+ then return $ C.Done Nothing ()
229+ else do
230+ -- Get the next chunk (if available) and the updated source
231+ (src', mbs) <- src C.$$& CL.head
232+
233+ -- If no chunk available, then there aren't enough bytes in the
234+ -- stream. Throw a ConnectionClosedByPeer
235+ bs <- maybe (liftIO $ throwIO ConnectionClosedByPeer) return mbs
236+
237+ let -- How many of the bytes in this chunk to send downstream
238+ toSend = min count (S.length bs)
239+ -- How many bytes will still remain to be sent downstream
240+ count' = count - toSend
241+ case () of
242+ ()
243+ -- The expected count is greater than the size of the
244+ -- chunk we just read. Send the entire chunk
245+ -- downstream, and then loop on this function for the
246+ -- next chunk.
247+ | count' > 0 -> do
248+ liftIO $ I.writeIORef ref (count', src')
249+ return $ C.HaveOutput (ibsIsolate ibs) (return ()) bs
250+
251+ -- The expected count is the total size of the chunk we
252+ -- just read. Send this chunk downstream, and then
253+ -- terminate the stream.
254+ | count == S.length bs -> do
255+ liftIO $ I.writeIORef ref (count', src')
256+ return $ C.HaveOutput (C.Done Nothing ()) (return ()) bs
257+
258+ -- Some of the bytes in this chunk should not be sent
259+ -- downstream. Split up the chunk into the sent and
260+ -- not-sent parts, add the not-sent parts onto the new
261+ -- source, and send the rest of the chunk downstream.
262+ | otherwise -> do
263+ let (x, y) = S.splitAt toSend bs
264+ liftIO $ I.writeIORef ref (count', C.HaveOutput src' (return ()) y)
265+ return $ C.HaveOutput (C.Done Nothing ()) (return ()) x
266+
267+-- | Extract the underlying @Source@ from an @IsolatedBSSource@, which will not
268+-- perform any more isolation.
269+ibsDone :: IsolatedBSSource -> IO (C.Source (ResourceT IO) ByteString)
270+ibsDone (IsolatedBSSource ref) = fmap snd $ I.readIORef ref
271+
213272serveConnection :: Settings
214273 -> T.Handle
215274 -> (SomeException -> IO ())
@@ -223,11 +282,11 @@ serveConnection settings th onException port app conn remoteHost' =
223282 where
224283 serveConnection' :: ResourceT IO ()
225284 serveConnection' = do
226− fromClient <- C.bufferSource $ connSource conn th
285+ let fromClient = connSource conn th
227286 serveConnection'' fromClient
228287
229288 serveConnection'' fromClient = do
230− (env, requireDone) <- parseRequest conn port remoteHost' fromClient
289+ (env, ibs) <- parseRequest conn port remoteHost' fromClient
231290 case settingsIntercept settings env of
232291 Nothing -> do
233292 -- Let the application run for as long as it wants
@@ -235,25 +294,22 @@ serveConnection settings th onException port app conn remoteHost' =
235294 res <- app env
236295
237296 -- flush the rest of the request body
238− requestBody env C.$$ CL.sinkNull
239− liftIO requireDone
297+ ibsIsolate ibs C.$$ CL.sinkNull
298+ fromClient' <- liftIO $ ibsDone ibs
240299
241300 liftIO $ T.resume th
242301 keepAlive <- sendResponse th env conn res
243− when keepAlive $ serveConnection'' fromClient
302+ when keepAlive $ serveConnection'' fromClient'
244303 Just intercept -> do
245304 liftIO $ T.pause th
246305 intercept fromClient conn
247306
248−-- Require that the entire request body has been consumed
249−type RequireDone = IO ()
250−
251307parseRequest :: Connection -> Port -> SockAddr
252− -> C.BufferedSource (ResourceT IO) S.ByteString
253− -> ResourceT IO (Request, RequireDone)
254−parseRequest conn port remoteHost' src = do
255− headers' <- src C.$$ takeHeaders
256− parseRequest' conn port headers' remoteHost' src
308+ -> C.Source (ResourceT IO) S.ByteString
309+ -> ResourceT IO (Request, IsolatedBSSource)
310+parseRequest conn port remoteHost' src1 = do
311+ (src2, headers') <- src1 C.$$& takeHeaders
312+ parseRequest' conn port headers' remoteHost' src2
257313
258314-- FIXME come up with good values here
259315bytesPerRead, maxTotalHeaderLength :: Int
@@ -289,8 +345,8 @@ parseRequest' :: Connection
289345 -> Port
290346 -> [ByteString]
291347 -> SockAddr
292− -> C.BufferedSource (ResourceT IO) S.ByteString
293− -> ResourceT IO (Request, RequireDone)
348+ -> C.Source (ResourceT IO) S.ByteString -- FIXME was buffered
349+ -> ResourceT IO (Request, IsolatedBSSource)
294350parseRequest' _ _ [] _ _ = throwIO $ NotEnoughLines []
295351parseRequest' conn port (firstLine:otherLines) remoteHost' src = do
296352 (method, rpath', gets, httpversion) <- parseFirst firstLine
@@ -307,64 +363,9 @@ parseRequest' conn port (firstLine:otherLines) remoteHost' src = do
307363 Nothing -> 0
308364 Just bs -> readInt bs
309365 let serverName' = takeUntil 58 host -- ':'
310− (rbody, requireDone) <-
311− if len0 == 0
312− then return (mempty, return ())
313− else do
314− -- We can't use the standard isolate, as its counter is not
315− -- kept in a mutable variable.
316− lenRef <- liftIO $ I.newIORef len0
317− let isolate = C.NeedInput push C.Closed
318− push bs = flip C.ConduitM (return ()) $ do
319− len <- liftIO $ I.readIORef lenRef
320− let (a, b) = S.splitAt len bs
321− len' = len - S.length a
322− liftIO $ I.writeIORef lenRef len'
323− return $ if len' == 0
324− then
325− let mleftover = if S.null b then Nothing else Just b
326− final = C.Finished mleftover
327− in if S.null a
328− then final
329− else C.HaveOutput final (return ()) a
330− else C.HaveOutput isolate (return ()) a
331−
332− -- Make sure that we don't connect to the source after the
333− -- isolate conduit closes.
334− --
335− -- Here's the issue: we fuse our buffered request body with
336− -- an isolate conduit which ensures no more than X bytes
337− -- are read. Suppose we read all X bytes, and then we call
338− -- requestBody again. What happens?
339− --
340− -- Previously, we would try to read one more chunk from the
341− -- buffered source. This is inherent to conduit: we
342− -- wouldn't know that the isolate Conduit isn't accepting
343− -- more data until after we've pushed some data to it. This
344− -- results in hanging, since there's no data available on
345− -- the wire.
346− --
347− -- Instead, we add a wrapper that checks if the request
348− -- body has already been depleted before making that first
349− -- pull.
350− --
351− -- Possible optimization: do away with the Conduit
352− -- entirely. However, this may be less efficient overall,
353− -- as we'd now have to check the BufferedSource status on
354− -- each call. Worth looking into.
355−
356− wrap src' = C.SourceM
357− (do
358− len <- liftIO $ I.readIORef lenRef
359− if len <= 0
360− then return C.Closed
361− else return src')
362− (return ())
363−
364− requireDone = do
365− len <- liftIO $ I.readIORef lenRef
366− when (len > 0) $ throwIO ConnectionClosedByPeer
367− return (wrap $ src C.$= isolate, requireDone)
366+ (rbody, ibs) <- liftIO $ do
367+ ibs <- fmap IsolatedBSSource $ I.newIORef (len0, src)
368+ return (ibsIsolate ibs, ibs)
368369 return (Request
369370 { requestMethod = method
370371 , httpVersion = httpversion
@@ -379,7 +380,7 @@ parseRequest' conn port (firstLine:otherLines) remoteHost' src = do
379380 , remoteHost = remoteHost'
380381 , requestBody = rbody
381382 , vault = mempty
382− }, requireDone)
383+ }, ibs)
383384
384385
385386takeUntil :: Word8 -> ByteString -> ByteString
@@ -533,7 +534,7 @@ sendResponse th req conn r = sendResponse' r
533534 T.tickle th
534535 return isPersist
535536 where
536− body = fmap (\x -> case x of
537+ body = fmap2 (\x -> case x of
537538 C.Flush -> flush
538539 C.Chunk builder -> builder) bodyFlush
539540 headers' = headers version s hs
@@ -543,7 +544,13 @@ sendResponse th req conn r = sendResponse' r
543544 chunk :: C.Conduit Builder (ResourceT IO) Builder
544545 chunk = C.NeedInput push close
545546 push x = C.HaveOutput chunk (return ()) (chunkedTransferEncoding x)
546− close = C.Open C.Closed (return ()) chunkedTransferTerminator
547+ close = C.HaveOutput (C.Done Nothing ()) (return ()) chunkedTransferTerminator
548+
549+fmap2 :: Functor m => (o1 -> o2) -> C.Pipe i o1 m r -> C.Pipe i o2 m r
550+fmap2 f (C.HaveOutput p c o) = C.HaveOutput (fmap2 f p) c (f o)
551+fmap2 f (C.NeedInput p c) = C.NeedInput (fmap2 f . p) (fmap2 f c)
552+fmap2 f (C.PipeM mp c) = C.PipeM (fmap (fmap2 f) mp) c
553+fmap2 _ (C.Done i x) = C.Done i x
547554
548555parseHeaderNoAttr :: ByteString -> H.Header
549556parseHeaderNoAttr s =
@@ -559,27 +566,27 @@ connSource :: Connection -> T.Handle -> C.Source (ResourceT IO) ByteString
559566connSource Connection { connRecv = recv } th =
560567 src
561568 where
562− src = C.SourceM (do
569+ src = C.PipeM (do
563570 bs <- liftIO recv
564571 if S.null bs
565− then return C.Closed
572+ then return $ C.Done Nothing ()
566573 else do
567574 when (S.length bs >= 2048) $ liftIO $ T.tickle th
568− return (C.Open src (return ()) bs))
575+ return (C.HaveOutput src (return ()) bs))
569576 (return ())
570577
571578-- | Use 'connSendAll' to send this data while respecting timeout rules.
572579connSink :: Connection -> T.Handle -> C.Sink B.ByteString (ResourceT IO) ()
573580connSink Connection { connSendAll = send } th =
574− C.Processing push close
581+ sink
575582 where
583+ sink = C.NeedInput push close
576584 close = liftIO (T.resume th)
577− push x = C.SinkM $ do
578− liftIO $ do
579− T.resume th
580− send x
581− T.pause th
582− return (C.Processing push close)
585+ push x = C.PipeM (liftIO $ do
586+ T.resume th
587+ send x
588+ T.pause th
589+ return sink) (liftIO $ T.resume th)
583590 -- We pause timeouts before passing control back to user code. This ensures
584591 -- that a timeout will only ever be executed when Warp is in control. We
585592 -- also make sure to resume the timeout after the completion of user code
@@ -599,7 +606,7 @@ data Settings = Settings
599606 , settingsHost :: HostPreference -- ^ Default value: HostIPv4
600607 , settingsOnException :: SomeException -> IO () -- ^ What to do with exceptions thrown by either the application or server. Default: ignore server-generated exceptions (see 'InvalidRequest') and print application-generated applications to stderr.
601608 , settingsTimeout :: Int -- ^ Timeout value in seconds. Default value: 30
602− , settingsIntercept :: Request -> Maybe (C.BufferedSource (ResourceT IO) S.ByteString -> Connection -> ResourceT IO ())
609+ , settingsIntercept :: Request -> Maybe (C.Source (ResourceT IO) S.ByteString -> Connection -> ResourceT IO ())
603610 , settingsManager :: Maybe Manager -- ^ Use an existing timeout manager instead of spawning a new one. If used, 'settingsTimeout' is ignored. Default is 'Nothing'
604611 }
605612
@@ -635,7 +642,7 @@ data THStatus = THStatus
635642
636643takeHeaders :: C.Sink ByteString (ResourceT IO) [ByteString]
637644takeHeaders =
638− C.Processing (push (THStatus 0 id id)) close
645+ C.NeedInput (push (THStatus 0 id id)) close
639646 where
640647 close = throwIO IncompleteHeaders
641648
@@ -650,7 +657,7 @@ takeHeaders =
650657 let len' = len + bsLen
651658 prepend' = prepend . S.append bs
652659 status = THStatus len' lines prepend'
653− in C.Processing (push status) close
660+ in C.NeedInput (push status) close
654661 -- Found a newline at position end.
655662 Just end ->
656663 let start = end + 1 -- start of next chunk
@@ -678,7 +685,7 @@ takeHeaders =
678685 then let bs' = SU.unsafeDrop start bs
679686 in push status bs'
680687 -- no more bytes in this chunk, ask for more
681− else C.Processing (push status) close
688+ else C.NeedInput (push status) close
682689 where
683690 bsLen = S.length bs
684691 mnl = S.elemIndex 10 bs