207 lines
5.2 KiB
Plaintext
207 lines
5.2 KiB
Plaintext
// ************************************
|
||
// @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)
|
||
}
|
||
}
|