diff --git a/config/server_ws.go b/config/server_ws.go index 85ccc63..15947cf 100644 --- a/config/server_ws.go +++ b/config/server_ws.go @@ -26,7 +26,7 @@ type ( // TLS 证书key地址 TLSKey string `json:",optional"` // check origin - CheckOrigin func(*http.Request) bool + CheckOrigin func(*http.Request) bool `json:",optional"` } ) diff --git a/server_ws.go b/server_ws.go index f7e12fc..eda6bab 100644 --- a/server_ws.go +++ b/server_ws.go @@ -11,8 +11,20 @@ import ( "strings" ) +type WsConfOption func(conf config.WSServerConf) + +func WithWSCheckOrigin(fn func(*http.Request) bool) WsConfOption { + return func(conf config.WSServerConf) { + conf.CheckOrigin = fn + } +} + // ListenWebsocket 开始监听Websocket -func (ngin *Engine) ListenWebsocket(conf config.WSServerConf) error { +func (ngin *Engine) ListenWebsocket(conf config.WSServerConf, opts ...WsConfOption) error { + for _, opt := range opts { + opt(conf) + } + err := ngin.setup() if err != nil { nlog.Errorf("%s failed to setup server, err:%v", ngin.LogPrefix(), err)