Fix WebSocket TOCTOU race conditions and concurrent write panics (#402)

backend/websocket/debate_spectator.go

@@ -231,7 +231,7 @@ func DebateWebsocketHandler(c *gin.Context) {
 	// Send initial poll snapshot
 	snapshot, err := loadPollSnapshot(debateID)
 	if err == nil && snapshot != nil {
-		conn.WriteJSON(snapshot)
+		client.WriteJSON(snapshot)
 	} else if err != nil {
 	}
 
@@ -254,7 +254,7 @@ func DebateWebsocketHandler(c *gin.Context) {
 		},
 		"timestamp": time.Now().Unix(),
 	}
-	conn.WriteJSON(presenceEvent)
+	client.WriteJSON(presenceEvent)
 
 	// Read pump
 	go readPump(client, hub)

backend/websocket/gamification.go

@@ -23,6 +23,13 @@ func (gc *GamificationClient) SafeWriteJSON(v interface{}) error {
 	return gc.Conn.WriteJSON(v)
 }
 
+// SafeWriteMessage safely writes raw WebSocket messages to the gamification client's connection
+func (gc *GamificationClient) SafeWriteMessage(messageType int, data []byte) error {
+	gc.writeMu.Lock()
+	defer gc.writeMu.Unlock()
+	return gc.Conn.WriteMessage(messageType, data)
+}
+
 // Global gamification hub for broadcasting events to all connected clients
 var (
 	gamificationClients = make(map[*GamificationClient]bool)

backend/websocket/gamification_handler.go

@@ -122,7 +122,7 @@ func GamificationWebSocketHandler(c *gin.Context) {
 
 		// Handle ping/pong for keepalive
 		if messageType == websocket.PingMessage {
-			if err := conn.WriteMessage(websocket.PongMessage, nil); err != nil {
+			if err := client.SafeWriteMessage(websocket.PongMessage, nil); err != nil {
 				log.Printf("Error writing pong: %v", err)
 				break
 			}

backend/websocket/team_websocket.go

@@ -229,6 +229,13 @@ func TeamWebsocketHandler(c *gin.Context) {
 		Team2Ready:   make(map[string]bool),
 	}
 
+	// Upgrade the connection
+	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		log.Println("Team WebSocket upgrade error:", err)
+		return
+	}
+
 	// Insert room if absent
 	teamRoomsMutex.Lock()
 	room, exists := teamRooms[roomKey]
@@ -238,15 +245,9 @@ func TeamWebsocketHandler(c *gin.Context) {
 	} else {
 		// discard prepared room; existing room will be used
 	}
+	room.Mutex.Lock()
 	teamRoomsMutex.Unlock()
 
-	// Upgrade the connection
-	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
-	if err != nil {
-		log.Println("Team WebSocket upgrade error:", err)
-		return
-	}
-
 	// CRITICAL: Validate userTeamID matches one of the debate teams before creating client
 	userTeamIDHex := userTeamID.Hex()
 	team1IDHex := debate.Team1ID.Hex()
@@ -255,6 +256,7 @@ func TeamWebsocketHandler(c *gin.Context) {
 	if userTeamIDHex != team1IDHex && userTeamIDHex != team2IDHex {
 		log.Printf("[TeamWebsocketHandler] ❌ ERROR: UserTeamID %s doesn't match Team1ID %s or Team2ID %s", userTeamIDHex, team1IDHex, team2IDHex)
 		c.JSON(http.StatusInternalServerError, gin.H{"error": "Team assignment error"})
+		room.Mutex.Unlock()
 		conn.Close()
 		return
 	}
@@ -278,7 +280,6 @@ func TeamWebsocketHandler(c *gin.Context) {
 		Tokens:       10, // Initial tokens
 	}
 
-	room.Mutex.Lock()
 	room.Clients[conn] = client
 	room.Mutex.Unlock()
 
@@ -354,13 +355,19 @@ func TeamWebsocketHandler(c *gin.Context) {
 			userID := client.UserID.Hex()
 			room.Mutex.Lock()
 			delete(room.Clients, conn)
+			clientCount := len(room.Clients)
+			room.Mutex.Unlock()
+
 			// If room is empty, delete it
-			if len(room.Clients) == 0 {
+			if clientCount == 0 {
 				teamRoomsMutex.Lock()
-				delete(teamRooms, roomKey)
+				room.Mutex.Lock()
+				if len(room.Clients) == 0 {
+					delete(teamRooms, roomKey)
+				}
+				room.Mutex.Unlock()
 				teamRoomsMutex.Unlock()
 			}
-			room.Mutex.Unlock()
 
 			// Notify remaining clients that this user has left
 			broadcastExcept(room, conn, map[string]any{

backend/websocket/websocket.go

@@ -269,24 +269,24 @@ func WebsocketHandler(c *gin.Context) {
 		return
 	}
 
+	// Upgrade the connection.
+	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		return
+	}
+
 	// Create the room if it doesn't exist.
 	roomsMutex.Lock()
 	if _, exists := rooms[roomID]; !exists {
 		rooms[roomID] = &Room{Clients: make(map[*websocket.Conn]*Client)}
 	}
 	room := rooms[roomID]
+	room.Mutex.Lock()
 	roomsMutex.Unlock()
 
-	// Upgrade the connection.
-	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
-	if err != nil {
-		return
-	}
-
 	// Check if this is a spectator connection (they want to receive video streams)
 	// Allow spectators to connect even if room has 2 debaters
 	isSpectator := strings.EqualFold(c.Query("spectator"), "true")
-	room.Mutex.Lock()
 	currentDebaters := 0
 	for _, existing := range room.Clients {
 		if !existing.IsSpectator {
@@ -300,7 +300,6 @@ func WebsocketHandler(c *gin.Context) {
 		conn.Close()
 		return
 	}
-	room.Mutex.Unlock()
 
 	if avatarURL == "" {
 		avatarURL = "https://api.dicebear.com/9.x/big-ears/svg?seed=Nolan"
@@ -332,11 +331,7 @@ func WebsocketHandler(c *gin.Context) {
 		client.ConnectionID = uuid.New().String()
 	}
 
-	// Mark as spectator if needed (we can add a field to Client struct for this)
-	// For now, we'll handle it through the message handlers
-
-	// Send current participants to the new client
-	room.Mutex.Lock()
+	// Add client to room
 	room.Clients[conn] = client
 	room.Mutex.Unlock()
 
@@ -414,14 +409,18 @@ func WebsocketHandler(c *gin.Context) {
 				delete(room.Clients, conn)
 			}
 			clientCount = len(room.Clients)
+			room.Mutex.Unlock()
 
 			// If room is empty, delete it.
 			if clientCount == 0 {
 				roomsMutex.Lock()
-				delete(rooms, roomID)
+				room.Mutex.Lock()
+				if len(room.Clients) == 0 {
+					delete(rooms, roomID)
+				}
+				room.Mutex.Unlock()
 				roomsMutex.Unlock()
 			}
-			room.Mutex.Unlock()
 
 			if exists && disconnectedClient.IsSpectator {
 				log.Printf("[ws] spectator disconnected: room=%s connectionId=%s user=%s", roomID, disconnectedClient.ConnectionID, disconnectedClient.Email)
@@ -642,7 +641,7 @@ func handlePhaseChange(room *Room, conn *websocket.Conn, message Message, roomID
 				"currentTurn": currentTurn,
 				"phase":       message.Phase,
 			}
-			if err := clientConn.WriteJSON(response); err != nil {
+			if err := client.SafeWriteJSON(response); err != nil {
 			}
 		}
 	}