Implementing WebSockets in Golang

Introduction

In the age of instant communication and dynamic content, real-time data transfer has become a crucial part of modern web applications. Whether it’s a live chat, notifications, collaborative tools, or real-time dashboards, sending and receiving data instantly is key. One of the most efficient ways to achieve this in web development is through WebSockets. In this article, we’ll dive deep into WebSockets, focusing on how to implement them in Golang to build real-time applications. We’ll cover everything from setting up a basic WebSocket server to handling multiple connections, broadcasting messages, ensuring security, and creating an end-to-end working system with a client.

What are WebSockets and Why Are They Important?

WebSockets provide a full-duplex communication channel over a single, long-lived connection, allowing servers to send data to clients in real time without waiting for a request from the client. This is a stark contrast to traditional HTTP, where the client must initiate every request, leading to inefficiencies in real-time communication.

WebSockets vs HTTP and Other Protocols

• HTTP: In a traditional HTTP request-response cycle, the client sends a request, and the server responds. This model is inefficient for real-time updates as it requires continuous polling or long-polling, both of which are resource-intensive.
• WebSockets: WebSockets establish a persistent connection between the client and server, allowing either side to send messages at any time. This reduces latency and increases efficiency.
• Server-Sent Events (SSE): SSE is another protocol for real-time updates, where the server pushes updates to the client over an HTTP connection. However, SSE is one-way (server-to-client) and not as flexible as WebSockets.

Use Cases for WebSockets

WebSockets are ideal for applications that require real-time updates, such as:

• Live Chat Applications
• Online Gaming
• Real-Time Notifications
• Collaborative Editing Tools
• Financial Tickers and Live Data Feeds

Setting Up a Basic WebSocket Server in Golang

Let’s start by setting up a basic WebSocket server in Golang. We’ll use the popular gorilla/websocket package, which provides a simple and effective way to work with WebSockets.

Step 1: Installing the Gorilla WebSocket Package

First, you need to install the gorilla/websocket package:

go get github.com/gorilla/websocket

Step 2: Creating the WebSocket Server

Now, let’s create a simple WebSocket server that listens for incoming connections and echoes back any message it receives.

package main

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

// Upgrader is used to upgrade HTTP connections to WebSocket connections.
var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
       return true
    },
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
    // Upgrade the HTTP connection to a WebSocket connection
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
       fmt.Println("Error upgrading:", err)
       return
    }
    defer conn.Close()
    // Listen for incoming messages
    for {
       // Read message from the client
       _, message, err := conn.ReadMessage()
       if err != nil {
          fmt.Println("Error reading message:", err)
          break
       }
       fmt.Printf("Received: %s\\n", message)
       // Echo the message back to the client
       if err := conn.WriteMessage(websocket.TextMessage, message); err != nil {
          fmt.Println("Error writing message:", err)
          break
       }
    }
}

func main() {
    http.HandleFunc("/ws", wsHandler)
    fmt.Println("WebSocket server started on :8080")
    err := http.ListenAndServe(":8080", nil)
    if err != nil {
       fmt.Println("Error starting server:", err)
    }
}

Step 3: Running the Server

Run the server using go run:

go run main.go

Now, your WebSocket server is running on http://localhost:8080/ws. You can connect to it using a WebSocket client and send messages to see them echoed back.

Creating a WebSocket Client

Now that we have the WebSocket server up and running, let’s create a simple HTML and JavaScript client that connects to our Golang WebSocket server.

Step 1: Building the HTML Client

Create an index.html file with the following content:

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>WebSocket Client</title>
</head>
<body>
<h1>WebSocket Client</h1>
<div>
    <input type="text" id="messageInput" placeholder="Enter your message">
    <button onclick="sendMessage()">Send</button>
</div>
<div id="messages"></div>

<script>
    let ws;

    function connect() {
        ws = new WebSocket("ws://localhost:8080/ws");

        ws.onopen = function() {
            console.log("Connected to WebSocket server");
        };

        ws.onmessage = function(event) {
            let messageDisplay = document.getElementById("messages");
            messageDisplay.innerHTML += `<p>${event.data}</p>`;
        };

        ws.onclose = function() {
            console.log("WebSocket connection closed, retrying...");
            setTimeout(connect, 1000); // Reconnect after 1 second
        };

        ws.onerror = function(error) {
            console.error("WebSocket error:", error);
        };
    }

    function sendMessage() {
        let input = document.getElementById("messageInput");
        let message = input.value;
        ws.send(message);
        input.value = "";
    }

    connect();
</script>
</body>
</html>

Step 2: Running the Client

Open index.html in your web browser. It will automatically connect to the WebSocket server running on ws://localhost:8080/ws.

Step 3: Testing the Connection

Type a message in the input field and click the “Send” button. The message will be sent to the server, which will then echo it back, and the message will be displayed in the “messages” section on the page.

Handling WebSocket Connections Concurrently

When building real-time applications, it’s common to have multiple clients connected to your WebSocket server simultaneously. Handling these connections efficiently is crucial for performance and scalability.

Using Goroutines for Concurrent Handling

Golang’s lightweight concurrency model, powered by Goroutines, makes it easy to handle multiple WebSocket connections simultaneously.

func wsHandler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
       fmt.Println("Error upgrading:", err)
       return
    }
    defer conn.Close()

    go handleConnection(conn)
}

func handleConnection(conn *websocket.Conn) {
    for {
       _, message, err := conn.ReadMessage()
       if err != nil {
          fmt.Println("Error reading message:", err)
          break
       }
       fmt.Printf("Received: %s\n", message)

       if err := conn.WriteMessage(websocket.TextMessage, message); err != nil {
          fmt.Println("Error writing message:", err)
          break
       }
    }
}

In this example, each new WebSocket connection is handled in a separate Goroutine. This allows multiple clients to interact with the server concurrently without blocking each other.

Broadcasting Messages to Multiple Clients

Broadcasting messages to multiple connected clients is a common requirement in real-time applications like chat rooms or live notifications. Let’s extend our WebSocket server to support broadcasting.

Implementing a Broadcast System

We’ll use a map to store all active WebSocket connections and a channel to distribute messages to all connected clients.

package main

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

var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
       return true
    },
}

var clients = make(map[*websocket.Conn]bool) // Connected clients
var broadcast = make(chan []byte)            // Broadcast channel
var mutex = &sync.Mutex{}                    // Protect clients map

func wsHandler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
       fmt.Println("Error upgrading:", err)
       return
    }
    defer conn.Close()

    mutex.Lock()
    clients[conn] = true
    mutex.Unlock()

    for {
       _, message, err := conn.ReadMessage()
       if err != nil {
          mutex.Lock()
          delete(clients, conn)
          mutex.Unlock()
          break
       }
       broadcast <- message
    }
}

func handleMessages() {
    for {
       // Grab the next message from the broadcast channel
       message := <-broadcast

       // Send the message to all connected clients
       mutex.Lock()
       for client := range clients {
          err := client.WriteMessage(websocket.TextMessage, message)
          if err != nil {
             client.Close()
             delete(clients, client)
          }
       }
       mutex.Unlock()
    }
}

func main() {
    http.HandleFunc("/ws", wsHandler)
    go handleMessages()
    fmt.Println("WebSocket server started on :8080")
    err := http.ListenAndServe(":8080", nil)
    if err != nil {
       fmt.Println("Error starting server:", err)
    }
}

Testing Broadcasting

With this implementation, when any client sends a message to the server, it is broadcast to all connected clients. You can test this by opening multiple instances of the index.html client and sending messages from each.

Ensuring Secure WebSocket Communication

Security is a critical aspect of any WebSocket implementation. Here are some best practices to ensure secure WebSocket communication.

Protecting Against Cross-Site WebSocket Hijacking (CSWSH)

CSWSH is a security vulnerability where an attacker can hijack a WebSocket connection from another origin. To protect against this, validate the origin header in your WebSocket server:

var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
       origin := r.Header.Get("Origin")
       return origin == "<http://yourdomain.com>"
    },
}

Using Secure WebSockets (wss://)

For production applications, always use wss:// instead of ws:// to encrypt the communication channel with SSL/TLS. This prevents man-in-the-middle attacks and ensures that the data exchanged between the client and server is secure.

Deploying a WebSocket-Based Golang Application

Once your WebSocket server is ready, the next step is deploying it to a production environment. Here’s how to do it.

Deploying on a VPS or Cloud Provider

1. Choose a VPS or cloud provider: DigitalOcean, AWS, or Google Cloud are popular choices.
2. Set up a reverse proxy: Use Nginx or Caddy to handle SSL termination and forward WebSocket connections to your Go server.
3. Ensure SSL/TLS is enabled: Obtain a certificate from Let’s Encrypt for your domain and configure Nginx or Caddy to use it.

Scaling WebSocket Servers

WebSocket servers can be scaled horizontally by running multiple instances and using a load balancer to distribute incoming connections. Ensure that your server can handle concurrent connections efficiently, and consider using a message broker like Redis to manage state across distributed WebSocket servers.

Conclusion

WebSockets are a powerful tool for building real-time applications, enabling instant communication between clients and servers. In this article, we’ve explored how to implement WebSockets in Golang, from setting up a basic server to handling concurrent connections, broadcasting messages, ensuring security, and building a full end-to-end working system with a client.

With the knowledge and examples provided, you’re well-equipped to start building your own real-time applications in Golang. Remember, the key to mastering WebSockets lies in understanding the underlying concepts and experimenting with different implementations. Happy coding!

Implementing WebSockets in Golang was originally published in WiseMonks on Medium, where people are continuing the conversation by highlighting and responding to this story.