{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Network.WebSockets.Skews
(
Server
, Args(..)
, RequestHandler
, start
, reinit
, threadId
, enqueRequestHandler
, enqueResponse
, replaceRequestHandlers
, setDefaultRequestHandler
, setDefaultResponse
, respondWith
, doNothing
, sendToClients
, forgetDefaultRequestHandler
, forgetDefaultResponse
, forgetReceivedRequests
, listeningPort
, listeningHost
, recentlyReceived
, countConnectedClients
) where
import Control.Applicative ((<|>))
import Control.Concurrent (ThreadId, forkIO)
import Control.Exception (handleJust)
import Control.Monad (mapM_)
import qualified Data.ByteString.Lazy as B
import qualified Data.IORef as IOR
import Data.Monoid (mempty)
import qualified Deque as Q
import qualified Network.WebSockets as WS
data Server =
Server
{ requestHandlerQueue :: !(IOR.IORef (Q.Deque RequestHandler))
, defaultRequestHandler :: !(IOR.IORef (Maybe RequestHandler))
, recentlyReceivedRef :: !(IOR.IORef (Q.Deque B.ByteString))
, clientConnections :: !(IOR.IORef [WS.Connection])
, threadId :: !ThreadId
, listeningPort :: !Int
, listeningHost :: !String
}
type RequestHandler = B.ByteString -> IO (Maybe B.ByteString)
respondWith :: B.ByteString -> RequestHandler
respondWith = const . return . Just
doNothing :: RequestHandler
doNothing = const $ return Nothing
data Args = Args
{ host :: String
, portNumber :: Int
}
start :: Args -> IO Server
start (Args listeningHost listeningPort) = do
requestHandlerQueue <- IOR.newIORef mempty
defaultRequestHandler <- IOR.newIORef Nothing
recentlyReceivedRef <- IOR.newIORef mempty
clientConnections <- IOR.newIORef mempty
threadId <- forkIO $ WS.runServer listeningHost listeningPort $ \pc -> do
c <- WS.acceptRequest pc
addClientConnection clientConnections c
let
loop = do
m <- WS.receive c
case m of
WS.DataMessage _ _ _ dat -> do
mbrh <- deque requestHandlerQueue
dmbrh <- IOR.readIORef defaultRequestHandler
let bs = WS.fromDataMessage dat
case mbrh <|> dmbrh of
Just rh -> maybe (return ()) (WS.sendBinaryData c) =<< rh bs
Nothing -> return ()
IOR.modifyIORef recentlyReceivedRef (Q.snoc bs)
loop
WS.ControlMessage (WS.Close _ _) ->
WS.sendClose c ("Bye" :: B.ByteString)
_other -> loop
ignoreConnectionClosed loop
return Server {..}
addClientConnection :: IOR.IORef [WS.Connection] -> WS.Connection -> IO ()
addClientConnection ccsr c =
IOR.atomicModifyIORef' ccsr (\ccs -> (c : ccs, ()))
enqueRequestHandler :: Server -> RequestHandler -> IO ()
enqueRequestHandler Server {..} rh =
IOR.atomicModifyIORef' requestHandlerQueue (\q -> (Q.snoc rh q, ()))
enqueResponse :: Server -> B.ByteString -> IO ()
enqueResponse s = enqueRequestHandler s . respondWith
setDefaultRequestHandler :: Server -> RequestHandler -> IO ()
setDefaultRequestHandler Server {..} =
IOR.atomicWriteIORef defaultRequestHandler . Just
replaceRequestHandlers :: Server -> [RequestHandler] -> IO ()
replaceRequestHandlers Server {..} =
IOR.atomicWriteIORef requestHandlerQueue . Q.fromList
setDefaultResponse :: Server -> B.ByteString -> IO ()
setDefaultResponse s = setDefaultRequestHandler s . respondWith
forgetDefaultRequestHandler :: Server -> IO ()
forgetDefaultRequestHandler Server {..} =
IOR.atomicWriteIORef defaultRequestHandler Nothing
forgetDefaultResponse :: Server -> IO ()
forgetDefaultResponse = forgetDefaultRequestHandler
forgetReceivedRequests :: Server -> IO ()
forgetReceivedRequests Server {..} =
IOR.atomicWriteIORef recentlyReceivedRef mempty
sendToClients :: Server -> B.ByteString -> IO ()
sendToClients Server {..} msg =
mapM_ (`WS.sendBinaryData` msg) =<< IOR.readIORef clientConnections
reinit :: Server -> IO ()
reinit Server {..} = do
IOR.atomicWriteIORef requestHandlerQueue mempty
IOR.atomicWriteIORef defaultRequestHandler Nothing
IOR.atomicWriteIORef recentlyReceivedRef mempty
mapM_ (ignoreConnectionClosed . (`WS.sendClose` ("Bye" :: B.ByteString)))
=<< IOR.readIORef clientConnections
IOR.atomicWriteIORef clientConnections []
recentlyReceived :: Server -> IO (Q.Deque B.ByteString)
recentlyReceived Server {..} = IOR.readIORef recentlyReceivedRef
deque :: IOR.IORef (Q.Deque a) -> IO (Maybe a)
deque qr = IOR.atomicModifyIORef' qr $ \q -> case Q.uncons q of
Just (x, qLeft) -> (qLeft, Just x)
_ -> (mempty, Nothing)
ignoreConnectionClosed :: IO () -> IO ()
ignoreConnectionClosed = handleJust selectClosed (const $ return ())
where
selectClosed :: WS.ConnectionException -> Maybe ()
selectClosed WS.ConnectionClosed = Just ()
selectClosed _ = Nothing
countConnectedClients :: Server -> IO Int
countConnectedClients Server {..} = length <$> IOR.readIORef clientConnections