Playing around with a simple chat server, where clients connect via telnet and send messages to server

chat_server.go

package main

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"os"
	"regexp"
	"strings"
)

// Represents each telnet connection to the chat server
type Client struct {
	server *Server
	conn   net.Conn
	name   string
	room   string

	// server sends messages to client over this channel
	ch chan string
}

// The chat server
type Server struct {

	// mapping of all existing chatrooms to a map of all the clients in that chatroom
	allClients map[string]map[*Client]bool

	//server receives messages from the various clients on these channels
	chatMessageCh chan Message
	joinRoomCh    chan Message
	leaveRoomCh   chan Message

	// If a client terminates a connection, that event is received on this channel
	leaveClientCh chan *Client
}

// TODO make different kind of message, LeaveRoomMessage, JoinRoomMessage, ChatMessage etc
// for now keeping things simple
// Using instances of Message, clients communicate with the Server over channels
type Message struct {
	sender *Client // originator of the message
	text   string  // only populated if its a plain chat message
	room   string  // only populated if its join room message
}

func (client *Client) bootup() (okay bool) {
	io.WriteString(client.conn, "Iniatializing chat client\n")
	io.WriteString(client.conn, "By default all conversations are in the 'global' chat room\n")
	io.WriteString(client.conn, "To join/create a chat room type the follwing command:\n")
	io.WriteString(client.conn, "     joinroom <single_letter_room_name>\n")
	io.WriteString(client.conn, "To leave the current chatroom\n")
	io.WriteString(client.conn, "     leaveroom\n")
	io.WriteString(client.conn, "But first, to get started, enter your alias and hit enter\n")

	b := bufio.NewReader(client.conn)
	name, _, err := b.ReadLine()
	client.name = strings.TrimSpace(string(name))
	if client.name == "" || err != nil {
		io.WriteString(client.conn, "Name cannot be empty, try again next time!\n")
		return false
	}
	io.WriteString(client.conn, fmt.Sprintf("Welcome %v:", client.name))
	io.WriteString(client.conn, "Feel free to start chatting in the 'global' room or join another room\n")
	return true
}

// Routes message to appropriate channel on server
func (c *Client) routeMessageToServer(text string) {

	re := regexp.MustCompile(`\s?joinroom\s+(\w+)\s?`)
	matches := re.FindStringSubmatch(text)
	if len(matches) == 2 {
		c.server.joinRoomCh <- Message{sender: c, room: matches[1]}
		return
	}

	leaveRoomPattern := `\s?leaveroom\s?`
	matched, err := regexp.MatchString(leaveRoomPattern, text)
	if matched || err != nil {
		c.server.leaveRoomCh <- Message{sender: c}
		return
	}

	c.server.chatMessageCh <- Message{sender: c, text: text}
}

// Client will listen for any text entered by user, and will route message to
// appropriate channel on server. It will also be listening on its channel for any message
// from server
func (c *Client) listenBlocking() {

	defer func() {
		// We should leave the room if connection is closed
		if c.conn != nil {
			c.conn.Close()
		}
		c.server.leaveClientCh <- c
	}()

	// For all text that client enters on console, route it to server
	go func() {
		r := bufio.NewReader(c.conn)
		for {

			if line, err := r.ReadString('\n'); err == nil {
				line = strings.TrimSpace(line)
				if line != "" {
					go c.routeMessageToServer(line)
				}
			} else {
				// If error is present, this might mean the connection has been terminated
				// close channel and get out
				fmt.Printf("Client: %v has terminated the connection\n", c.name)
				close(c.ch)
				break
			}
		}
	}()

	// For any message recieved from server, send it to clients console
	for msg := range c.ch {
		if _, err := io.WriteString(c.conn, msg); err != nil {
			// This could mean the connection was terminated, we should return
			return
		}
	}
}

func (s *Server) stop() {
	// TODO see if any cleanup needs to be done, revisit
}

func (s *Server) removeClientFromRoom(c *Client, room string) {
	msg := fmt.Sprintf("Room leaving - room : %v, client:%v\n", room, c.name)
	fmt.Printf("%s", msg)

	if s.allClients[room] == nil {
		return
	}
	for client, _ := range s.allClients[room] {
		go func(client *Client) {
			if c != client {
				client.ch <- msg
			}
		}(client)
	}
	delete(s.allClients[room], c)
}

func (s *Server) addClientToRoom(c *Client, room string) {
	msg := fmt.Sprintf("Room joining - room: %v, client: %v\n", room, c.name)
	fmt.Printf("%s", msg)

	if s.allClients[room] == nil {
		s.allClients[room] = make(map[*Client]bool)
	}
	s.allClients[room][c] = true
	for client, _ := range s.allClients[room] {
		go func(client *Client) {
			client.ch <- msg
		}(client)
	}
}

func (s *Server) propagateChatMessage(message Message) {

	msgText := fmt.Sprintf("Chat message - room: %v, client: %v, message: %v\n",
		message.sender.room, message.sender.name, message.text,
	)
	fmt.Printf("%s", msgText)

	sender := message.sender
	clientsInRoom := s.allClients[sender.room]
	if clientsInRoom == nil {
		return
	}
	for client, _ := range clientsInRoom {
		go func(sender *Client, receiver *Client, msgText string) {
			if sender != receiver {
				receiver.ch <- msgText
			}
		}(sender, client, msgText)
	}
}

// The server will listen on all its channels for any kind of message from any client
func (s *Server) listenBlocking() {

	fmt.Println("Chat server starting, will wait for clients to join")
	for {
		select {

		case message := <-s.chatMessageCh:
			go s.propagateChatMessage(message)

		case message := <-s.joinRoomCh:
			client := message.sender

			if client.room == message.room {
				client.ch <- fmt.Sprintf("You are already in the room:%v\n", client.room)
				continue
			}
			if client.room != "" {
				// This means that is is NOT the initial joining of the client
				// He is already a part of some room
				go s.removeClientFromRoom(client, client.room)
			}
			go s.addClientToRoom(client, message.room)
			client.room = message.room

		case message := <-s.leaveRoomCh:
			client := message.sender
			if client.room == "global" {
				client.ch <- "To leave the global chat room you would have to leave chat!\n"
				continue
			}
			go s.removeClientFromRoom(client, client.room)
			go s.addClientToRoom(client, "global")
			client.room = "global"

		case client := <-s.leaveClientCh:
			go s.removeClientFromRoom(client, client.room)
			fmt.Printf("Exiting chat - client: %v\n", client.name)
		}
	}
}

func main() {

	// TODO for now just listen on 6000, should be in a config file
	ln, err := net.Listen("tcp", ":6000")
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	server := &Server{
		allClients:    make(map[string]map[*Client]bool),
		chatMessageCh: make(chan Message),
		joinRoomCh:    make(chan Message),
		leaveRoomCh:   make(chan Message),
		leaveClientCh: make(chan *Client),
	}
	go server.listenBlocking()
	//defer server.stop()
	for {
		conn, err := ln.Accept()
		if err != nil {
			fmt.Println(err)
			continue
		}

		client := &Client{
			conn:   conn,
			server: server,
			ch:     make(chan string),
		}
		go func(c *Client) {
			if okay := c.bootup(); !okay {
				// Something went wrong, and client didn't enter name
				// too bad, close the connection
				c.conn.Close()
				return
			}
			go c.listenBlocking()
			joinRoomMessage := Message{sender: c, room: "global"}
			server.joinRoomCh <- joinRoomMessage
		}(client)
	}

}