diff --git a/router/option.go b/router/option.go index b79ce4a..ecee4bf 100644 --- a/router/option.go +++ b/router/option.go @@ -25,6 +25,17 @@ type serverOption struct { loggerCfg *middleware.AccessConfig // 日志配置 initContextData gin.HandlerFunc // 初始化一些请求数据 rateLimitInstance abstract.IRateLimit // 服务流控实例 + demoCheck gin.HandlerFunc // demo检查 +} + +// WithDemoCheck 设置demo检查 +func WithDemoCheck(f gin.HandlerFunc) SetServerOptionFunc { + return func(so *serverOption) { + if nil == f { + return + } + so.demoCheck = f + } } // WithRateLimitInstance 设置流控实例, 配置为 nil, 代表禁用 diff --git a/router/server.go b/router/server.go index 16e4f33..d0fcd50 100644 --- a/router/server.go +++ b/router/server.go @@ -143,6 +143,10 @@ func (s *Server) getGlobalMiddlewareList(option *serverOption) { AllowFiles: true, })) } + // demo 检查 + if nil != option.demoCheck { + s.globalMiddlewareList = append(s.globalMiddlewareList, option.demoCheck) + } if len(option.globalMiddlewareList) > 0 { // 自定义全局中间件追加 s.globalMiddlewareList = append(s.globalMiddlewareList, option.globalMiddlewareList...)