basic snap modified from jaspervdj

server.hs

{-# LANGUAGE OverloadedStrings #-}
import Data.Char (isPunctuation, isSpace)
import Data.Monoid (mappend)
import Data.Text (Text)
import Control.Exception (fromException)
import Control.Monad (forM_)
import Control.Concurrent (MVar, newMVar, modifyMVar_, readMVar)
import Control.Monad.IO.Class (liftIO)
import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Text.Encoding as T
import Data.List (foldl')
import Data.Monoid

import qualified Network.WebSockets as WS

import Snap.Http.Server.Config
import Snap.Http.Server (httpServe)
import qualified Snap.Core as Snap
import qualified Snap.Internal.Http.Types as Snap
import qualified Snap.Types.Headers as Headers

simpleConfig :: Config m a
simpleConfig = foldl' (\accum new -> new accum) emptyConfig base where
    base = [hostName, accessLog, errorLog, locale, port, ip, cert, key, compr, verbose]
    
    hostName = setHostname (bsFromString "localhost")
    accessLog = setAccessLog (ConfigFileLog "log/access.log")
    errorLog = setErrorLog (ConfigFileLog "log/error.log")
    locale = setLocale "US"
    port = setSSLPort 9160
    ip = setSSLBind (bsFromString "127.0.0.1")
    cert = setSSLCert "server.crt"
    key = setSSLKey "server.key"
    compr = setCompression True
    verbose = setVerbose True
    
    bsFromString = T.encodeUtf8 . T.pack

main :: IO ()
main = mainSSL

-- Snap without TLS
mainHTTP :: IO ()
mainHTTP = do
    state <- newMVar newServerState
    httpServe (setPort 9160 mempty) $ runWebSocketsSnap $ application state

-- Snap With TLS
mainSSL :: IO ()
mainSSL = do
    state <- newMVar newServerState
    httpServe simpleConfig $ runWebSocketsSnap $ application state


-- | The following function escapes from the current 'Snap.Snap' handler, and
-- continues processing the 'WS.WebSockets' action. The action to be executed
-- takes the 'WS.Request' as a parameter, because snap has already read this
-- from the socket.
runWebSocketsSnap :: WS.Protocol p
                  => (WS.Request -> WS.WebSockets p ())
                  -> Snap.Snap ()
runWebSocketsSnap = runWebSocketsSnapWith WS.defaultWebSocketsOptions

-- | Variant of 'runWebSocketsSnap' which allows custom options
runWebSocketsSnapWith :: WS.Protocol p
                      => WS.WebSocketsOptions
                      -> (WS.Request -> WS.WebSockets p ())
                      -> Snap.Snap ()
runWebSocketsSnapWith options ws = do
    rq <- Snap.getRequest
    Snap.escapeHttp $ \tickle writeEnd ->
        let options' = options
                { WS.onPong = tickle 30 >> WS.onPong options
                }

        in (WS.runWebSocketsWith options' (fromSnapRequest rq) ws writeEnd)

-- | Convert a snap request to a websockets request
fromSnapRequest :: Snap.Request -> WS.RequestHttpPart
fromSnapRequest rq = WS.RequestHttpPart
    { WS.requestHttpPath    = Snap.rqURI rq
    , WS.requestHttpHeaders = Headers.toList (Snap.rqHeaders rq)
    }







type Client = (Text, WS.Sink WS.Hybi00)
type ServerState = [Client]

newServerState :: ServerState
newServerState = []

numClients :: ServerState -> Int
numClients = length


clientExists :: Client -> ServerState -> Bool
clientExists client = any ((== fst client) . fst)


addClient :: Client -> ServerState -> ServerState
addClient client clients = client : clients


removeClient :: Client -> ServerState -> ServerState
removeClient client = filter ((/= fst client) . fst)


broadcast :: Text -> ServerState -> IO ()
broadcast message clients = do
    T.putStrLn message
    forM_ clients $ \(_, sink) -> WS.sendSink sink $ WS.textData message



application :: MVar ServerState -> WS.Request -> WS.WebSockets WS.Hybi00 ()
application state rq = do
    WS.acceptRequest rq
    WS.getVersion >>= liftIO . putStrLn . ("Client version: " ++)
    sink <- WS.getSink
    msg <- WS.receiveData
    clients <- liftIO $ readMVar state
    case msg of
        _   | not (prefix `T.isPrefixOf` msg) -> WS.sendTextData ("Wrong announcement" :: Text)
            | any ($ fst client) [T.null, T.any isPunctuation, T.any isSpace] ->
                    WS.sendTextData ("Name cannot " `mappend`
                        "contain punctuation or whitespace, and " `mappend`
                        "cannot be empty" :: Text)
            | clientExists client clients -> WS.sendTextData ("User already exists" :: Text)
            | otherwise -> do
                liftIO $ modifyMVar_ state $ \s -> do
                    let s' = addClient client s
                    WS.sendSink sink $ WS.textData $
                        "Welcome! Users: " `mappend`
                        T.intercalate ", " (map fst s)
                    broadcast (fst client `mappend` " joined") s'
                    return s'
                talk state client
            where
                prefix = "Hi! I am "
                client = (T.drop (T.length prefix) msg, sink)


talk :: WS.Protocol p => MVar ServerState -> Client -> WS.WebSockets p ()
talk state client@(user, _) = flip WS.catchWsError catchDisconnect $ do
    msg <- WS.receiveData
    liftIO $ readMVar state >>= broadcast
        (user `mappend` ": " `mappend` msg)
    talk state client
    where
        catchDisconnect e = case fromException e of
            Just WS.ConnectionClosed -> liftIO $ modifyMVar_ state $ \s -> do
                let s' = removeClient client s
                broadcast (user `mappend` " disconnected") s'
                putStrLn $ "Message : " ++ show e
                return s'
            _ -> return ()