goProject/trunk/framework/websocketServer/hookHandler.go

// ************************************
// @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)
	}
}