simple chat server example

gistfile1.hs

module Main where

import Data.Char (ord)

import Network.Socket hiding (recv)
import qualified Data.ByteString as S
import qualified Data.ByteString.Char8 as C8
import qualified Data.ByteString.UTF8 as U8
import Network.Socket.ByteString (recv, sendAll)

import Control.Monad (unless, when, liftM)
import Control.Monad.Trans (liftIO)
import Control.Monad.Reader (ReaderT(..), runReaderT, ask, asks, local)
import Control.Applicative ( (<$>), (<*>), pure )

import Control.Concurrent (forkIO)
import Control.Concurrent.MVar (newMVar, newEmptyMVar, takeMVar, putMVar, MVar(..))
import Control.Concurrent.Chan
import Data.IORef (newIORef, readIORef, writeIORef, IORef(..))

import System.Environment (getArgs)

data ChanData = NewMember S.ByteString -- name
              | MemberLeft S.ByteString -- name
              | Message S.ByteString S.ByteString -- Message author content
              | ServerWillClose

data ServerState = ServerState { closing :: IORef Bool
                               , clientsCount :: MVar Integer
                               , serverChannel :: Chan ChanData
                               , quitNow :: MVar ()}
                   
type ServerRun = ReaderT ServerState IO

data ClientState = ClientState { clientName :: S.ByteString
                               , clientConn :: Socket
                               , clientChannel :: Chan ChanData }

data RunState = RunState { serverState :: ServerState
                         , clientState :: ClientState }

type ClientRun = ReaderT RunState IO


main :: IO ()
main = withSocketsDo $
       do (port : _) <- getArgs
          sock <- listenSockOnPort $ read port
          st <- defaultServerState
          forkIO $ runReaderT (acceptSockWith sock handleClient) st
          takeMVar $ quitNow st
          sClose sock

defaultServerState :: IO ServerState
defaultServerState = ServerState <$> newIORef False <*> newMVar 0 <*> newChan <*> newEmptyMVar

handleClient :: ClientRun ()
handleClient = do
  (RunState sst cst) <- ask
  name <- clientPromptName
  let newSt = RunState sst (cst {clientName = name})
  broadcast $ NewMember name
  liftIO $ do
    forkIO $ runReaderT writeToClient newSt
    runReaderT listenClient newSt

writeToClient :: ClientRun ()
writeToClient = do
  clChan <- asks (clientChannel . clientState)
  me <- asks (clientName . clientState)
  msg <- liftIO $ readChan clChan
  case msg of
    NewMember name | name /= me -> clientWriteLine $ S.concat [U8.fromString "New member: ", name]
    MemberLeft name | name /= me -> clientWriteLine $ S.concat [U8.fromString "Member left: ", name]
    Message author str | author /= me -> clientWriteLine $ S.concat [author, C8.pack ": ", str]
    ServerWillClose -> clientWriteLine $ U8.fromString "Server will close soon"
    _ -> return ()
  writeToClient


listenClient :: ClientRun ()
listenClient = do
  conn <- asks (clientConn . clientState)
  msg <- clientRecvLine
  let close = S.take 5 msg == C8.pack "close"
  let eof = S.null msg
  when close $ do
    broadcast ServerWillClose
    setClosing True
  when (not close && not eof) $ do 
    broadcastMessage msg
    listenClient
  when eof $ do
    name <- asks (clientName . clientState)
    broadcast $ MemberLeft name
  

setClosing :: Bool -> ClientRun ()
setClosing v = do
  cc <- asks (closing . serverState)
  liftIO $ writeIORef cc v

clientRecvLine :: ClientRun S.ByteString
clientRecvLine = do
  sock <- asks (clientConn . clientState)
  str <- liftIO $ recv sock 1024
  return $ S.takeWhile ((/=) 0X0d) str

clientWriteLine :: S.ByteString -> ClientRun ()
clientWriteLine str = do
  sock <- asks (clientConn . clientState)
  liftIO $ sendAll sock $ C8.snoc str '\n'

clientPromptName :: ClientRun S.ByteString
clientPromptName = do
  clientWriteLine $ U8.fromString "Please enter your name"
  clientRecvLine

broadcast :: ChanData -> ClientRun ()
broadcast d = do
  chan <- asks (serverChannel . serverState)
  liftIO $ (writeChan chan d >> readChan chan >> return ())

broadcastMessage :: S.ByteString -> ClientRun ()
broadcastMessage str = do
  name <- asks (clientName . clientState)
  broadcast $ Message name str

listenSockOnPort :: Int -> IO (Socket)
listenSockOnPort n = do
  (serverAddr : _) <- getAddrInfo
                      (Just (defaultHints {addrFlags = [AI_PASSIVE,AI_NUMERICSERV]})) -- How to bind to socket
                      Nothing  -- listen to all interfaces
                      (Just $ show n)
  let (family, addr) = tupolev2 addrFamily addrAddress $ serverAddr
  sock <- socket family Stream defaultProtocol
  setSocketOption sock ReuseAddr 1
  bindSocket sock addr
  listen sock 1
  return sock

acceptSockWith :: Socket -> ClientRun () -> ServerRun ()
acceptSockWith sock action = do
  (conn, _) <- liftIO $ accept sock
  sst <- ask
  cc <- liftIO $ modClientsCount succ $ clientsCount sst
  liftIO . forkIO $ do
    clChan <- dupChan $ serverChannel sst
    runReaderT action $ RunState sst $ ClientState (C8.pack $ show cc) conn clChan
    cl <- modClientsCount pred $ clientsCount sst
    sClose conn
    must_close <- liftIO $ readIORef $ closing sst
    when (cl == 0 && must_close) $ putMVar (quitNow sst) ()
  acceptSockWith sock action
  where 
    modClientsCount :: (Integer -> Integer) -> MVar Integer -> IO Integer
    modClientsCount op mc = do
      c <- liftM op $ takeMVar mc
      putMVar mc c
      putStrLn $ "There are " ++ (show c) ++ " clients online"
      return c

tupolev2 :: (a -> b) -> (a -> c) -> a -> (b, c)
tupolev2 f g = \x -> (f x, g x)