From 7ee4b893cdc3e2bea8c2be39e254e4f3b13c5695 Mon Sep 17 00:00:00 2001 From: kercylan98 Date: Mon, 25 Dec 2023 15:17:03 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20server=20=E5=8C=85=E6=96=B0=E5=A2=9E=20?= =?UTF-8?q?WithWebsocketConnInitializer=20=E5=87=BD=E6=95=B0=EF=BC=8C?= =?UTF-8?q?=E6=94=AF=E6=8C=81=E5=AF=B9=20websocket=20=E8=BF=9E=E6=8E=A5?= =?UTF-8?q?=E6=89=93=E5=BC=80=E5=90=8E=E8=BF=9B=E8=A1=8C=E5=88=9D=E5=A7=8B?= =?UTF-8?q?=E5=8C=96=E8=AE=BE=E7=BD=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- server/options.go | 53 +++++++++++++++++++++++++++++------------------ server/server.go | 5 +++++ 2 files changed, 38 insertions(+), 20 deletions(-) diff --git a/server/options.go b/server/options.go index 29f441f..dbe14c7 100644 --- 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 index 6688cea..1b8daa8 100644 --- 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 {