feat: server 包新增 WithWebsocketConnInitializer 函数,支持对 websocket 连接打开后进行初始化设置

This commit is contained in:
kercylan98 2023-12-25 15:17:03 +08:00
parent 2639412f96
commit 7ee4b893cd
2 changed files with 38 additions and 20 deletions

View File

@ -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"
@ -51,6 +52,18 @@ type runtime struct {
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 的方式创建服务器

View File

@ -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 {