@@ -10,13 +10,10 @@ module Network.Wai.Handler.WebSockets
1010import Data.ByteString (ByteString)
1111import qualified Data.ByteString.Char8 as BC
1212import qualified Data.CaseInsensitive as CI
13−import Data.Conduit
14−import Data.IORef (newIORef, readIORef, writeIORef)
1513import Network.HTTP.Types (status500)
1614import qualified Network.Wai as Wai
1715import qualified Network.WebSockets as WS
1816import qualified Network.WebSockets.Connection as WS
19−import System.IO.Streams (InputStream, OutputStream)
2017import qualified System.IO.Streams as Streams
2118
2219--------------------------------------------------------------------------------
@@ -29,10 +26,10 @@ websocketsOr :: WS.ConnectionOptions
2926 -> WS.ServerApp
3027 -> Wai.Application
3128 -> Wai.Application
32−websocketsOr opts app backup req =
29+websocketsOr opts app backup req sendResponse =
3330 case websocketsApp opts app req of
34− Nothing -> backup req
35− Just res -> return res
31+ Nothing -> backup req sendResponse
32+ Just res -> sendResponse res
3633
3734--------------------------------------------------------------------------------
3835websocketsApp :: WS.ConnectionOptions
@@ -61,13 +58,15 @@ getRequestHead req = WS.RequestHead
6158runWebSockets :: WS.ConnectionOptions
6259 -> WS.RequestHead
6360 -> (WS.PendingConnection -> IO a)
64− -> Source IO ByteString
65− -> Sink ByteString IO ()
61+ -> IO ByteString
62+ -> (ByteString -> IO ())
6663 -> IO a
6764runWebSockets opts req app src sink = do
6865
69− is <- srcToInput src
70− os <- sinkToOutput sink >>= Streams.builderStream
66+ is <- Streams.makeInputStream $ do
67+ bs <- src
68+ return $ if BC.null bs then Nothing else Just bs
69+ os <- Streams.makeOutputStream (maybe (return ()) sink) >>= Streams.builderStream
7170
7271 let pc = WS.PendingConnection
7372 { WS.pendingOptions = opts
@@ -78,20 +77,3 @@ runWebSockets opts req app src sink = do
7877 }
7978
8079 app pc
81−
82−srcToInput :: Source IO ByteString -> IO (InputStream ByteString)
83−srcToInput src0 = do
84− (rsrc0, ()) <- src0 $$+ return ()
85− ref <- newIORef rsrc0
86− Streams.makeInputStream $ do
87− rsrc <- readIORef ref
88− (rsrc', mbs) <- rsrc $$++ await
89− writeIORef ref rsrc'
90− return mbs
91−
92−sinkToOutput :: Sink ByteString IO () -> IO (OutputStream ByteString)
93−sinkToOutput sink =
94− Streams.makeOutputStream output
95− where
96− output Nothing = return ()
97− output (Just bs) = yield bs $$ sink