kefu_cloud/server/internal/ws/ws.go

package ws

import (
	"encoding/json"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	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
	Kind      string
	SessionID *uint
	Send      chan []byte
}

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
	register   chan *Client
	unregister chan *Client
	mu         sync.RWMutex
}

var DefaultHub = NewHub()

func NewHub() *Hub {
	return &Hub{
		clients:    make(map[*Client]bool),
		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 {
		case client := <-h.register:
			h.mu.Lock()
			h.clients[client] = true
			h.mu.Unlock()

		case client := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				close(client.Send)
			}
			h.mu.Unlock()
		}
	}
}

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 {
			continue
		}
		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)
		}
	}
}

// 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{
		UserID:   userID,
		TenantID: tenantID,
		Role:     role,
		Kind:     "agent",
		Send:     make(chan []byte, 256),
	}
	log.Printf("WebSocket 连接: user=%d tenant=%d", userID, tenantID)
	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)
}