From 5ab67999d3d875e2b6df34ee8e4b4efce89c94ee Mon Sep 17 00:00:00 2001
From: lc-official <lc-official@qq.com>
Date: Mon, 11 Aug 2025 11:54:58 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8Dwebsocket?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 util/ws/websocket.go | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/util/ws/websocket.go b/util/ws/websocket.go
index f36c5d3..e9aed61 100644
--- a/util/ws/websocket.go
+++ b/util/ws/websocket.go
@@ -3,6 +3,7 @@ package ws
 import (
 	"context"
 	"github.com/gorilla/websocket"
+	"log"
 )
 
 type Packet struct {
@@ -27,16 +28,21 @@ func Binary(buf []byte) *Packet {
 type Conn struct {
 	WriteChan chan *Packet
 	ReadChan  chan []byte
-	ctx       context.Context
+	Context   context.Context
+	cancel    context.CancelFunc
 	conn      *websocket.Conn
 	close     bool
 }
 
 func Wrap(conn *websocket.Conn) *Conn {
+	ctx, cancel := context.WithCancel(context.Background())
 	return &Conn{
 		WriteChan: make(chan *Packet),
 		ReadChan:  make(chan []byte),
 		conn:      conn,
+		Context:   ctx,
+		cancel:    cancel,
+		close:     false,
 	}
 }
 
@@ -44,6 +50,9 @@ func (conn *Conn) ReadLoop() {
 	for !conn.close {
 		_, message, err := conn.conn.ReadMessage()
 		if err != nil {
+			if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseAbnormalClosure) {
+				log.Println("websocket read error:", err)
+			}
 			conn.Close()
 			break
 		}
@@ -63,16 +72,14 @@ func (conn *Conn) WriteLoop() {
 			if err != nil {
 				break
 			}
-		default:
-			if conn.close {
-				return
-			}
+		case <-conn.Context.Done():
+			return
 		}
 	}
 }
 
 func (conn *Conn) Close() error {
-	conn.ReadChan <- nil
+	conn.cancel()
 	conn.close = true
 	return conn.conn.Close()
 }