--------------------------------------------------------------------------------
-- | This provides a simple stand-alone server for 'WebSockets' applications.
-- Note that in production you want to use a real webserver such as snap or
-- warp.
{-# LANGUAGE OverloadedStrings #-}
module Network.WebSockets.Server
    ( ServerApp
    , runServer
    , ServerOptions (..)
    , defaultServerOptions
    , runServerWithOptions
    , runServerWith
    , makeListenSocket
    , makePendingConnection
    , makePendingConnectionFromStream

    , PongTimeout
    ) where


--------------------------------------------------------------------------------
import           Control.Concurrent            (threadDelay)
import qualified Control.Concurrent.Async      as Async
import           Control.Exception             (Exception, allowInterrupt,
                                                bracket, bracketOnError,
                                                finally, mask_, throwIO)
import           Control.Monad                 (forever, void, when)
import qualified Data.IORef                    as IORef
import           Data.Maybe                    (isJust)
import           Network.Socket                (Socket)
import qualified Network.Socket                as S
import qualified System.Clock                  as Clock


--------------------------------------------------------------------------------
import           Network.WebSockets.Connection
import           Network.WebSockets.Http
import qualified Network.WebSockets.Stream     as Stream
import           Network.WebSockets.Types


--------------------------------------------------------------------------------
-- | WebSockets application that can be ran by a server. Once this 'IO' action
-- finishes, the underlying socket is closed automatically.
type ServerApp = PendingConnection -> IO ()


--------------------------------------------------------------------------------
-- | Provides a simple server. This function blocks forever.  Note that this
-- is merely provided for quick-and-dirty or internal applications, but for real
-- applications, you should use a real server.
--
-- For example:
--
-- * Performance is reasonable under load, but:
-- * No protection against DoS attacks is provided.
-- * No logging is performed.
-- * ...
--
-- Glue for using this package with real servers is provided by:
--
-- * <https://hackage.haskell.org/package/wai-websockets>
--
-- * <https://hackage.haskell.org/package/websockets-snap>
runServer :: String     -- ^ Address to bind
          -> Int        -- ^ Port to listen on
          -> ServerApp  -- ^ Application
          -> IO ()      -- ^ Never returns
runServer :: String -> Int -> ServerApp -> IO ()
runServer String
host Int
port ServerApp
app = String -> Int -> ConnectionOptions -> ServerApp -> IO ()
runServerWith String
host Int
port ConnectionOptions
defaultConnectionOptions ServerApp
app


--------------------------------------------------------------------------------
-- | A version of 'runServer' which allows you to customize some options.
runServerWith :: String -> Int -> ConnectionOptions -> ServerApp -> IO ()
runServerWith :: String -> Int -> ConnectionOptions -> ServerApp -> IO ()
runServerWith String
host Int
port ConnectionOptions
opts = ServerOptions -> ServerApp -> IO ()
forall a. ServerOptions -> ServerApp -> IO a
runServerWithOptions ServerOptions
defaultServerOptions
    { serverHost              = host
    , serverPort              = port
    , serverConnectionOptions = opts
    }
{-# DEPRECATED runServerWith "Use 'runServerWithOptions' instead" #-}


--------------------------------------------------------------------------------
data ServerOptions = ServerOptions
    { ServerOptions -> String
serverHost              :: String
    , ServerOptions -> Int
serverPort              :: Int
    , ServerOptions -> ConnectionOptions
serverConnectionOptions :: ConnectionOptions
    -- | Require a pong from the client every N seconds; otherwise kill the
    -- connection.  If you use this, you should also use 'withPingThread' to
    -- send a ping at a smaller interval; for example N/2.
    , ServerOptions -> Maybe Int
serverRequirePong       :: Maybe Int
    }


--------------------------------------------------------------------------------
defaultServerOptions :: ServerOptions
defaultServerOptions :: ServerOptions
defaultServerOptions = ServerOptions
    { serverHost :: String
serverHost              = String
"127.0.0.1"
    , serverPort :: Int
serverPort              = Int
8080
    , serverConnectionOptions :: ConnectionOptions
serverConnectionOptions = ConnectionOptions
defaultConnectionOptions
    , serverRequirePong :: Maybe Int
serverRequirePong       = Maybe Int
forall a. Maybe a
Nothing
    }


--------------------------------------------------------------------------------
-- | Customizable version of 'runServer'.  Never returns until killed.
--
-- Please use the 'defaultServerOptions' combined with record updates to set the
-- fields you want.  This way your code is unlikely to break on future changes.
runServerWithOptions :: ServerOptions -> ServerApp -> IO a
runServerWithOptions :: forall a. ServerOptions -> ServerApp -> IO a
runServerWithOptions ServerOptions
opts ServerApp
app = IO a -> IO a
forall a. IO a -> IO a
S.withSocketsDo (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$
    IO Socket -> (Socket -> IO ()) -> (Socket -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
    (String -> Int -> IO Socket
makeListenSocket String
host Int
port)
    Socket -> IO ()
S.close ((Socket -> IO a) -> IO a) -> (Socket -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Socket
sock -> IO a -> IO a
forall a. IO a -> IO a
mask_ (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ IO () -> IO a
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO a) -> IO () -> IO a
forall a b. (a -> b) -> a -> b
$ do
        IO ()
allowInterrupt
        (Socket
conn, SockAddr
_) <- Socket -> IO (Socket, SockAddr)
S.accept Socket
sock

        -- This IORef holds a time at which the thread may be killed.  This time
        -- can be extended by calling 'tickle'.
        IORef Int64
killRef <- Int64 -> IO (IORef Int64)
forall a. a -> IO (IORef a)
IORef.newIORef (Int64 -> IO (IORef Int64)) -> IO Int64 -> IO (IORef Int64)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
killDelay) (Int64 -> Int64) -> IO Int64 -> IO Int64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Int64
getSecs
        let tickle :: IO ()
tickle = IORef Int64 -> Int64 -> IO ()
forall a. IORef a -> a -> IO ()
IORef.writeIORef IORef Int64
killRef (Int64 -> IO ()) -> IO Int64 -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
killDelay) (Int64 -> Int64) -> IO Int64 -> IO Int64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Int64
getSecs

        -- Update the connection options to call 'tickle' whenever a pong is
        -- received.
        let connOpts' :: ConnectionOptions
connOpts'
                | Bool -> Bool
not Bool
useKiller = ConnectionOptions
connOpts
                | Bool
otherwise     = ConnectionOptions
connOpts
                    { connectionOnPong = tickle >> connectionOnPong connOpts
                    }

        -- Run the application.
        Async ()
appAsync  <- ((forall a. IO a -> IO a) -> IO ()) -> IO (Async ())
forall a. ((forall a. IO a -> IO a) -> IO a) -> IO (Async a)
Async.asyncWithUnmask (((forall a. IO a -> IO a) -> IO ()) -> IO (Async ()))
-> ((forall a. IO a -> IO a) -> IO ()) -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask ->
            (IO () -> IO ()
forall a. IO a -> IO a
unmask (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                Socket -> ConnectionOptions -> ServerApp -> IO ()
runApp Socket
conn ConnectionOptions
connOpts' ServerApp
app) IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally`
            (Socket -> IO ()
S.close Socket
conn)

        -- Install the killer if required.
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
useKiller (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO (Async ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Async ()) -> IO ()) -> IO (Async ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
Async.async (IORef Int64 -> Async () -> IO ()
forall {a}. IORef Int64 -> Async a -> IO ()
killer IORef Int64
killRef Async ()
appAsync)
  where
    host :: String
host     = ServerOptions -> String
serverHost ServerOptions
opts
    port :: Int
port     = ServerOptions -> Int
serverPort ServerOptions
opts
    connOpts :: ConnectionOptions
connOpts = ServerOptions -> ConnectionOptions
serverConnectionOptions ServerOptions
opts

    -- Get the current number of seconds on some clock.
    getSecs :: IO Int64
getSecs = TimeSpec -> Int64
Clock.sec (TimeSpec -> Int64) -> IO TimeSpec -> IO Int64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Clock -> IO TimeSpec
Clock.getTime Clock
Clock.Monotonic

    -- Parse the 'serverRequirePong' options.
    useKiller :: Bool
useKiller = Maybe Int -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Int -> Bool) -> Maybe Int -> Bool
forall a b. (a -> b) -> a -> b
$ ServerOptions -> Maybe Int
serverRequirePong ServerOptions
opts
    killDelay :: Int64
killDelay = Int64 -> (Int -> Int64) -> Maybe Int -> Int64
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int64
0 Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ServerOptions -> Maybe Int
serverRequirePong ServerOptions
opts)

    -- Thread that reads the killRef, and kills the application if enough time
    -- has passed.
    killer :: IORef Int64 -> Async a -> IO ()
killer IORef Int64
killRef Async a
appAsync = do
        Int64
killAt   <- IORef Int64 -> IO Int64
forall a. IORef a -> IO a
IORef.readIORef IORef Int64
killRef
        Int64
now      <- IO Int64
getSecs
        Maybe (Either SomeException a)
appState <- Async a -> IO (Maybe (Either SomeException a))
forall a. Async a -> IO (Maybe (Either SomeException a))
Async.poll Async a
appAsync
        case Maybe (Either SomeException a)
appState of
            -- Already finished/killed/crashed, we can give up.
            Just Either SomeException a
_ -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            -- Should not be killed yet.  Wait and try again.
            Maybe (Either SomeException a)
Nothing | Int64
now Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
< Int64
killAt -> do
                Int -> IO ()
threadDelay (Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
killDelay Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000)
                IORef Int64 -> Async a -> IO ()
killer IORef Int64
killRef Async a
appAsync
            -- Time to kill.
            Maybe (Either SomeException a)
_ -> Async a -> PongTimeout -> IO ()
forall e a. Exception e => Async a -> e -> IO ()
Async.cancelWith Async a
appAsync PongTimeout
PongTimeout


--------------------------------------------------------------------------------
-- | Create a standardized socket on which you can listen for incomming
-- connections. Should only be used for a quick and dirty solution! Should be
-- preceded by the call 'Network.Socket.withSocketsDo'.
makeListenSocket :: String -> Int -> IO Socket
makeListenSocket :: String -> Int -> IO Socket
makeListenSocket String
host Int
port = do
  AddrInfo
addr:[AddrInfo]
_ <- Maybe AddrInfo -> Maybe String -> Maybe String -> IO [AddrInfo]
S.getAddrInfo (AddrInfo -> Maybe AddrInfo
forall a. a -> Maybe a
Just AddrInfo
hints) (String -> Maybe String
forall a. a -> Maybe a
Just String
host) (String -> Maybe String
forall a. a -> Maybe a
Just (Int -> String
forall a. Show a => a -> String
show Int
port))
  IO Socket
-> (Socket -> IO ()) -> (Socket -> IO Socket) -> IO Socket
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracketOnError
    (Family -> SocketType -> ProtocolNumber -> IO Socket
S.socket (AddrInfo -> Family
S.addrFamily AddrInfo
addr) SocketType
S.Stream ProtocolNumber
S.defaultProtocol)
    Socket -> IO ()
S.close
    (\Socket
sock -> do
        ()
_     <- Socket -> SocketOption -> Int -> IO ()
S.setSocketOption Socket
sock SocketOption
S.ReuseAddr Int
1
        ()
_     <- Socket -> SocketOption -> Int -> IO ()
S.setSocketOption Socket
sock SocketOption
S.NoDelay   Int
1
        Socket -> SockAddr -> IO ()
S.bind Socket
sock (AddrInfo -> SockAddr
S.addrAddress AddrInfo
addr)
        Socket -> Int -> IO ()
S.listen Socket
sock Int
5
        Socket -> IO Socket
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Socket
sock
        )
  where
    hints :: AddrInfo
hints = AddrInfo
S.defaultHints { S.addrSocketType = S.Stream }


--------------------------------------------------------------------------------
runApp :: Socket
       -> ConnectionOptions
       -> ServerApp
       -> IO ()
runApp :: Socket -> ConnectionOptions -> ServerApp -> IO ()
runApp Socket
socket ConnectionOptions
opts ServerApp
app =
    IO PendingConnection -> ServerApp -> ServerApp -> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
        (Socket -> ConnectionOptions -> IO PendingConnection
makePendingConnection Socket
socket ConnectionOptions
opts)
        (Stream -> IO ()
Stream.close (Stream -> IO ()) -> (PendingConnection -> Stream) -> ServerApp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PendingConnection -> Stream
pendingStream)
        ServerApp
app


--------------------------------------------------------------------------------
-- | Turns a socket, connected to some client, into a 'PendingConnection'. The
-- 'PendingConnection' should be closed using 'Stream.close' later.
makePendingConnection
    :: Socket -> ConnectionOptions -> IO PendingConnection
makePendingConnection :: Socket -> ConnectionOptions -> IO PendingConnection
makePendingConnection Socket
socket ConnectionOptions
opts = do
    Stream
stream <- Socket -> IO Stream
Stream.makeSocketStream Socket
socket
    Stream -> ConnectionOptions -> IO PendingConnection
makePendingConnectionFromStream Stream
stream ConnectionOptions
opts


-- | More general version of 'makePendingConnection' for 'Stream.Stream'
-- instead of a 'Socket'.
makePendingConnectionFromStream
    :: Stream.Stream -> ConnectionOptions -> IO PendingConnection
makePendingConnectionFromStream :: Stream -> ConnectionOptions -> IO PendingConnection
makePendingConnectionFromStream Stream
stream ConnectionOptions
opts = do
    -- TODO: we probably want to send a 40x if the request is bad?
    Maybe RequestHead
mbRequest <- Stream -> Parser RequestHead -> IO (Maybe RequestHead)
forall a. Stream -> Parser a -> IO (Maybe a)
Stream.parse Stream
stream (Bool -> Parser RequestHead
decodeRequestHead Bool
False)
    case Maybe RequestHead
mbRequest of
        Maybe RequestHead
Nothing      -> ConnectionException -> IO PendingConnection
forall e a. Exception e => e -> IO a
throwIO ConnectionException
ConnectionClosed
        Just RequestHead
request -> PendingConnection -> IO PendingConnection
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return PendingConnection
            { pendingOptions :: ConnectionOptions
pendingOptions  = ConnectionOptions
opts
            , pendingRequest :: RequestHead
pendingRequest  = RequestHead
request
            , pendingOnAccept :: Connection -> IO ()
pendingOnAccept = \Connection
_ -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            , pendingStream :: Stream
pendingStream   = Stream
stream
            }


--------------------------------------------------------------------------------
-- | Internally used exception type used to kill connections if there
-- is a pong timeout.
data PongTimeout = PongTimeout deriving Int -> PongTimeout -> ShowS
[PongTimeout] -> ShowS
PongTimeout -> String
(Int -> PongTimeout -> ShowS)
-> (PongTimeout -> String)
-> ([PongTimeout] -> ShowS)
-> Show PongTimeout
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PongTimeout -> ShowS
showsPrec :: Int -> PongTimeout -> ShowS
$cshow :: PongTimeout -> String
show :: PongTimeout -> String
$cshowList :: [PongTimeout] -> ShowS
showList :: [PongTimeout] -> ShowS
Show


--------------------------------------------------------------------------------
instance Exception PongTimeout