Back

Mastering WebSockets With Go

Tutorial on how to use WebSockets to build real-time APIs in Go

by Percy Bolmér, November 22, 2022

Image by Percy Bolmér. Gopher by Takuya Ueda, Original Go Gopher by Renée French (CC BY 3.0)

If we think about it, regular HTTP APIs are dumb, like, really dumb. We can fetch data by sending a request for the data. If we have to keep data fresh on a website, we will have to continuously request the data, so-called Polling.

All Images in this article is made by Percy Bolmér. Gopher by Takuya Ueda, Original Go Gopher by Renée French (CC BY 3.0)

This is like having a kid in the backseat asking “Are we there yet”, instead of having the driver say “We are there now”. This is the way we started to use when designing websites, Silly isn’t it?

Thankfully, developers have solved this with technologies like WebSockets, WebRTC, gRPC, HTTP2 Stream, ServerSent Events, and other bi-directional communications.

WebSockets is one of the oldest ways to communicate in a bi-directional way and is widely used today. It is supported by most browsers and is relatively easy to use.

In this tutorial, we will cover what WebSockets are and how they work, how to use them in Go to communicate between servers and clients. We will also explore some regular pitfalls that I’ve seen in WebSocket APIs, and how to solve them.

During the tutorial, we will be building a chat application where you can enter different chat rooms. The WebSocket server will be built using Go, and the client connecting in vanilla JavaScript. The patterns we learn and apply could easily be adapted when connecting using a Websocket Client written in Go, Java, React, or any other language.

This article is also recorded and can be viewed on my YouTube channel.

What Are WebSockets & Why You Should Care

How a WebSocket is initialized in easy terms

The WebSocket standard is defined in RFC 645.

WebSockets uses HTTP to send an initial request to the server. This is a regular HTTP request, but it contains a special HTTP header Connection: Upgrade. This tells the server that the client is trying to upgrade the HTTP requests TCP connection into a long-running WebSocket. If the server responds with an HTTP 101 Switching Protocols then the connection will be kept alive, making it possible for the Client and Server to send messages bidirectional, full-duplex.

Once this connection is agreed upon, we can send and receive data from both parties. There isn’t more to WebSockets, that is probably what you need to understand about them to get going.

If you want to understand more about what is going on under the hood during the setup, I can recommend the RFC.

You might wonder if a real-time solution is needed. So here is a few areas where WebSockets are often used.

Chat Applications — Applications that need to receive and relay messages to other clients, this is a perfect match for WebSockets.
Games — If you develop a game that has multiplayer and is web-based then WebSockets are truly a match made in heaven. You can push data from the clients and broadcast it to all other players.
Feeds — For applications that need a feed of data, the updated data can easily be pushed to any client using WebSockets.
Real-time Data — Basically anytime you have real-time data, WebSockets is a great solution.

Beginning The Foundation For The Application

The Project Setup — A Go Backend and a JavaSript Client

We will begin by setting up a simple HTTP server that also hosts our web application using a file server. I want to refrain from using any web framework such as React, so we will use native JavaScript. Usually, the steps are very similar when connecting to the WebSocket, so you should have no trouble porting it into whatever framework you use.

Begin by initializing a new module

go mod init programmingpercy.tech/websockets-go

We then create a new file main.go which will be our starting point of the application.

We will first set up the application to serve an API and host the HTML/JS code. Once we have that up, we will start on the actual WebSocket implementation so that it is easier to follow along.

Let us fill main.go with a simple code to host the website that we will soon build. We will only serve the frontend directory that we will create and store the frontend code inside.

package main

import (
	"log"
	"net/http"
)

func main() {
	setupAPI()

	// Serve on port :8080, fudge yeah hardcoded port
	log.Fatal(http.ListenAndServe(":8080", nil))
}

// setupAPI will start all Routes and their Handlers
func setupAPI() {
	// Serve the ./frontend directory at Route /
	http.Handle("/", http.FileServer(http.Dir("./frontend")))
}

main.go — The first version that simply hosts the Frontend

Now let us add the front end, it will be a simple raw HTML/JS/CSS file that shows our amazingly looking chat application. It consists of one form chatroom-selection which our users can use to enter a certain chatroom and a second form chatroom-message that is used to send messages via WebSocket.

This is just simple HTML and JavaScript, but no WebSocket implementation is yet implemented. The only thing that can be worthwhile mentioning is window[“WebSocket”] which is a global you can use to check if the client’s browser does support WebSocket. If this is undefined then we will alert the user that their browser isn’t supported.

Create a folder named frontend and a file named index.html. Then fill the index.html with the following Gist. I won’t be covering the HTML and JS parts, I expect you to be familiar with them.

<!DOCTYPE html>
<html lang="en">

<head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <meta http-equiv="X-UA-Compatible" content="ie=edge" />
    <title>PP - Websockets</title>
</head>

<body>
    <div class="center">
        <h1>Amazing Chat Application</h1>
        <h3 id="chat-header">Currently in chat: general</h3>

        <!--
        Here is a form that allows us to select what Chatroom to be in
        -->
        <form id="chatroom-selection">
            <label for="chatroom">Chatroom:</label>
            <input type="text" id="chatroom" name="chatroom"><br><br>
            <input type="submit" value="Change chatroom">
        </form>

        <br>
        <!--
        Textarea to show messages from users
        -->
        <textarea class="messagearea" id="chatmessages" readonly name="chatmessages" rows="4" cols="50"
            placeholder="Welcome to the general chatroom, here messages from others will appear"></textarea>

        <br>
        <!--
        Chatroom-message form is used to send messages
        -->
        <form id="chatroom-message">
            <label for="message">Message:</label>
            <input type="text" id="message" name="message"><br><br>
            <input type="submit" value="Send message">
        </form>
    </div>

    <!--
        Javascript that is used to Connect to Websocket and Handle New messages
    -->
    <script type="text/javascript">

        // selectedchat is by default General.
        var selectedchat = "general";

        /**
         * changeChatRoom will update the value of selectedchat
         * and also notify the server that it changes chatroom
         * */
        function changeChatRoom() {
            // Change Header to reflect the Changed chatroom
            var newchat = document.getElementById("chatroom");
            if (newchat != null && newchat.value != selectedchat) {
                console.log(newchat);
            }
            return false;
        }
        /**
         * sendMessage will send a new message onto the Websocket
         * */
        function sendMessage() {
            var newmessage = document.getElementById("message");
            if (newmessage != null) {
                console.log(newmessage);
            }
            return false;
        }
        /**
         * Once the website loads, we want to apply listeners and connect to websocket
         * */
        window.onload = function () {
            // Apply our listener functions to the submit event on both forms
            // we do it this way to avoid redirects
            document.getElementById("chatroom-selection").onsubmit = changeChatRoom;
            document.getElementById("chatroom-message").onsubmit = sendMessage;

            // Check if the browser supports WebSocket
            if (window["WebSocket"]) {
                console.log("supports websockets");
            } else {
                alert("Not supporting websockets");
            }
        };
    </script>

    <style type="text/css">
        body {
            overflow: hidden;
            padding: 0;
            margin: 0;
            width: 100%;
            height: 100%;
            background: rgb(66, 56, 56);
        }

        .center {
            margin: auto;
            width: 50%;
            border: 3px solid green;
            padding: 10px;
        }
    </style>

</body>

</html>

frontend/index.html — Simple website without WebSockets yet

If you run the application with go run main.go in the terminal and visit localhost:8080 you should be greeted with an amazing website that has everything we need to start implementing the WebSockets.

localhost:8080 — Amazing chat application

Right now, sending messages and changing the chatroom does nothing but print to the console, but that is what we will implement.

Connecting WebSocket Between Clients & Server

To get started we will add to the frontend so that it tries to connect to our WebSocket API. This is easy in JavaScript and can be done with one line of code.

In JavaScript, there is a built-in WebSocket library you can use without importing anything. We can create the client with new WebSocket(URL) however first we need to create the URL. The URL is composed of the protocol just like a regular HTTP URL, followed by the path. It is a standard to place WebSockets on a /ws endpoint.

There are two protocols when we use WebSockets, it is ws and there is wss. It works just like HTTP and HTTPS, the extra S stands for secure and will apply a SSL encryption on the traffic.

It is highly recommended to use that but requires a certificate, we will apply this later.

Let us add a line that connects to ws://localhost/ws in the windows.onload function.

       /**
         * Once the website loads, we want to apply listeners and connect to websocket
         * */
        window.onload = function () {
            // Apply our listener functions to the submit event on both forms
            // we do it this way to avoid redirects
            document.getElementById("chatroom-selection").onsubmit = changeChatRoom;
            document.getElementById("chatroom-message").onsubmit = sendMessage;

            // Check if the browser supports WebSocket
            if (window["WebSocket"]) {
                console.log("supports websockets");
                // Connect to websocket
                conn = new WebSocket("ws://" + document.location.host + "/ws");
            } else {
                alert("Not supporting websockets");
            }
        };

index.html — Added connecting to the backend

You can go ahead and rerun the program by now, and when you visit the website you should see an error printed to the console that we cannot connect. This is simply because our backend does not accept connections yet.

Let us update the backend code to accept WebSocket connections.

We will begin by building our Manager which is used to serve the connections and upgrade regular HTTP requests into a WebSocket connection, the manager will also be responsible for keeping track of all Clients.

We will use Gorillas WebSocket library to handle connections, this is done by creating an Upgrader which takes in an HTTP request and upgrades the TCP connection. We will assign the buffer size to the Upgrader which will be applied to all new clients.

The manager will expose a regular HTTP HandlerFunc named serveWS which we will host on /ws endpoint. At this point, we will upgrade the connection and then simply close it again, but we can verify that we can connect this way. Create a file named manager.go and fill the code from the gist into it.

package main

import (
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

var (
	/**
	websocketUpgrader is used to upgrade incomming HTTP requests into a persitent websocket connection
	*/
	websocketUpgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
)

// Manager is used to hold references to all Clients Registered, and Broadcasting etc
type Manager struct {
}

// NewManager is used to initalize all the values inside the manager
func NewManager() *Manager {
	return &Manager{}
}

// serveWS is a HTTP Handler that the has the Manager that allows connections
func (m *Manager) serveWS(w http.ResponseWriter, r *http.Request) {

	log.Println("New connection")
	// Begin by upgrading the HTTP request
	conn, err := websocketUpgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	// We wont do anything yet so close connection again
	conn.Close()
}

manager.go — The manager starting point, can accept and upgrade HTTP requests

We also need to add the serveWS into the /ws endpoint so that the front end can connect. We will initiate a new manager and add the HTTP handler in the setupAPI function inside main.go.

// setupAPI will start all Routes and their Handlers
func setupAPI() {
	// Create a Manager instance used to handle WebSocket Connections
	manager := NewManager()

	// Serve the ./frontend directory at Route /
	http.Handle("/", http.FileServer(http.Dir("./frontend")))
	http.HandleFunc("/ws", manager.serveWS)
}

main.go — Exposing the Managers serveWS function as a endpoint

You can run the software by running

go run *.go

If you proceed with visiting the website you should notice that there is no longer an error being printed in the console, the connection is now accepted.

Clients & Management

The manager who is responsible for all clients

We could add all the client logic into the serveWS function, but it might become very large. I recommend creating a Client struct used to handle single connections, the struct is responsible for all logic related to clients, and managed by the Manager.

The client will also be responsible for reading / writing messages in a concurrently safe way. The WebSocket connection in Go only allows one concurrent writer, so we can handle this by using an unbuffered channel. This is a technique recommended by the developers of the Gorilla library themselves.

Before we implement messages, let us make sure we create the Client struct and give the Manager the ability to Add and Delete clients.

I’ve created a new file called client.go which will be pretty small for now and hold any logic related to our clients.

I will create a new type named ClientList which is simply a map that can be used to look up a Client. I also like to have each client hold a reference to the manager, this allows us to more easily manage state even from clients.

package main

import "github.com/gorilla/websocket"

// ClientList is a map used to help manage a map of clients
type ClientList map[*Client]bool

// Client is a websocket client, basically a frontend visitor
type Client struct {
	// the websocket connection
	connection *websocket.Conn

	// manager is the manager used to manage the client
	manager *Manager
}

// NewClient is used to initialize a new Client with all required values initialized
func NewClient(conn *websocket.Conn, manager *Manager) *Client {
	return &Client{
		connection: conn,
		manager:    manager,
	}
}

client.go — The first draft of the client

It is time to update the manager to hold the newly created ClientList. Since many people can connect concurrently we also want the manager to implement the sync.RWMutex so we can lock it before adding clients.

We will also update the NewManager function to initialize a ClientList.

The function serveWS will be updated to Create a new Client with the connection and add it to the manager.

We will also update the manager with an addClient function that inserts the clients and a removeClient that deletes them. The removal will make sure to gracefully close the connection.

package main

import (
	"log"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
)

var (
	/**
	websocketUpgrader is used to upgrade incomming HTTP requests into a persitent websocket connection
	*/
	websocketUpgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
)

// Manager is used to hold references to all Clients Registered, and Broadcasting etc
type Manager struct {
	clients ClientList

	// Using a syncMutex here to be able to lcok state before editing clients
	// Could also use Channels to block
	sync.RWMutex
}

// NewManager is used to initalize all the values inside the manager
func NewManager() *Manager {
	return &Manager{
		clients: make(ClientList),
	}
}

// serveWS is a HTTP Handler that the has the Manager that allows connections
func (m *Manager) serveWS(w http.ResponseWriter, r *http.Request) {

	log.Println("New connection")
	// Begin by upgrading the HTTP request
	conn, err := websocketUpgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	// Create New Client
	client := NewClient(conn, m)
	// Add the newly created client to the manager
	m.addClient(client)
}

// addClient will add clients to our clientList
func (m *Manager) addClient(client *Client) {
	// Lock so we can manipulate
	m.Lock()
	defer m.Unlock()

	// Add Client
	m.clients[client] = true
}

// removeClient will remove the client and clean up
func (m *Manager) removeClient(client *Client) {
	m.Lock()
	defer m.Unlock()

	// Check if Client exists, then delete it
	if _, ok := m.clients[client]; ok {
		// close connection
		client.connection.Close()
		// remove
		delete(m.clients, client)
	}
}

manager.go — The manager with ability to add and remove Clients

Right now we have everything in place to accept new clients and add them. We cant remove clients properly yet but we will soon.

We will have to implement so that our clients can Read and Write messages.

Reading And Writing Messages

Writing Messages Concurrently in a safe way

Reading and writing messages might seem like an easy task, and it is. There is however a small pitfall that many people miss. The WebSocket connection is only allowed to have one concurrent writer, we can fix this by having an unbuffered channel act as a locker.

We will update the serveWS function inside manager.go to start up two goroutines per client once they are created. For now, we will comment out the writing until fully implemented.

// serveWS is a HTTP Handler that the has the Manager that allows connections
func (m *Manager) serveWS(w http.ResponseWriter, r *http.Request) {

	log.Println("New connection")
	// Begin by upgrading the HTTP request
	conn, err := websocketUpgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	// Create New Client
	client := NewClient(conn, m)
	// Add the newly created client to the manager
	m.addClient(client)
  // Start the read / write processes
	go client.readMessages()
	// go client.writeMessages()
}

manager.go — Updating the serveWS to start a Read/write goroutine for each client

We will begin by adding the reading process since it is a bit easier.

Reading a message from the Socket is done by using ReadMessage which returns a messagetype, the payload, and an error.

The message type is used to explain what type of message is being sent if it’s a Ping, pong, data or binary message, etc. All types can be read about in the RFC.

The error will be returned in case something goes wrong, it will also return an error once the connection is closed. So we will want to check for certain close messages to print them, but for a regular close, we won’t log.

// readMessages will start the client to read messages and handle them
// appropriatly.
// This is suppose to be ran as a goroutine
func (c *Client) readMessages() {
	defer func() {
		// Graceful Close the Connection once this
		// function is done
		c.manager.removeClient(c)
	}()
	// Loop Forever
	for {
		// ReadMessage is used to read the next message in queue
		// in the connection
		messageType, payload, err := c.connection.ReadMessage()

		if err != nil {
			// If Connection is closed, we will Recieve an error here
			// We only want to log Strange errors, but not simple Disconnection
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("error reading message: %v", err)
			}
			break // Break the loop to close conn & Cleanup
		}
		log.Println("MessageType: ", messageType)
		log.Println("Payload: ", string(payload))
	}
}

client.go — Added a read message function

We can update the frontend code and try sending a few messages to verify that it works as intended.

Inside index.html we have a function named sendMessage which right now prints out the message in the console. We can simply update this to instead push the message out onto the WebSocket. Sending messages with JavaScript is as simple as using the conn.send function.

        /**
         * sendMessage will send a new message onto the Websocket
         * */
        function sendMessage() {
            var newmessage = document.getElementById("message");
            if (newmessage != null) {
                conn.send(newmessage.value);
            }
            return false;
        }

index.html — Sending the message

Restart your program and enter a message in the UI and press the Send Message button you should see the message type and the payload sent inside the stdout.

Right now we can only send messages but nothing is done with the messages, it is time to let us add the ability to write messages.

Remember that I wrote that we can only write with one concurrent process to the WebSocket? This can be solved in many ways, one way that Gorilla themself recommend is using an unbuffered channel to block concurrent writes. When any process wants to write on the client’s connection, it will instead write the message to the unbuffered channel, which will block if there is any other process currently writing. This makes us able to avoid any concurrency problems.

We will update the client struct to hold this channel, and the constructor function to initialize it for us.

// Client is a websocket client, basically a frontend visitor
type Client struct {
	// the websocket connection
	connection *websocket.Conn

	// manager is the manager used to manage the client
	manager *Manager
	// egress is used to avoid concurrent writes on the WebSocket
	egress chan []byte
}

// NewClient is used to initialize a new Client with all required values initialized
func NewClient(conn *websocket.Conn, manager *Manager) *Client {
	return &Client{
		connection: conn,
		manager:    manager,
		egress:     make(chan []byte),
	}
}

client.go — Adding the egress channel which acts as a gateway

The writeMessages function is pretty similar to the readMessages. However in this case we won’t receive an Err telling us the connection is closed. We will be the ones who send a CloseMessage to the front-end client once the egress channel is closed.

In go, we can see if a channel is closed by accepting 2 output parameters, the second being a boolean to indicate that the channel is closed.

We will be using the connections WriteMessage function which accepts the messagetype as the first input parameter and the payload as the second.


// writeMessages is a process that listens for new messages to output to the Client
func (c *Client) writeMessages() {
	defer func() {
		// Graceful close if this triggers a closing
		c.manager.removeClient(c)
	}()

	for {
		select {
		case message, ok := <-c.egress:
			// Ok will be false Incase the egress channel is closed
			if !ok {
				// Manager has closed this connection channel, so communicate that to frontend
				if err := c.connection.WriteMessage(websocket.CloseMessage, nil); err != nil {
					// Log that the connection is closed and the reason
					log.Println("connection closed: ", err)
				}
				// Return to close the goroutine
				return
			}
			// Write a Regular text message to the connection
			if err := c.connection.WriteMessage(websocket.TextMessage, message); err != nil {
				log.Println(err)
			}
			log.Println("sent message")
		}

	}
}

client.go — The function to handle any messages that are suppose to be sent

If you are familiar with Go, you might have noticed that we used a for select here which is redundant for now. We will later in this tutorial add more cases to the selection.

Make sure you uncomment go client.writeMessages in the serveWS function.

Right now any messages that are pushed on the egress will be sent to the client. No process writes messages to the egress currently, but we can make a quick hack to test that it is working as expected.

We will make every message received in the readMessages broadcast to all other clients. We will do this by simply outputting all the input messages onto each client’s egress.


// readMessages will start the client to read messages and handle them
// appropriatly.
// This is suppose to be ran as a goroutine
func (c *Client) readMessages() {
	defer func() {
		// Graceful Close the Connection once this
		// function is done
		c.manager.removeClient(c)
	}()
	// Loop Forever
	for {
		// ReadMessage is used to read the next message in queue
		// in the connection
		messageType, payload, err := c.connection.ReadMessage()

		if err != nil {
			// If Connection is closed, we will Recieve an error here
			// We only want to log Strange errors, but simple Disconnection
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("error reading message: %v", err)
			}
			break // Break the loop to close conn & Cleanup
		}
		log.Println("MessageType: ", messageType)
		log.Println("Payload: ", string(payload))

		// Hack to test that WriteMessages works as intended
		// Will be replaced soon
		for wsclient := range c.manager.clients {
			wsclient.egress <- payload
		}
	}
}

client.go — Added a small broadcast to each message that is recieved

We have only added the for loop at line 29, we will remove that later. This is only to test that the whole Reading and Writing are working as intended.

It is time to update the front end to handle incoming messages. JavaScript handles WebSocket events by firing some events that we can apply listeners to.

The events are all explained in the docs. We can cover them quickly.

Close — fires when a WebSocket closes.
Error — fires when a WebSocket is closed due to an error.
Message — fires when a WebSocket receives a new message
Open — fires when a WebSocket connection is opened

Depending on what you want to do in your front end you can assign these events handlers. We are interested in the message event, so we will add a listener which just prints them to the console for now.

Once the connection is opened we will add a simple function to print the event sent. This event object contains a bunch of data such as the timestamp sent and the message type. We will be wanting the payload which is contained in the data field.

        /**
         * Once the website loads, we want to apply listeners and connect to websocket
         * */
        window.onload = function () {
            // Apply our listener functions to the submit event on both forms
            // we do it this way to avoid redirects
            document.getElementById("chatroom-selection").onsubmit = changeChatRoom;
            document.getElementById("chatroom-message").onsubmit = sendMessage;

            // Check if the browser supports WebSocket
            if (window["WebSocket"]) {
                console.log("supports websockets");
                // Connect to websocket
                conn = new WebSocket("ws://" + document.location.host + "/ws");
            
                // Add a listener to the onmessage event
                conn.onmessage = function(evt) {
                    console.log(evt);
                }
            
            } else {
                alert("Not supporting websockets");
            }
        };

index.html — Adding a onmessage listener to handle incoming messages

You can now try rebooting the software and visiting the website and sending a few messages. You should see in the console that the events are being sent and received.

This means both the reading and writing works for now.

Scaling Using An Event Approach

Structuring messages sent on the WebSocket

We can connect, and we can send and receive messages now. This is all great and we have a basic setup.

Now, this might work if you only want to send one kind of message. I usually find that creating an Event / Type based approach makes scaling the WebSocket much easier.

What this means is that we create a default format in which we send each message. In this format, we have a certain field that explains what kind of message type it is, and then a payload.

What, does this sound familiar?

It should because it is basically what the WebSockets does right now. The exception is that we will send our messages as a JSON object which our applications can use to route to the correct action/function to perform.

This is a way I find easy to use, easy to scale and to make the WebSocket be leveraged across many use cases. It is sort of an RPC solution.

We begin by adding the Event class in the JavaScript file so that we can parse incoming messages. We then pass these Events into a routeEvent function that checks the value of the field type and passes the event into the real handler.

In the onmessage listener, we will expect JSON formatted data that fits into the Event class.

We will also create a function called sendEvent which will take in an event name, and the payload. It creates the event based on the input and sends it as a JSON.

Whenever a user is sending a message using the sendMessage it will call upon the sendEvent.

The following gist shows the JavaScript part for dealing with this.

<!DOCTYPE html>
<html lang="en">

<head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <meta http-equiv="X-UA-Compatible" content="ie=edge" />
    <title>PP - Websockets</title>
</head>

<body>
    <div class="center">
        <h1>Amazing Chat Application</h1>
        <h3 id="chat-header">Currently in chat: general</h3>

        <!--
        Here is a form that allows us to select what Chatroom to be in
        -->
        <form id="chatroom-selection">
            <label for="chatroom">Chatroom:</label>
            <input type="text" id="chatroom" name="chatroom"><br><br>
            <input type="submit" value="Change chatroom">
        </form>

        <br>
        <!--
        Textarea to show messages from users
        -->
        <textarea class="messagearea" id="chatmessages" readonly name="chatmessages" rows="4" cols="50"
            placeholder="Welcome to the general chatroom, here messages from others will appear"></textarea>

        <br>
        <!--
        Chatroom-message form is used to send messages
        -->
        <form id="chatroom-message">
            <label for="message">Message:</label>
            <input type="text" id="message" name="message"><br><br>
            <input type="submit" value="Send message">
        </form>
    </div>

    <!--
        Javascript that is used to Connect to Websocket and Handle New messages
    -->
    <script type="text/javascript">

        // selectedchat is by default General.
        var selectedchat = "general";

        /**
         * Event is used to wrap all messages Send and Recieved
         * on the Websocket
         * The type is used as a RPC
         * */
        class Event {
            // Each Event needs a Type
            // The payload is not required
            constructor(type, payload) {
                this.type = type;
                this.payload = payload;
            }
        }
        /**
         * routeEvent is a proxy function that routes
         * events into their correct Handler
         * based on the type field
         * */
        function routeEvent(event) {

            if (event.type === undefined) {
                alert("no 'type' field in event");
            }
            switch (event.type) {
                case "new_message":
                    console.log("new message");
                    break;
                default:
                    alert("unsupported message type");
                    break;
            }

        }

        /**
         * changeChatRoom will update the value of selectedchat
         * and also notify the server that it changes chatroom
         * */
        function changeChatRoom() {
            // Change Header to reflect the Changed chatroom
            var newchat = document.getElementById("chatroom");
            if (newchat != null && newchat.value != selectedchat) {
                console.log(newchat);
            }
            return false;
        }
        /**
         * sendMessage will send a new message onto the Chat
         * */
        function sendMessage() {
            var newmessage = document.getElementById("message");
            if (newmessage != null) {
                sendEvent("send_message", newmessage.value)
            }
            return false;
        }

        /**
         * sendEvent
         * eventname - the event name to send on
         * payload - the data payload
         * */
        function sendEvent(eventName, payload) {
            // Create a event Object with a event named send_message
            const event = new Event(eventName, payload);
            // Format as JSON and send
            conn.send(JSON.stringify(event));
        }
        /**
         * Once the website loads, we want to apply listeners and connect to websocket
         * */
        window.onload = function () {
            // Apply our listener functions to the submit event on both forms
            // we do it this way to avoid redirects
            document.getElementById("chatroom-selection").onsubmit = changeChatRoom;
            document.getElementById("chatroom-message").onsubmit = sendMessage;

            // Check if the browser supports WebSocket
            if (window["WebSocket"]) {
                console.log("supports websockets");
                // Connect to websocket
                conn = new WebSocket("ws://" + document.location.host + "/ws");

                // Add a listener to the onmessage event
                conn.onmessage = function (evt) {
                    console.log(evt);
                    // parse websocket message as JSON
                    const eventData = JSON.parse(evt.data);
                    // Assign JSON data to new Event Object
                    const event = Object.assign(new Event, eventData);
                    // Let router manage message
                    routeEvent(event);
                }

            } else {
                alert("Not supporting websockets");
            }
        };
    </script>

    <style type="text/css">
        body {
            overflow: hidden;
            padding: 0;
            margin: 0;
            width: 100%;
            height: 100%;
            background: rgb(66, 56, 56);
        }

        .center {
            margin: auto;
            width: 50%;
            border: 3px solid green;
            padding: 10px;
        }
    </style>

</body>

</html>

index.html — The JavaScript tag is updated to handle Events.

Now that the website has the logic in place to accept Events and send them, we will need to make the backend handle them also.

Begin by making a file called event.go which will contain all logic for events.

We will want to have the Event struct in the back-end, and it should be a replica of the Event class from JavaScript.


package main

import "encoding/json"

// Event is the Messages sent over the websocket
// Used to differ between different actions
type Event struct {
	// Type is the message type sent
	Type string `json:"type"`
	// Payload is the data Based on the Type
	Payload json.RawMessage `json:"payload"`
}

event.go — The websocket Event struct

Notice that the data type of the payload is a json.RawMessage is because we want users to be able to send whatever payload they want. It is up to the event handler to know how the payload data is structured.

When a message is received on the backend, we will use the type field to route it into the appropriate EventHandler, and eventhandler is a function signature. So it is easy to add new functionality by just creating new functions that fulfill the signature pattern.

The EventHandler signature will accept the Event and the Client the message came from. It will also return an error. We accept the Client because some handlers might want to return a response or send some other Event back to the client once it is completed.

We will also add a SendMessageEvent which is the format expected inside the payload of the event.

package main

import "encoding/json"

// Event is the Messages sent over the websocket
// Used to differ between different actions
type Event struct {
	// Type is the message type sent
	Type string `json:"type"`
	// Payload is the data Based on the Type
	Payload json.RawMessage `json:"payload"`
}

// EventHandler is a function signature that is used to affect messages on the socket and triggered
// depending on the type
type EventHandler func(event Event, c *Client) error
	
const (
	// EventSendMessage is the event name for new chat messages sent
	EventSendMessage = "send_message"
)

// SendMessageEvent is the payload sent in the
// send_message event
type SendMessageEvent struct {
	Message string `json:"message"`
	From    string `json:"from"`
}

event.go — Added the EventHandler signature and the SendMessageEvent

I like having the Manager store the map of EventHandlers. This allows me to easily add stuff, in a real application the manager could contain a database repository, etc. We will add it and add a new function named setupEventHandlers that is used to add the needed ones.

A nice way to easily scale having a bunch of handlers is then storing these EventHandlers in a Map, and using the Type as the key. So instead of using a switch to route the event, we will keep a map that holds all the handlers.

We add a routeEvent function that takes in the incoming event and selects the correct handler from the map.

If you have a keen eye, you might have noticed that the routeEvent itself is a EventHandler and could be used that way if wanted.

...

var (
	ErrEventNotSupported = errors.New("this event type is not supported")
)

// Manager is used to hold references to all Clients Registered, and Broadcasting etc
type Manager struct {
	clients ClientList

	// Using a syncMutex here to be able to lcok state before editing clients
	// Could also use Channels to block
	sync.RWMutex
	// handlers are functions that are used to handle Events
	handlers map[string]EventHandler
}

// NewManager is used to initalize all the values inside the manager
func NewManager() *Manager {
	m := &Manager{
		clients:  make(ClientList),
		handlers: make(map[string]EventHandler),
	}
	m.setupEventHandlers()
	return m
}

// setupEventHandlers configures and adds all handlers
func (m *Manager) setupEventHandlers() {
	m.handlers[EventSendMessage] = func(e Event, c *Client) error {
		fmt.Println(e)
		return nil
	}
}

// routeEvent is used to make sure the correct event goes into the correct handler
func (m *Manager) routeEvent(event Event, c *Client) error {
	// Check if Handler is present in Map
	if handler, ok := m.handlers[event.Type]; ok {
		// Execute the handler and return any err
		if err := handler(event, c); err != nil {
			return err
		}
		return nil
	} else {
		return ErrEventNotSupported
	}
}
...

manager.go — Added EventHandlers to the manager

The final piece before we have the whole event infrastructure in place is to change the Client. The client’s readMessages should marshal the incoming JSON into an Event and then use the Manager to route it.

We will also modify the Clients egress channel to not send raw bytes, but instead Event. This also means we need to change the writeMessages to marshal the data before sending it.

package main

import (
	"encoding/json"
	"log"

	"github.com/gorilla/websocket"
)

// ClientList is a map used to help manage a map of clients
type ClientList map[*Client]bool

// Client is a websocket client, basically a frontend visitor
type Client struct {
	// the websocket connection
	connection *websocket.Conn

	// manager is the manager used to manage the client
	manager *Manager
	// egress is used to avoid concurrent writes on the WebSocket
	egress chan Event
}

// NewClient is used to initialize a new Client with all required values initialized
func NewClient(conn *websocket.Conn, manager *Manager) *Client {
	return &Client{
		connection: conn,
		manager:    manager,
		egress:     make(chan Event),
	}
}

// readMessages will start the client to read messages and handle them
// appropriatly.
// This is suppose to be ran as a goroutine
func (c *Client) readMessages() {
	defer func() {
		// Graceful Close the Connection once this
		// function is done
		c.manager.removeClient(c)
	}()
	// Loop Forever
	for {
		// ReadMessage is used to read the next message in queue
		// in the connection
		_, payload, err := c.connection.ReadMessage()

		if err != nil {
			// If Connection is closed, we will Recieve an error here
			// We only want to log Strange errors, but simple Disconnection
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("error reading message: %v", err)
			}
			break // Break the loop to close conn & Cleanup
		}
		// Marshal incoming data into a Event struct
		var request Event
		if err := json.Unmarshal(payload, &request); err != nil {
			log.Printf("error marshalling message: %v", err)
			break // Breaking the connection here might be harsh xD
		}
		// Route the Event
		if err := c.manager.routeEvent(request, c); err != nil {
			log.Println("Error handeling Message: ", err)
		}
	}
}

// writeMessages is a process that listens for new messages to output to the Client
func (c *Client) writeMessages() {
	defer func() {
		// Graceful close if this triggers a closing
		c.manager.removeClient(c)
	}()

	for {
		select {
		case message, ok := <-c.egress:
			// Ok will be false Incase the egress channel is closed
			if !ok {
				// Manager has closed this connection channel, so communicate that to frontend
				if err := c.connection.WriteMessage(websocket.CloseMessage, nil); err != nil {
					// Log that the connection is closed and the reason
					log.Println("connection closed: ", err)
				}
				// Return to close the goroutine
				return
			}
			data, err := json.Marshal(message)
			if err != nil {
				log.Println(err)
				return // closes the connection, should we really
			}
			// Write a Regular text message to the connection
			if err := c.connection.WriteMessage(websocket.TextMessage, data); err != nil {
				log.Println(err)
			}
			log.Println("sent message")
		}

	}
}

client.go — Added the usage of Event instead of raw bytes

You can try restarting the backend using go run *.go and send a message. You should see that something like {send_message [34 49 50 51 34]} is being printed. The payload should be printed as bytes since the current Handler does not parse the raw bytes.

Before we implement that I’d like to make sure we cover a few more WebSocket-related topics.

The HeartBeats — Ping & Pong

Keeping Connections alive with Heartbeats

WebSockets allow both the Server and the Client to send a Ping frame. The Ping is used to check if the other part of the connection is still alive.

But not only do we check if we other connection is alive, but we also keep it alive. A WebSocket that is idle will / can close because it has been idle for too long, Ping & Pong allows us to easily keep the channel alive avoiding long-running connections with low traffic to close unexpectedly.

Whenever a Ping has been sent, the other party has to respond with a Pong. If no response is sent, you can assume that the other party is no longer alive.

It’s logical right, you don’t keep talking to somebody that doesn’t respond.

To implement it we will do it from the Server code. This means our API server will send Pings to each client frequently and wait for a Pong, and if it doesn’t we will remove that client.

Let us begin by defining the timers to use. Inside client.go we will create a pongWait and a pingInterval variable. PongWait is the seconds between the pongs we allow, and it will be reset by each pong from the client. If this time is exceeded we will close the connection, let us say that 10 seconds wait is reasonable.

pingInterval is how often we send pings to the clients. Note that this has to be LOWER than the pongWait. If we have a PingInterval that sends slower than pongWait, pongWait would cancel.

Etc, If we send a Ping every 15 seconds, but only allow the server to wait 10 seconds between pongs, then the connection would be closed after 10 seconds.


var (
	// pongWait is how long we will await a pong response from client
	pongWait = 10 * time.Second
	// pingInterval has to be less than pongWait, We cant multiply by 0.9 to get 90% of time
	// Because that can make decimals, so instead *9 / 10 to get 90%
	// The reason why it has to be less than PingRequency is becuase otherwise it will send a new Ping before getting response
	pingInterval = (pongWait * 9) / 10
)

client.go — Added the Timing variables, pingInterval Algorithm is borrowed from Gorilla

Now we will need to make the server send ping messages to each client. This will be done inside the writeMessages function of clients. We will create a timer that triggers based on the pingInterval, once triggered we will send a message of the type PingMessage with an empty payload.

The reason why we do this in the same function is that we want to remember that the connection does not allow concurrent writes. We could have another process send a Ping on the egress and add a messageType field on the Event struct, but I find that solution to be a bit more complex.

Since we run this in the same function we prevent it from concurrently writing because it will either read from the egress, or the timer, not both at the same time.

// writeMessages is a process that listens for new messages to output to the Client
func (c *Client) writeMessages() {
	// Create a ticker that triggers a ping at given interval
	ticker := time.NewTicker(pingInterval)
	defer func() {
		ticker.Stop()
		// Graceful close if this triggers a closing
		c.manager.removeClient(c)
	}()

	for {
		select {
		case message, ok := <-c.egress:
			// Ok will be false Incase the egress channel is closed
			if !ok {
				// Manager has closed this connection channel, so communicate that to frontend
				if err := c.connection.WriteMessage(websocket.CloseMessage, nil); err != nil {
					// Log that the connection is closed and the reason
					log.Println("connection closed: ", err)
				}
				// Return to close the goroutine
				return
			}

			data, err := json.Marshal(message)
			if err != nil {
				log.Println(err)
				return // closes the connection, should we really
			}
			// Write a Regular text message to the connection
			if err := c.connection.WriteMessage(websocket.TextMessage, data); err != nil {
				log.Println(err)
			}
			log.Println("sent message")
		case <-ticker.C:
			log.Println("ping")
			// Send the Ping
			if err := c.connection.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				log.Println("writemsg: ", err)
				return // return to break this goroutine triggeing cleanup
			}
		}

	}
}

client.go — The writeMessages now sends out frequent Pings

We are sending Pings, we don’t need to update the frontend code to respond. This is because the RFC spec says that any PingMessage should trigger a PongMessage to respond. The Browsers that support WebSocket have automatically built in so that the clients will respond to the ping messages.

So the Server is sending Pings to the Clients. The client responds with a Pong message, but what now?

We need to configure a PongHandler on the server. The PongHandler is a function that will be triggered on PongMessage. We will update readMessages to set an initial PongWait timer once started which will start counting down how long it will keep the connection alive.

The gorilla package allows us to easily set this using the SetReadDeadLine function. We will grab the current time, add PongWait to it, and set that to the connection.

We will also make a new function called pongHandler which will reset the timer using SetReadDeadLine every time a PongMessage is received by the client.

// readMessages will start the client to read messages and handle them
// appropriatly.
// This is suppose to be ran as a goroutine
func (c *Client) readMessages() {
	defer func() {
		// Graceful Close the Connection once this
		// function is done
		c.manager.removeClient(c)
	}()

	// Configure Wait time for Pong response, use Current time + pongWait
	// This has to be done here to set the first initial timer.
	if err := c.connection.SetReadDeadline(time.Now().Add(pongWait)); err != nil {
		log.Println(err)
		return
	}
	// Configure how to handle Pong responses
	c.connection.SetPongHandler(c.pongHandler)

	// Loop Forever
	for {
		// ReadMessage is used to read the next message in queue
		// in the connection
		_, payload, err := c.connection.ReadMessage()

		if err != nil {
			// If Connection is closed, we will Recieve an error here
			// We only want to log Strange errors, but simple Disconnection
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("error reading message: %v", err)
			}
			break // Break the loop to close conn & Cleanup
		}
		// Marshal incoming data into a Event struct
		var request Event
		if err := json.Unmarshal(payload, &request); err != nil {
			log.Printf("error marshalling message: %v", err)
			break // Breaking the connection here might be harsh xD
		}
		// Route the Event
		if err := c.manager.routeEvent(request, c); err != nil {
			log.Println("Error handeling Message: ", err)
		}
	}
}

// pongHandler is used to handle PongMessages for the Client
func (c *Client) pongHandler(pongMsg string) error {
	// Current time + Pong Wait time
	log.Println("pong")
	return c.connection.SetReadDeadline(time.Now().Add(pongWait))
}

client.go — Adding a PongHandler to reset time between pongs

Great, now we keep connections alive allowing the website to be long-running without disconnecting.

Try restarting your software and see that the logs are printing Pong and Pong on the Server.

We have most of the stuff implemented, it is time for some security.

Limiting Message Size

One rule of security is to always expect malicious usage. If people can, they will. So one thing that is good to always do is to limit the maximum size of a message that is allowed to be processed on the server.

This is to avoid a malicious user sending mega frames to DDOS or doing other bad stuff on your server.

Gorilla makes this very easy to be configured on the backend using the SetReadLimit which accepts the number of bytes that is allowed. If the message is bigger than the limit it will close the connection.

You must know the size of your messages so you don’t limit users that are using the application correctly.

In the chat, we are building we could impose a character limit on the front end and then specify a max size that matches the biggest message.

I will set a limit that each message can be max 512 bytes.

// readMessages will start the client to read messages and handle them
// appropriatly.
// This is suppose to be ran as a goroutine
func (c *Client) readMessages() {
	defer func() {
		// Graceful Close the Connection once this
		// function is done
		c.manager.removeClient(c)
	}()
	// Set Max Size of Messages in Bytes
	c.connection.SetReadLimit(512)
	// Configure Wait time for Pong response, use Current time + pongWait
	// This has to be done here to set the first initial timer.
	if err := c.connection.SetReadDeadline(time.Now().Add(pongWait)); err != nil {
		log.Println(err)
		return
	}
	// Configure how to handle Pong responses
	c.connection.SetPongHandler(c.pongHandler)
...

client.go — Setting Max read limit prevents mega frames

If you restart and try sending a long message, the connection will close.

Origin Check

Checking the originating location is important

As it currently stands we allow connections from anywhere to connect to our API. This is not very good UNLESS that’s what you want.

Usually, you have a frontend hosted on some Server, and that domain is the only allowed origin to connect. This is done to prevent Cross-Site Request Forgery.

To handle Origin checks we can build a function that accepts an HTTP request and see if the origin is allowed using a simple string check.

This function has to follow the signature func(r *http.Request) bool since the upgrader that upgrades the regular HTTP request into an HTTP connection has a field to accept a function like that. Before allowing the connection to be upgraded, it will execute our function on the request to verify the origin.


var (
	/**
	websocketUpgrader is used to upgrade incomming HTTP requests into a persitent websocket connection
	*/
	websocketUpgrader = websocket.Upgrader{
		// Apply the Origin Checker
		CheckOrigin:     checkOrigin,
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
)

var (
	ErrEventNotSupported = errors.New("this event type is not supported")
)

// checkOrigin will check origin and return true if its allowed
func checkOrigin(r *http.Request) bool {

	// Grab the request origin
	origin := r.Header.Get("Origin")

	switch origin {
	case "http://localhost:8080":
		return true
	default:
		return false
	}
}

manager.go — Added Origin check to the HTTP upgrader

If you want to test it, you could alter the port in the switch statement into something other than 8080 and try visiting the UI. You should see that it then crashes and with a message request origin not allowed.

Authentication

One important part of APIs is that we only should allow users that can authenticate.

WebSocket doesn’t come with any authentication utilities built in. This is not an issue though.

We will authenticate the users BEFORE they are allowed to establish the WebSocket connection, in the serveWS function.

There are two common ways of doing this. They both have some complexity to them but isn’t a deal breaker. Long ago you could pass a regular Basic authentication by adding user:password in the Websocket Connection URL. This has been deprecated since a while back.

There are two recommended solutions

A regular HTTP request to Authenticate returns an OneTimePassword (OTP) which can be used to connect to a WebSocket connection.
Connect WebSocket, but don’t accept any messages until a special Authentication message with credentials has been sent.

I prefer option number one, mainly for the fact that we don’t want bots to spam connections.

So the flow will be

The user authenticates using regular HTTP, an OTP ticket is returned to the user.
The user connects to the WebSocket using the OTP in the URL.

For this to work out, we will create a simple OTP solution. Note this solution is very basic and can be done better using official OTP packages etc, this is just to showcase the idea.

We will make a RetentionMap which is a simple Map holding OTPs. Any OTP that is older than 5 seconds will get thrown away.

We will also have to make a new login endpoint that accepts regular HTTP requests and authenticates a user. In our example, the authentication will be a simple string check. In a real production application, you should replace the authentication with a real solution. Covering authentication is a whole article of its own.

The serveWS needs to be updated so that it verifies the OTP once the user calls it, and we also need to make sure that the front end sends the OTP along the connection request.

Let us dig in by changing the front end first.

We want to create a simple login form and render it, along with a text displaying if we are connected or not. So we will begin by updating the index.html body.

<div class="center">
        <h1>Amazing Chat Application</h1>
        <h3 id="chat-header">Currently in chat: general</h3>
        <h3 id="connection-header">Connected to Websocket: false</h3>

        <!--
        Here is a form that allows us to select what Chatroom to be in
        -->
        <form id="chatroom-selection">
            <label for="chatroom">Chatroom:</label>
            <input type="text" id="chatroom" name="chatroom"><br><br>
            <input type="submit" value="Change chatroom">
        </form>

        <br>
        <!--
        Textarea to show messages from users
        -->
        <textarea class="messagearea" id="chatmessages" readonly name="chatmessages" rows="4" cols="50"
            placeholder="Welcome to the general chatroom, here messages from others will appear"></textarea>

        <br>
        <!--
        Chatroom-message form is used to send messages
        -->
        <form id="chatroom-message">
            <label for="message">Message:</label>
            <input type="text" id="message" name="message"><br><br>
            <input type="submit" value="Send message">
        </form>

        <!--
        login form is used to login
        -->
        <div style="border: 3px solid black;margin-top: 30px;">
            <form id="login-form">
                <label for="username">username:</label>
                <input type="text" id="username" name="username"><br>
                <label for="password">password:</label>
                <input type="password" id="password" name="password"><br><br>
                <input type="submit" value="Login">
            </form>
        </div>

    </div>

index.html — Added a login form in the body

Next, we will remove the WebSocket connection from happening in the document onload event. Because we don’t want to try to connect before a user is signed in.

We will create a connectWebsocket function that accepts an OTP input, that is appended as a GET parameter. The reason why we don’t add it as an HTTP header or a POST parameter is that it is not supported by the WebSocket client available in the Browsers.

We will also update the onload event to assign a handler to the loginform. This handler will send a request to /login and wait for an OTP to be returned, once that is returned it will trigger a WebSocket connection. Failures to authenticate will just send an alert.

Using onopen and onclose we can print out the correct connection status to the user. Update the script section in index.html to have the following functions.

        /**
         * login will send a login request to the server and then 
         * connect websocket
         * */
        function login() {
            let formData = {
                "username": document.getElementById("username").value,
                "password": document.getElementById("password").value
            }
            // Send the request
            fetch("login", {
                method: 'post',
                body: JSON.stringify(formData),
                mode: 'cors',
            }).then((response) => {
                if (response.ok) {
                    return response.json();
                } else {
                    throw 'unauthorized';
                }
            }).then((data) => {
                // Now we have a OTP, send a Request to Connect to WebSocket
                connectWebsocket(data.otp);
            }).catch((e) => { alert(e) });
            return false;
        }
        /**
         * ConnectWebsocket will connect to websocket and add listeners
         * */
        function connectWebsocket(otp) {
            // Check if the browser supports WebSocket
            if (window["WebSocket"]) {
                console.log("supports websockets");
                // Connect to websocket using OTP as a GET parameter
                conn = new WebSocket("ws://" + document.location.host + "/ws?otp="+ otp);

                // Onopen
                conn.onopen = function (evt) {
                    document.getElementById("connection-header").innerHTML = "Connected to Websocket: true";
                }

                conn.onclose = function(evt) {
                    // Set disconnected
                    document.getElementById("connection-header").innerHTML = "Connected to Websocket: false";
                }

                // Add a listener to the onmessage event
                conn.onmessage = function (evt) {
                    console.log(evt);
                    // parse websocket message as JSON
                    const eventData = JSON.parse(evt.data);
                    // Assign JSON data to new Event Object
                    const event = Object.assign(new Event, eventData);
                    // Let router manage message
                    routeEvent(event);
                }

            } else {
                alert("Not supporting websockets");
            }
        }
        /**
         * Once the website loads
         * */
        window.onload = function () {
            // Apply our listener functions to the submit event on both forms
            // we do it this way to avoid redirects
            document.getElementById("chatroom-selection").onsubmit = changeChatRoom;
            document.getElementById("chatroom-message").onsubmit = sendMessage;
            document.getElementById("login-form").onsubmit = login;


        };

index.html — Adding websocket supporting OTP in the script section

You could try the front end now, and you should see an alert when trying to log in.

After applying these changes to the frontend it’s time to allow the backend the verify the OTP. There are many ways to create an OTP, and some libs out there to help you with it. To keep this tutorial simple, I’ve written a very basic helper class that generates OTPs for us, removes them once they expire, and helps us verify them. There are much better ways to handle OTPs.

I’ve created a new file named otp.go which contains the following gist.

package main

import (
	"context"
	"time"

	"github.com/google/uuid"
)

type OTP struct {
	Key     string
	Created time.Time
}

type RetentionMap map[string]OTP

// NewRetentionMap will create a new retentionmap and start the retention given the set period
func NewRetentionMap(ctx context.Context, retentionPeriod time.Duration) RetentionMap {
	rm := make(RetentionMap)

	go rm.Retention(ctx, retentionPeriod)

	return rm
}

// NewOTP creates and adds a new otp to the map
func (rm RetentionMap) NewOTP() OTP {
	o := OTP{
		Key:     uuid.NewString(),
		Created: time.Now(),
	}

	rm[o.Key] = o
	return o
}

// VerifyOTP will make sure a OTP exists
// and return true if so
// It will also delete the key so it cant be reused
func (rm RetentionMap) VerifyOTP(otp string) bool {
	// Verify OTP is existing
	if _, ok := rm[otp]; !ok {
		// otp does not exist
		return false
	}
	delete(rm, otp)
	return true
}

// Retention will make sure old OTPs are removed
// Is Blocking, so run as a Goroutine
func (rm RetentionMap) Retention(ctx context.Context, retentionPeriod time.Duration) {
	ticker := time.NewTicker(400 * time.Millisecond)
	for {
		select {
		case <-ticker.C:
			for _, otp := range rm {
				// Add Retention to Created and check if it is expired
				if otp.Created.Add(retentionPeriod).Before(time.Now()) {
					delete(rm, otp.Key)
				}
			}
		case <-ctx.Done():
			return

		}
	}
}

otp.go — A retention map which removes expired keys

We need to update the manager to maintain a RetentionMap which we can use to verify OTPs in serveWS and to create new ones when users sign in using /login. We set the retention period to 5 seconds, and we also need to accept a context so we can cancel the underlying goroutine.


// Manager is used to hold references to all Clients Registered, and Broadcasting etc
type Manager struct {
	clients ClientList

	// Using a syncMutex here to be able to lcok state before editing clients
	// Could also use Channels to block
	sync.RWMutex
	// handlers are functions that are used to handle Events
	handlers map[string]EventHandler
	// otps is a map of allowed OTP to accept connections from
	otps RetentionMap
}

// NewManager is used to initalize all the values inside the manager
func NewManager(ctx context.Context) *Manager {
	m := &Manager{
		clients:  make(ClientList),
		handlers: make(map[string]EventHandler),
		// Create a new retentionMap that removes Otps older than 5 seconds
		otps: NewRetentionMap(ctx, 5*time.Second),
	}
	m.setupEventHandlers()
	return m
}

manager.go — Updated the struct to have a RetentionMap

Next, we need to implement the handler to run on /login, and it will be a simple handler. You should replace the authentication parts with a real login verification system. Our handler will accept a JSON formatted payload with username and password.

If the username is percy and the password 123 we will generate a new OTP and return that, if it does not match we will return an Unauthorized HTTP Status.

We also update the serveWS to accept an otp GET parameter.

// loginHandler is used to verify an user authentication and return a one time password
func (m *Manager) loginHandler(w http.ResponseWriter, r *http.Request) {

	type userLoginRequest struct {
		Username string `json:"username"`
		Password string `json:"password"`
	}

	var req userLoginRequest
	err := json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Authenticate user / Verify Access token, what ever auth method you use
	if req.Username == "percy" && req.Password == "123" {
		// format to return otp in to the frontend
		type response struct {
			OTP string `json:"otp"`
		}

		// add a new OTP
		otp := m.otps.NewOTP()

		resp := response{
			OTP: otp.Key,
		}

		data, err := json.Marshal(resp)
		if err != nil {
			log.Println(err)
			return
		}
		// Return a response to the Authenticated user with the OTP
		w.WriteHeader(http.StatusOK)
		w.Write(data)
		return
	}

	// Failure to auth
	w.WriteHeader(http.StatusUnauthorized)
}

// serveWS is a HTTP Handler that the has the Manager that allows connections
func (m *Manager) serveWS(w http.ResponseWriter, r *http.Request) {

	// Grab the OTP in the Get param
	otp := r.URL.Query().Get("otp")
	if otp == "" {
		// Tell the user its not authorized
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	// Verify OTP is existing
	if !m.otps.VerifyOTP(otp) {
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	log.Println("New connection")
	// Begin by upgrading the HTTP request
	conn, err := websocketUpgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	// Create New Client
	client := NewClient(conn, m)
	// Add the newly created client to the manager
	m.addClient(client)

	go client.readMessages()
	go client.writeMessages()
}

manager.go — Updated ServeWS and Login to handle OTP

Finally, we need to update main.go to host the login endpoint and pass in a Context to the Manager.

package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
)

func main() {

	// Create a root ctx and a CancelFunc which can be used to cancel retentionMap goroutine
	rootCtx := context.Background()
	ctx, cancel := context.WithCancel(rootCtx)

	defer cancel()

	setupAPI(ctx)

	// Serve on port :8080, fudge yeah hardcoded port
	log.Fatal(http.ListenAndServe(":8080", nil))

}

// setupAPI will start all Routes and their Handlers
func setupAPI(ctx context.Context) {

	// Create a Manager instance used to handle WebSocket Connections
	manager := NewManager(ctx)

	// Serve the ./frontend directory at Route /
	http.Handle("/", http.FileServer(http.Dir("./frontend")))
	http.HandleFunc("/login", manager.loginHandler)
	http.HandleFunc("/ws", manager.serveWS)

	http.HandleFunc("/debug", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, len(manager.clients))
	})
}

main.go — Adding a login endpoint and a Cancellation Context to manager

Once you have all this in place, you should now be able to use the front end, but only after you have used the login form successfully.

Try it out, pressing Send Message will not do anything. But after you log in you can view the WebSocket that it is getting messages.

We will only print the events to the console, but we will get there. Just one final security aspect to cover.

Encrypting Traffic Using HTTPS & WSS

Let us make one important thing very clear, right now we are using clear text traffic, and a very important part of going live into production is using HTTPS.

To make WebSockets use HTTPS instead, we can simply upgrade the Protocol from ws to wss. WSS is an acronym that stands for WebSockets Secure.

Open up index.html and replace the connection part in connectWebsocket to use WSS.

        function connectWebsocket(otp) {
            // Check if the browser supports WebSocket
            if (window["WebSocket"]) {
                console.log("supports websockets");
                // Connect to websocket using OTP as a GET parameter
                conn = new WebSocket("wss://" + document.location.host + "/ws?otp="+ otp);
...

index.html — Added the wss protocol to connection string

If you try the UI out now, it won’t connect because the backend does not support HTTPS. We can fix this by adding a certificate and key to the backend.

Don’t worry if you don’t own one, we will self-sign a certificate to use during this tutorial.

I’ve created a small script that creates a self-signed certificate using OpenSSL. You can see installation notes on their Github.

Create a file named gencert.bash, if you are using Windows you can run the commands manually.

#!/bin/bash

echo "creating server.key"
openssl genrsa -out server.key 2048
openssl ecparam -genkey -name secp384r1 -out server.key
echo "creating server.crt"
openssl req -new -x509 -sha256 -key server.key -out server.crt -batch -days 3650

gencert.bash — Create a self signed certificate

Execute the commands or run the bash script.

bash gencert.bash

You will see two new files, server.key and server.crt. You should never share these files. Store them in a better spot, so your developer doesn’t accidentally push them to GitHub (trust me this happens, people have bots for finding these mistakes)

The certificate we create should only be used for development purposes

Once you have those in place, we will have to update the main.go file to host the HTTP server using the certificate to encrypt the traffic. This is done by using ListenAndServeTLS instead of ListenAndServe. It works the same, but also takes in a path to a certificate file and a key file.

func main() {

	// Create a root ctx and a CancelFunc which can be used to cancel retentionMap goroutine
	rootCtx := context.Background()
	ctx, cancel := context.WithCancel(rootCtx)

	defer cancel()

	setupAPI(ctx)

	// Serve on port :8080, fudge yeah hardcoded port
	err := http.ListenAndServeTLS(":8080", "server.crt", "server.key", nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}

}

main.go — Using ListenAndServeTLS instead to use HTTPs

Don’t forget to update originChecker to allow HTTPS domain.

// checkOrigin will check origin and return true if its allowed
func checkOrigin(r *http.Request) bool {

	// Grab the request origin
	origin := r.Header.Get("Origin")

	switch origin {
		// Update this to HTTPS
	case "https://localhost:8080":
		return true
	default:
		return false
	}
}

manager.go — The origin has to be updated to support https

Restart the server using go run *.go, and this time, visit the https site instead.

You might have to accept the domain as insecure depending on your browser

You might see an error being printed as follows

2022/09/25 16:52:57 http: TLS handshake error from [::1]:51544: remote error: tls: unknown certificate

This is a remote error, this means it is being sent from the client to the server. This is telling you that the browser does not recognize your certificate provider (you) since it is self-signed. Don’t worry about it since it is a self-signed certificate only to be used in development.

If you are using a real certificate you won’t see that error.

Congratulations, you now are using HTTPS and your Websocket is using WSS.

Implementing A Few Event Handlers

Before we finish this tutorial, I want us to implement the actual Event handlers to make the chat properly work.

We have only implemented the framework for everything regarding WebSockets. It is time to implement some business logic in terms of Handlers.

I won’t cover more architectural principles or information about WebSockets going forward, we will only get some hands-on practice by finalizing, it won’t be much. Hopefully, you will see how easy it is to add further handlers and logic to the WebSocket API using this Event approach.

Let us begin by updating the manager.go to accept a real function in the setupEventHandlers.

// setupEventHandlers configures and adds all handlers
func (m *Manager) setupEventHandlers() {
	m.handlers[EventSendMessage] = SendMessageHandler
}

manager.go — Added SendMessageHandler as input to the event handler instead.

We want to implement the SendMessageHandler, which should accept a payload in the incoming event, marshal it, and then output it to all other clients.

Inside event.go we can add the following.

// NewMessageEvent is returned when responding to send_message
type NewMessageEvent struct {
	SendMessageEvent
	Sent time.Time `json:"sent"`
}

// SendMessageHandler will send out a message to all other participants in the chat
func SendMessageHandler(event Event, c *Client) error {
	// Marshal Payload into wanted format
	var chatevent SendMessageEvent
	if err := json.Unmarshal(event.Payload, &chatevent); err != nil {
		return fmt.Errorf("bad payload in request: %v", err)
	}

	// Prepare an Outgoing Message to others
	var broadMessage NewMessageEvent

	broadMessage.Sent = time.Now()
	broadMessage.Message = chatevent.Message
	broadMessage.From = chatevent.From

	data, err := json.Marshal(broadMessage)
	if err != nil {
		return fmt.Errorf("failed to marshal broadcast message: %v", err)
	}

	// Place payload into an Event
	var outgoingEvent Event
	outgoingEvent.Payload = data
	outgoingEvent.Type = EventNewMessage
	// Broadcast to all other Clients
	for client := range c.manager.clients {
		client.egress <- outgoingEvent
	}

	return nil

}

event.go — Adding a real handler to use to broadcast messages

That is all we need to do on the backend. We have to clean up the Frontend so that the javascript sends the Payload in the wanted format. So let’s add the same class but in JavaScript and send it in the event.

At the top of the Script section in index.html add the Class instance for both the Event types. They have to match the structs in the event.go so that the JSON format is the same.

        /**
         * SendMessageEvent is used to send messages to other clients
         * */
        class SendMessageEvent {
            constructor(message, from) {
                this.message = message;
                this.from = from;
            }
        }
        /**
         * NewMessageEvent is messages comming from clients
         * */
        class NewMessageEvent {
            constructor(message, from, sent) {
                this.message = message;
                this.from = from;
                this.sent = sent;
            }
        }

index.html — inside the script tags we define two classes

Then we have to update the sendMessage function that is triggered when somebody sends a new message. We have to make it send the correct payload type.

This should be a SendMessageEvent payload since that is what the handler in the server expects.

        /**
         * sendMessage will send a new message onto the Chat
         * */
        function sendMessage() {
            var newmessage = document.getElementById("message");
            if (newmessage != null) {
                let outgoingEvent = new SendMessageEvent(newmessage.value, "percy");
                sendEvent("send_message", outgoingEvent)
            }
            return false;
        }

index.html — sendMessage now sends the correct Payload

Finally, once a message is received on a client we should print it into the text area instead of the console. Let us update the routeEvent to expect a NewMessageEvent and pass it into a function that appends the message to the textarea.


        /**
         * routeEvent is a proxy function that routes
         * events into their correct Handler
         * based on the type field
         * */
        function routeEvent(event) {

            if (event.type === undefined) {
                alert("no 'type' field in event");
            }
            switch (event.type) {
                case "new_message":
                    // Format payload
                    const messageEvent = Object.assign(new NewMessageEvent, event.payload);
                    appendChatMessage(messageEvent);
                    break;
                default:
                    alert("unsupported message type");
                    break;
            }

        }
        /**
         * appendChatMessage takes in new messages and adds them to the chat
         * */
        function appendChatMessage(messageEvent) {
            var date = new Date(messageEvent.sent);
            // format message
            const formattedMsg = `${date.toLocaleString()}: ${messageEvent.message}`;
            // Append Message
            textarea = document.getElementById("chatmessages");
            textarea.innerHTML = textarea.innerHTML + "\n" + formattedMsg;
            textarea.scrollTop = textarea.scrollHeight;
        }

index.html — added so clients print the message once recieved

You should now be able to send messages between clients, you can easily try this. Open the UI on two browser tabs, log in and start chatting with yourself, but don’t stay up all night!

We can easily fix it so that we can manage different Chatrooms so that we don’t spew out all messages to everybody.

Let us begin by adding a new ChangeRoomEvent in index.html, and update the chat that the user has switched chatroom.

        /**
         * ChangeChatRoomEvent is used to switch chatroom
         * */
        class ChangeChatRoomEvent {
            constructor(name) {
                this.name = name;
            }
        }
        /**
         * changeChatRoom will update the value of selectedchat
         * and also notify the server that it changes chatroom
         * */
        function changeChatRoom() {
            // Change Header to reflect the Changed chatroom
            var newchat = document.getElementById("chatroom");
            if (newchat != null && newchat.value != selectedchat) {
                selectedchat = newchat.value;
                header = document.getElementById("chat-header").innerHTML = "Currently in chat: " + selectedchat;

                let changeEvent = new ChangeChatRoomEvent(selectedchat);
                sendEvent("change_room", changeEvent);
                textarea = document.getElementById("chatmessages");
                textarea.innerHTML = `You changed room into: ${selectedchat}`;
            }
            return false;
        }

index.html — Added changeroom event and logic

Add a new ChangeEvent in manager.go to the setupEventHandlers to handle this new event.

// setupEventHandlers configures and adds all handlers
func (m *Manager) setupEventHandlers() {
	m.handlers[EventSendMessage] = SendMessageHandler
	m.handlers[EventChangeRoom] = ChatRoomHandler
}

manager.go — Added the ChatRoomEvent & ChatRoomHandler

We can add a chatroom field to the Client struct so that we can know what chatroom the user has selected.

// Client is a websocket client, basically a frontend visitor
type Client struct {
	// the websocket connection
	connection *websocket.Conn

	// manager is the manager used to manage the client
	manager *Manager
	// egress is used to avoid concurrent writes on the WebSocket
	egress chan Event
  // chatroom is used to know what room user is in
	chatroom string
}

client.go — The chatroom field is added

Inside event.go we will add the ChatRoomHandler that will simply overwrite the new chatroom field in the client.

We will also make sure SendMessageHandler checks that the other clients are in the same room before sending the event.

package main

import (
	"encoding/json"
	"fmt"
	"time"
)

// Event is the Messages sent over the websocket
// Used to differ between different actions
type Event struct {
	// Type is the message type sent
	Type string `json:"type"`
	// Payload is the data Based on the Type
	Payload json.RawMessage `json:"payload"`
}

// EventHandler is a function signature that is used to affect messages on the socket and triggered
// depending on the type
type EventHandler func(event Event, c *Client) error

const (
	// EventSendMessage is the event name for new chat messages sent
	EventSendMessage = "send_message"
	// EventNewMessage is a response to send_message
	EventNewMessage = "new_message"
	// EventChangeRoom is event when switching rooms
	EventChangeRoom = "change_room"
)

// SendMessageEvent is the payload sent in the
// send_message event
type SendMessageEvent struct {
	Message string `json:"message"`
	From    string `json:"from"`
}

// NewMessageEvent is returned when responding to send_message
type NewMessageEvent struct {
	SendMessageEvent
	Sent time.Time `json:"sent"`
}

// SendMessageHandler will send out a message to all other participants in the chat
func SendMessageHandler(event Event, c *Client) error {
	// Marshal Payload into wanted format
	var chatevent SendMessageEvent
	if err := json.Unmarshal(event.Payload, &chatevent); err != nil {
		return fmt.Errorf("bad payload in request: %v", err)
	}

	// Prepare an Outgoing Message to others
	var broadMessage NewMessageEvent

	broadMessage.Sent = time.Now()
	broadMessage.Message = chatevent.Message
	broadMessage.From = chatevent.From

	data, err := json.Marshal(broadMessage)
	if err != nil {
		return fmt.Errorf("failed to marshal broadcast message: %v", err)
	}

	// Place payload into an Event
	var outgoingEvent Event
	outgoingEvent.Payload = data
	outgoingEvent.Type = EventNewMessage
	// Broadcast to all other Clients
	for client := range c.manager.clients {
		// Only send to clients inside the same chatroom
		if client.chatroom == c.chatroom {
			client.egress <- outgoingEvent
		}

	}
	return nil
}

type ChangeRoomEvent struct {
	Name string `json:"name"`
}

// ChatRoomHandler will handle switching of chatrooms between clients
func ChatRoomHandler(event Event, c *Client) error {
	// Marshal Payload into wanted format
	var changeRoomEvent ChangeRoomEvent
	if err := json.Unmarshal(event.Payload, &changeRoomEvent); err != nil {
		return fmt.Errorf("bad payload in request: %v", err)
	}

	// Add Client to chat room
	c.chatroom = changeRoomEvent.Name

	return nil
}

event.go — The ChatRoomHandler is added

Great, we know what a superb chat app that allows users to switch chatrooms.

You should visit the UI and give it a go!

Conclusion

In this tutorial, we built a whole framework for a Websocket server.

We have a server that accepts WebSockets in a Secure, Scalable, and managed way.

We have covered the following aspects

How to Connect WebSockets
How to effectively read and write messages to the WebSockets.
How to structure a go backend API with WebSockets
How to use an Event-based design for an easily managed WebSocket API.
How to keep connections alive using a heart beating technique named PingPong
How to avoid users from exploiting the WebSocket by limiting message size to avoid Jumbo frames.
How to limit the allowed origins the WebSocket allows
How to authenticate when using WebSockets, by implementing an OTP ticketing system
How to add HTTPS and WSS to the WebSockets.

I strongly believe this tutorial covers everything you need to learn before getting started with your WebSocket API.

If you have any questions, ideas, or feedback, I strongly encourage you to reach out on any of my social media listed below.

I hope you enjoyed this article, I know I did.

Mastering WebSockets With Go

Previous

Next