refactor: fix writePump goroutine leak and data race, add room/client limits for DoS prevention, add WebSocket auto-reconnection

Author: shijiashuai

changelog/2026-03-09_stability-and-security.md

@@ -0,0 +1,33 @@
+# Stability & Security Refactoring
+
+Date: 2026-03-09
+
+## Critical Bug Fixes
+
+### writePump goroutine leak and data race (hub.go)
+- `HandleWS` used a `defer` to close resources, but the ordering caused a race:
+  - `writePump` goroutine called `c.conn.Close()` on write error
+  - The read goroutine's defer also called `c.Close()`
+  - Two goroutines closing the same connection concurrently = data race
+- `writePump` could also leak if `close(client.send)` was never reached
+- **Fix**: Explicit sequential cleanup after read loop exits:
+  1. `removeClient` — prevents new messages from being routed to send channel
+  2. `close(client.send)` — terminates writePump's range loop
+  3. `c.Close()` — only the read goroutine owns conn lifecycle now
+- `writePump` no longer calls `c.conn.Close()` on write error; instead drains remaining messages and returns
+
+### No room/client limits — DoS vector (hub.go)
+- Any client could create unlimited rooms and join unlimited times, exhausting server memory
+- **Fix**: Added `MaxRooms = 1000` and `MaxClientsPerRoom = 50` limits in `addClient()`
+- Rejections are logged for monitoring
+
+## Improvements
+
+### WebSocket auto-reconnection (app.js)
+- On unexpected disconnect, the client just set state to 'idle' — user had to manually rejoin
+- **Fix**: On unexpected disconnect (not manual leave), automatically attempt reconnection after 2 seconds if a roomId exists
+- Shows "连接断开，正在重连…" status message during reconnection
+
+### Files Modified
+- `internal/signal/hub.go` — goroutine lifecycle fix, room limits, writePump drain
+- `web/app.js` — WebSocket auto-reconnection

internal/signal/hub.go

@@ -10,6 +10,12 @@ import (
 	"github.com/gorilla/websocket"
 )
 
+// Limits to prevent resource exhaustion
+const (
+	MaxRooms          = 1000
+	MaxClientsPerRoom = 50
+)
+
 type Hub struct {
 	mu    sync.RWMutex
 	rooms map[string]map[string]*Client
@@ -86,11 +92,8 @@ func (h *Hub) HandleWS(w http.ResponseWriter, r *http.Request) {
 	log.Printf("signal: ws connected from %s", r.RemoteAddr)
 	client := &Client{conn: c, send: make(chan Message, 32)}
 	go client.writePump()
-	defer func() {
-		h.removeClient(client)
-		close(client.send)
-		c.Close()
-	}()
+
+	// readPump: blocks until read error (disconnect / protocol error)
 	for {
 		var msg Message
 		if err := c.ReadJSON(&msg); err != nil {
@@ -115,6 +118,14 @@ func (h *Hub) HandleWS(w http.ResponseWriter, r *http.Request) {
 			log.Printf("signal: unknown msg type=%s room=%s from=%s", msg.Type, msg.Room, msg.From)
 		}
 	}
+
+	// Cleanup: order matters to avoid goroutine leak and data race.
+	// 1. Remove from hub first (prevents new messages being sent to client.send)
+	h.removeClient(client)
+	// 2. Close send channel (terminates writePump's range loop)
+	close(client.send)
+	// 3. Close WebSocket (writePump may still be draining; conn.Close is safe to call concurrently)
+	c.Close()
 }
 
 func (h *Hub) addClient(c *Client) {
@@ -125,11 +136,21 @@ func (h *Hub) addClient(c *Client) {
 	}
 	m, ok := h.rooms[c.room]
 	if !ok {
+		// Enforce max rooms limit
+		if len(h.rooms) >= MaxRooms {
+			log.Printf("signal: room limit reached (%d), rejecting join room=%s id=%s", MaxRooms, c.room, c.id)
+			return
+		}
 		m = make(map[string]*Client)
 		h.rooms[c.room] = m
 	}
+	// Enforce per-room client limit
+	if len(m) >= MaxClientsPerRoom {
+		log.Printf("signal: room %s full (%d clients), rejecting id=%s", c.room, MaxClientsPerRoom, c.id)
+		return
+	}
 	m[c.id] = c
-	log.Printf("signal: join room=%s id=%s", c.room, c.id)
+	log.Printf("signal: join room=%s id=%s (room size: %d, total rooms: %d)", c.room, c.id, len(m), len(h.rooms))
 	broadcastMembers(c.room, m)
 }
 
@@ -197,8 +218,11 @@ func (c *Client) writePump() {
 	for msg := range c.send {
 		if err := c.conn.WriteJSON(msg); err != nil {
 			log.Printf("signal: write message error room=%s id=%s: %v", c.room, c.id, err)
-			c.conn.Close()
-			break
+			// Don't close conn here — the read goroutine owns conn lifecycle.
+			// Drain remaining messages from send channel so close(send) doesn't block.
+			for range c.send {
+			}
+			return
 		}
 	}
 }

web/app.js

@@ -190,13 +190,29 @@ function connectWS() {
   ws.onclose = () => {
     stopPing()
     ws = null
-    closePeerConnection()
-    roomId = null
-    if (remoteInput) remoteInput.value = ''
-    if (manualClose) { manualClose = false; setError('') }
-    else { setError('信令服务器连接已关闭') }
-    setState('idle')
-    renderMembers([])
+    if (manualClose) {
+      manualClose = false
+      setError('')
+      closePeerConnection()
+      roomId = null
+      if (remoteInput) remoteInput.value = ''
+      setState('idle')
+      renderMembers([])
+    } else {
+      // Unexpected disconnect — attempt reconnection
+      setError('连接断开，正在重连…')
+      if (roomId) {
+        setTimeout(() => {
+          if (roomId && state !== 'idle') {
+            connectWS()
+          }
+        }, 2000)
+      } else {
+        closePeerConnection()
+        setState('idle')
+        renderMembers([])
+      }
+    }
   }
 }