@@ -21,62 +21,79 @@ import qualified Network.WebSockets.Connection as WS
2121import System.IO.Streams (InputStream, OutputStream)
2222import qualified System.IO.Streams as Streams
2323
24−
24+--------------------------------------------------------------------------------
2525-- | For use with 'settingsIntercept' from the Warp web server.
26−intercept :: WS.Protocol p
27− => (WS.Request -> WS.WebSockets p ())
26+intercept :: WS.ServerApp
2827 -> Wai.Request
2928 -> Maybe (Source (ResourceT IO) ByteString -> Warp.Connection -> ResourceT IO ())
30−intercept = interceptWith WS.defaultWebSocketsOptions
29+intercept = interceptWith WS.defaultConnectionOptions
3130
31+--------------------------------------------------------------------------------
3232-- | Variation of 'intercept' which allows custom options.
33−interceptWith :: WS.Protocol p
34− => WS.WebSocketsOptions
35− -> (WS.Request -> WS.WebSockets p ())
33+interceptWith :: WS.ConnectionOptions
34+ -> WS.ServerApp
3635 -> Wai.Request
3736 -> Maybe (Source (ResourceT IO) ByteString -> Warp.Connection -> ResourceT IO ())
38−interceptWith opts app req = case lookup "upgrade" $ Wai.requestHeaders req of
37+interceptWith opts app req = case lookup "upgrade" (Wai.requestHeaders req) of
3938 Just s
40− | S.map toLower s == "websocket" -> Just $ runWebSockets opts req' app
39+ | BC.map toLower s == "websocket" -> Just $ runWebSockets opts req' app
4140 | otherwise -> Nothing
4241 _ -> Nothing
43− where
44− req' = WS.RequestHttpPart (Wai.rawPathInfo req) (Wai.requestHeaders req)
45− (Wai.isSecure req)
42+ where
43+ req' = WS.RequestHead (Wai.rawPathInfo req) (Wai.requestHeaders req) (Wai.isSecure req)
4644
47−-- | Internal function to run the WebSocket iteratee using the conduit library
48−runWebSockets :: WS.Protocol p
49− => WS.WebSocketsOptions
50− -> WS.RequestHttpPart
51− -> (WS.Request -> WS.WebSockets p ())
45+--------------------------------------------------------------------------------
46+---- | Internal function to run the WebSocket io-streams using the conduit library
47+runWebSockets :: WS.ConnectionOptions
48+ -> WS.RequestHead
49+ -> WS.ServerApp
5250 -> Source (ResourceT IO) ByteString
5351 -> Warp.Connection
5452 -> ResourceT IO ()
5553runWebSockets opts req app source conn = do
56− step <- liftIO $ E.runIteratee $ WS.runWebSocketsWith opts req app send
57− source $$ sink (E.returnI step)
58− where
59− send = iterConnection conn
6054
61− sink iter = await >>= maybe (close iter) (push iter)
55+ (is, os) <- lift $ connectionToStreams conn
56+
57+ let pc = WS.PendingConnection
58+ { WS.pendingOptions = opts
59+ , WS.pendingRequest = req
60+ , WS.pendingOnAccept = forkPingThread
61+ , WS.pendingIn = is
62+ , WS.pendingOut = os
63+ }
64+
65+ lift $ app pc
66+
67+--------------------------------------------------------------------------------
68+-- | Start a ping thread in the background
69+forkPingThread :: WS.Connection -> IO ()
70+forkPingThread conn = do
71+ _ <- forkIO pingThread
72+ return ()
73+ where
74+ pingThread = handle ignore $ forever $ do
75+ WS.sendPing conn (BC.pack "ping")
76+ threadDelay $ 30 * 1000 * 1000
6277
63− push iter bs = do
64− step <- liftIO $ E.runIteratee $ E.enumList 1 [bs] E.$$ iter
65− case step of
66− E.Continue _ -> sink $ E.returnI step
67− E.Yield out inp -> maybe (return ()) leftover (streamToMaybe inp) >> return out
68− E.Error e -> liftIO $ monadThrow e
69− close iter = do
70− _ <- liftIO $ E.runIteratee $ E.enumEOF E.$$ iter
71− return ()
78+ ignore :: SomeException -> IO ()
79+ ignore _ = return ()
7280
73−iterConnection :: Warp.Connection -> E.Iteratee ByteString IO ()
74−iterConnection c = E.continue go
75− where
76− go (E.Chunks []) = E.continue go
77− go (E.Chunks cs) = E.tryIO (Warp.connSendMany c cs) >> E.continue go
78− go E.EOF = E.continue go
81+------------------------------------------------------------------------------
82+-- | Converts a 'Connection' to an 'InputStream' \/ 'OutputStream' pair. Note that,
83+-- as is usually the case in @io-streams@, writing a 'Nothing' to the generated
84+-- 'OutputStream' does not cause the underlying 'Connection' to be closed.
85+connectionToStreams :: Warp.Connection
86+ -> IO (InputStream ByteString, OutputStream Builder)
87+connectionToStreams connection = do
88+ is <- Streams.makeInputStream input
89+ os <- Streams.makeOutputStream output
90+ return $! (is, os)
91+
92+ where
93+ input = do
94+ s <- Warp.connRecv connection
95+ return $! if BC.null s then Nothing else Just s
7996
80−streamToMaybe :: E.Stream S.ByteString -> Maybe S.ByteString
81−streamToMaybe E.EOF = Nothing
82−streamToMaybe (E.Chunks bs) = Just $ S.concat bs
97+ output Nothing = return $! ()
98+ output (Just s') = if BC.null s then return $! () else Warp.connSendAll connection s
99+ where s = Builder.toByteString s'