{-|
Module: IHP.WebSocket
Description: Building blocks for websocket applications
Copyright: (c) digitally induced GmbH, 2020
-}
{-# LANGUAGE AllowAmbiguousTypes #-}
module IHP.WebSocket
( WSApp (..)
, startWSApp
, setState
, getState
, receiveData
, receiveDataMessage
, sendTextData
, sendJSON
)
where

import IHP.Prelude
import qualified Network.WebSockets as Websocket
import IHP.ApplicationContext
import IHP.Controller.RequestContext
import qualified Data.UUID as UUID
import qualified Data.Maybe as Maybe
import qualified Control.Exception as Exception
import IHP.Controller.Context
import qualified Data.Aeson as Aeson

import qualified IHP.Log.Types as Log
import qualified IHP.Log as Log

import Control.Concurrent.Chan
import Control.Concurrent
import System.Timeout
import Data.Function (fix)
import qualified Network.WebSockets.Connection as WebSocket

class WSApp state where
    initialState :: state

    run :: (?state :: IORef state, ?context :: ControllerContext, ?applicationContext :: ApplicationContext, ?modelContext :: ModelContext, ?connection :: Websocket.Connection) => IO ()
    run = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

    onPing :: (?state :: IORef state, ?context :: ControllerContext, ?applicationContext :: ApplicationContext, ?modelContext :: ModelContext, ?connection :: Websocket.Connection) => IO ()
    onPing = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

    onClose :: (?state :: IORef state, ?context :: ControllerContext, ?applicationContext :: ApplicationContext, ?modelContext :: ModelContext, ?connection :: Websocket.Connection) => IO ()
    onClose = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

startWSApp :: forall state. (WSApp state, ?applicationContext :: ApplicationContext, ?requestContext :: RequestContext, ?context :: ControllerContext, ?modelContext :: ModelContext) => Websocket.Connection -> IO ()
startWSApp :: forall state.
(WSApp state, ?applicationContext::ApplicationContext,
 ?requestContext::RequestContext, ?context::ControllerContext,
 ?modelContext::ModelContext) =>
Connection -> IO ()
startWSApp Connection
connection' = do
    IORef state
state <- forall a. a -> IO (IORef a)
newIORef (forall state. WSApp state => state
initialState @state)
    IORef UTCTime
lastPongAt <- IO UTCTime
getCurrentTime forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a. a -> IO (IORef a)
newIORef


    let connection :: Connection
connection = IORef UTCTime -> Connection -> Connection
installPongHandler IORef UTCTime
lastPongAt Connection
connection'
    let ?state = IORef state
state
    let ?connection = Connection
connection
    let pingHandler :: IO ()
pingHandler = do
            Int
seconds <- IORef UTCTime -> IO Int
secondsSinceLastPong IORef UTCTime
lastPongAt
            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
seconds forall a. Ord a => a -> a -> Bool
> Int
pingWaitTime forall a. Num a => a -> a -> a
* Int
2) (forall e a. Exception e => e -> IO a
throwIO PongTimeout
PongTimeout)
            forall state.
(WSApp state, ?state::IORef state, ?context::ControllerContext,
 ?applicationContext::ApplicationContext,
 ?modelContext::ModelContext, ?connection::Connection) =>
IO ()
onPing @state

    Either SomeException ()
result <- forall e a. Exception e => IO a -> IO (Either e a)
Exception.try ((forall a. Connection -> Int -> IO () -> IO a -> IO a
WebSocket.withPingThread Connection
connection Int
pingWaitTime IO ()
pingHandler (forall state.
(WSApp state, ?state::IORef state, ?context::ControllerContext,
 ?applicationContext::ApplicationContext,
 ?modelContext::ModelContext, ?connection::Connection) =>
IO ()
run @state)) forall a b. IO a -> IO b -> IO a
`Exception.finally` forall state.
(WSApp state, ?state::IORef state, ?context::ControllerContext,
 ?applicationContext::ApplicationContext,
 ?modelContext::ModelContext, ?connection::Connection) =>
IO ()
onClose @state)
    case Either SomeException ()
result of
        Left (e :: SomeException
e@Exception.SomeException{}) ->
            case forall e. Exception e => SomeException -> Maybe e
Exception.fromException SomeException
e of
                (Just ConnectionException
Websocket.ConnectionClosed) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                (Just (Websocket.CloseRequest {})) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                (Just ConnectionException
other) -> forall a. Text -> a
error (Text
"Unhandled Websocket exception: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> Text
show ConnectionException
other)
                Maybe ConnectionException
Nothing -> forall context string.
(?context::context, LoggingProvider context, ToLogStr string) =>
string -> IO ()
Log.error (forall a. Show a => a -> Text
tshow SomeException
e)
        Right ()
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

setState :: (?state :: IORef state) => state -> IO ()
setState :: forall state. (?state::IORef state) => state -> IO ()
setState state
newState = forall a. IORef a -> a -> IO ()
writeIORef ?state::IORef state
?state state
newState

getState :: (?state :: IORef state) => IO state
getState :: forall state. (?state::IORef state) => IO state
getState = forall a. IORef a -> IO a
readIORef ?state::IORef state
?state

receiveData :: (?connection :: Websocket.Connection, Websocket.WebSocketsData a) => IO a
receiveData :: forall a. (?connection::Connection, WebSocketsData a) => IO a
receiveData = forall a. WebSocketsData a => Connection -> IO a
Websocket.receiveData ?connection::Connection
?connection

receiveDataMessage :: (?connection :: Websocket.Connection) => IO Websocket.DataMessage
receiveDataMessage :: (?connection::Connection) => IO DataMessage
receiveDataMessage = Connection -> IO DataMessage
Websocket.receiveDataMessage ?connection::Connection
?connection

sendTextData :: (?connection :: Websocket.Connection, Websocket.WebSocketsData text) => text -> IO ()
sendTextData :: forall text.
(?connection::Connection, WebSocketsData text) =>
text -> IO ()
sendTextData text
text = forall a. WebSocketsData a => Connection -> a -> IO ()
Websocket.sendTextData ?connection::Connection
?connection text
text

-- | Json encode a payload and send it over the websocket wire
--
-- __Example:__
--
-- > message <- Aeson.decode <$> receiveData @LByteString
-- >
-- > case message of
-- >     Just decodedMessage -> handleMessage decodedMessage
-- >     Nothing -> sendJSON FailedToDecodeMessageError
--
sendJSON :: (?connection :: Websocket.Connection, Aeson.ToJSON value) => value -> IO ()
sendJSON :: forall value.
(?connection::Connection, ToJSON value) =>
value -> IO ()
sendJSON value
payload = forall text.
(?connection::Connection, WebSocketsData text) =>
text -> IO ()
sendTextData (forall a. ToJSON a => a -> ByteString
Aeson.encode value
payload)

instance Websocket.WebSocketsData UUID where
    fromDataMessage :: DataMessage -> UUID
fromDataMessage (Websocket.Text ByteString
byteString Maybe Text
_) = ByteString -> Maybe UUID
UUID.fromLazyASCIIBytes ByteString
byteString forall {t1} {t2}. t1 -> (t1 -> t2) -> t2
|> forall a. HasCallStack => Maybe a -> a
Maybe.fromJust
    fromDataMessage (Websocket.Binary ByteString
byteString) = ByteString -> Maybe UUID
UUID.fromLazyASCIIBytes ByteString
byteString forall {t1} {t2}. t1 -> (t1 -> t2) -> t2
|> forall a. HasCallStack => Maybe a -> a
Maybe.fromJust
    fromLazyByteString :: ByteString -> UUID
fromLazyByteString ByteString
byteString = ByteString -> Maybe UUID
UUID.fromLazyASCIIBytes ByteString
byteString forall {t1} {t2}. t1 -> (t1 -> t2) -> t2
|> forall a. HasCallStack => Maybe a -> a
Maybe.fromJust
    toLazyByteString :: UUID -> ByteString
toLazyByteString = UUID -> ByteString
UUID.toLazyASCIIBytes

data PongTimeout
    = PongTimeout
    deriving (Int -> PongTimeout -> ShowS
[PongTimeout] -> ShowS
PongTimeout -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PongTimeout] -> ShowS
$cshowList :: [PongTimeout] -> ShowS
show :: PongTimeout -> String
$cshow :: PongTimeout -> String
showsPrec :: Int -> PongTimeout -> ShowS
$cshowsPrec :: Int -> PongTimeout -> ShowS
Show)

instance Exception PongTimeout

pingWaitTime :: Int
pingWaitTime :: Int
pingWaitTime = Int
30

installPongHandler :: IORef UTCTime -> WebSocket.Connection -> WebSocket.Connection
installPongHandler :: IORef UTCTime -> Connection -> Connection
installPongHandler IORef UTCTime
lastPongAt Connection
connection =
    Connection
connection { connectionOptions :: ConnectionOptions
WebSocket.connectionOptions = Connection
connection.connectionOptions { connectionOnPong :: IO ()
WebSocket.connectionOnPong = IORef UTCTime -> IO ()
connectionOnPong IORef UTCTime
lastPongAt }  }

connectionOnPong :: IORef UTCTime -> IO ()
connectionOnPong :: IORef UTCTime -> IO ()
connectionOnPong IORef UTCTime
lastPongAt = do
    UTCTime
now <- IO UTCTime
getCurrentTime
    forall a. IORef a -> a -> IO ()
writeIORef IORef UTCTime
lastPongAt UTCTime
now

secondsSinceLastPong :: IORef UTCTime -> IO Int
secondsSinceLastPong :: IORef UTCTime -> IO Int
secondsSinceLastPong IORef UTCTime
lastPongAt = do
    UTCTime
now <- IO UTCTime
getCurrentTime
    UTCTime
last <- forall a. IORef a -> IO a
readIORef IORef UTCTime
lastPongAt
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. (RealFrac a, Integral b) => a -> b
ceiling forall a b. (a -> b) -> a -> b
$ NominalDiffTime -> Pico
nominalDiffTimeToSeconds forall a b. (a -> b) -> a -> b
$ UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
now UTCTime
last