Using Middleware for Authentication in WebSocket Connections in Go

In modern web applications, real-time communication is an essential feature. WebSockets serve this purpose well by facilitating full-duplex communication channels over a single TCP connection. When working with WebSockets in Go, adding authentication is a crucial step to ensure that only authorized users can establish connections. Middleware provides a graceful approach to repeatedly apply pre-request checks like authentication.

What is Middleware?

Middleware is a way of adding layers to request handling. Each layer can perform its own processing and then call the next middleware in stack. This modular design makes it easy to address cross-cutting concerns such as logging, authentication, and more.

Integrating Middleware for Authentication

When using the Go language to handle WebSocket connections, you'll typically utilize a package like gorilla/websocket. We'll create a middleware function that checks authentication status before a WebSocket handshake is completed. Let's dive into the implementation.

Step 1: Install Required Packages

You'll need the following Go package to manage WebSocket connections. Run this command:

go get -u github.com/gorilla/websocket

Step 2: Setup WebSocket and Middleware

Let's create an authentication middleware for handling WebSocket connection requests. This function will check for a valid token before letting the request proceed to the WebSocket upgrader.

package main

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

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

func authMiddleware(next http.HandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        token := r.Header.Get("Authorization")
        if token != "expected-secret-token" {
            http.Error(w, "Unauthorized", http.StatusUnauthorized)
            return
        }
        next.ServeHTTP(w, r) // Call the next handler
    }
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        http.Error(w, "Could not open websocket connection", http.StatusBadRequest)
        return
    }
    defer conn.Close()

    // Handle WebSocket connection
    for {
        _, message, err := conn.ReadMessage()
        if err != nil {
            break
        }
        // Echo the message back
        if err = conn.WriteMessage(websocket.TextMessage, message); err != nil {
            break
        }
    }
}

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