Go กับ WebSocket Security - เพิ่มความปลอดภัยด้วย Session และ Authentication!

ทำไมต้องใช้ Session และ Authentication กับ WebSocket?

WebSocket ไม่มีการส่งข้อมูล Credential (เช่น Token หรือ Cookie) โดยอัตโนมัติเหมือน HTTP ดังนั้นเราต้องมีวิธีตรวจสอบสิทธิ์ของผู้ใช้ทุกครั้งที่มีการเชื่อมต่อ เพื่อให้มั่นใจว่าเฉพาะผู้ใช้ที่ได้รับอนุญาตเท่านั้นที่จะสามารถสื่อสารกับเซิร์ฟเวอร์ได้

วิธีการเพิ่ม Authentication ให้กับ WebSocket

1.ตรวจสอบ Token หรือ Cookie ก่อนอนุญาตให้เชื่อมต่อ

2.เก็บสถานะการเชื่อมต่อไว้ใน Session เพื่อระบุตัวตนผู้ใช้

3.ตรวจสอบสิทธิ์ของผู้ใช้ในทุกข้อความที่ส่งมา

ขั้นตอนการเพิ่ม Authentication ให้ WebSocket

1. ตรวจสอบ Token ก่อนอนุญาตการเชื่อมต่อ

ในขั้นตอนการอัปเกรดการเชื่อมต่อ WebSocket สามารถตรวจสอบ Token ได้จาก Header ของ Request

ในตัวอย่างนี้:

ใช้ authMiddleware เพื่อตรวจสอบ Token ใน Header ของ Request ก่อนอนุญาตการเชื่อมต่อ

หาก Token ไม่ถูกต้อง จะส่งสถานะ 401 (Unauthorized) กลับไป

ตัวอย่างโค้ด:

package main

import (
    "fmt"
    "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 != "valid-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 {
        fmt.Println("Error upgrading connection:", err)
        return
    }
    defer conn.Close()

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

func main() {
    http.HandleFunc("/ws", authMiddleware(handleConnections))

    fmt.Println("WebSocket server started at :8080/ws")
    http.ListenAndServe(":8080", nil)
}

2. ใช้ Session เพื่อเก็บข้อมูลผู้ใช้

สามารถใช้ map หรือ Session Manager เพื่อเก็บสถานะการเชื่อมต่อของผู้ใช้แต่ละคน

ในตัวอย่างนี้ :

ใช้ Token เป็น Key ใน sessions เพื่อระบุการเชื่อมต่อของผู้ใช้แต่ละคน

เมื่อการเชื่อมต่อถูกปิด จะลบ Token ออกจาก sessions

ตัวอย่าง :

var sessions = make(map[string]*websocket.Conn)

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

    sessions[token] = conn
    defer delete(sessions, token)

    for {
        _, msg, err := conn.ReadMessage()
        if err != nil {
            fmt.Println("Error reading message:", err)
            break
        }
        fmt.Printf("Received from %s: %s\n", token, msg)
    }
}

3. ตรวจสอบสิทธิ์ในทุกข้อความที่ส่งมา

หากต้องการตรวจสอบสิทธิ์ในแต่ละข้อความที่ส่งมา สามารถเพิ่มตรรกะการตรวจสอบในฟังก์ชัน ReadMessage

ตัวอย่าง :

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

    sessions[token] = conn
    defer delete(sessions, token)

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

        // ตัวอย่างการตรวจสอบสิทธิ์ในข้อความ
        if string(msg) == "admin-only-action" && token != "admin-token" {
            conn.WriteMessage(websocket.TextMessage, []byte("Unauthorized action"))
            continue
        }

        fmt.Printf("Message from %s: %s\n", token, msg)
    }
}

กรณีใช้งานจริง : การส่งข้อมูลเฉพาะกลุ่มผู้ใช้

สามารถส่งข้อมูลเฉพาะกลุ่ม เช่น กลุ่ม Admin หรือผู้ใช้เฉพาะคน โดยใช้ Token หรือ Role เป็นตัวกรอง

ตัวอย่าง :

func broadcastToAdmins(message string) {
    for token, conn := range sessions {
        if token == "admin-token" { // ตัวอย่างการกรองเฉพาะ Admin
            conn.WriteMessage(websocket.TextMessage, []byte(message))
        }
    }
}

สรุปง่ายๆ

• ใช้ Middleware เพื่อตรวจสอบ Token ก่อนอนุญาตการเชื่อมต่อ
• เก็บสถานะการเชื่อมต่อของผู้ใช้ใน Session เพื่อใช้งานต่อ
• เพิ่มตรรกะการตรวจสอบสิทธิ์ในทุกข้อความที่ส่งมาเพื่อเพิ่มความปลอดภัย