WebSocket chat that fails

Server.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}
module Server (talk, newServer, Server) where

import           Control.Concurrent.Async
import           Control.Concurrent.STM
import           Control.Exception
import           Control.Monad
import           Data.Foldable            (traverse_)
import           Data.Functor             (void)
import           Data.Map
import           Data.Text                (Text, pack, unpack)
import qualified Network.WebSockets       as WS

type ClientName = String
data Client = Client
  { clientName     :: ClientName
  , clientConn     :: WS.Connection
  , clientKicked   :: TVar (Maybe String)
  , clientSendChan :: TChan Message
  }

data Message
  = Notice String
  | Tell ClientName String
  | Broadcast ClientName String
  | Command String

newClient :: ClientName -> WS.Connection -> STM Client
newClient name conn = do
  c <- newTChan
  k <- newTVar Nothing
  return Client { clientName     = name
                , clientConn     = conn
                , clientKicked   = k
                , clientSendChan = c
                }

sendMessage :: Client -> Message -> STM ()
sendMessage Client{..} msg =
  writeTChan clientSendChan msg

data Server = Server
  { clients :: TVar (Map ClientName Client)
  }

newServer :: IO Server
newServer = do
  c <- newTVarIO empty
  return Server { clients = c }

broadcast :: Server -> Message -> STM ()
broadcast Server{..} msg = do
  clientMap <- readTVar clients
  traverse_
    (\client -> sendMessage client msg) $
    elems clientMap

talk :: Server -> WS.Connection -> IO ()
talk server@Server{..} conn = do
  send "What is your name?"
  name <- unpack <$> WS.receiveData conn
  if Prelude.null name
    then talk server conn
    else
    do
      ok <- checkAddClient server name conn
      case ok of
        Nothing -> do
          send "This name is already in use."
          talk server conn
        Just client ->
          runClient server client `finally` removeClient server name
  where
    send :: Text -> IO ()
    send = WS.sendTextData conn

runClient :: Server -> Client -> IO ()
runClient server@Server{..} client@Client{..} = do
  void $ race serve receive
  where
    receive = forever $ do
      msg <- unpack <$> WS.receiveData clientConn
      atomically $ sendMessage client (Command msg)

    serve = join . atomically $ do
      k <- readTVar clientKicked
      case k of
        Just reason -> return $
          WS.sendTextData clientConn $
          ("You have been kicked: " :: Text) <> pack reason
        Nothing -> do
          msg <- readTChan clientSendChan
          return $ do
            continue <- handleMessage server client msg
            when continue $ serve

handleMessage :: Server -> Client -> Message -> IO Bool
handleMessage server Client{..} message =
  case message of
    Notice msg         -> output $ "*** " <> pack  msg
    Tell name msg      -> output $ "*" <> pack name <> "*: " <> pack  msg
    Broadcast name msg -> output $ "<" <> pack name <> ">: " <> pack msg
    Command msg ->
      case words msg of
        "/kick" : who : why  -> do
          atomically $ kick server who (unwords why)
          return True
        _ -> do
          atomically $ broadcast server $ Broadcast clientName msg
          return True
  where
    output :: Text -> IO Bool
    output t = WS.sendTextData clientConn t >> return True

kick :: Server -> String -> String -> STM ()
kick Server{..} name reason = do
  clientMap <- readTVar clients
  case Data.Map.lookup name clientMap of
    Nothing -> return ()
    Just client -> do
      let k = clientKicked client
      kicked <- readTVar k
      case kicked of
        Just _  -> return ()
        Nothing -> writeTVar k $ Just reason

checkAddClient :: Server -> ClientName -> WS.Connection -> IO (Maybe Client)
checkAddClient server@Server{..} name conn = atomically $ do
  clientMap <- readTVar clients
  if member name clientMap
    then return Nothing
    else do client <- newClient name conn
            writeTVar clients $ insert name client clientMap
            broadcast server  $ Notice (name <> " has connected")
            return (Just client)

removeClient :: Server -> ClientName -> IO ()
removeClient server@Server{..} name = atomically $ do
  modifyTVar' clients $ delete name
  broadcast server $ Notice (name <> " left")