Golang实现HTTP请求速率限制的插件库Tollbooth的使用

Tollbooth是一个通用的中间件，用于对HTTP请求进行速率限制。

版本说明

• v1.0.0: 保持旧API但所有第三方模块移至自己的仓库
• v2.x.x: 全新的API，代码清理、线程安全和自动过期数据结构
• v3.x.x: 修复了golang.org/x/time/rate的正确使用方式
• v4.x.x: 支持浮点数的每秒最大请求数
• v5.x.x: 添加了go.mod和go.sum
• v6.x.x: 替换go-cache为github.com/go-pkgz/expirable-cache
• v7.x.x: 替换time/rate为嵌入式time/rate以支持更多速率限制头
• v8.x.x: 修复RemoteIP漏洞，新增HTTPMiddleware函数

五分钟教程

package main

import (
    "net/http"
    "github.com/didip/tollbooth/v8"
    "github.com/didip/tollbooth/v8/limiter"
)

func HelloHandler(w http.ResponseWriter, req *http.Request) {
    w.Write([]byte("Hello, World!"))
}

func main() {
    // 创建每个处理程序的请求限制器
    lmt := tollbooth.NewLimiter(1, nil)

    // 版本>=8的新功能，必须明确定义如何选择IP地址
    lmt.SetIPLookup(limiter.IPLookup{
        Name:           "X-Real-IP",
        IndexFromRight: 0,
    })

    // 版本>=8的新功能，HTTPMiddleware是与标准路由器兼容的替代方案
    http.Handle("/", tollbooth.HTTPMiddleware(lmt)(http.HandlerFunc(HelloHandler)))
    // 旧语法：
    // http.Handle("/", tollbooth.LimitFuncHandler(lmt, HelloHandler))

    http.ListenAndServe(":12345", nil)
}

主要特性

1. 多种限制条件：可以基于远程IP、路径、方法、自定义头和基本认证用户名进行限制

import (
    "time"
    "github.com/didip/tollbooth/v8"
    "github.com/didip/tollbooth/v8/limiter"
)

lmt := tollbooth.NewLimiter(1, nil)

// 或创建带有可过期令牌桶的限制器
lmt = tollbooth.NewLimiter(1, &limiter.ExpirableOptions{DefaultExpirationTTL: time.Hour})

// 版本>=8的新功能，必须明确定义如何选择IP地址
lmt.SetIPLookup(limiter.IPLookup{
    Name:           "X-Real-IP",
    IndexFromRight: 0,
})

// 仅限制GET和POST请求
lmt.SetMethods([]string{"GET", "POST"})

// 基于基本认证用户名限制
lmt.SetBasicAuthUsers([]string{"bob", "jane", "didip", "vip"})
lmt.RemoveBasicAuthUsers([]string{"vip"})

// 限制包含特定值的请求头
lmt.SetHeader("X-Access-Token", []string{"abc123", "xyz098"})
lmt.RemoveHeader("X-Access-Token")
lmt.RemoveHeaderEntries("X-Access-Token", []string{"limitless-token"})

1. 自定义过期时间：可以设置各种元素的过期TTL

import "time"

lmt := tollbooth.NewLimiter(1, nil)

// 设置令牌桶的自定义过期TTL
lmt.SetTokenBucketExpirationTTL(time.Hour)

// 设置基本认证用户的过期TTL
lmt.SetBasicAuthExpirationTTL(time.Hour)

// 设置头条目的过期TTL
lmt.SetHeaderEntryExpirationTTL(time.Hour)

1. 自定义拒绝响应：可以设置自定义消息或拒绝回调函数

lmt := tollbooth.NewLimiter(1, nil)

lmt.SetIPLookup(limiter.IPLookup{
    Name:           "X-Forwarded-For",
    IndexFromRight: 0,
})

// 设置自定义消息
lmt.SetMessage("You have reached maximum request limit.")

// 设置自定义内容类型
lmt.SetMessageContentType("text/plain; charset=utf-8")

// 设置拒绝时的自定义函数
lmt.SetOnLimitReached(func(w http.ResponseWriter, r *http.Request) { 
    fmt.Println("A request was rejected") 
})

1. 响应头信息：在被拒绝时会返回以下HTTP响应头：

   • X-Rate-Limit-Limit
   • X-Rate-Limit-Duration
   • X-Rate-Limit-Request-Forwarded-For
   • X-Rate-Limit-Request-Remote-Addr

2. 使用Token Bucket算法：不需要外部存储

其他Web框架集成

Tollbooth可以与其他Web框架集成，社区提供了以下适配器：

• Chi
• Echo
• FastHTTP
• Gin
• GoRestful
• HTTPRouter
• Iris
• Negroni

完整示例

package main

import (
    "fmt"
    "net/http"
    "time"
    "github.com/didip/tollbooth/v8"
    "github.com/didip/tollbooth/v8/limiter"
)

func main() {
    // 创建限制器：1请求/秒，令牌桶1小时后过期
    lmt := tollbooth.NewLimiter(1, &limiter.ExpirableOptions{
        DefaultExpirationTTL: time.Hour,
    })

    // 设置IP查找方式
    lmt.SetIPLookup(limiter.IPLookup{
        Name:           "X-Forwarded-For",
        IndexFromRight: 0,
    })

    // 仅限制GET请求
    lmt.SetMethods([]string{"GET"})

    // 设置自定义拒绝消息
    lmt.SetMessage("您已达到请求速率限制，请稍后再试。")
    lmt.SetMessageContentType("text/plain; charset=utf-8")

    // 自定义拒绝处理函数
    lmt.SetOnLimitReached(func(w http.ResponseWriter, r *http.Request) {
        fmt.Printf("请求被限制: %s %s\n", r.Method, r.URL.Path)
    })

    // 创建处理函数
    helloHandler := func(w http.ResponseWriter, req *http.Request) {
        w.Write([]byte("Hello, World!"))
    }

    // 应用速率限制中间件
    http.Handle("/", tollbooth.HTTPMiddleware(lmt)(http.HandlerFunc(helloHandler)))

    fmt.Println("服务器启动在 :8080")
    http.ListenAndServe(":8080", nil)
}

这个示例展示了Tollbooth的主要功能，包括速率限制设置、IP查找配置、方法限制、自定义消息和拒绝回调等。