如果你运行的是 HTTP 服务,并且希望限制 HTTP 的访问频率,那么可以使用一些稳定的工具,例如。 github.com/didip/tollb… 。 但是,如果你正在构建一个非常简单的东西,那么你自己实现它并不难。

我们可以使用一个实验性的Go包 x/time/rate

在本教程中,我们将创建一个简单的中间件来实现基于 IP 限制 HTTP 的访问频率。

纯 HTTP 服务

让我们从构建一个简单的 HTTP 服务开始,它有非常简单的终端。但是它的访问频率可能非常高,这就是为什么我们要增加一个频率限制。

package main

import (
    "log"
    "net/http"
)

func main() {
    mux := http.NewServeMux()
    mux.HandleFunc("/", okHandler)

    if err := http.ListenAndServe(":8888", mux); err != nil {
        log.Fatalf("unable to start server: %s", err.Error())
    }
}

func okHandler(w http.ResponseWriter, r *http.Request) {
    // Some very expensive database call
    w.Write([]byte("alles gut"))
}

我们通过 main.go 启动服务,监听 :8888 端口,然后我们就有了一个终端 / 。

golang.org/x/time/rate

我们将使用 x/time/rate 的 Go 包,它提供了一个令牌桶速率限制器算法。rate#Limiter 控制事件发生的频率。它实现了一个大小为 b 的“令牌桶”,最初是满的,然后以每秒 r 个令牌的速率重新填充。通常来讲,在任何足够大的时间间隔内,限制器将速率限制为每秒 r 个令牌,最大突发大小为 b 个事件。

因为我们要实现每个 IP 地址的频率限制器,所以我们还需要维护一个限制器的映射。

package main

import (
    "sync"

    "golang.org/x/time/rate"
)

// IPRateLimiter .
type IPRateLimiter struct {
    ips map[string]*rate.Limiter
    mu  *sync.RWMutex
    r   rate.Limit
    b   int
}

// NewIPRateLimiter .
func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
    i := &IPRateLimiter{
        ips: make(map[string]*rate.Limiter),
        mu:  &sync.RWMutex{},
        r:   r,
        b:   b,
    }

    return i
}

// AddIP creates a new rate limiter and adds it to the ips map,
// using the IP address as the key
func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {
    i.mu.Lock()
    defer i.mu.Unlock()

    limiter := rate.NewLimiter(i.r, i.b)

    i.ips[ip] = limiter

    return limiter
}

// GetLimiter returns the rate limiter for the provided IP address if it exists.
// Otherwise calls AddIP to add IP address to the map
func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
    i.mu.Lock()
    limiter, exists := i.ips[ip]

    if !exists {
        i.mu.Unlock()
        return i.AddIP(ip)
    }

    i.mu.Unlock()

    return limiter
}

NewIPRateLimiter 创建一个 IP 限制器实例,HTTP 服务必须调用这个实例的 GetLimiter 来获取指定 IP 的限制器(从映射中获取或生成一个新的)。

中间件

让我们升级我们的 HTTP 服务并向所有终端添加中间件,这样如果 IP 已经达到限制,它将响应 429 Too Many Requests ,否则,它将继续处理请求。

在 limitMiddleware 函数中,每当中间件接收到 HTTP 请求时,我们都调用全局限制器的 Allow() 方法。如果存储桶中没有剩余的令牌,Allow() 将返回 false,并向用户发送 429 Too Many Requests 响应。否则,调用Allow() 将只使用存储桶中的一个令牌,然后将请求传递给链中的下一个处理程序。

package main

import (
    "log"
    "net/http"
)

var limiter = NewIPRateLimiter(1, 5)

func main() {
    mux := http.NewServeMux()
    mux.HandleFunc("/", okHandler)

    if err := http.ListenAndServe(":8888", limitMiddleware(mux)); err != nil {
        log.Fatalf("unable to start server: %s", err.Error())
    }
}

func limitMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        limiter := limiter.GetLimiter(r.RemoteAddr)
        if !limiter.Allow() {
            http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
            return
        }

        next.ServeHTTP(w, r)
    })
}

func okHandler(w http.ResponseWriter, r *http.Request) {
    // Some very expensive database call
    w.Write([]byte("alles gut"))
}

编译 & 执行

go get golang.org/x/time/rate
go build -o server .
./server

测试

我喜欢使用一个非常好的工具进行 HTTP 负载测试,叫做 vegeta(它也是用Go编写的)。

brew install vegeta

我们需要创建一个简单的配置文件,说明我们想要生成什么请求。

GET http://localhost:8888/

然后运行攻击 10 秒,每个时间单位 100 个请求。

vegeta attack -duration=10s -rate=100 -targets=vegeta.conf | vegeta report

因此,你将看到一些请求返回了 200 ,但大多数请求返回了 429 。

原文链接:https://dev.to/plutov/rate-limiting-http-requests-in-go-based-on-ip-address-542g