From f95191c8afda278e9fc8d9627801c93c9709d4ad Mon Sep 17 00:00:00 2001
From: UUBulb <35923940+uubulb@users.noreply.github.com>
Date: Wed, 17 Jul 2024 10:06:23 +0800
Subject: [PATCH] fix: concurrent write to single WebSocket connection (#392)

---
 model/monitor.go            | 2 +-
 pkg/websocketx/safe_conn.go | 6 ++++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/model/monitor.go b/model/monitor.go
index 459a036..69c367c 100644
--- a/model/monitor.go
+++ b/model/monitor.go
@@ -120,7 +120,7 @@ func (m *Monitor) AfterFind(tx *gorm.DB) error {
 
 // IsServiceSentinelNeeded 判断该任务类型是否需要进行服务监控 需要则返回true
 func IsServiceSentinelNeeded(t uint64) bool {
-	return t != TaskTypeCommand && t != TaskTypeTerminal && t != TaskTypeUpgrade
+	return t != TaskTypeCommand && t != TaskTypeTerminalGRPC && t != TaskTypeUpgrade
 }
 
 func (m *Monitor) InitSkipServers() error {
diff --git a/pkg/websocketx/safe_conn.go b/pkg/websocketx/safe_conn.go
index 1d29019..512bddf 100644
--- a/pkg/websocketx/safe_conn.go
+++ b/pkg/websocketx/safe_conn.go
@@ -28,6 +28,12 @@ func (conn *Conn) Write(data []byte) (int, error) {
 	return len(data), nil
 }
 
+func (conn *Conn) WriteMessage(messageType int, data []byte) error {
+	conn.writeLock.Lock()
+	defer conn.writeLock.Unlock()
+	return conn.Conn.WriteMessage(messageType, data)
+}
+
 func (conn *Conn) Read(data []byte) (int, error) {
 	if len(conn.dataBuf) > 0 {
 		n := copy(data, conn.dataBuf)