module Network.Wai.Handler.WebSockets
( intercept
, interceptWith
) where
import Control.Monad.IO.Class (liftIO)
import Data.ByteString (ByteString)
import Data.Char (toLower)
import qualified Data.ByteString.Char8 as S
import Data.Conduit
import qualified Data.Enumerator as E
import qualified Network.Wai as Wai
import qualified Network.Wai.Handler.Warp as Warp
import qualified Network.WebSockets as WS
intercept :: WS.Protocol p
=> (WS.Request -> WS.WebSockets p ())
-> Wai.Request
-> Maybe (Source (ResourceT IO) ByteString -> Warp.Connection -> ResourceT IO ())
intercept = interceptWith WS.defaultWebSocketsOptions
interceptWith :: WS.Protocol p
=> WS.WebSocketsOptions
-> (WS.Request -> WS.WebSockets p ())
-> Wai.Request
-> Maybe (Source (ResourceT IO) ByteString -> Warp.Connection -> ResourceT IO ())
interceptWith opts app req = case lookup "upgrade" $ Wai.requestHeaders req of
Just s
| S.map toLower s == "websocket" -> Just $ runWebSockets opts req' app
| otherwise -> Nothing
_ -> Nothing
where
req' = WS.RequestHttpPart (Wai.rawPathInfo req) (Wai.requestHeaders req)
(Wai.isSecure req)
runWebSockets :: WS.Protocol p
=> WS.WebSocketsOptions
-> WS.RequestHttpPart
-> (WS.Request -> WS.WebSockets p ())
-> Source (ResourceT IO) ByteString
-> Warp.Connection
-> ResourceT IO ()
runWebSockets opts req app source conn = do
step <- liftIO $ E.runIteratee $ WS.runWebSocketsWith opts req app send
source $$ sink (E.returnI step)
where
send = iterConnection conn
sink iter = await >>= maybe (close iter) (push iter)
push iter bs = do
step <- liftIO $ E.runIteratee $ E.enumList 1 [bs] E.$$ iter
case step of
E.Continue _ -> sink $ E.returnI step
E.Yield out inp -> maybe (return ()) leftover (streamToMaybe inp) >> return out
E.Error e -> liftIO $ monadThrow e
close iter = do
_ <- liftIO $ E.runIteratee $ E.enumEOF E.$$ iter
return ()
iterConnection :: Warp.Connection -> E.Iteratee ByteString IO ()
iterConnection c = E.continue go
where
go (E.Chunks []) = E.continue go
go (E.Chunks cs) = E.tryIO (Warp.connSendMany c cs) >> E.continue go
go E.EOF = E.continue go
streamToMaybe :: E.Stream S.ByteString -> Maybe S.ByteString
streamToMaybe E.EOF = Nothing
streamToMaybe (E.Chunks bs) = Just $ S.concat bs