@@ -1,6 +1,7 @@
11module Network.Wai.Handler.Warp.Conduit where
22
33import Control.Exception
4+import Control.Monad (unless)
45import Control.Monad.IO.Class (MonadIO, liftIO)
56import Control.Monad.Trans.Class (lift)
67import Data.ByteString (ByteString)
@@ -27,47 +28,44 @@ newtype IsolatedBSSource = IsolatedBSSource (I.IORef (Int, ResumableSource (Reso
2728ibsIsolate :: IsolatedBSSource -> Source (ResourceT IO) ByteString
2829ibsIsolate ibs@(IsolatedBSSource ref) = do
2930 (count, src) <- liftIO $ I.readIORef ref
30− if count == 0
31− -- No more bytes wanted downstream, so we're done.
32− then return ()
33− else do
34− -- Get the next chunk (if available) and the updated source
35− (src', mbs) <- lift $ src $$++ CL.head
36−
37− -- If no chunk available, then there aren't enough bytes in the
38− -- stream. Throw a ConnectionClosedByPeer
39− bs <- maybe (liftIO $ throwIO ConnectionClosedByPeer) return mbs
40−
41− let -- How many of the bytes in this chunk to send downstream
42− toSend = min count (S.length bs)
43− -- How many bytes will still remain to be sent downstream
44− count' = count - toSend
45− case () of
46− ()
47− -- The expected count is greater than the size of the
48− -- chunk we just read. Send the entire chunk
49− -- downstream, and then loop on this function for the
50− -- next chunk.
51− | count' > 0 -> do
52− liftIO $ I.writeIORef ref (count', src')
53− yield bs
54− ibsIsolate ibs
55−
56− -- The expected count is the total size of the chunk we
57− -- just read. Send this chunk downstream, and then
58− -- terminate the stream.
59− | count == S.length bs -> do
60− liftIO $ I.writeIORef ref (count', src')
61− yield bs
62−
63− -- Some of the bytes in this chunk should not be sent
64− -- downstream. Split up the chunk into the sent and
65− -- not-sent parts, add the not-sent parts onto the new
66− -- source, and send the rest of the chunk downstream.
67− | otherwise -> do
68− let (x, y) = S.splitAt toSend bs
69− liftIO $ I.writeIORef ref (count', fmapResume (yield y >>) src')
70− yield x
31+ unless (count == 0) $ do
32+ -- Get the next chunk (if available) and the updated source
33+ (src', mbs) <- lift $ src $$++ CL.head
34+
35+ -- If no chunk available, then there aren't enough bytes in the
36+ -- stream. Throw a ConnectionClosedByPeer
37+ bs <- maybe (liftIO $ throwIO ConnectionClosedByPeer) return mbs
38+
39+ let -- How many of the bytes in this chunk to send downstream
40+ toSend = min count (S.length bs)
41+ -- How many bytes will still remain to be sent downstream
42+ count' = count - toSend
43+ case () of
44+ ()
45+ -- The expected count is greater than the size of the
46+ -- chunk we just read. Send the entire chunk
47+ -- downstream, and then loop on this function for the
48+ -- next chunk.
49+ | count' > 0 -> do
50+ liftIO $ I.writeIORef ref (count', src')
51+ yield bs
52+ ibsIsolate ibs
53+
54+ -- The expected count is the total size of the chunk we
55+ -- just read. Send this chunk downstream, and then
56+ -- terminate the stream.
57+ | count == S.length bs -> do
58+ liftIO $ I.writeIORef ref (count', src')
59+ yield bs
60+
61+ -- Some of the bytes in this chunk should not be sent
62+ -- downstream. Split up the chunk into the sent and
63+ -- not-sent parts, add the not-sent parts onto the new
64+ -- source, and send the rest of the chunk downstream.
65+ | otherwise -> do
66+ let (x, y) = S.splitAt toSend bs
67+ liftIO $ I.writeIORef ref (count', fmapResume (yield y >>) src')
68+ yield x
7169
7270-- | Extract the underlying @Source@ from an @IsolatedBSSource@, which will not
7371-- perform any more isolation.