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)
}

Explanation

• websocket.Upgrader is used to upgrade an HTTP connection to a WebSocket connection.
• The authMiddleware function checks if the incoming request contains a valid authorization token.
• If the token is valid, it calls the next handler; otherwise, it returns a 401 Unauthorized error.
• The wsHandler manages the WebSocket traffic once authorized.

Testing the Secure WebSocket Server

Start the server using:

go run main.go

Then, connect to the WebSocket server with a tool like websocat:

websocat ws://localhost:8080/ws -H "Authorization: expected-secret-token"

If the token is correct, the server will establish a WebSocket connection and echo back any messages sent to it.

Conclusion

Using middleware for WebSocket authentication is efficient, as it cleanly separates concerns without duplicating code across multiple handlers. Integrated securely within your Go application using Gorilla WebSocket and middleware, you can ensure verified client connections, enhancing both security and usability.