@@ -1,46 +1,48 @@
11{-# LANGUAGE OverloadedStrings #-}
22module Network.Wai.Handler.WebSockets
3− ( intercept
4− , interceptWith
3+ ( websocketsApp
4+ , websocketsOr
55 ) where
66
7−import Control.Monad.IO.Class (liftIO)
8−import Blaze.ByteString.Builder (Builder)
9−import qualified Blaze.ByteString.Builder as Builder
107import Data.ByteString (ByteString)
118import qualified Data.ByteString.Char8 as BC
12−import Data.Char (toLower)
9+import qualified Data.CaseInsensitive as CI
1310import Data.Conduit
11+import Data.IORef (newIORef, readIORef, writeIORef)
12+import Network.HTTP.Types (status500)
1413import qualified Network.Wai as Wai
15−import qualified Network.Wai.Handler.Warp as Warp
1614import qualified Network.WebSockets as WS
1715import qualified Network.WebSockets.Connection as WS
1816import System.IO.Streams (InputStream, OutputStream)
1917import qualified System.IO.Streams as Streams
2018
2119--------------------------------------------------------------------------------
22−-- | For use with 'settingsIntercept' from the Warp web server.
23−intercept :: WS.ServerApp
24− -> Wai.Request
25− -> Maybe (Source IO ByteString -> Warp.Connection -> IO ())
26−intercept = interceptWith WS.defaultConnectionOptions
20+websocketsOr :: WS.ConnectionOptions
21+ -> WS.ServerApp
22+ -> Wai.Application
23+ -> Wai.Application
24+websocketsOr opts app backup req =
25+ case websocketsApp opts app req of
26+ Nothing -> backup req
27+ Just res -> return res
2728
2829--------------------------------------------------------------------------------
29−-- | Variation of 'intercept' which allows custom options.
30−interceptWith :: WS.ConnectionOptions
30+websocketsApp :: WS.ConnectionOptions
3131 -> WS.ServerApp
3232 -> Wai.Request
33− -> Maybe (Source IO ByteString -> Warp.Connection -> IO ())
34−interceptWith opts app req = case lookup "upgrade" (Wai.requestHeaders req) of
35− Just s
36− | BC.map toLower s == "websocket" -> Just $ runWebSockets opts req' app
37− | otherwise -> Nothing
38− _ -> Nothing
39− where
40− req' = WS.RequestHead
41− (Wai.rawPathInfo req `BC.append` Wai.rawQueryString req)
42− (Wai.requestHeaders req)
43− (Wai.isSecure req)
33+ -> Maybe Wai.Response
34+websocketsApp opts app req
35+ | fmap CI.mk (lookup "upgrade" $ Wai.requestHeaders req) == Just "websocket" =
36+ Just $ flip Wai.responseRaw backup $ \src sink ->
37+ runWebSockets opts req' app src sink
38+ | otherwise = Nothing
39+ where
40+ req' = WS.RequestHead
41+ (Wai.rawPathInfo req `BC.append` Wai.rawQueryString req)
42+ (Wai.requestHeaders req)
43+ (Wai.isSecure req)
44+ backup = Wai.responseLBS status500 [("Content-Type", "text/plain")]
45+ "The web application attempted to send a WebSockets response, but WebSockets are not supported by your WAI handler."
4446
4547--------------------------------------------------------------------------------
4648---- | Internal function to run the WebSocket io-streams using the conduit library
@@ -48,13 +50,14 @@ runWebSockets :: WS.ConnectionOptions
4850 -> WS.RequestHead
4951 -> WS.ServerApp
5052 -> Source IO ByteString
51− -> Warp.Connection
53+ -> Sink ByteString IO ()
5254 -> IO ()
53−runWebSockets opts req app _ conn = do
55+runWebSockets opts req app src sink = do
5456
55− (is, os) <- liftIO $ connectionToStreams conn
57+ is <- srcToInput src
58+ os <- sinkToOutput sink >>= Streams.builderStream
5659
57− let pc = WS.PendingConnection
60+ let pc = WS.PendingConnection
5861 { WS.pendingOptions = opts
5962 , WS.pendingRequest = req
6063 , WS.pendingOnAccept = \_ -> return ()
@@ -62,24 +65,21 @@ runWebSockets opts req app _ conn = do
6265 , WS.pendingOut = os
6366 }
6467
65− liftIO $ app pc
68+ app pc
6669
67−------------------------------------------------------------------------------
68−-- | Converts a 'Connection' to an 'InputStream' \/ 'OutputStream' pair. Note that,
69−-- as is usually the case in @io-streams@, writing a 'Nothing' to the generated
70−-- 'OutputStream' does not cause the underlying 'Connection' to be closed.
71−connectionToStreams :: Warp.Connection
72− -> IO (InputStream ByteString, OutputStream Builder)
73−connectionToStreams connection = do
74− is <- Streams.makeInputStream input
75− os <- Streams.makeOutputStream output
76− return $! (is, os)
77−
78− where
79− input = do
80− s <- Warp.connRecv connection
81− return $! if BC.null s then Nothing else Just s
70+srcToInput :: Source IO ByteString -> IO (InputStream ByteString)
71+srcToInput src0 = do
72+ (rsrc0, ()) <- src0 $$+ return ()
73+ ref <- newIORef rsrc0
74+ Streams.makeInputStream $ do
75+ rsrc <- readIORef ref
76+ (rsrc', mbs) <- rsrc $$++ await
77+ writeIORef ref rsrc'
78+ return mbs
8279
83− output Nothing = return $! ()
84− output (Just s') = if BC.null s then return $! () else Warp.connSendAll connection s
85− where s = Builder.toByteString s'
80+sinkToOutput :: Sink ByteString IO () -> IO (OutputStream ByteString)
81+sinkToOutput sink =
82+ Streams.makeOutputStream output
83+ where
84+ output Nothing = return ()
85+ output (Just bs) = yield bs $$ sink