// ************************************ // @package: websocketServer // @description: // @author: // @revision history: // @create date: 2022-02-16 18:13:45 // ************************************ package websocketServer import ( "time" "github.com/gorilla/websocket" webServer "Framework/webServer" logUtil "goutil/logUtil" ) // 事件回调函数 type EventCallbackFuncs struct { // websocket连接事件 OnConnFunc func(ctx *Context) // websocket关闭事件 OnCloseFunc func(ctx *Context) // websocket接收事件 OnMsgFunc func(ctx *Context, msgType int, msgData []byte) } // 用户自定义数据结构 type userDatas struct { // WsServer 或 WssServer 指针 server interface{} // 事件回调函数 eventCallback *EventCallbackFuncs } // iServerMgr // @description: WsServer/WssServer内部管理接口,以便hookHandler中统一访问 type iServerMgr interface { // 升级为websocket upgrade(ctx *Context) (conn *websocket.Conn, err error) // 将连接从连接池删除 delConn(conn *websocket.Conn) // 获取接收到Ping消息时,是否自动回复Pong信息 GetAutoPong() bool } // 回调勾子,将http/https升级为websocket func hookHandler(webServerCtx *webServer.Context) { defer func() { if r := recover(); r != nil { logUtil.LogUnknownError(r) } }() userDataI := webServerCtx.GetUserData() if userDataI == nil { return } userData, ok := userDataI.(*userDatas) if !ok { logUtil.ErrorLog("userData type error") return } // 通信结束信号 cls := make(chan struct{}) ctx := &Context{ webServerCtx: webServerCtx, // web_server环境 cls: cls, // 关闭连接信号 } var serverMgr iServerMgr var conn *websocket.Conn var err error // 转为iServerMgr switch userData.server.(type) { case *WsServer: if svr, ok := userData.server.(*WsServer); ok { serverMgr = svr } else { logUtil.ErrorLog("server type not WsServer") return } case *WssServer: if svr, ok := userData.server.(*WssServer); ok { serverMgr = svr } else { logUtil.ErrorLog("server type not WssServer") return } default: logUtil.ErrorLog("server type not WsServer or WssServer") return } // 升级为websocket conn, err = serverMgr.upgrade(ctx) if err != nil { if err.Error() == "connManager(disableNewConn)" { // 禁用新连接。正常功能,直接返回 return } logUtil.ErrorLog("websocket Upgrade failed: %v", err) return } // 将连接从连接池删除 defer serverMgr.delConn(conn) // 关闭连接 defer conn.Close() // 默认情况下,ReadMessage不会读取到ping/pong/close消息(内部有专门的处理函数) // 设置心跳包处理函数 conn.SetPingHandler(func(msg string) error { // 只要收到消息,都需要更新最近一次收到心跳包的时间 ctx.heartbeat = time.Now() if serverMgr.GetAutoPong() { // 自动回应一个Pong心跳 go ctx.SendMessage(websocket.PongMessage, []byte(msg)) } // 接收消息回调 if userData.eventCallback.OnMsgFunc != nil { userData.eventCallback.OnMsgFunc(ctx, websocket.PingMessage, []byte(msg)) } return nil }) // 设置关闭包处理函数 conn.SetCloseHandler(func(code int, text string) error { // 所有向cls写入,都使用select超时结构,以保证这儿不会一直阻塞,确保此协程能退出 select { case <-time.After(time.Millisecond * 10): case cls <- struct{}{}: } return nil }) // 设置最近一次收到心跳包的时间 ctx.heartbeat = time.Now() // 新连接回调 if userData.eventCallback.OnConnFunc != nil { userData.eventCallback.OnConnFunc(ctx) } // 开启读协程 go func() { defer func() { if r := recover(); r != nil { logUtil.LogUnknownError(r) // 有异常出现(可能是用户回调中出现异常);执行到这儿,需要关闭连接 // 所有向cls写入,都使用select超时结构,以保证这儿不会一直阻塞,确保此协程能退出 select { case <-time.After(time.Millisecond * 10): case cls <- struct{}{}: } } }() for { // 注意:ReadMessage不会读取到心跳包和关闭包数据;心跳包/关闭包需要设置专门的处理函数 // 但内部对心跳包/关闭包的处理,也是由ReadMessage函数触发的(也就是不调用ReadMessage函数,可能也不会触发对心跳包/关闭包的处理); // 经测试和内部代码确认:调用心跳包/关闭包处理函数时,ReadMessage不会返回;心跳包/关闭包处理函数调用完毕后ReadMessage才可能返回 mt, msg, err := conn.ReadMessage() if err != nil { // 所有向cls写入,都使用select超时结构,以保证这儿不会一直阻塞,确保此协程能退出 select { case <-time.After(time.Millisecond * 10): case cls <- struct{}{}: } break } // 只要收到消息,都需要更新最近一次收到心跳包的时间 ctx.heartbeat = time.Now() // 接收消息回调 if userData.eventCallback.OnMsgFunc != nil { userData.eventCallback.OnMsgFunc(ctx, mt, msg) } } }() // 等待退出 <-cls // 设置已关闭标志 ctx.isClosed = true // 关闭回调 if userData.eventCallback.OnCloseFunc != nil { userData.eventCallback.OnCloseFunc(ctx) } }