88fcec9 · typeable/wai

21import System.IO.Streams (InputStream, OutputStream)

22import qualified System.IO.Streams as Streams

23

24−

24+--------------------------------------------------------------------------------

25-- | For use with 'settingsIntercept' from the Warp web server.

26−intercept :: WS.Protocol p

26+intercept :: WS.ServerApp

27− => (WS.Request -> WS.WebSockets p ())

28 -> Wai.Request

27 -> Wai.Request

29 -> Maybe (Source (ResourceT IO) ByteString -> Warp.Connection -> ResourceT IO ())

28 -> Maybe (Source (ResourceT IO) ByteString -> Warp.Connection -> ResourceT IO ())

30−intercept = interceptWith WS.defaultWebSocketsOptions

29+intercept = interceptWith WS.defaultConnectionOptions

31

30

31+--------------------------------------------------------------------------------

32-- | Variation of 'intercept' which allows custom options.

33−interceptWith :: WS.Protocol p

33+interceptWith :: WS.ConnectionOptions

34− => WS.WebSocketsOptions

34+ -> WS.ServerApp

35− -> (WS.Request -> WS.WebSockets p ())

36 -> Wai.Request

35 -> Wai.Request

37 -> Maybe (Source (ResourceT IO) ByteString -> Warp.Connection -> ResourceT IO ())

36 -> 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

39 Just s

38 Just s

40− | S.map toLower s == "websocket" -> Just $ runWebSockets opts req' app

39+ | BC.map toLower s == "websocket" -> Just $ runWebSockets opts req' app

41 | otherwise -> Nothing

40 | otherwise -> Nothing

42 _ -> Nothing

41 _ -> Nothing

43− where

42+ where

44− req' = WS.RequestHttpPart (Wai.rawPathInfo req) (Wai.requestHeaders req)

43+ req' = WS.RequestHead (Wai.rawPathInfo req) (Wai.requestHeaders req) (Wai.isSecure req)

45− (Wai.isSecure req)

46

44

45+--------------------------------------------------------------------------------

47−-- | Internal function to run the WebSocket iteratee using the conduit library

46+---- | Internal function to run the WebSocket io-streams using the conduit library

48−runWebSockets :: WS.Protocol p

47+runWebSockets :: WS.ConnectionOptions

49− => WS.WebSocketsOptions

48+ -> WS.RequestHead

50− -> WS.RequestHttpPart

49+ -> WS.ServerApp

51− -> (WS.Request -> WS.WebSockets p ())

52 -> Source (ResourceT IO) ByteString

50 -> Source (ResourceT IO) ByteString

53 -> Warp.Connection

51 -> Warp.Connection

54 -> ResourceT IO ()

52 -> ResourceT IO ()

55runWebSockets opts req app source conn = do

53runWebSockets 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

60

54

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

62

77

63− push iter bs = do

78+ ignore :: SomeException -> IO ()

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 ()

79+ ignore _ = return ()

72

80

73−iterConnection :: Warp.Connection -> E.Iteratee ByteString IO ()

81+------------------------------------------------------------------------------

74−iterConnection c = E.continue go

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+

75− where

92+ where

76− go (E.Chunks []) = E.continue go

93+ input = do

77− go (E.Chunks cs) = E.tryIO (Warp.connSendMany c cs) >> E.continue go

94+ s <- Warp.connRecv connection

78− go E.EOF = E.continue go

95+ return $! if BC.null s then Nothing else Just s

79

96

80−streamToMaybe :: E.Stream S.ByteString -> Maybe S.ByteString

97+ output Nothing = return $! ()

81−streamToMaybe E.EOF = Nothing

98+ output (Just s') = if BC.null s then return $! () else Warp.connSendAll connection s

82−streamToMaybe (E.Chunks bs) = Just $ S.concat bs

99+ where s = Builder.toByteString s'