Go 中基于 IP 地址的 HTTP 请求限流

Go语言精选

共 6673字,需浏览 14分钟

 ·

2021-03-27 06:38

点击上方蓝色“Go语言中文网”关注,每天一起学 Go

如果你在运行 HTTP 服务并且想对 endpoints 进行限速,你可以使用维护良好的工具,例如 github.com/didip/tollbooth[1]。但是如果你在构建一些非常简单的东西,自己实现并不困难。

我们可以使用已经存在的试验性的 Go 包 x/time/rate

在本教程中,我们将创建一个基于用户 IP 地址进行速率限制的简单的中间件。

「干净的」HTTP 服务

让我们从构建一个简单的 HTTP 服务开始,该服务具有非常简单的 endpiont。这可能是个非常「重」的 endpoint,因此我们想在这里添加速率限制。

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 并拥有一个 endpoint /

golang.org/x/time/rate

我们将使用 Go 中 x/time/rate 包,该包提供了令牌桶限速算法。rate#Limiter[2] 控制事件发生的频率。它实现了一个大小为 b 的「令牌桶」,初始化时是满的,并且以每秒 r 个令牌的速率重新填充。非正式地,在任意足够长的时间间隔中,限速器将速率限制在每秒 r 个令牌,最大突发事件为 b 个。

既然我们想要实现基于 IP 地址的限速器,我们还需要维护一个限速器的 map。

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 的限速器(从 map 中获取,或者生成一个新的)。

中间件

让我们升级我们的 HTTP 服务,将中间件添加到所有的 endpoints 中,因此,如果某 IP 达到了限制速率,服务将会返回 429 Too Many Requests,否则服务将处理该请求。

在 limitMiddleware 函数中,每次中间件收到 HTTP 请求,我们都会调用全局限速器的 Allow() 方法。如果令牌桶中没有剩余的令牌,Allow() 将返回 false,我们返回给用户 429 Too Many Requests 响应。否则,调用 Allow() 将消耗桶中的一个令牌,我们将控制权传递给调用链的下一个处理器。

package main

import (
    "log"
    "net/http"
)

var limiter = NewIPRateLimiter(15)

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

测试

有一个非常棒的工具称作 vegeta,我喜欢在 HTTP 负载测试中使用(它也是用 Go 编写的)

brew install vegeta

我们需要创建一个简单的配置文件,声明我们想要发送的请求。

GET http://localhost:8888/

然后,以每个时间单元 100 个请求的速率攻击 10 秒。

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

结果,你将看到一些请求返回 200,但大多数返回 429。


via: https://pliutau.com/rate-limit-http-requests/

作者:ALEX PLIUTAU[3] 译者:DoubleLuck[4] 校对:unknwon[5]

本文由 GCTT[6] 原创编译,Go 中文网[7] 荣誉推出

参考资料

[1]

github.com/didip/tollbooth: https://github.com/didip/tollbooth

[2]

rate#Limiter: https://godoc.org/golang.org/x/time/rate#Limiter

[3]

ALEX PLIUTAU: https://pliutau.com/

[4]

DoubleLuck: https://github.com/DoubleLuck

[5]

unknwon: https://github.com/unknwon

[6]

GCTT: https://github.com/studygolang/GCTT

[7]

Go 中文网: https://studygolang.com/



推荐阅读


福利

我为大家整理了一份从入门到进阶的Go学习资料礼包,包含学习建议:入门看什么,进阶看什么。关注公众号 「polarisxu」,回复 ebook 获取;还可以回复「进群」,和数万 Gopher 交流学习。

浏览 48
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报