golang-cookbook

Appearance

在 Go 的 Gin 框架中，限制请求（如限制请求速率、请求大小或请求频率）是非常常见的需求。下面是几种在 Gin 中实现限制功能的方法。

1. 限制请求大小（Limit Request Body Size） ​

可以使用 Gin 中的 MaxMultipartMemory 来限制上传文件的大小，或者通过中间件限制请求的 body 大小。

go
package main

import (
    "fmt"
    "github.com/gin-gonic/gin"
)

func main() {
    r := gin.Default()

    // 设置请求体最大大小为 10MB
    r.MaxMultipartMemory = 10 << 20 // 10MB

    r.POST("/upload", func(c *gin.Context) {
        // 处理上传逻辑
        c.String(200, "File uploaded successfully")
    })

    r.Run(":8080")
}

2. 限制请求速率（Rate Limiting） ​

要实现基于时间的请求速率限制，可以使用如 github.com/gin-contrib/limiter 或 golang.org/x/time/rate 这样的库来创建自定义中间件。

使用 github.com/gin-contrib/limiter： ​

limiter 是一个非常常见的库，能够帮助你实现简单的请求速率限制。以下是如何使用它：

go
package main

import (
    "github.com/gin-gonic/gin"
    "github.com/gin-contrib/limiter"
    "github.com/gin-contrib/limiter/devices/redis"
    "time"
)

func main() {
    r := gin.Default()

    // 限制每秒最多 5 次请求
    r.Use(limiter.New(&limiter.Config{
        Limiter: redis.NewLimiter("localhost:6379", "rate-limit", time.Second*1, 5),
    }))

    r.GET("/", func(c *gin.Context) {
        c.JSON(200, gin.H{
            "message": "Hello, World!",
        })
    })

    r.Run(":8080")
}

使用 golang.org/x/time/rate： ​

你也可以通过 rate 包来手动实现速率限制。这里是一个简单的例子，展示了如何使用 rate.NewLimiter 来限制请求速率。

go
package main

import (
    "fmt"
    "github.com/gin-gonic/gin"
    "golang.org/x/time/rate"
    "time"
)

func main() {
    r := gin.Default()

    // 限制每秒 1 次请求
    limiter := rate.NewLimiter(rate.Every(time.Second), 1)

    r.GET("/", func(c *gin.Context) {
        if limiter.Allow() {
            c.JSON(200, gin.H{
                "message": "Request accepted",
            })
        } else {
            c.JSON(429, gin.H{
                "message": "Too many requests",
            })
        }
    })

    r.Run(":8080")
}

3. 限制 IP 请求频率（IP Rate Limiting） ​

如果你希望限制每个 IP 地址的请求频率，可以基于 gin-contrib/limiter 或自定义存储来实现 IP 限流。

使用 gin-contrib/limiter 限制每个 IP 的请求： ​

go
package main

import (
    "fmt"
    "github.com/gin-gonic/gin"
    "github.com/gin-contrib/limiter"
    "github.com/gin-contrib/limiter/devices/memory"
    "time"
)

func main() {
    r := gin.Default()

    // 限制每个 IP 每秒最多 1 次请求
    r.Use(limiter.New(&limiter.Config{
        Limiter: memory.NewLimiter(time.Second*1, 1),
        KeyFunc: limiter.IPKeyFunc, // 使用 IP 作为限制的 Key
    }))

    r.GET("/", func(c *gin.Context) {
        c.JSON(200, gin.H{
            "message": "Hello, World!",
        })
    })

    r.Run(":8080")
}

4. 限制请求头（Headers Limit） ​

你还可以限制请求头的大小，特别是在处理 POST 请求时，限制特定头部内容的长度或数量。Gin 本身没有直接提供这样的限制，但你可以通过中间件手动检查请求头。

go
package main

import (
    "github.com/gin-gonic/gin"
    "fmt"
)

func main() {
    r := gin.Default()

    // 自定义中间件限制请求头大小
    r.Use(func(c *gin.Context) {
        maxHeaderSize := 1024 * 10 // 设置最大请求头大小为 10KB
        c.Request.Header.Set("Max-Request-Size", fmt.Sprintf("%d", maxHeaderSize))
        
        // 如果请求头大小超出限制，拒绝请求
        if c.Request.ContentLength > maxHeaderSize {
            c.JSON(400, gin.H{
                "message": "Request headers too large",
            })
            c.Abort()
            return
        }

        c.Next()
    })

    r.GET("/", func(c *gin.Context) {
        c.JSON(200, gin.H{
            "message": "Hello, World!",
        })
    })

    r.Run(":8080")
}

总结 ​

  1. 请求体大小限制：通过 MaxMultipartMemory 来限制上传文件的大小，或在中间件中检查请求体大小。
  2. 请求速率限制：使用 gin-contrib/limiter 或 golang.org/x/time/rate 实现速率限制。
  3. IP 请求频率限制：可以通过 gin-contrib/limiter 库或自定义存储方式，限制每个 IP 地址的请求频率。
  4. 请求头限制：通过中间件自定义检查请求头的大小，限制头部内容。

根据不同的需求，你可以选择合适的方式来限制请求，确保应用程序的安全性和性能。