Building a Lightweight WebSocket Server in Go

Introduction
Prerequisites
Setting up Environment
Creating a WebSocket Server
Handling WebSocket Connections
Broadcasting Messages
Conclusion

Introduction

In this tutorial, we will learn how to build a lightweight WebSocket server using Go programming language. By the end of this tutorial, you will have a working WebSocket server that can handle multiple client connections and broadcast messages to all connected clients.

WebSocket is a communication protocol that provides full-duplex communication over a single, long-lived, connection. This makes it suitable for applications that require real-time updates, such as chat applications, gaming, and collaborative editing tools.

Prerequisites

To follow along with this tutorial, you should have a basic understanding of Go programming language and familiarity with concepts such as functions, variables, and structs. You will also need to have Go installed on your system. If you haven’t done so, head over to the official Go website and follow the installation instructions for your operating system.

Setting up Environment

Before we dive into building our WebSocket server, let’s create a new Go module and set up the necessary packages and dependencies.

Create a new directory for your project:

```shell
mkdir websocket-server
cd websocket-server
```

Initialize a new Go module:

```shell
go mod init github.com/your-username/websocket-server
```

Install the Gorilla WebSocket library, which provides a WebSocket implementation for Go:
```
```shell
go get github.com/gorilla/websocket
```
```
With our environment set up, we can now proceed to create our WebSocket server.

Creating a WebSocket Server

Create a new file called main.go and open it in a text editor.

Import the required packages:

```go
package main

import (
    "fmt"
    "log"
    "net/http"
    "github.com/gorilla/websocket"
)
```

Define a struct to represent our WebSocket server:

```go
type Server struct {
    clients   map[*websocket.Conn]bool
    broadcast chan []byte
}
```

Define a method to handle incoming WebSocket connections:

```go
func (s *Server) handleConnections(w http.ResponseWriter, r *http.Request) {
    upgrader := websocket.Upgrader{
        CheckOrigin: func(r *http.Request) bool {
            return true
        },
    }

    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Fatal(err)
    }

    s.clients[conn] = true

    for {
        var msg []byte
        err := conn.ReadMessage(&msg)
        if err != nil {
            log.Printf("Error: %v", err)
            delete(s.clients, conn)
            break
        }
        s.broadcast <- msg
    }
}
```

Define a method to handle broadcasting messages to all connected clients:

```go
func (s *Server) handleMessages() {
    for {
        msg := <-s.broadcast
        for client := range s.clients {
            err := client.WriteMessage(websocket.TextMessage, msg)
            if err != nil {
                log.Printf("Error: %v", err)
                client.Close()
                delete(s.clients, client)
            }
        }
    }
}
```

Modify the main function to initialize the WebSocket server:

```go
func main() {
    server := Server{
        clients:   make(map[*websocket.Conn]bool),
        broadcast: make(chan []byte),
    }

    http.HandleFunc("/ws", server.handleConnections)

    go server.handleMessages()

    log.Println("WebSocket server started on localhost:8000")
    err := http.ListenAndServe(":8000", nil)
    if err != nil {
        log.Fatal(err)
    }
}
```

Save the file and build the project:
```
```shell
go build
```
```
Congratulations! You have just created a basic WebSocket server in Go. Let’s now move on to handling WebSocket connections.

Handling WebSocket Connections

Our WebSocket server is currently set up to handle incoming connections and broadcast messages to all connected clients. However, we don’t have any web page to test our server yet. We’ll create a simple HTML file that connects to the WebSocket server and sends/receives messages.

Create a new directory called static in the root of your project.
Inside the static directory, create a new file called index.html and open it in a text editor.

Add the following HTML code to the index.html file:

```html
<!DOCTYPE html>
<html>
<head>
    <title>WebSocket Example</title>
</head>
<body>
    <input type="text" id="message">
    <button onclick="sendMessage()">Send</button>

    <ul id="messages"></ul>

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

        socket.onmessage = function(event) {
            var messageList = document.getElementById("messages");
            var li = document.createElement("li");
            li.appendChild(document.createTextNode(event.data));
            messageList.appendChild(li);
        };

        function sendMessage() {
            var messageInput = document.getElementById("message");
            var message = messageInput.value;

            socket.send(message);

            messageInput.value = "";
        }
    </script>
</body>
</html>
```

Save the file.

Broadcasting Messages

With our WebSocket server and web page set up, let’s test our application by broadcasting messages to all connected clients.

Start the WebSocket server by running the compiled executable:
```
```shell
./websocket-server
```
```
Open a web browser and navigate to http://localhost:8000/static/index.html.
Open multiple browser tabs or windows and connect all of them to the WebSocket server.
Enter a message in any of the browser tabs/windows and click the “Send” button.
```
You should see the message appear on all connected clients' screens.
```
Congratulations! You have successfully built a lightweight WebSocket server in Go. You can now extend the functionality of the server by implementing additional features such as user authentication, chat rooms, and message persistence.

Conclusion

In this tutorial, we have learned how to build a lightweight WebSocket server using Go programming language. We started by setting up the environment and installing the necessary packages. Then, we created a struct to represent our WebSocket server and defined methods to handle incoming connections and broadcast messages. Finally, we tested our server by creating a simple HTML file that connects to the server and sends/receives messages.

WebSockets provide a powerful mechanism for building real-time applications that require bidirectional communication between clients and servers. With the knowledge gained from this tutorial, you can now explore more advanced concepts and build more sophisticated WebSocket applications.

I hope you found this tutorial helpful! If you have any questions or feedback, please leave a comment below.

Happy coding!

Published: 19 May 2022