Websocket - Chat in real-time with Golang

I. How to the system work?

I will start with 2 channels: Alice and Bob. When Alice enters the message chat, it will send data to the web socket server with channel Alice. We will receive the data and write the message data to Bob's channel -> Bob will read the message. The opposite will be the same.

II. Implement

1. Back-end

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

var mapWsConn = make(map[string]*websocket.Conn)

func main() {
	http.HandleFunc("/chat", LoadPageChat)
	http.HandleFunc("/ws", InitWebsocket)

	log.Fatal(http.ListenAndServe(":3000", nil))
}

The main goroutine will init 2 API: load page chat and init web socket server. The server will run in port 3000

func LoadPageChat(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Access-Control-Allow-Origin", "*")

	path, err := os.Getwd()
	if err != nil {
		fmt.Fprintf(w, "%s", "error")
		return
	}

	content, err := os.ReadFile(path + "/chat-using-websocket/chat.html")
	if err != nil {
		fmt.Fprintf(w, "%s", "error")
		return
	}

	fmt.Fprintf(w, "%s", content)
}

The function LoadPageChat will read the file chat.html and handle the user interface. The API will return HTML for front-end to display for the user.

func InitWebsocket(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Access-Control-Allow-Origin", "*")

	channel := r.URL.Query().Get("channel")
	if r.Header.Get("Origin") != "http://"+r.Host {
		fmt.Fprintf(w, "%s", "error")
		return
	}

	if _, ok := mapWsConn[channel]; !ok {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			fmt.Fprintf(w, "%s", "error")
			return
		}

		mapWsConn[channel] = conn
	}

	for {
		var msg map[string]string
		err := mapWsConn[channel].ReadJSON(&msg)
		if err != nil {
			fmt.Println("Error reading JSON:", err)
			break
		}
		fmt.Printf("Received: %s\n", msg)

		otherConn := getConn(channel)
		if otherConn == nil {
			continue
		}

		err = otherConn.WriteJSON(msg)
		if err != nil {
			fmt.Println("Error writing JSON:", err)
			break
		}
	}
}

func getConn(channel string) *websocket.Conn {
	for key, conn := range mapWsConn {
		if key != channel {
			return conn
		}
	}

	return nil
}

The function will init web socket's connection. We will consume the channel to get the message and send it to another channel.

Example: Alice sends the message "Hi" to channel Alice -> the channel Alice will receive "Hi" -> write the message "Hi" to channel Bob

2. Front-end

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>Chat application</title>
</head>
<body>
<div class="container">
    <div class="chat">
        <div class="santaSays">
            <div class="text-box-santa">
                <div class="text">
                    <p>Hi there, my child!</p>
                    <p>What can I help you with?</p>
                </div>
            </div>
        </div>
        <div class="userSays">
            <div class="text">
                <p>Hello, Santa!</p>
                <p>I'd like to know when you'll bring my gift?</p>
            </div>
        </div>
    </div>
    <hr>
    <div class="message-box">
        <div class="message-input">
            <input id="inputText" type="text" placeholder="What can I help you with?">
        </div>
        <div class="send-btn">
            <i class="fa-solid fa-paper-plane plane"></i>
        </div>
    </div>
</div>
</body>
</html>

This is the chat application interface.

<script>
    const chat = document.querySelector(".chat");
    const inputText = document.getElementById("inputText");
    let ws;

    if (window.WebSocket === undefined) {
        console.log("Your browser does not support WebSockets")
    } else {
        ws = initWS();
    }

    function initWS() {
        let socket = new WebSocket("ws://" + window.location.host + "/ws" + window.location.search)

        socket.onopen = function() {
            console.log("Socket is open")
        };

        // receive data from server
        socket.onmessage = function (e) {
            let pS = document.createElement("p");
            pS.innerHTML = JSON.parse(e.data).message;
            pS.classList.add("santaMessage");
            chat.appendChild(pS);
            chat.scrollTop = chat.scrollHeight;
        }

        // close socket
        socket.onclose = function () {
            console.log("Socket closed")
        }

        return socket;
    }

    inputText.addEventListener("keyup", (e) => {
        if (e.key === "Enter") {
            let pU = document.createElement("p");
            pU.innerHTML = inputText.value;
            pU.classList.add("userMessage");
            chat.appendChild(pU);
            chat.scrollTop = chat.scrollHeight;

            ws.send(JSON.stringify({message: inputText.value}));
            inputText.value = "";
        }
    })
</script>