tam

server/internal/websocket/websocket.go

@@ -5,12 +5,18 @@ import (
 	"log"
 	"strconv"
 	"sync"
+	"time"
 
 	"github.com/gofiber/websocket/v2"
 	"gorm.io/gorm"
 	internalDb "server/internal/db"
 )
 
+const (
+	wsPingInterval = 30 * time.Second
+	wsPongWait     = 60 * time.Second
+)
+
 type SocketMsg struct {
 	Event string         `json:"event"`
 	Data  map[string]any `json:"data"`
@@ -32,6 +38,12 @@ func (c *SocketClient) WriteJSON(v any) error {
 	return c.Conn.WriteJSON(v)
 }
 
+func (c *SocketClient) WriteRaw(msg []byte) error {
+	c.writeMu.Lock()
+	defer c.writeMu.Unlock()
+	return c.Conn.WriteMessage(websocket.TextMessage, msg)
+}
+
 
 type WsHub struct {
 	mu              sync.RWMutex
@@ -133,6 +145,11 @@ func (h *WsHub) Unregister(c *SocketClient) {
 	defer h.mu.Unlock()
 
 	if c.Role == "student" {
+		current, exists := h.students[c.StudentID]
+		if !exists || current != c {
+			log.Printf("[WS] Student %d stale disconnect ignored (replaced by newer connection)", c.StudentID)
+			return
+		}
 		delete(h.students, c.StudentID)
 		log.Printf("[WS] Student %d disconnected", c.StudentID)
 
@@ -146,7 +163,6 @@ func (h *WsHub) Unregister(c *SocketClient) {
 					})
 				}
 			}
-			delete(h.subscribers, c.StudentID)
 		}
 	} else if c.Role == "teacher" {
 		delete(h.teachers, c.Addr)
@@ -248,12 +264,15 @@ func (h *WsHub) Unsubscribe(teacherAddr string, studentID int64) {
 // Chuyển tiếp frame ảnh từ Student đến các Teacher đã subscribe
 func (h *WsHub) RelayFrame(studentID int64, event string, data map[string]any) {
 	h.mu.RLock()
-	defer h.mu.RUnlock()
-
 	teachersList, exists := h.subscribers[studentID]
 	if !exists || len(teachersList) == 0 {
+		h.mu.RUnlock()
 		return
 	}
+	// Snapshot subscriber addresses while holding read lock
+	addrs := make([]string, len(teachersList))
+	copy(addrs, teachersList)
+	h.mu.RUnlock()
 
 	relayEvent := "teacher:screenshot-stream-frame"
 	if event == "webcam_stream_frame" {
@@ -267,10 +286,18 @@ func (h *WsHub) RelayFrame(studentID int64, event string, data map[string]any) {
 			"imageBuffer": data["imageBuffer"],
 		},
 	}
+	msgBytes, err := json.Marshal(msg)
+	if err != nil {
+		return
+	}
 
-	for _, addr := range teachersList {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	for _, addr := range addrs {
 		if t, found := h.teachers[addr]; found {
-			_ = t.WriteJSON(msg)
+			if err := t.WriteRaw(msgBytes); err != nil {
+				log.Printf("[WS] Relay to teacher %s failed: %v", addr, err)
+			}
 		}
 	}
 }
@@ -308,11 +335,37 @@ func WebSocketHandler(db *gorm.DB) func(*websocket.Conn) {
 			c.Close()
 		}()
 
+		c.SetReadDeadline(time.Now().Add(wsPongWait))
+		c.SetPongHandler(func(string) error {
+			return c.SetReadDeadline(time.Now().Add(wsPongWait))
+		})
+
+		pingDone := make(chan struct{})
+		defer close(pingDone)
+		go func() {
+			ticker := time.NewTicker(wsPingInterval)
+			defer ticker.Stop()
+			for {
+				select {
+				case <-pingDone:
+					return
+				case <-ticker.C:
+					client.writeMu.Lock()
+					err := c.WriteMessage(websocket.PingMessage, nil)
+					client.writeMu.Unlock()
+					if err != nil {
+						return
+					}
+				}
+			}
+		}()
+
 		for {
 			_, msgBytes, err := c.ReadMessage()
 			if err != nil {
 				break
 			}
+			_ = c.SetReadDeadline(time.Now().Add(wsPongWait))
 
 			var msg SocketMsg
 			if err := json.Unmarshal(msgBytes, &msg); err != nil {