@@ -3,24 +3,44 @@ module Network.Wai.Handler.WebSockets
33 ( intercept
44 ) where
55
6−import Network.Wai (Request, requestHeaders, rawPathInfo, requestHeaders)
7−import Network.Socket (Socket)
8−import Data.Char (toLower)
9−import Data.Enumerator (Iteratee)
6+import Control.Monad.IO.Class (liftIO)
107import Data.ByteString (ByteString)
8+import Data.Char (toLower)
9+import Network.Socket (Socket)
10+import Network.Wai (Request, requestHeaders, rawPathInfo, requestHeaders)
1111import qualified Data.ByteString.Char8 as S
12−import Network.WebSockets (WebSockets, runWebSockets, RequestHttpPart(RequestHttpPart), Protocol)
12+import qualified Data.Conduit as C
13+import qualified Data.Enumerator as E
1314import qualified Network.WebSockets as WS
14−import Network.Socket.Enumerator (iterSocket)
15+import qualified Network.WebSockets.Internal as WS
1516
16−intercept :: Protocol p
17− => (WS.Request -> WebSockets p ())
17+-- | For use with 'settingsIntercept' from the Warp web server.
18+intercept :: WS.Protocol p
19+ => (WS.Request -> WS.WebSockets p ())
1820 -> Request
19− -> Maybe (Socket -> Iteratee ByteString IO ())
20−intercept app req =
21− case lookup "upgrade" $ requestHeaders req of
22− Just s | S.map toLower s=="websocket" -> Just $ \socket -> do
23− runWebSockets req' app (iterSocket socket)
24− _ -> Nothing
21+ -> Maybe (C.BufferedSource IO ByteString -> Socket -> C.ResourceT IO ())
22+intercept app req = case lookup "upgrade" $ requestHeaders req of
23+ Just s
24+ | S.map toLower s == "websocket" -> Just $ runWebSockets req' app
25+ | otherwise -> Nothing
26+ _ -> Nothing
27+ where
28+ req' = WS.RequestHttpPart (rawPathInfo req) (requestHeaders req)
29+
30+-- | Internal function to run the WebSocket iteratee using the conduit library
31+runWebSockets :: WS.Protocol p
32+ => WS.RequestHttpPart
33+ -> (WS.Request -> WS.WebSockets p ())
34+ -> C.BufferedSource IO ByteString
35+ -> Socket
36+ -> C.ResourceT IO ()
37+runWebSockets req app source sock = do
38+ step <- liftIO $ E.runIteratee $ WS.runWebSockets req app send
39+ source C.$$ C.sinkState (E.returnI step) push close
2540 where
26− req' = RequestHttpPart (rawPathInfo req) (requestHeaders req)
41+ send = WS.iterSocket sock
42+ close = const $ return ()
43+
44+ push iter bs = do
45+ step <- liftIO $ E.runIteratee $ E.enumList 1 [bs] E.$$ iter
46+ return (E.returnI step, C.Processing)