feat: 实现网关服务的长连接功能

backend/gateway/internal/ws/client.go

@@ -0,0 +1,144 @@
+package ws
+
+import (
+	"encoding/json"
+	"log"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+const (
+	// 写入超时时间
+	writeWait = 10 * time.Second
+
+	// 读取超时时间
+	pongWait = 60 * time.Second
+
+	// 发送 ping 的时间间隔，必须小于 pongWait
+	pingPeriod = (pongWait * 9) / 10
+
+	// 最大消息大小
+	maxMessageSize = 4096
+)
+
+type Client struct {
+	hub  *Hub
+	conn *websocket.Conn
+	send chan *Message
+}
+
+func NewClient(hub *Hub, conn *websocket.Conn) *Client {
+	return &Client{
+		hub:  hub,
+		conn: conn,
+		send: make(chan *Message, 256),
+	}
+}
+
+func (c *Client) readPump() {
+	defer func() {
+		c.hub.unregister <- c
+		c.conn.Close()
+	}()
+
+	c.conn.SetReadLimit(maxMessageSize)
+	c.conn.SetReadDeadline(time.Now().Add(pongWait))
+	c.conn.SetPongHandler(func(string) error {
+		c.conn.SetReadDeadline(time.Now().Add(pongWait))
+		return nil
+	})
+
+	for {
+		_, message, err := c.conn.ReadMessage()
+		if err != nil {
+			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
+				log.Printf("error: %v", err)
+			}
+			break
+		}
+
+		var msg Message
+		if err := json.Unmarshal(message, &msg); err != nil {
+			log.Printf("error unmarshaling message: %v", err)
+			continue
+		}
+
+		// 处理消息
+		switch msg.Type {
+		case MessageTypePing:
+			// 回复 pong
+			pongMsg := &Message{Type: MessageTypePong}
+			c.send <- pongMsg
+		case MessageTypeText:
+			// 广播文本消息
+			c.hub.broadcast <- &msg
+		case MessageTypeCommand:
+			// 处理命令
+			c.handleCommand(&msg)
+		}
+	}
+}
+
+func (c *Client) writePump() {
+	ticker := time.NewTicker(pingPeriod)
+	defer func() {
+		ticker.Stop()
+		c.conn.Close()
+	}()
+
+	for {
+		select {
+		case message, ok := <-c.send:
+			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if !ok {
+				// Hub 关闭了通道
+				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
+				return
+			}
+
+			w, err := c.conn.NextWriter(websocket.TextMessage)
+			if err != nil {
+				return
+			}
+
+			msgBytes, err := json.Marshal(message)
+			if err != nil {
+				log.Printf("error marshaling message: %v", err)
+				return
+			}
+
+			w.Write(msgBytes)
+
+			if err := w.Close(); err != nil {
+				return
+			}
+
+		case <-ticker.C:
+			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
+			pingMsg := &Message{Type: MessageTypePing}
+			pingBytes, err := json.Marshal(pingMsg)
+			if err != nil {
+				return
+			}
+			if err := c.conn.WriteMessage(websocket.TextMessage, pingBytes); err != nil {
+				return
+			}
+		}
+	}
+}
+
+func (c *Client) handleCommand(msg *Message) {
+	// 处理命令逻辑
+	// 这里可以根据命令类型执行不同的操作
+	log.Printf("Received command: %v", msg.Data)
+
+	// 示例：回复命令执行结果
+	response := &Message{
+		Type:    MessageTypeText,
+		Content: "Command executed successfully",
+		Data:    msg.Data,
+	}
+
+	c.send <- response
+}