Go and WebSocket Security - Enhancing the Security of Your System!

In this episode, we will explore how to increase the security of your WebSocket Cluster by using Encryption and Authorization to prevent attacks and protect the privacy of your data.

Why Does a WebSocket Cluster Need Security?

1. Prevent Eavesdropping : Data transmitted through WebSocket must be secure from malicious entities.

2. Access Control : Only authorized users should be able to use the WebSocket.

3. Prevent DDoS Attacks : Reduce the risk of attacks targeting the WebSocket Cluster.

Enhancing Security for WebSocket Cluster

1. Use WSS (WebSocket Secure)
WSS is WebSocket operating over TLS (Transport Layer Security), which helps encrypt the data sent over the network.
Setting up NGINX for WSS :

server {
    listen 443 ssl;
    server_name yourdomain.com;

    ssl_certificate /path/to/your/certificate.pem;
    ssl_certificate_key /path/to/your/key.pem;

    location /ws/ {
        proxy_pass http://localhost:8080; # ชี้ไปยัง WebSocket Server
        proxy_http_version 1.1;
        proxy_set_header Upgrade $http_upgrade;
        proxy_set_header Connection "upgrade";
        proxy_set_header Host $host;
    }
}

2. Add Authentication Token
Before establishing a WebSocket connection, verify the Token to authenticate the user.
Example of Token verification in Go:
In this example:
The authMiddleware checks the Token before allowing the connection.

func authMiddleware(next http.HandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        token := r.URL.Query().Get("token")
        if token != "valid-token" { // ตัวอย่างตรวจสอบ Token แบบง่าย
            http.Error(w, "Unauthorized", http.StatusUnauthorized)
            return
        }
        next(w, r)
    }
}

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

    log.Println("Client connected with valid token")
    for {
        _, msg, err := conn.ReadMessage()
        if err != nil {
            log.Println("Error reading message:", err)
            break
        }
        log.Printf("Received: %s", msg)
    }
}

func main() {
    http.HandleFunc("/ws", authMiddleware(handleConnections))
    log.Println("WebSocket server with token authentication started at :8080")
    log.Fatal(http.ListenAndServe(":8080", nil))
}

3. Use Role-Based Access Control (RBAC)
Define permissions for each type of user, such as Admin and Member, to limit actions.
Example of using RBAC:
In this example:
If the user is not an Admin, they will not be able to perform certain actions.

func handleConnections(w http.ResponseWriter, r *http.Request) {
    token := r.URL.Query().Get("token")
    role := "member" // สมมติว่าได้ Role จาก Token

    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("Error upgrading connection:", err)
        return
    }
    defer conn.Close()

    for {
        _, msg, err := conn.ReadMessage()
        if err != nil {
            log.Println("Error reading message:", err)
            break
        }

        if string(msg) == "admin-action" && role != "admin" {
            conn.WriteMessage(websocket.TextMessage, []byte("Unauthorized action"))
            continue
        }

        log.Printf("Message from %s: %s", role, msg)
    }
}

4. Prevent DDoS Attacks with Rate Limiting
Rate Limiting helps restrict the number of requests each user can send within a specified time frame.
Example of Rate Limiting usage:

var rateLimiter = make(map[string]int)

func rateLimitMiddleware(next http.HandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        ip := r.RemoteAddr
        rateLimiter[ip]++

        if rateLimiter[ip] > 10 { // อนุญาต 10 คำขอต่อ 1 นาที
            http.Error(w, "Too many requests", http.StatusTooManyRequests)
            return
        }

        next(w, r)
    }
}

Further Improvements

• Use JWT (JSON Web Token) : For secure authentication.
• Encrypt Sensitive Data : Use AES or RSA for data requiring high security.
• Monitor Usage : Use Prometheus and Grafana to monitor usage statistics.

Summary

• Use WSS (WebSocket Secure) to encrypt connections.
• Authenticate users with Tokens or JWT.
• Implement Role-Based Access Control (RBAC) to manage permissions.
• Prevent DDoS with Rate Limiting.