修复会话安全与实时消息

2026-07-14 15:06:01 +08:00
parent 064af29d3b
commit 8cff2a5824
29 changed files with 1961 additions and 546 deletions
@@ -11,31 +11,29 @@ import (
 )

 var upgrader = websocket.Upgrader{
-	CheckOrigin: func(r *http.Request) bool { return true },
+	CheckOrigin:  func(r *http.Request) bool { return true },
+	Subprotocols: []string{"kefu-v1", "kefu-visitor-v1"},
 }

 type Client struct {
-	Conn     *websocket.Conn
-	UserID   uint
-	TenantID uint
-	Role     string
-	Send     chan []byte
+	Conn      *websocket.Conn
+	UserID    uint
+	TenantID  uint
+	Role      string
+	Kind      string
+	SessionID *uint
+	Send      chan []byte
 }

-type Message struct {
-	Type      string `json:"type"`
-	SessionID uint   `json:"session_id"`
-	Content   string `json:"content,omitempty"`
-	FromID    uint   `json:"from_id,omitempty"`
-	FromName  string `json:"from_name,omitempty"`
-	TenantID  uint   `json:"tenant_id"`
-	Seq       int    `json:"seq,omitempty"`
-	Timestamp int64  `json:"timestamp"`
+type Event struct {
+	Type      string      `json:"type"`
+	SessionID uint        `json:"session_id,omitempty"`
+	Data      interface{} `json:"data,omitempty"`
+	Timestamp int64       `json:"timestamp"`
 }

 type Hub struct {
 	clients    map[*Client]bool
-	broadcast  chan []byte
 	register   chan *Client
 	unregister chan *Client
 	mu         sync.RWMutex
@@ -46,12 +44,20 @@ var DefaultHub = NewHub()
 func NewHub() *Hub {
 	return &Hub{
 		clients:    make(map[*Client]bool),
-		broadcast:  make(chan []byte, 256),
 		register:   make(chan *Client),
 		unregister: make(chan *Client),
 	}
 }

+func NewEvent(eventType string, sessionID uint, data interface{}) ([]byte, error) {
+	return json.Marshal(Event{
+		Type:      eventType,
+		SessionID: sessionID,
+		Data:      data,
+		Timestamp: time.Now().UnixMilli(),
+	})
+}
+
 func (h *Hub) Run() {
 	for {
 		select {
@@ -67,93 +73,117 @@ func (h *Hub) Run() {
 				close(client.Send)
 			}
 			h.mu.Unlock()
-
-		case msg := <-h.broadcast:
-			h.mu.RLock()
-			for client := range h.clients {
-				select {
-				case client.Send <- msg:
-				default:
-					close(client.Send)
-					delete(h.clients, client)
-				}
-			}
-			h.mu.RUnlock()
 		}
 	}
 }

-func (h *Hub) BroadcastToTenant(tenantID uint, msg []byte) {
+func (h *Hub) send(client *Client, message []byte) {
+	select {
+	case client.Send <- message:
+	default:
+	}
+}
+
+// BroadcastToSession 仅把消息推送给当前客服、主管/管理员及该会话的访客。
+func (h *Hub) BroadcastToSession(tenantID, sessionID uint, agentID *uint, message []byte) {
 	h.mu.RLock()
 	defer h.mu.RUnlock()
+
 	for client := range h.clients {
-		if client.TenantID == tenantID {
-			select {
-			case client.Send <- msg:
-			default:
-			}
-		}
-	}
-}
-
-func HandleWebSocket(c *Client) {
-	conn := c.Conn
-
-	go func() {
-		ticker := time.NewTicker(30 * time.Second)
-		defer ticker.Stop()
-		for {
-			select {
-			case msg, ok := <-c.Send:
-				if !ok {
-					conn.WriteMessage(websocket.CloseMessage, []byte{})
-					return
-				}
-				conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
-				if err := conn.WriteMessage(websocket.TextMessage, msg); err != nil {
-					return
-				}
-			case <-ticker.C:
-				conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
-				if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
-					return
-				}
-			}
-		}
-	}()
-
-	for {
-		_, msgBytes, err := conn.ReadMessage()
-		if err != nil {
-			break
-		}
-
-		var msg Message
-		if err := json.Unmarshal(msgBytes, &msg); err != nil {
+		if client.TenantID != tenantID {
 			continue
 		}
-		msg.TenantID = c.TenantID
-		msg.FromID = c.UserID
-		msg.Timestamp = time.Now().UnixMilli()
-
-		reply, _ := json.Marshal(msg)
-		DefaultHub.BroadcastToTenant(c.TenantID, reply)
+		if client.Kind == "visitor" {
+			if client.SessionID != nil && *client.SessionID == sessionID {
+				h.send(client, message)
+			}
+			continue
+		}
+		if client.Role == "admin" || client.Role == "supervisor" ||
+			(agentID != nil && client.Role == "agent" && client.UserID == *agentID) {
+			h.send(client, message)
+		}
 	}
 }

-func Upgrade(w http.ResponseWriter, r *http.Request, userID, tenantID uint, role string) (*Client, error) {
+// BroadcastToTenantStaff 只通知租户内的工作人员，不向访客泄露其他会话事件。
+func (h *Hub) BroadcastToTenantStaff(tenantID uint, message []byte) {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	for client := range h.clients {
+		if client.TenantID == tenantID && client.Kind == "agent" {
+			h.send(client, message)
+		}
+	}
+}
+
+func HandleWebSocket(client *Client) {
+	defer func() {
+		DefaultHub.unregister <- client
+		client.Conn.Close()
+	}()
+
+	go writePump(client)
+	for {
+		if _, _, err := client.Conn.ReadMessage(); err != nil {
+			return
+		}
+	}
+}
+
+func writePump(client *Client) {
+	ticker := time.NewTicker(30 * time.Second)
+	defer ticker.Stop()
+	for {
+		select {
+		case message, ok := <-client.Send:
+			if !ok {
+				client.Conn.WriteMessage(websocket.CloseMessage, []byte{})
+				return
+			}
+			client.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+			if err := client.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
+				return
+			}
+		case <-ticker.C:
+			client.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+			if err := client.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				return
+			}
+		}
+	}
+}
+
+func upgrade(w http.ResponseWriter, r *http.Request, client *Client) (*Client, error) {
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		return nil, err
 	}
+	client.Conn = conn
+	DefaultHub.register <- client
+	return client, nil
+}
+
+func UpgradeAgent(w http.ResponseWriter, r *http.Request, userID, tenantID uint, role string) (*Client, error) {
 	client := &Client{
-		Conn:     conn,
 		UserID:   userID,
 		TenantID: tenantID,
 		Role:     role,
+		Kind:     "agent",
 		Send:     make(chan []byte, 256),
 	}
-	DefaultHub.register <- client
 	log.Printf("WebSocket 连接: user=%d tenant=%d", userID, tenantID)
-	return client, nil
+	return upgrade(w, r, client)
+}
+
+func UpgradeVisitor(w http.ResponseWriter, r *http.Request, tenantID, sessionID uint) (*Client, error) {
+	client := &Client{
+		TenantID:  tenantID,
+		Kind:      "visitor",
+		SessionID: &sessionID,
+		Send:      make(chan []byte, 256),
+	}
+	log.Printf("访客 WebSocket 连接: session=%d tenant=%d", sessionID, tenantID)
+	return upgrade(w, r, client)
 }
@@ -0,0 +1,41 @@
+package ws
+
+import (
+	"testing"
+	"time"
+)
+
+func TestBroadcastToSessionRestrictsRecipients(t *testing.T) {
+	hub := NewHub()
+	go hub.Run()
+
+	sessionID := uint(12)
+	agentID := uint(7)
+	assignedAgent := &Client{TenantID: 1, UserID: agentID, Role: "agent", Kind: "agent", Send: make(chan []byte, 1)}
+	otherAgent := &Client{TenantID: 1, UserID: 8, Role: "agent", Kind: "agent", Send: make(chan []byte, 1)}
+	supervisor := &Client{TenantID: 1, UserID: 9, Role: "supervisor", Kind: "agent", Send: make(chan []byte, 1)}
+	visitor := &Client{TenantID: 1, Kind: "visitor", SessionID: &sessionID, Send: make(chan []byte, 1)}
+	otherSessionID := uint(13)
+	otherVisitor := &Client{TenantID: 1, Kind: "visitor", SessionID: &otherSessionID, Send: make(chan []byte, 1)}
+	otherTenant := &Client{TenantID: 2, UserID: 7, Role: "agent", Kind: "agent", Send: make(chan []byte, 1)}
+
+	for _, client := range []*Client{assignedAgent, otherAgent, supervisor, visitor, otherVisitor, otherTenant} {
+		hub.register <- client
+	}
+
+	hub.BroadcastToSession(1, sessionID, &agentID, []byte(`{"type":"message"}`))
+	for _, client := range []*Client{assignedAgent, supervisor, visitor} {
+		select {
+		case <-client.Send:
+		case <-time.After(time.Second):
+			t.Fatalf("应收到会话消息的客户端未收到：%+v", client)
+		}
+	}
+	for _, client := range []*Client{otherAgent, otherVisitor, otherTenant} {
+		select {
+		case <-client.Send:
+			t.Fatalf("不应收到会话消息的客户端收到消息：%+v", client)
+		default:
+		}
+	}
+}