From 40860ee0f9c2887de1fefec750fdb4e4377f73e7 Mon Sep 17 00:00:00 2001
From: NoahLan <6995syu@163.com>
Date: Wed, 21 Jun 2023 15:20:03 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8DwsConf=E7=9A=84func?=
 =?UTF-8?q?=E5=8F=82=E6=95=B0=E9=97=AE=E9=A2=98=EF=BC=8C=E6=B7=BB=E5=8A=A0?=
 =?UTF-8?q?checkorigin=E7=9A=84=E6=B3=A8=E5=85=A5=E6=96=B9=E5=BC=8F?=
 =?UTF-8?q?=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 config/server_ws.go |  2 +-
 server_ws.go        | 14 +++++++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

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)