Skip to content
This repository
branch: master
file 98 lines (86 sloc) 3.682 kb
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97
{-# LANGUAGE OverloadedStrings #-}
module Network.Wai.Handler.WebSockets
    ( websocketsApp
    , websocketsOr
    , isWebSocketsReq
    , getRequestHead
    , runWebSockets
    ) where

import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BC
import qualified Data.CaseInsensitive as CI
import Data.Conduit
import Data.IORef (newIORef, readIORef, writeIORef)
import Network.HTTP.Types (status500)
import qualified Network.Wai as Wai
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

--------------------------------------------------------------------------------
isWebSocketsReq :: Wai.Request -> Bool
isWebSocketsReq req =
    fmap CI.mk (lookup "upgrade" $ Wai.requestHeaders req) == Just "websocket"

--------------------------------------------------------------------------------
websocketsOr :: WS.ConnectionOptions
             -> WS.ServerApp
             -> Wai.Application
             -> Wai.Application
websocketsOr opts app backup req =
    case websocketsApp opts app req of
        Nothing -> backup req
        Just res -> return res

--------------------------------------------------------------------------------
websocketsApp :: WS.ConnectionOptions
              -> WS.ServerApp
              -> Wai.Request
              -> Maybe Wai.Response
websocketsApp opts app req
    | isWebSocketsReq req =
        Just $ flip Wai.responseRaw backup $ \src sink ->
            runWebSockets opts req' app src sink
    | otherwise = Nothing
  where
    req' = getRequestHead req
    backup = Wai.responseLBS status500 [("Content-Type", "text/plain")]
                "The web application attempted to send a WebSockets response, but WebSockets are not supported by your WAI handler."

--------------------------------------------------------------------------------
getRequestHead :: Wai.Request -> WS.RequestHead
getRequestHead req = WS.RequestHead
    (Wai.rawPathInfo req `BC.append` Wai.rawQueryString req)
    (Wai.requestHeaders req)
    (Wai.isSecure req)

--------------------------------------------------------------------------------
---- | Internal function to run the WebSocket io-streams using the conduit library
runWebSockets :: WS.ConnectionOptions
              -> WS.RequestHead
              -> (WS.PendingConnection -> IO a)
              -> Source IO ByteString
              -> Sink ByteString IO ()
              -> IO a
runWebSockets opts req app src sink = do

    is <- srcToInput src
    os <- sinkToOutput sink >>= Streams.builderStream

    let pc = WS.PendingConnection
                { WS.pendingOptions = opts
                , WS.pendingRequest = req
                , WS.pendingOnAccept = \_ -> return ()
                , WS.pendingIn = is
                , WS.pendingOut = os
                }

    app pc

srcToInput :: Source IO ByteString -> IO (InputStream ByteString)
srcToInput src0 = do
    (rsrc0, ()) <- src0 $$+ return ()
    ref <- newIORef rsrc0
    Streams.makeInputStream $ do
        rsrc <- readIORef ref
        (rsrc', mbs) <- rsrc $$++ await
        writeIORef ref rsrc'
        return mbs

sinkToOutput :: Sink ByteString IO () -> IO (OutputStream ByteString)
sinkToOutput sink =
    Streams.makeOutputStream output
  where
    output Nothing = return ()
    output (Just bs) = yield bs $$ sink
Something went wrong with that request. Please try again.