@@ -7,34 +7,34 @@ module Network.Wai.Handler.WebSockets
77import Control.Monad.IO.Class (liftIO)
88import Data.ByteString (ByteString)
99import Data.Char (toLower)
10−import Network.Socket (Socket)
11−import Network.Wai (Request, requestHeaders, rawPathInfo, requestHeaders)
1210import qualified Data.ByteString.Char8 as S
1311import qualified Data.Conduit as C
1412import qualified Data.Enumerator as E
13+import qualified Network.Wai as Wai
14+import qualified Network.Wai.Handler.Warp as Warp
1515import qualified Network.WebSockets as WS
16−import qualified Network.WebSockets.Internal as WS
1716
1817-- | For use with 'settingsIntercept' from the Warp web server.
1918intercept :: WS.Protocol p
2019 => (WS.Request -> WS.WebSockets p ())
21− -> Request
22− -> Maybe (C.BufferedSource IO ByteString -> Socket -> C.ResourceT IO ())
20+ -> Wai.Request
21+ -> Maybe (C.BufferedSource IO ByteString -> Warp.Connection -> C.ResourceT IO ())
2322intercept = interceptWith WS.defaultWebSocketsOptions
2423
2524-- | Variation of 'intercept' which allows custom options.
2625interceptWith :: WS.Protocol p
2726 => WS.WebSocketsOptions
2827 -> (WS.Request -> WS.WebSockets p ())
29− -> Request
30− -> Maybe (C.BufferedSource IO ByteString -> Socket -> C.ResourceT IO ())
31−interceptWith opts app req = case lookup "upgrade" $ requestHeaders req of
28+ -> Wai.Request
29+ -> Maybe (C.BufferedSource IO ByteString -> Warp.Connection -> C.ResourceT IO ())
30+interceptWith opts app req = case lookup "upgrade" $ Wai.requestHeaders req of
3231 Just s
3332 | S.map toLower s == "websocket" -> Just $ runWebSockets opts req' app
3433 | otherwise -> Nothing
3534 _ -> Nothing
3635 where
37− req' = WS.RequestHttpPart (rawPathInfo req) (requestHeaders req)
36+ req' = WS.RequestHttpPart (Wai.rawPathInfo req) (Wai.requestHeaders req)
37+ (Wai.isSecure req)
3838
3939-- | Internal function to run the WebSocket iteratee using the conduit library
4040runWebSockets :: WS.Protocol p
@@ -42,15 +42,24 @@ runWebSockets :: WS.Protocol p
4242 -> WS.RequestHttpPart
4343 -> (WS.Request -> WS.WebSockets p ())
4444 -> C.BufferedSource IO ByteString
45− -> Socket
45+ -> Warp.Connection
4646 -> C.ResourceT IO ()
47−runWebSockets opts req app source sock = do
47+runWebSockets opts req app source conn = do
4848 step <- liftIO $ E.runIteratee $ WS.runWebSocketsWith opts req app send
4949 source C.$$ C.sinkState (E.returnI step) push close
5050 where
51− send = WS.iterSocket sock
52− close = const $ return ()
51+ send = iterConnection conn
5352
5453 push iter bs = do
5554 step <- liftIO $ E.runIteratee $ E.enumList 1 [bs] E.$$ iter
56− return (E.returnI step, C.Processing)
55+ return $ C.StateProcessing $ E.returnI step
56+ close iter = do
57+ _ <- liftIO $ E.runIteratee $ E.enumEOF E.$$ iter
58+ return ()
59+
60+iterConnection :: Warp.Connection -> E.Iteratee ByteString IO ()
61+iterConnection c = E.continue go
62+ where
63+ go (E.Chunks []) = E.continue go
64+ go (E.Chunks cs) = E.tryIO (Warp.connSendMany c cs) >> E.continue go
65+ go E.EOF = E.continue go