From b2f81ead9fae7e65e2e15bc72c26c0ee86f0e2ea Mon Sep 17 00:00:00 2001
From: TsMask <340112800@qq.com>
Date: Fri, 18 Apr 2025 17:12:22 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20ws=E8=BF=9E=E6=8E=A5=E5=BC=82=E5=B8=B8?=
 =?UTF-8?q?=E4=B8=AD=E6=96=AD=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/modules/ws/service/ws.go         | 11 +++++++++--
 src/modules/ws/service/ws_receive.go |  8 ++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/modules/ws/service/ws.go b/src/modules/ws/service/ws.go
index 264c1aab..5a0caea1 100644
--- a/src/modules/ws/service/ws.go
+++ b/src/modules/ws/service/ws.go
@@ -10,6 +10,7 @@ import (
 	"be.ems/src/framework/resp"
 	"be.ems/src/framework/utils/generate"
 	"be.ems/src/modules/ws/model"
+
 	"github.com/gorilla/websocket"
 )
 
@@ -30,7 +31,7 @@ func (s *WS) UpgraderWs(w http.ResponseWriter, r *http.Request) *websocket.Conn
 	wsUpgrader := websocket.Upgrader{
 		Subprotocols: []string{"omc-ws"},
 		// 设置消息发送缓冲区大小（byte），如果这个值设置得太小，可能会导致服务端在发送大型消息时遇到问题
-		WriteBufferSize: 1024,
+		WriteBufferSize: 4096,
 		// 消息包启用压缩
 		EnableCompression: true,
 		// ws握手超时时间
@@ -195,11 +196,18 @@ func (s *WS) ClientWriteListen(wsClient *model.WSClient) {
 	wsClient.MsgChan <- msgByte
 	// 消息发送监听
 	for msg := range wsClient.MsgChan {
+		// PONG句柄
+		if string(msg) == "ws:pong" {
+			wsClient.LastHeartbeat = time.Now().UnixMilli()
+			wsClient.Conn.WriteMessage(websocket.PongMessage, []byte{})
+			continue
+		}
 		// 关闭句柄
 		if string(msg) == "ws:close" {
 			wsClient.Conn.WriteMessage(websocket.CloseMessage, []byte{})
 			return
 		}
+
 		// 发送消息
 		err := wsClient.Conn.WriteMessage(websocket.TextMessage, msg)
 		if err != nil {
@@ -207,6 +215,5 @@ func (s *WS) ClientWriteListen(wsClient *model.WSClient) {
 			s.ClientClose(wsClient.ID)
 			return
 		}
-		wsClient.LastHeartbeat = time.Now().UnixMilli()
 	}
 }
diff --git a/src/modules/ws/service/ws_receive.go b/src/modules/ws/service/ws_receive.go
index 6bb8f571..4a1720f1 100644
--- a/src/modules/ws/service/ws_receive.go
+++ b/src/modules/ws/service/ws_receive.go
@@ -51,6 +51,14 @@ func (s *WSReceive) Commont(client *model.WSClient, reqMsg model.WSRequest) {
 	case "close":
 		s.close(client)
 		return
+	case "ping", "PING":
+		resByte, _ := json.Marshal(resp.Ok(map[string]any{
+			"requestId": reqMsg.RequestID,
+			"data":      "PONG",
+		}))
+		client.MsgChan <- resByte
+		client.MsgChan <- []byte("ws:pong")
+		return
 	case "ps":
 		resByte, err = processor.GetProcessData(reqMsg.RequestID, reqMsg.Data)
 	case "net":