feat: server 包新增 WithWebsocketConnInitializer 函数，支持对 websocket 连接打开后进…

…行初始化设置
kercylan98 · Dec 25, 2023 · 7ee4b89 · 7ee4b89
1 parent 2639412
commit 7ee4b89
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 20 deletions.
diff --git a/server/options.go b/server/options.go
@@ -6,6 +6,7 @@ import (
 	"github.com/kercylan98/minotaur/utils/log"
 	"github.com/kercylan98/minotaur/utils/timer"
 	"google.golang.org/grpc"
+	"net/http"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -31,26 +32,38 @@ type option struct {
 }
 
 type runtime struct {
-	deadlockDetect               time.Duration       // 是否开启死锁检测
-	supportMessageTypes          map[int]bool        // websocket 模式下支持的消息类型
-	certFile, keyFile            string              // TLS文件
-	tickerPool                   *timer.Pool         // 定时器池
-	ticker                       *timer.Ticker       // 定时器
-	tickerAutonomy               bool                // 定时器是否独立运行
-	connTickerSize               int                 // 连接定时器大小
-	websocketReadDeadline        time.Duration       // websocket 连接超时时间
-	websocketCompression         int                 // websocket 压缩等级
-	websocketWriteCompression    bool                // websocket 写入压缩
-	limitLife                    time.Duration       // 限制最大生命周期
-	packetWarnSize               int                 // 数据包大小警告
-	messageStatisticsDuration    time.Duration       // 消息统计时长
-	messageStatisticsLimit       int                 // 消息统计数量
-	messageStatistics            []*atomic.Int64     // 消息统计数量
-	messageStatisticsLock        *sync.RWMutex       // 消息统计锁
-	dispatcherBufferSize         int                 // 消息分发器缓冲区大小
-	connWriteBufferSize          int                 // 连接写入缓冲区大小
-	disableAutomaticReleaseShunt bool                // 是否禁用自动释放分流渠道
-	websocketUpgrader            *websocket.Upgrader // websocket 升级器
+	deadlockDetect               time.Duration                                                                       // 是否开启死锁检测
+	supportMessageTypes          map[int]bool                                                                        // websocket 模式下支持的消息类型
+	certFile, keyFile            string                                                                              // TLS文件
+	tickerPool                   *timer.Pool                                                                         // 定时器池
+	ticker                       *timer.Ticker                                                                       // 定时器
+	tickerAutonomy               bool                                                                                // 定时器是否独立运行
+	connTickerSize               int                                                                                 // 连接定时器大小
+	websocketReadDeadline        time.Duration                                                                       // websocket 连接超时时间
+	websocketCompression         int                                                                                 // websocket 压缩等级
+	websocketWriteCompression    bool                                                                                // websocket 写入压缩
+	limitLife                    time.Duration                                                                       // 限制最大生命周期
+	packetWarnSize               int                                                                                 // 数据包大小警告
+	messageStatisticsDuration    time.Duration                                                                       // 消息统计时长
+	messageStatisticsLimit       int                                                                                 // 消息统计数量
+	messageStatistics            []*atomic.Int64                                                                     // 消息统计数量
+	messageStatisticsLock        *sync.RWMutex                                                                       // 消息统计锁
+	dispatcherBufferSize         int                                                                                 // 消息分发器缓冲区大小
+	connWriteBufferSize          int                                                                                 // 连接写入缓冲区大小
+	disableAutomaticReleaseShunt bool                                                                                // 是否禁用自动释放分流渠道
+	websocketUpgrader            *websocket.Upgrader                                                                 // websocket 升级器
+	websocketConnInitializer     func(writer http.ResponseWriter, request *http.Request, conn *websocket.Conn) error // websocket 连接初始化
+}
+
+// WithWebsocketConnInitializer 通过 websocket 连接初始化的方式创建服务器，当 initializer 返回错误时，服务器将不会处理该连接的后续逻辑
+//   - 该选项仅在创建 NetworkWebsocket 服务器时有效
+func WithWebsocketConnInitializer(initializer func(writer http.ResponseWriter, request *http.Request, conn *websocket.Conn) error) Option {
+	return func(srv *Server) {
+		if srv.network != NetworkWebsocket {
+			return
+		}
+		srv.websocketConnInitializer = initializer
+	}
 }
 
 // WithWebsocketUpgrade 通过指定 websocket.Upgrader 的方式创建服务器

diff --git a/server/server.go b/server/server.go
@@ -258,6 +258,11 @@ func (slf *Server) Run(addr string) error {
 			if err != nil {
 				return
 			}
+			if slf.websocketConnInitializer != nil {
+				if err = slf.websocketConnInitializer(writer, request, ws); err != nil {
+					return
+				}
+			}
 			if len(ip) == 0 {
 				addr := ws.RemoteAddr().String()
 				if index := strings.LastIndex(addr, ":"); index != -1 {