修改bug

ws/iconn.go

@@ -1,6 +1,7 @@
 package ws
 
 type IConn interface {
+	Addr() string
 	NotifyClose()
 	SendMsg(data []byte) error
 	Name() string

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()

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...)

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