WebSockets in Go

WebSockets are a powerful communication protocol that enables real-time, full-duplex communication between a client and a server over a single, long-lived connection. Unlike traditional HTTP requests, which follow a request-response model, WebSockets facilitate bidirectional communication, allowing data to be transmitted from both the client and the server without the need for continuous polling. In this chapter, we'll explore how to implement WebSockets in Go, from the basics to more advanced concepts.

Basics of WebSockets

What are WebSockets?

WebSockets provide a standardized way for browsers and servers to communicate in real-time. They are particularly useful for applications that require low-latency communication, such as chat applications, online gaming, or live data feeds.

How do WebSockets work?

Unlike traditional HTTP, which follows a request-response model, WebSockets start with an HTTP handshake, after which the connection is upgraded to a WebSocket connection. Once established, both the client and server can send messages asynchronously over the same connection.

Setting Up a WebSocket Server in Go

Let’s start by setting up a basic WebSocket server in Go using the gorilla/websocket package, a popular WebSocket library in the Go ecosystem.

				
					package main

import (
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

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

func handler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			log.Println(err)
			return
		}
		if err := conn.WriteMessage(messageType, p); err != nil {
			log.Println(err)
			return
		}
	}
}

func main() {
	http.HandleFunc("/", handler)
	log.Fatal(http.ListenAndServe(":8080", nil))
}

We import the gorilla/websocket package, which provides utilities for working with WebSockets.
We define an HTTP handler function that upgrades the incoming HTTP connection to a WebSocket connection.
Inside the handler, we continuously read messages from the client and echo them back.

Advanced WebSocket Concepts

Broadcasting Messages

In real-world applications, we often need to broadcast messages from one client to multiple clients. Let’s extend our WebSocket server to support broadcasting.

				
					type Client struct {
	conn *websocket.Conn
}

var (
	clients    = make(map[*Client]bool)
	broadcast  = make(chan []byte)
	register   = make(chan *Client)
	unregister = make(chan *Client)
)

func handleMessages() {
	for {
		message := <-broadcast
		for client := range clients {
			if err := client.conn.WriteMessage(websocket.TextMessage, message); err != nil {
				log.Println(err)
				return
			}
		}
	}
}

func main() {
	go handleMessages()

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Println(err)
			return
		}
		defer conn.Close()

		client := &Client{conn: conn}
		register <- client
		defer func() { unregister <- client }()

		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				log.Println(err)
				return
			}
			broadcast <- message
		}
	})

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

We define a Client struct to represent WebSocket clients.
We use channels to manage client registration, unregistration, and message broadcasting.
The handleMessages function continuously listens for messages on the broadcast channel and sends them to all connected clients.
In the main function, we spawn a goroutine to handle message broadcasting.
Inside the HTTP handler, we register each client, handle incoming messages, and broadcast them to all clients.

Securing WebSocket Connections with TLS

In many cases, transmitting data over WebSocket connections requires an extra layer of security, especially when dealing with sensitive information. Transport Layer Security (TLS) provides encryption and authentication to ensure data confidentiality and integrity. Let’s enhance our WebSocket server to use TLS for secure communication.

				
					package main

import (
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

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

func handler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			log.Println(err)
			return
		}
		if err := conn.WriteMessage(messageType, p); err != nil {
			log.Println(err)
			return
		}
	}
}

func main() {
	http.HandleFunc("/", handler)

	// Serve using TLS
	certFile := "server.crt"
	keyFile := "server.key"
	log.Fatal(http.ListenAndServeTLS(":8080", certFile, keyFile, nil))
}

We’ve updated the main function to use http.ListenAndServeTLS instead of http.ListenAndServe, enabling TLS support.
certFile and keyFile specify the paths to the TLS certificate and key files. You’ll need to generate these files or obtain them from a certificate authority.

Handling Binary Data

While our examples so far have focused on sending text-based messages, WebSockets also support binary data transmission. Let’s modify our WebSocket server to handle binary messages.

				
					func handler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	for {
		messageType, data, err := conn.ReadMessage()
		if err != nil {
			log.Println(err)
			return
		}
		if err := conn.WriteMessage(messageType, data); err != nil {
			log.Println(err)
			return
		}
	}
}

We’ve made no changes to the handler function itself, but it’s important to note that the messageType returned by conn.ReadMessage() indicates whether the message is text or binary. We use the same messageType when writing the message back to ensure consistency.

Advanced Client Authentication

In real-world applications, you may need to implement more advanced authentication mechanisms to secure WebSocket connections. This could involve using JSON Web Tokens (JWT), OAuth, or other authentication protocols. Let’s briefly discuss how you might integrate JWT-based authentication into our WebSocket server.

				
					// Example JWT-based authentication middleware
func authenticate(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		tokenString := r.Header.Get("Authorization")
		if tokenString == "" {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		// Verify JWT token
		token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
			// Validate the token signing method etc.
		})
		if err != nil {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		// If token is valid, proceed to the next handler
		next.ServeHTTP(w, r)
	})
}

// Usage
func main() {
	http.HandleFunc("/", handler)
	http.Handle("/secured", authenticate(http.HandlerFunc(securedHandler)))

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

We’ve defined a middleware function authenticate that extracts the JWT token from the request header, verifies its validity, and either proceeds to the next handler or returns an “Unauthorized” error.
The authenticate middleware is applied to a specific route (/secured), indicating that only authenticated clients can access that endpoint.
Inside the main function, we use http.Handle to associate the authenticate middleware with the /secured route.

By incorporating TLS support, handling binary data, and implementing advanced client authentication, we've elevated our WebSocket server to meet the demands of real-world applications. These advanced concepts provide a solid foundation for building secure, high-performance WebSocket-based communication systems in Go. Experiment with these techniques and adapt them to your specific use cases to create robust and scalable applications. Happy coding !❤️