A simple TCP chat room with go(lang).

TCPChatRoom.go

/*
	Simple TCP chat room

	start with arg
		-l ADDRESS:PORT

	you can use nc as client, by
		nc ADDRESS PORT

-------------------------------------
|	Author:		Richard Chen		|
|	Website:	https://iochen.com/	|
|	GitHub:		@iochen				|
|	Twitter:	@realRichardChen	|
|	Telegram:	@ioware				|
|	LICENSE:	MIT LICENSE			|
-------------------------------------
*/

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"strconv"
	"strings"
	"time"
)

type Message struct {
	Username string
	Data     []byte
}

// pass message through channels
var RoomMsg = make(map[int]chan *Message)

// send UserNumber[i] times to broadcast
var UserNumber = make(map[int]int)

// if not exist, then make one
var RoomExist = make(map[int]bool)

func main() {
	flagListen := flag.String("l", "localhost:8080", "listen address")
	flag.Parse()

	l, err := net.Listen("tcp4", *flagListen)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer l.Close()

	log.Printf("Listening on: %s\n", *flagListen)

	for {
		c, err := l.Accept()
		if err != nil {
			fmt.Println(err)
			return
		}
		go handleConnection(c)
	}
}

func handleConnection(c net.Conn) {
	defer c.Close()

	// ask for room ID
	_, err := c.Write([]byte("Please enter your room number: "))
	if err != nil {
		log.Println(err)
		return
	}
	ri, err := bufio.NewReader(c).ReadString('\n')
	if err == io.EOF {
		return
	} else if err != nil {
		log.Println(err)
		return
	}
	ri = strings.TrimSpace(ri)
	roomID, err := strconv.Atoi(string(ri))
	if err != nil {
		log.Println(err)
		return
	}

	// create channel if not exist
	if !RoomExist[roomID] {
		RoomMsg[roomID] = make(chan *Message)
	}
	// set to exist
	RoomExist[roomID] = true

	// ask for username
	_, err = c.Write([]byte("Please enter your username: "))
	username, err := bufio.NewReader(c).ReadString('\n')
	username = strings.TrimSpace(username)
	if err == io.EOF {
		return
	} else if err != nil {
		log.Println(err)
		return
	}
	// add user && remove user before return
	UserNumber[roomID]++
	defer func() {
		UserNumber[roomID]--
	}()

	// receive message
	go ReceiveMsg(roomID, username, c)

	for {
		select {
		case msg := <-RoomMsg[roomID]:
			// if username matches, the client has had the message,
			// so there's no need to print again
			if msg.Username != username {
				// clear current line with '\r'
				_, err := c.Write(append([]byte("\r"+msg.Username+": "), msg.Data...))
				if err != nil {
					log.Println(err)
					return
				}

				// print "USERNAME: "
				_, err = c.Write([]byte(username + ": "))
				if err != nil {
					log.Println(err)
					return
				}
			}
		}
	}

}

func ReceiveMsg(roomID int, username string, c net.Conn) {
	for {
		// if don't sleep, "\r" will clear "USERNAME: "
		// there are better solutions, like counting receivers after broadcasting,
		// but that will make it more complex
		time.Sleep(1e4)

		_, err := c.Write([]byte(username + ": "))
		if err != nil {
			log.Println(err)
			return
		}
		netData, err := bufio.NewReader(c).ReadBytes('\n')
		if err == io.EOF {
			return
		} else if err != nil {
			log.Println(err)
			return
		}

		message := &Message{
			Username: username,
			Data:     netData,
		}

		// send UserNumber[roomID] times to broadcast
		for i := 0; i < UserNumber[roomID]; i++ {
			RoomMsg[roomID] <- message
		}
	}
}