From d152df3a1d302be6e0d43a79330b1f448f247cff Mon Sep 17 00:00:00 2001 From: Shikong <919411476@qq.com> Date: Tue, 25 Oct 2022 15:01:57 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=20waf=20=E6=B7=BB=E5=8A=A0?= =?UTF-8?q?=20=E8=AE=BF=E9=97=AE=E9=A2=91=E7=8E=87=E9=99=90=E5=88=B6?= =?UTF-8?q?=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- common/config/config.go | 6 +++ config.yaml | 5 ++- controller/test/test.go | 22 ++++++++++ docs/docs.go | 46 +++++++++++++++++++++ docs/swagger.json | 46 +++++++++++++++++++++ docs/swagger.yaml | 28 +++++++++++++ model/generic/generate/config.yml | 1 + services/services.go | 2 + services/waf/waf.go | 67 +++++++++++++++++++++++++++++++ 9 files changed, 222 insertions(+), 1 deletion(-) create mode 100644 services/waf/waf.go diff --git a/common/config/config.go b/common/config/config.go index e0b8d51..d8b1745 100644 --- a/common/config/config.go +++ b/common/config/config.go @@ -11,6 +11,7 @@ type BasicConfig struct { Jwt JwtConfig `yaml:"jwt"` NanoId NanoIdConfig `yaml:"nanoId"` Captcha CaptchaConfig `yaml:"captcha"` + Waf Waf `yaml:"waf"` } type NanoIdConfig struct { @@ -69,3 +70,8 @@ type CaptchaConfig struct { //Expire 失效时间 (秒) Expire int64 `json:"expired"` } + +type Waf struct { + //RateLimit 访问频率 + RateLimit int64 `yaml:"rateLimit"` +} diff --git a/config.yaml b/config.yaml index 30ac001..934307f 100644 --- a/config.yaml +++ b/config.yaml @@ -41,4 +41,7 @@ captcha: noiseCount: 100 showLineOptions: 4 source: 1234567890qwertyuioplkjhgfdsazxcvbnm - expire: 300 \ No newline at end of file + expire: 300 + +waf: + rateLimit: 100 \ No newline at end of file diff --git a/controller/test/test.go b/controller/test/test.go index 675073c..62edb0a 100644 --- a/controller/test/test.go +++ b/controller/test/test.go @@ -9,6 +9,7 @@ import ( "gofiber.study.skcks.cn/global" "gofiber.study.skcks.cn/model/dto" "gofiber.study.skcks.cn/model/vo" + "gofiber.study.skcks.cn/services/waf" "time" ) @@ -180,3 +181,24 @@ func (c *Controller) VerifyCaptchaTest() { return ctx.JSON(response.NewResponse("验证成功")) }) } + +// WafRateLimitTest waf 访问频率限制 测试 +// +// @Summary waf 访问频率限制 测试 +// @Description waf 访问频率限制 测试 +// @Tags Test +// @Accept json +// @Produce json +// @Success 200 {object} response.Response{data=string} +// @Failure default {object} errorx.CodeErrorResponse +// @Router /test/waf/access [get] +func (c *Controller) WafRateLimitTest() { + c.Router.Add(fiber.MethodGet, "/waf/access", func(ctx *fiber.Ctx) error { + access := waf.Service.Access(ctx.IP()) + if !access { + return ctx.JSON(errorx.NewErrorWithCode(fiber.StatusTooManyRequests, "访问频率超过限制")) + } + + return ctx.JSON(response.NewResponse("正常访问")) + }) +} diff --git a/docs/docs.go b/docs/docs.go index 19816e9..ac72ec7 100644 --- a/docs/docs.go +++ b/docs/docs.go @@ -717,6 +717,47 @@ const docTemplate = `{ } } }, + "/test/waf/access": { + "get": { + "description": "waf 访问频率限制 测试", + "consumes": [ + "application/json" + ], + "produces": [ + "application/json" + ], + "tags": [ + "Test" + ], + "summary": "waf 访问频率限制 测试", + "responses": { + "200": { + "description": "OK", + "schema": { + "allOf": [ + { + "$ref": "#/definitions/response.Response" + }, + { + "type": "object", + "properties": { + "data": { + "type": "string" + } + } + } + ] + } + }, + "default": { + "description": "", + "schema": { + "$ref": "#/definitions/errorx.CodeErrorResponse" + } + } + } + } + }, "/user/account": { "get": { "security": [ @@ -928,6 +969,11 @@ const docTemplate = `{ "type": "string", "example": "data:image/png;base64,iVBO..." }, + "expire": { + "description": "过期时间 unix", + "type": "integer", + "example": 10000000 + }, "id": { "description": "验证码 id", "type": "string", diff --git a/docs/swagger.json b/docs/swagger.json index 1df356a..a7d00a9 100644 --- a/docs/swagger.json +++ b/docs/swagger.json @@ -709,6 +709,47 @@ } } }, + "/test/waf/access": { + "get": { + "description": "waf 访问频率限制 测试", + "consumes": [ + "application/json" + ], + "produces": [ + "application/json" + ], + "tags": [ + "Test" + ], + "summary": "waf 访问频率限制 测试", + "responses": { + "200": { + "description": "OK", + "schema": { + "allOf": [ + { + "$ref": "#/definitions/response.Response" + }, + { + "type": "object", + "properties": { + "data": { + "type": "string" + } + } + } + ] + } + }, + "default": { + "description": "", + "schema": { + "$ref": "#/definitions/errorx.CodeErrorResponse" + } + } + } + } + }, "/user/account": { "get": { "security": [ @@ -920,6 +961,11 @@ "type": "string", "example": "data:image/png;base64,iVBO..." }, + "expire": { + "description": "过期时间 unix", + "type": "integer", + "example": 10000000 + }, "id": { "description": "验证码 id", "type": "string", diff --git a/docs/swagger.yaml b/docs/swagger.yaml index d74cb2e..9fbed69 100644 --- a/docs/swagger.yaml +++ b/docs/swagger.yaml @@ -106,6 +106,10 @@ definitions: description: 验证码图片 base64 example: data:image/png;base64,iVBO... type: string + expire: + description: 过期时间 unix + example: 10000000 + type: integer id: description: 验证码 id example: abcdefg123456 @@ -556,6 +560,30 @@ paths: summary: sonyFlake id 生成测试 tags: - Test + /test/waf/access: + get: + consumes: + - application/json + description: waf 访问频率限制 测试 + produces: + - application/json + responses: + "200": + description: OK + schema: + allOf: + - $ref: '#/definitions/response.Response' + - properties: + data: + type: string + type: object + default: + description: "" + schema: + $ref: '#/definitions/errorx.CodeErrorResponse' + summary: waf 访问频率限制 测试 + tags: + - Test /user/account: get: consumes: diff --git a/model/generic/generate/config.yml b/model/generic/generate/config.yml index 3370c8c..49a84e9 100644 --- a/model/generic/generate/config.yml +++ b/model/generic/generate/config.yml @@ -11,6 +11,7 @@ targets: column_mapper: snake # how column name map to class or struct field name include_tables: # tables included, you can use ** - user + - waf table_prefix: "" multiple_files: true # generate multiple files or one template: | # template for code file, it has higher perior than template_path diff --git a/services/services.go b/services/services.go index b703a29..cad8563 100644 --- a/services/services.go +++ b/services/services.go @@ -3,9 +3,11 @@ package services import ( "gofiber.study.skcks.cn/services/auth" "gofiber.study.skcks.cn/services/user" + "gofiber.study.skcks.cn/services/waf" ) func Init() { auth.InitService() user.InitService() + waf.InitService() } diff --git a/services/waf/waf.go b/services/waf/waf.go new file mode 100644 index 0000000..ebfaef0 --- /dev/null +++ b/services/waf/waf.go @@ -0,0 +1,67 @@ +package waf + +import ( + "context" + "github.com/go-redis/redis/v8" + "gofiber.study.skcks.cn/common/config" + "gofiber.study.skcks.cn/common/logger" + "gofiber.study.skcks.cn/global" + "time" + "xorm.io/xorm" +) + +var Service *Waf + +const ( + StoreName = "waf" + Separator = ":" +) + +type Waf struct { + db *xorm.EngineGroup + store *redis.Client + config config.Waf +} + +func InitService() { + Service = &Waf{ + db: global.DataSources, + store: global.Redis, + config: global.Config.Waf, + } + // + //ctx := context.Background() + //table, err := Service.store.HGetAll(ctx, StoreName).Result() + //utils.MainAppExec(func() { + // if err != nil { + // logger.Log.Errorf("[waf] 初始化出错 %s", err) + // } + // + // logger.Log.Debugf("[waf] \n%v", table) + //}) +} + +func (w *Waf) Access(ip string) bool { + key := StoreName + Separator + "access" + Separator + ip + ctx := context.Background() + num, err := w.store.LLen(ctx, key).Result() + if err != nil { + logger.Log.Errorf("[waf] access 出错 %s", err) + return false + } + + if num < w.config.RateLimit { + w.store.LPush(ctx, key, time.Now().Unix()) + return true + } else { + last, _ := w.store.LIndex(ctx, key, -1).Int64() + if time.Now().Unix()-last < 60 { + logger.Log.Infof("[waf] ip:%s 访问频率超过限制 %d", ip, w.config.RateLimit) + return false + } else { + w.store.LPush(ctx, key, time.Now().Unix()) + w.store.LTrim(ctx, key, 0, w.config.RateLimit-1) + return true + } + } +}