Skip to content
This repository

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
branch: master
Fetching contributors…

Octocat-spinner-32-eaf2f5

Cannot retrieve contributors at this time

file 250 lines (208 sloc) 8.104 kb
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249
{-
Adapted from haskell-chat-sever-example which is
Copyright (c) 2012, Joseph Adams

Modifications (c) 2012, Simon Marlow
-}

{-# LANGUAGE RecordWildCards #-}
module Main where

import Control.Concurrent
import Control.Concurrent.STM
import qualified Data.Map as Map
import Data.Map (Map)
import System.IO
import Control.Exception
import Network
import Control.Monad
import Text.Printf
import qualified Data.Foldable as F

{-
Notes

- protocol:
Server: "Name?"
Client: <string>
-- if <string> is already in use, ask for another name
-- Commands:
-- /tell <string> message... (single-user tell)
-- /quit (exit)
-- /kick <string> (kick another user)
-- message... (broadcast to all connected users)

- a client needs to both listen for commands from the socket and
listen for activity from other clients. Therefore we're going to
need at least two threads per client (for listening to multiple
things). Easiest is to use STM for in-process communication, and to
have a receiving thread that listens on the socket and forwards to a
TChan.

- Handle all errors properly, be async-exception safe

- Consistency:
- if two clients simultaneously kick a third client, only one will be
successful

Ideas for exercises:

- add a "/names" command to list the currently connected users

- add a timeout to the "What is your name?" question

- broadcast is inefficient because it uses an unbounded transaction.
Use a broadcast channel instead (TChan + dupTChan).

- add a "/nick" command to change the current client's name.
Careful! The name is stored as an immutable value in the Client
record. What is the best way to make it modifiable?

- flood prevention: prevent a client from issuing more than a
certain number of messages in a given time.
-}

port :: Int
port = 44444

main :: IO ()
main = withSocketsDo $ do
  server <- newServer
  sock <- listenOn (PortNumber (fromIntegral port))
  printf "Listening on port %d\n" port
  forever $ do
      (handle, host, port) <- accept sock
      printf "Accepted connection from %s: %s\n" host (show port)
      forkIO (talk server handle `finally` hClose handle)


-- ---------------------------------------------------------------------------
-- Data structures and initialisation

type ClientName = String

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

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

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

newClient :: ClientName -> Handle -> STM Client
newClient name handle = do
  c <- newTChan
  k <- newTVar Nothing
  return Client { clientName = name
                , clientHandle = handle
                , clientSendChan = c
                , clientKicked = k
                }

data Client = Client
  { clientName :: ClientName
  , clientHandle :: Handle
  , clientKicked :: TVar (Maybe String)
  , clientSendChan :: TChan Message
  }


-- -----------------------------------------------------------------------------
-- Basic operations

broadcast :: Server -> Message -> STM ()
broadcast Server{..} msg =
    readTVar clients >>= F.mapM_ (\client -> sendMessage client msg)

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

tell :: Server -> ClientName -> ClientName -> String -> STM ()
tell Server{..} from who msg = do
    clientmap <- readTVar clients
    case Map.lookup who clientmap of
        Nothing -> return ()
        Just client -> sendMessage client $ Tell from msg

kick :: Server -> Client -> ClientName -> STM (IO ())
kick Server{..} client@Client{clientHandle=handle} who = do
    clientmap <- readTVar clients
    case Map.lookup who clientmap of
        Nothing ->
           return $ hPutStrLn handle (who ++ " is not connected")
        Just victim -> do
           writeTVar (clientKicked victim) $ Just ("by " ++ clientName client)
           return $ hPutStrLn handle ("you kicked " ++ who)

-- -----------------------------------------------------------------------------
-- The main server

talk :: Server -> Handle -> IO ()
talk server@Server{..} handle = do
    hSetNewlineMode handle universalNewlineMode
        -- Swallow carriage returns sent by telnet clients
    hSetBuffering handle LineBuffering
    readName
  where
    readName = do
      hPutStrLn handle "What is your name?"
      name <- hGetLine handle
      if null name
         then readName
         else mask $ \restore -> do
                ok <- checkAddClient server name handle
                case ok of
                  Nothing -> restore $ do
                     hPrintf handle
                        "The name %s is in use, please choose another\n" name
                     readName
                  Just client ->
                     restore (runClient server client)
                       `finally` removeClient server name

checkAddClient :: Server -> ClientName -> Handle -> IO (Maybe Client)
checkAddClient server@Server{..} name handle = atomically $ do
    clientmap <- readTVar clients
    if Map.member name clientmap
       then return Nothing
       else do client <- newClient name handle
               modifyTVar' clients $ Map.insert name client
               broadcast server $ Notice $ name ++ " has connected"
               return (Just client)

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


runClient :: Server -> Client -> IO ()
runClient server@Server{..} client@Client{..}
 = concurrently send receive
 where
    send = join $ atomically $ do
        k <- readTVar clientKicked
        case k of
            Just reason -> return $
                hPutStrLn clientHandle $ "You have been kicked: " ++ reason
            Nothing -> do
                msg <- readTChan clientSendChan
                return $ do
                    continue <- handleMessage server client msg
                    when continue $ send

    receive = do
       msg <- hGetLine clientHandle
       atomically $ sendMessage client $ Command msg
       receive

handleMessage :: Server -> Client -> Message -> IO Bool
handleMessage server client@Client{..} message =
  case message of
     Notice msg -> output $ "*** " ++ msg
     Tell name msg -> output $ "*" ++ name ++ "*: " ++ msg
     Broadcast name msg -> output $ "<" ++ name ++ ">: " ++ msg
     Command msg ->
       case words msg of
           ["/kick", who] -> do
               join $ atomically $ kick server client who
               return True
           "/tell" : who : what -> do
               atomically $ tell server clientName who (unwords what)
               return True
           ["/quit"] ->
               return False
           ('/':_):_ -> do
               hPutStrLn clientHandle $ "Unrecognised command: " ++ msg
               return True
           _ -> do
               atomically $ broadcast server $ Broadcast clientName msg
               return True
 where
   output s = do hPutStrLn clientHandle s; return True

-- ----------------------------------------------------------------------------
-- Utils

-- | @concurrently left right@ runs @left@ and @right@ in separate
-- threads. When either @left@ or @right@ completes, whether
-- successfully or by throwing an exception, the other thread is
-- killed with @killThread@, and the call to @concurrently@ returns.

concurrently :: IO () -> IO () -> IO ()
concurrently left right = do
    done <- newEmptyMVar
    mask $ \restore -> do
        let
            spawn x = forkIO $ restore x `finally` tryPutMVar done ()
            stop threads = mapM_ killThread threads
        --
        tids <- mapM spawn [left,right]
        restore (takeMVar done) `onException` stop tids
        stop tids
Something went wrong with that request. Please try again.