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