添加 waf

添加 访问频率限制测试

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"`
+}

config.yaml

@@ -42,3 +42,6 @@ captcha:
   showLineOptions: 4
   source: 1234567890qwertyuioplkjhgfdsazxcvbnm
   expire: 300
+
+waf:
+  rateLimit: 100

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("正常访问"))
+	})
+}

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",

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",

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:

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

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()
 }

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
+		}
+	}
+}