We’ll be building a basic chat box using Go and gorilla/websocket.

Go WebSockets

Begin by creating the webserver

Create a websocket upgrader to upgrade http connections to ws.

package main

import "net/http"
import "github.com/gorilla/websocket"

var upgrader = websocket.Upgrader{}

// upgrades the connection to ws using the upgrader  
func wsPage(res http.ResponseWriter, req *http.Request){

}

// serves the chat box
func homePage(res http.ResponseWriter, req *http.Request){
    http.ServeFile(res, req, "index.html")    
}

func main(){
    http.HandleFunc("/ws", wsPage)
    http.HandleFunc("/", homePage)    
    http.ListenAndServe(":8080", nil)
}

Here’s the index.html for the chat box. Read this for reference on how client side websocket applications work.

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title></title>
</head>
<body>

    <input type="text" placeholder="message" id="textbox">
    <button id="button">Send</button>
    <div id="box"></div>

    <script>
        var socket = new WebSocket("ws://localhost:8080/ws");

        var button = document.getElementById("button");
        button.addEventListener("click", function(event){
            var text = document.getElementById("textbox").value;        
            socket.send(text);
        });

        socket.onopen = function(event){
            console.log("Socket opened successfully");
        }

        socket.onmessage = function(event){
            var box = document.createElement("div");
            box.innerHTML = event.data;
            document.getElementById("box").appendChild(box);
        }
        
        window.onbeforeunload = function(event){
            socket.close();
        }
    </script>
</body>
</html>

Hub and Client

The Hub stores a map of all the connected clients and also broadcasts messages to them all.

For those unfamiliar with channels read this. Basically, they allow goroutines to communicate by passing data through them.

type Hub struct {
    clients map[*Client]bool
    broadcast     chan []byte
    addClient     chan *Client
    removeClient  chan *Client
}

// initialize a new hub
var hub = Hub{
    broadcast:     make(chan []byte),
    addClient:     make(chan *Client),
    removeClient:  make(chan *Client),
    clients:       make(map[*Client]bool),
}

// Runs forever as a goroutine
func (hub *Hub) start() {
    for {
        // one of these fires when a channel 
        // receives data
        select {
        case conn := <-hub.addClient:
            // add a new client
            hub.clients[conn] = true
        case conn := <-hub.removeClient:
            // remove a client
            if _, ok := hub.clients[conn]; ok {
                delete(hub.clients, conn)
                close(conn.send)
            }
        case message := <-hub.broadcast:
            // broadcast a message to all clients
            for conn := range hub.clients {
                select {
                case conn.send <- message:
                default:
                    close(conn.send)
                    delete(hub.clients, conn)
                }
            }
        }
    }
}

The Client stores the client’s websocket connection and sends/receives messages.

type Client struct {
    ws *websocket.Conn
    // Hub passes broadcast messages to this channel
    send chan []byte
}

// Hub broadcasts a new message and this fires 
func (c *Client) write() {
    // make sure to close the connection incase the loop exits
    defer func() {
        c.ws.Close()
    }()

    for {
        select {
        case message, ok := <-c.send:
            if !ok {
                c.ws.WriteMessage(websocket.CloseMessage, []byte{})
                return
            }

            c.ws.WriteMessage(websocket.TextMessage, message)
        }
    }
}

// New message received so pass it to the Hub 
func (c *Client) read() {
    defer func() {
        hub.removeClient <- c
        c.ws.Close()
    }()

    for {
        _, message, err := c.ws.ReadMessage()
        if err != nil {
            hub.removeClient <- c
            c.ws.Close()
            break
        }

        hub.broadcast <- message
    }
}

We should now change the wsPage() handler to create a new client after upgrading the connection and storing it in the Hub.

func wsPage(res http.ResponseWriter, req *http.Request) {
    conn, err := upgrader.Upgrade(res, req, nil)
    if err != nil {
        http.NotFound(res, req)
        return
    }

    client := &Client{
        ws:   conn,
        send: make(chan []byte),
    }

    hub.addClient <- client

    go client.write()
    go client.read()
}

And done, here’s the full source code or view it on Github:

package main

import "net/http"
import "github.com/gorilla/websocket"

var upgrader = websocket.Upgrader{}

type Hub struct {
    clients map[*Client]bool
    broadcast     chan []byte
    addClient     chan *Client
    removeClient  chan *Client
}

var hub = Hub{
    broadcast:     make(chan []byte),
    addClient:     make(chan *Client),
    removeClient:  make(chan *Client),
    clients:       make(map[*Client]bool),
}

func (hub *Hub) start() {
    for {
        select {
        case conn := <-hub.addClient:
            hub.clients[conn] = true
        case conn := <-hub.removeClient:
            if _, ok := hub.clients[conn]; ok {
                delete(hub.clients, conn)
                close(conn.send)
            }
        case message := <-hub.broadcast:
            for conn := range hub.clients {
                select {
                case conn.send <- message:
                default:
                    close(conn.send)
                    delete(hub.clients, conn)
                }
            }
        }
    }
}

type Client struct {
    ws *websocket.Conn
    send chan []byte
}

func (c *Client) write() {
    defer func() {
        c.ws.Close()
    }()

    for {
        select {
        case message, ok := <-c.send:
            if !ok {
                c.ws.WriteMessage(websocket.CloseMessage, []byte{})
                return
            }

            c.ws.WriteMessage(websocket.TextMessage, message)
        }
    }
}

func (c *Client) read() {
    defer func() {
        hub.removeClient <- c
        c.ws.Close()
    }()

    for {
        _, message, err := c.ws.ReadMessage()
        if err != nil {
            hub.removeClient <- c
            c.ws.Close()
            break
        }

        hub.broadcast <- message
    }
}

func wsPage(res http.ResponseWriter, req *http.Request) {
    conn, err := upgrader.Upgrade(res, req, nil)

    if err != nil {
        http.NotFound(res, req)
        return
    }

    client := &Client{
        ws:   conn,
        send: make(chan []byte),
    }

    hub.addClient <- client

    go client.write()
    go client.read()
}

func homePage(res http.ResponseWriter, req *http.Request){
    http.ServeFile(res, req, "index.html")    
}

func main(){
    go hub.start()
    http.HandleFunc("/ws", wsPage)
    http.HandleFunc("/", homePage)    
    http.ListenAndServe(":8080", nil)
}