{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Yesod.WebSockets
    ( -- * Core API
      WebSocketsT
    , webSockets
    , webSocketsWith
    , webSocketsOptions
    , webSocketsOptionsWith
    , receiveData
    , receiveDataE
    , receiveDataMessageE
    , sendPing
    , sendPingE
    , sendClose
    , sendCloseE
    , sendTextData
    , sendTextDataE
    , sendBinaryData
    , sendBinaryDataE
    , sendDataMessageE
      -- * Conduit API
    , sourceWS
    , sinkWSText
    , sinkWSBinary
      -- * Async helpers
    , race
    , race_
    , concurrently
    , concurrently_
      -- * Re-exports from websockets
    , WS.defaultConnectionOptions
    , WS.ConnectionOptions (..)
    ) where

import           Control.Monad                  (forever, when)
import           Control.Monad.Reader           (ReaderT, runReaderT, MonadReader, ask)
import           Conduit
import qualified Network.Wai.Handler.WebSockets as WaiWS
import qualified Network.WebSockets             as WS
import qualified Yesod.Core                     as Y
import           UnliftIO (SomeException, tryAny, MonadIO, liftIO, MonadUnliftIO, withRunInIO, race, race_, concurrently, concurrently_)

-- | A transformer for a WebSockets handler.
--
-- Since 0.1.0
type WebSocketsT = ReaderT WS.Connection

-- | Attempt to run a WebSockets handler. This function first checks if the
-- client initiated a WebSockets connection and, if so, runs the provided
-- application, short-circuiting the rest of your handler. If the client did
-- not request a WebSockets connection, the rest of your handler will be called
-- instead.
--
-- Since 0.1.0
webSockets
  :: (MonadUnliftIO m, Y.MonadHandler m)
  => WebSocketsT m ()
  -> m ()
webSockets :: WebSocketsT m () -> m ()
webSockets = ConnectionOptions -> WebSocketsT m () -> m ()
forall (m :: * -> *).
(MonadUnliftIO m, MonadHandler m) =>
ConnectionOptions -> WebSocketsT m () -> m ()
webSocketsOptions ConnectionOptions
WS.defaultConnectionOptions

-- | Varient of 'webSockets' which allows you to specify
-- the WS.ConnectionOptions setttings when upgrading to a websocket connection.
--
-- Since 0.2.5
webSocketsOptions
  :: (MonadUnliftIO m, Y.MonadHandler m)
  => WS.ConnectionOptions
  -> WebSocketsT m ()
  -> m ()
webSocketsOptions :: ConnectionOptions -> WebSocketsT m () -> m ()
webSocketsOptions ConnectionOptions
opts = ConnectionOptions
-> (RequestHead -> m (Maybe AcceptRequest))
-> WebSocketsT m ()
-> m ()
forall (m :: * -> *).
(MonadUnliftIO m, MonadHandler m) =>
ConnectionOptions
-> (RequestHead -> m (Maybe AcceptRequest))
-> WebSocketsT m ()
-> m ()
webSocketsOptionsWith ConnectionOptions
opts ((RequestHead -> m (Maybe AcceptRequest))
 -> WebSocketsT m () -> m ())
-> (RequestHead -> m (Maybe AcceptRequest))
-> WebSocketsT m ()
-> m ()
forall a b. (a -> b) -> a -> b
$ m (Maybe AcceptRequest) -> RequestHead -> m (Maybe AcceptRequest)
forall a b. a -> b -> a
const (m (Maybe AcceptRequest) -> RequestHead -> m (Maybe AcceptRequest))
-> m (Maybe AcceptRequest)
-> RequestHead
-> m (Maybe AcceptRequest)
forall a b. (a -> b) -> a -> b
$ Maybe AcceptRequest -> m (Maybe AcceptRequest)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe AcceptRequest -> m (Maybe AcceptRequest))
-> Maybe AcceptRequest -> m (Maybe AcceptRequest)
forall a b. (a -> b) -> a -> b
$ AcceptRequest -> Maybe AcceptRequest
forall a. a -> Maybe a
Just (AcceptRequest -> Maybe AcceptRequest)
-> AcceptRequest -> Maybe AcceptRequest
forall a b. (a -> b) -> a -> b
$ Maybe ByteString -> Headers -> AcceptRequest
WS.AcceptRequest Maybe ByteString
forall a. Maybe a
Nothing []

-- | Varient of 'webSockets' which allows you to specify the 'WS.AcceptRequest'
-- setttings when upgrading to a websocket connection.
--
-- Since 0.2.4
webSocketsWith :: (MonadUnliftIO m, Y.MonadHandler m)
               => (WS.RequestHead -> m (Maybe WS.AcceptRequest))
               -- ^ A Nothing indicates that the websocket upgrade request should not happen
               -- and instead the rest of the handler will be called instead.  This allows
               -- you to use 'WS.getRequestSubprotocols' and only accept the request if
               -- a compatible subprotocol is given.  Also, the action runs before upgrading
               -- the request to websockets, so you can also use short-circuiting handler
               -- actions such as 'Y.invalidArgs'.
               -> WebSocketsT m ()
               -> m ()
webSocketsWith :: (RequestHead -> m (Maybe AcceptRequest))
-> WebSocketsT m () -> m ()
webSocketsWith = ConnectionOptions
-> (RequestHead -> m (Maybe AcceptRequest))
-> WebSocketsT m ()
-> m ()
forall (m :: * -> *).
(MonadUnliftIO m, MonadHandler m) =>
ConnectionOptions
-> (RequestHead -> m (Maybe AcceptRequest))
-> WebSocketsT m ()
-> m ()
webSocketsOptionsWith ConnectionOptions
WS.defaultConnectionOptions

-- | Varient of 'webSockets' which allows you to specify both
-- the WS.ConnectionOptions and the 'WS.AcceptRequest'
-- setttings when upgrading to a websocket connection.
--
-- Since 0.2.5
webSocketsOptionsWith :: (MonadUnliftIO m, Y.MonadHandler m)
               => WS.ConnectionOptions
               -- ^ Custom websockets options
               -> (WS.RequestHead -> m (Maybe WS.AcceptRequest))
               -- ^ A Nothing indicates that the websocket upgrade request should not happen
               -- and instead the rest of the handler will be called instead.  This allows
               -- you to use 'WS.getRequestSubprotocols' and only accept the request if
               -- a compatible subprotocol is given.  Also, the action runs before upgrading
               -- the request to websockets, so you can also use short-circuiting handler
               -- actions such as 'Y.invalidArgs'.
               -> WebSocketsT m ()
               -> m ()
webSocketsOptionsWith :: ConnectionOptions
-> (RequestHead -> m (Maybe AcceptRequest))
-> WebSocketsT m ()
-> m ()
webSocketsOptionsWith ConnectionOptions
wsConnOpts RequestHead -> m (Maybe AcceptRequest)
buildAr WebSocketsT m ()
inner = do
    Request
req <- m Request
forall (m :: * -> *). MonadHandler m => m Request
Y.waiRequest
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Request -> Bool
WaiWS.isWebSocketsReq Request
req) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        let rhead :: RequestHead
rhead = Request -> RequestHead
WaiWS.getRequestHead Request
req
        Maybe AcceptRequest
mar <- RequestHead -> m (Maybe AcceptRequest)
buildAr RequestHead
rhead
        case Maybe AcceptRequest
mar of
            Maybe AcceptRequest
Nothing -> () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            Just AcceptRequest
ar ->
                (IO ByteString -> (ByteString -> IO ()) -> m ()) -> m ()
forall (m :: * -> *) a.
(MonadHandler m, MonadUnliftIO m) =>
(IO ByteString -> (ByteString -> IO ()) -> m ()) -> m a
Y.sendRawResponseNoConduit
                  ((IO ByteString -> (ByteString -> IO ()) -> m ()) -> m ())
-> (IO ByteString -> (ByteString -> IO ()) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \IO ByteString
src ByteString -> IO ()
sink -> ((forall a. m a -> IO a) -> IO ()) -> m ()
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO ()) -> m ())
-> ((forall a. m a -> IO a) -> IO ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
runInIO -> ConnectionOptions
-> RequestHead
-> (PendingConnection -> IO ())
-> IO ByteString
-> (ByteString -> IO ())
-> IO ()
forall a.
ConnectionOptions
-> RequestHead
-> (PendingConnection -> IO a)
-> IO ByteString
-> (ByteString -> IO ())
-> IO a
WaiWS.runWebSockets
                    ConnectionOptions
wsConnOpts
                    RequestHead
rhead
                    (\PendingConnection
pconn -> do
                        Connection
conn <- PendingConnection -> AcceptRequest -> IO Connection
WS.acceptRequestWith PendingConnection
pconn AcceptRequest
ar
                        let app :: IO ()
app = m () -> IO ()
forall a. m a -> IO a
runInIO (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ WebSocketsT m () -> Connection -> m ()
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT WebSocketsT m ()
inner Connection
conn
#if MIN_VERSION_websockets(0,12,6)
                        Connection -> Int -> IO () -> IO () -> IO ()
forall a. Connection -> Int -> IO () -> IO a -> IO a
WS.withPingThread Connection
conn Int
30 (() -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO ()
app
#else
                        WS.forkPingThread conn 30
                        app
#endif
                    )
                    IO ByteString
src
                    ByteString -> IO ()
sink

-- | Wrapper for capturing exceptions
wrapWSE :: (MonadIO m, MonadReader WS.Connection m)
        => (WS.Connection -> a -> IO ())
        -> a
        -> m (Either SomeException ())
wrapWSE :: (Connection -> a -> IO ()) -> a -> m (Either SomeException ())
wrapWSE Connection -> a -> IO ()
ws a
x = do
  Connection
conn <- m Connection
forall r (m :: * -> *). MonadReader r m => m r
ask
  IO (Either SomeException ()) -> m (Either SomeException ())
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either SomeException ()) -> m (Either SomeException ()))
-> IO (Either SomeException ()) -> m (Either SomeException ())
forall a b. (a -> b) -> a -> b
$ IO () -> IO (Either SomeException ())
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny (IO () -> IO (Either SomeException ()))
-> IO () -> IO (Either SomeException ())
forall a b. (a -> b) -> a -> b
$ Connection -> a -> IO ()
ws Connection
conn a
x

wrapWS :: (MonadIO m, MonadReader WS.Connection m)
       => (WS.Connection -> a -> IO ())
       -> a
       -> m ()
wrapWS :: (Connection -> a -> IO ()) -> a -> m ()
wrapWS Connection -> a -> IO ()
ws a
x = do
  Connection
conn <- m Connection
forall r (m :: * -> *). MonadReader r m => m r
ask
  IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Connection -> a -> IO ()
ws Connection
conn a
x

-- | Receive a piece of data from the client.
--
-- Since 0.1.0
receiveData
  :: (MonadIO m, MonadReader WS.Connection m, WS.WebSocketsData a)
  => m a
receiveData :: m a
receiveData = do
  Connection
conn <- m Connection
forall r (m :: * -> *). MonadReader r m => m r
ask
  IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Connection -> IO a
forall a. WebSocketsData a => Connection -> IO a
WS.receiveData Connection
conn

-- | Receive a piece of data from the client.
-- Capture SomeException as the result or operation
-- Since 0.2.2
receiveDataE
  :: (MonadIO m, MonadReader WS.Connection m, WS.WebSocketsData a)
  => m (Either SomeException a)
receiveDataE :: m (Either SomeException a)
receiveDataE = do
  Connection
conn <- m Connection
forall r (m :: * -> *). MonadReader r m => m r
ask
  IO (Either SomeException a) -> m (Either SomeException a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either SomeException a) -> m (Either SomeException a))
-> IO (Either SomeException a) -> m (Either SomeException a)
forall a b. (a -> b) -> a -> b
$ IO a -> IO (Either SomeException a)
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny (IO a -> IO (Either SomeException a))
-> IO a -> IO (Either SomeException a)
forall a b. (a -> b) -> a -> b
$ Connection -> IO a
forall a. WebSocketsData a => Connection -> IO a
WS.receiveData Connection
conn

-- | Receive an application message.
-- Capture SomeException as the result or operation
-- Since 0.2.3
receiveDataMessageE
  :: (MonadIO m, MonadReader WS.Connection m)
  => m (Either SomeException WS.DataMessage)
receiveDataMessageE :: m (Either SomeException DataMessage)
receiveDataMessageE = do
  Connection
conn <- m Connection
forall r (m :: * -> *). MonadReader r m => m r
ask
  IO (Either SomeException DataMessage)
-> m (Either SomeException DataMessage)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either SomeException DataMessage)
 -> m (Either SomeException DataMessage))
-> IO (Either SomeException DataMessage)
-> m (Either SomeException DataMessage)
forall a b. (a -> b) -> a -> b
$ IO DataMessage -> IO (Either SomeException DataMessage)
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny (IO DataMessage -> IO (Either SomeException DataMessage))
-> IO DataMessage -> IO (Either SomeException DataMessage)
forall a b. (a -> b) -> a -> b
$ Connection -> IO DataMessage
WS.receiveDataMessage Connection
conn

-- | Send a textual message to the client.
--
-- Since 0.1.0
sendTextData
  :: (MonadIO m, WS.WebSocketsData a, MonadReader WS.Connection m)
  => a
  -> m ()
sendTextData :: a -> m ()
sendTextData = (Connection -> a -> IO ()) -> a -> m ()
forall (m :: * -> *) a.
(MonadIO m, MonadReader Connection m) =>
(Connection -> a -> IO ()) -> a -> m ()
wrapWS Connection -> a -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendTextData

-- | Send a textual message to the client.
-- Capture SomeException as the result or operation
-- and can be used like
-- `either handle_exception return =<< sendTextDataE ("Welcome" :: Text)`
-- Since 0.2.2
sendTextDataE
  :: (MonadIO m, WS.WebSocketsData a, MonadReader WS.Connection m)
  => a
  -> m (Either SomeException ())
sendTextDataE :: a -> m (Either SomeException ())
sendTextDataE = (Connection -> a -> IO ()) -> a -> m (Either SomeException ())
forall (m :: * -> *) a.
(MonadIO m, MonadReader Connection m) =>
(Connection -> a -> IO ()) -> a -> m (Either SomeException ())
wrapWSE Connection -> a -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendTextData

-- | Send a binary message to the client.
--
-- Since 0.1.0
sendBinaryData
  :: (MonadIO m, WS.WebSocketsData a, MonadReader WS.Connection m)
  => a
  -> m ()
sendBinaryData :: a -> m ()
sendBinaryData = (Connection -> a -> IO ()) -> a -> m ()
forall (m :: * -> *) a.
(MonadIO m, MonadReader Connection m) =>
(Connection -> a -> IO ()) -> a -> m ()
wrapWS Connection -> a -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendBinaryData

-- | Send a binary message to the client.
-- Capture SomeException as the result of operation
-- Since 0.2.2
sendBinaryDataE
  :: (MonadIO m, WS.WebSocketsData a, MonadReader WS.Connection m)
  => a
  -> m (Either SomeException ())
sendBinaryDataE :: a -> m (Either SomeException ())
sendBinaryDataE = (Connection -> a -> IO ()) -> a -> m (Either SomeException ())
forall (m :: * -> *) a.
(MonadIO m, MonadReader Connection m) =>
(Connection -> a -> IO ()) -> a -> m (Either SomeException ())
wrapWSE Connection -> a -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendBinaryData

-- | Send a ping message to the client.
--
-- Since 0.2.2
sendPing
  :: (MonadIO m, WS.WebSocketsData a, MonadReader WS.Connection m)
  => a
  -> m ()
sendPing :: a -> m ()
sendPing = (Connection -> a -> IO ()) -> a -> m ()
forall (m :: * -> *) a.
(MonadIO m, MonadReader Connection m) =>
(Connection -> a -> IO ()) -> a -> m ()
wrapWS Connection -> a -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendPing

-- | Send a ping message to the client.
-- Capture SomeException as the result of operation
-- Since 0.2.2
sendPingE
  :: (MonadIO m, WS.WebSocketsData a, MonadReader WS.Connection m)
  => a
  -> m (Either SomeException ())
sendPingE :: a -> m (Either SomeException ())
sendPingE = (Connection -> a -> IO ()) -> a -> m (Either SomeException ())
forall (m :: * -> *) a.
(MonadIO m, MonadReader Connection m) =>
(Connection -> a -> IO ()) -> a -> m (Either SomeException ())
wrapWSE Connection -> a -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendPing

-- | Send a DataMessage to the client.
-- Capture SomeException as the result of operation
-- Since 0.2.3
sendDataMessageE
  :: (MonadIO m, MonadReader WS.Connection m)
  => WS.DataMessage
  -> m (Either SomeException ())
sendDataMessageE :: DataMessage -> m (Either SomeException ())
sendDataMessageE DataMessage
x = do
  Connection
conn <- m Connection
forall r (m :: * -> *). MonadReader r m => m r
ask
  IO (Either SomeException ()) -> m (Either SomeException ())
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either SomeException ()) -> m (Either SomeException ()))
-> IO (Either SomeException ()) -> m (Either SomeException ())
forall a b. (a -> b) -> a -> b
$ IO () -> IO (Either SomeException ())
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny (IO () -> IO (Either SomeException ()))
-> IO () -> IO (Either SomeException ())
forall a b. (a -> b) -> a -> b
$ Connection -> DataMessage -> IO ()
WS.sendDataMessage Connection
conn DataMessage
x

-- | Send a close request to the client.
--
-- Since 0.2.2
sendClose
  :: (MonadIO m, WS.WebSocketsData a, MonadReader WS.Connection m)
  => a
  -> m ()
sendClose :: a -> m ()
sendClose = (Connection -> a -> IO ()) -> a -> m ()
forall (m :: * -> *) a.
(MonadIO m, MonadReader Connection m) =>
(Connection -> a -> IO ()) -> a -> m ()
wrapWS Connection -> a -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendClose

-- | Send a close request to the client.
-- Capture SomeException as the result of operation
-- Since 0.2.2
sendCloseE
  :: (MonadIO m, WS.WebSocketsData a, MonadReader WS.Connection m)
  => a
  -> m (Either SomeException ())
sendCloseE :: a -> m (Either SomeException ())
sendCloseE = (Connection -> a -> IO ()) -> a -> m (Either SomeException ())
forall (m :: * -> *) a.
(MonadIO m, MonadReader Connection m) =>
(Connection -> a -> IO ()) -> a -> m (Either SomeException ())
wrapWSE Connection -> a -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendClose

-- | A @Source@ of WebSockets data from the user.
--
-- Since 0.1.0
sourceWS
  :: (MonadIO m, WS.WebSocketsData a, MonadReader WS.Connection m)
  => ConduitT i a m ()
sourceWS :: ConduitT i a m ()
sourceWS = ConduitT i a m () -> ConduitT i a m ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (ConduitT i a m () -> ConduitT i a m ())
-> ConduitT i a m () -> ConduitT i a m ()
forall a b. (a -> b) -> a -> b
$ m a -> ConduitT i a m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m a
forall (m :: * -> *) a.
(MonadIO m, MonadReader Connection m, WebSocketsData a) =>
m a
receiveData ConduitT i a m a -> (a -> ConduitT i a m ()) -> ConduitT i a m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> ConduitT i a m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield

-- | A @Sink@ for sending textual data to the user.
--
-- Since 0.1.0
sinkWSText
  :: (MonadIO m, WS.WebSocketsData a, MonadReader WS.Connection m)
  => ConduitT a o m ()
sinkWSText :: ConduitT a o m ()
sinkWSText = (a -> m ()) -> ConduitT a o m ()
forall (m :: * -> *) a o.
Monad m =>
(a -> m ()) -> ConduitT a o m ()
mapM_C a -> m ()
forall (m :: * -> *) a.
(MonadIO m, WebSocketsData a, MonadReader Connection m) =>
a -> m ()
sendTextData

-- | A @Sink@ for sending binary data to the user.
--
-- Since 0.1.0
sinkWSBinary
  :: (MonadIO m, WS.WebSocketsData a, MonadReader WS.Connection m)
  => ConduitT a o m ()
sinkWSBinary :: ConduitT a o m ()
sinkWSBinary = (a -> m ()) -> ConduitT a o m ()
forall (m :: * -> *) a o.
Monad m =>
(a -> m ()) -> ConduitT a o m ()
mapM_C a -> m ()
forall (m :: * -> *) a.
(MonadIO m, WebSocketsData a, MonadReader Connection m) =>
a -> m ()
sendBinaryData