Day 66 — WebSocket

Learning Objectives

• Build a WebSocket server from the wire format (no library)
• Understand the upgrade handshake — why some WAF rules break WebSocket
• Implement real-time event broadcasting

Key Concepts

WebSocket Upgrade:
  Client sends HTTP Upgrade request:
    GET /ws HTTP/1.1
    Host: example.com
    Upgrade: websocket
    Connection: Upgrade
    Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
    Sec-WebSocket-Version: 13

  Server responds with 101 Switching Protocols:
    HTTP/1.1 101 Switching Protocols
    Upgrade: websocket
    Connection: Upgrade
    Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=
    (Accept = base64(SHA1(key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")))

  After this: raw WebSocket frames (not HTTP)

WebSocket Frame:
  [1B: FIN+RSV+Opcode][1B: MASK+Payload-len][0-8B: extended len][0-4B: mask][payload]

Exercise: WebSocket Server (No Library)

package main

import (
    "bufio"
    "crypto/sha1"
    "encoding/base64"
    "encoding/binary"
    "fmt"
    "io"
    "net"
    "net/http"
    "sync"
)

const wsGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

func computeAcceptKey(key string) string {
    h := sha1.New()
    h.Write([]byte(key + wsGUID))
    return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

type WSConn struct {
    conn net.Conn
    mu   sync.Mutex
}

// ReadFrame reads one WebSocket frame
func (ws *WSConn) ReadFrame() (opcode byte, payload []byte, err error) {
    b := make([]byte, 2)
    if _, err := io.ReadFull(ws.conn, b); err != nil {
        return 0, nil, err
    }

    opcode = b[0] & 0x0F
    masked := (b[1] & 0x80) != 0
    length := int64(b[1] & 0x7F)

    switch length {
    case 126:
        ext := make([]byte, 2)
        io.ReadFull(ws.conn, ext)
        length = int64(binary.BigEndian.Uint16(ext))
    case 127:
        ext := make([]byte, 8)
        io.ReadFull(ws.conn, ext)
        length = int64(binary.BigEndian.Uint64(ext))
    }

    if length > 10*1024*1024 { // 10MB limit
        return 0, nil, fmt.Errorf("frame too large: %d bytes", length)
    }

    var mask [4]byte
    if masked {
        io.ReadFull(ws.conn, mask[:])
    }

    payload = make([]byte, length)
    io.ReadFull(ws.conn, payload)

    if masked {
        for i := range payload {
            payload[i] ^= mask[i%4]
        }
    }

    return opcode, payload, nil
}

// WriteFrame writes a text frame to the WebSocket connection
func (ws *WSConn) WriteFrame(opcode byte, payload []byte) error {
    ws.mu.Lock()
    defer ws.mu.Unlock()

    header := make([]byte, 2, 10)
    header[0] = 0x80 | opcode // FIN=1, opcode
    length := len(payload)
    switch {
    case length <= 125:
        header[1] = byte(length)
    case length <= 65535:
        header[1] = 126
        ext := make([]byte, 2)
        binary.BigEndian.PutUint16(ext, uint16(length))
        header = append(header, ext...)
    default:
        header[1] = 127
        ext := make([]byte, 8)
        binary.BigEndian.PutUint64(ext, uint64(length))
        header = append(header, ext...)
    }

    ws.conn.Write(header)
    _, err := ws.conn.Write(payload)
    return err
}

// Hub broadcasts messages to all connected WebSocket clients
type Hub struct {
    mu      sync.RWMutex
    clients map[*WSConn]bool
}

func (h *Hub) Add(ws *WSConn) {
    h.mu.Lock()
    h.clients[ws] = true
    h.mu.Unlock()
}

func (h *Hub) Remove(ws *WSConn) {
    h.mu.Lock()
    delete(h.clients, ws)
    h.mu.Unlock()
}

func (h *Hub) Broadcast(msg []byte) {
    h.mu.RLock()
    defer h.mu.RUnlock()
    for ws := range h.clients {
        go ws.WriteFrame(0x01, msg) // 0x01 = text frame
    }
}

var hub = &Hub{clients: make(map[*WSConn]bool)}

func wsHandler(w http.ResponseWriter, r *http.Request) {
    // Validate WebSocket upgrade headers
    if r.Header.Get("Upgrade") != "websocket" {
        http.Error(w, "not a websocket upgrade", 400)
        return
    }

    key := r.Header.Get("Sec-WebSocket-Key")
    if key == "" {
        http.Error(w, "missing Sec-WebSocket-Key", 400)
        return
    }

    // Hijack the underlying TCP connection
    hijacker, ok := w.(http.Hijacker)
    if !ok {
        http.Error(w, "hijack not supported", 500)
        return
    }
    conn, rw, err := hijacker.Hijack()
    if err != nil {
        http.Error(w, err.Error(), 500)
        return
    }

    // Send 101 Switching Protocols
    fmt.Fprintf(conn, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: %s\r\n\r\n",
        computeAcceptKey(key))

    ws := &WSConn{conn: conn}
    _ = rw
    hub.Add(ws)
    defer hub.Remove(ws)

    fmt.Printf("[WS] client connected: %s\n", conn.RemoteAddr())

    // Read loop
    for {
        opcode, payload, err := ws.ReadFrame()
        if err != nil {
            fmt.Printf("[WS] client disconnected: %v\n", err)
            return
        }
        switch opcode {
        case 0x01: // text
            fmt.Printf("[WS] message: %s\n", payload)
            hub.Broadcast(payload) // broadcast to all clients
        case 0x08: // close
            return
        case 0x09: // ping
            ws.WriteFrame(0x0A, payload) // pong
        }
    }
}

func main() {
    http.HandleFunc("/ws", wsHandler)
    http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
        fmt.Fprint(w, `<html><body>
<script>
var ws = new WebSocket("ws://localhost:8080/ws");
ws.onmessage = function(e) { document.body.innerHTML += "<p>" + e.data + "</p>"; };
ws.onopen = function() { ws.send("hello from browser"); };
</script></body></html>`)
    })

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

Real-World Context

WAF rules that block <script> tags will break legitimate WebSocket applications that send HTML content. Understanding the WebSocket handshake (HTTP Upgrade → raw framing) explains why WAFs need separate WebSocket-aware rules. Your WAF (Week 10) must not inspect WebSocket frames as HTTP bodies.