From 828f3af1fa9659503993080ebf8a20957161a9d6 Mon Sep 17 00:00:00 2001
From: liuxiaobo <1224730913@qq.com>
Date: Thu, 29 May 2025 20:30:18 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9bug?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ws/iconn.go    |  1 +
 ws/wsClient.go | 11 +++++++++--
 ws/wsConn.go   |  4 ++++
 ws/ws_test.go  | 10 +++++++++-
 4 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/ws/iconn.go b/ws/iconn.go
index 130a411..6f981ad 100644
--- a/ws/iconn.go
+++ b/ws/iconn.go
@@ -1,6 +1,7 @@
 package ws
 
 type IConn interface {
+	Addr() string
 	NotifyClose()
 	SendMsg(data []byte) error
 	Name() string
diff --git a/ws/wsClient.go b/ws/wsClient.go
index bf25924..95065b8 100644
--- a/ws/wsClient.go
+++ b/ws/wsClient.go
@@ -11,15 +11,20 @@ import (
 	"github.com/gorilla/websocket"
 )
 
+type IOnFunc interface {
+	OnMessage(msg []byte) error
+}
+
 type Client struct {
 	conn     *websocket.Conn
 	sendChan chan *wsMessage
 	ctx      context.Context
 	cancel   context.CancelFunc
 	wg       sync.WaitGroup
+	onFunc   IOnFunc
 }
 
-func NewClient(url string) (*Client, error) {
+func NewClient(url string, onFunc IOnFunc) (*Client, error) {
 	dialer := websocket.DefaultDialer
 	dialer.HandshakeTimeout = 30 * time.Second
 	conn, _, err := dialer.Dial(url, http.Header{"User-Agent": {"MyClient/1.0"}})
@@ -34,6 +39,7 @@ func NewClient(url string) (*Client, error) {
 		sendChan: make(chan *wsMessage, 100),
 		ctx:      ctx,
 		cancel:   cancel,
+		onFunc:   onFunc,
 	}, nil
 }
 
@@ -64,7 +70,8 @@ func (c *Client) readLoop() {
 			case websocket.PongMessage:
 
 			case websocket.TextMessage, websocket.BinaryMessage:
-				log.DebugF("收到消息，类型:%v 内容:%v", messageType, string(message))
+				_ = c.onFunc.OnMessage(message)
+
 			case websocket.CloseMessage:
 				log.Debug("收到关闭帧")
 				c.Stop()
diff --git a/ws/wsConn.go b/ws/wsConn.go
index 2e61726..1bd9cb5 100644
--- a/ws/wsConn.go
+++ b/ws/wsConn.go
@@ -193,6 +193,10 @@ func (c *wsConnect) Name() string {
 	return fmt.Sprintf("用户:%v, 地址:%v", c.userId, c.wsConn.RemoteAddr())
 }
 
+func (c *wsConnect) Addr() string {
+	return c.wsConn.RemoteAddr().String()
+}
+
 func (c *wsConnect) Log(format string, v ...interface{}) string {
 	s := fmt.Sprintf("连接:%v, id:%v ", c.wsConn.RemoteAddr().String(), c.id)
 	return s + fmt.Sprintf(format, v...)
diff --git a/ws/ws_test.go b/ws/ws_test.go
index 519f01a..1ade146 100644
--- a/ws/ws_test.go
+++ b/ws/ws_test.go
@@ -26,8 +26,16 @@ func wsServer() {
 	s.Run()
 }
 
+type onFunc struct {
+}
+
+func (onFunc) OnMessage(msg []byte) error {
+	log.DebugF("收到消息，内容:%v", string(msg))
+	return nil
+}
+
 func wsClient() {
-	client, err := NewClient(fmt.Sprintf("ws://%v", addr))
+	client, err := NewClient(fmt.Sprintf("ws://%v", addr), onFunc{})
 	if err != nil {
 		log.Fatal(err.Error())
 		return