module Network.Wai.Handler.WebSockets
( intercept
, interceptWith
) where
import Control.Monad.IO.Class (liftIO)
import Blaze.ByteString.Builder (Builder)
import qualified Blaze.ByteString.Builder as Builder
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BC
import Data.Char (toLower)
import Data.Conduit
import qualified Network.Wai as Wai
import qualified Network.Wai.Handler.Warp as Warp
import qualified Network.WebSockets as WS
import qualified Network.WebSockets.Connection as WS
import System.IO.Streams (InputStream, OutputStream)
import qualified System.IO.Streams as Streams
intercept :: WS.ServerApp
-> Wai.Request
-> Maybe (Source IO ByteString -> Warp.Connection -> IO ())
intercept = interceptWith WS.defaultConnectionOptions
interceptWith :: WS.ConnectionOptions
-> WS.ServerApp
-> Wai.Request
-> Maybe (Source IO ByteString -> Warp.Connection -> IO ())
interceptWith opts app req = case lookup "upgrade" (Wai.requestHeaders req) of
Just s
| BC.map toLower s == "websocket" -> Just $ runWebSockets opts req' app
| otherwise -> Nothing
_ -> Nothing
where
req' = WS.RequestHead (Wai.rawPathInfo req) (Wai.requestHeaders req) (Wai.isSecure req)
runWebSockets :: WS.ConnectionOptions
-> WS.RequestHead
-> WS.ServerApp
-> Source IO ByteString
-> Warp.Connection
-> IO ()
runWebSockets opts req app _ conn = do
(is, os) <- liftIO $ connectionToStreams conn
let pc = WS.PendingConnection
{ WS.pendingOptions = opts
, WS.pendingRequest = req
, WS.pendingOnAccept = \_ -> return ()
, WS.pendingIn = is
, WS.pendingOut = os
}
liftIO $ app pc
connectionToStreams :: Warp.Connection
-> IO (InputStream ByteString, OutputStream Builder)
connectionToStreams connection = do
is <- Streams.makeInputStream input
os <- Streams.makeOutputStream output
return $! (is, os)
where
input = do
s <- Warp.connRecv connection
return $! if BC.null s then Nothing else Just s
output Nothing = return $! ()
output (Just s') = if BC.null s then return $! () else Warp.connSendAll connection s
where s = Builder.toByteString s'