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

如果你在運行 HTTP 服務(wù)并且想對 endpoints 進(jìn)行限速，你可以使用維護(hù)良好的工具，例如 github.com/didip/tollbooth[1]。但是如果你在構(gòu)建一些非常簡單的東西，自己實現(xiàn)并不困難。

我們可以使用已經(jīng)存在的試驗性的 Go 包 x/time/rate。

在本教程中，我們將創(chuàng)建一個基于用戶 IP 地址進(jìn)行速率限制的簡單的中間件。

「干凈的」HTTP 服務(wù)

讓我們從構(gòu)建一個簡單的 HTTP 服務(wù)開始，該服務(wù)具有非常簡單的 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 中，我們啟動服務(wù)，該服務(wù)監(jiān)聽 :8888 并擁有一個 endpoint /。

golang.org/x/time/rate

我們將使用 Go 中 x/time/rate 包，該包提供了令牌桶限速算法。rate#Limiter[2] 控制事件發(fā)生的頻率。它實現(xiàn)了一個大小為 b 的「令牌桶」，初始化時是滿的，并且以每秒 r 個令牌的速率重新填充。非正式地，在任意足夠長的時間間隔中，限速器將速率限制在每秒 r 個令牌，最大突發(fā)事件為 b 個。

既然我們想要實現(xiàn)基于 IP 地址的限速器，我們還需要維護(hù)一個限速器的 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 創(chuàng)建了一個 IP 限速器的實例，HTTP 服務(wù)需要調(diào)用該實例的 GetLimiter 方法去獲取特定 IP 的限速器（從 map 中獲取，或者生成一個新的）。

中間件

讓我們升級我們的 HTTP 服務(wù)，將中間件添加到所有的 endpoints 中，因此，如果某 IP 達(dá)到了限制速率，服務(wù)將會返回 429 Too Many Requests，否則服務(wù)將處理該請求。

在 limitMiddleware 函數(shù)中，每次中間件收到 HTTP 請求，我們都會調(diào)用全局限速器的 Allow() 方法。如果令牌桶中沒有剩余的令牌，Allow() 將返回 false，我們返回給用戶 429 Too Many Requests 響應(yīng)。否則，調(diào)用 Allow() 將消耗桶中的一個令牌，我們將控制權(quán)傳遞給調(diào)用鏈的下一個處理器。

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

構(gòu)建 & 運行

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

測試

有一個非常棒的工具稱作 vegeta，我喜歡在 HTTP 負(fù)載測試中使用（它也是用 Go 編寫的）

brew install vegeta

我們需要創(chuàng)建一個簡單的配置文件，聲明我們想要發(fā)送的請求。

GET http://localhost:8888/

然后，以每個時間單元 100 個請求的速率攻擊 10 秒。

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

結(jié)果，你將看到一些請求返回 200，但大多數(shù)返回 429。

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

作者：ALEX PLIUTAU[3] 譯者：DoubleLuck[4] 校對：unknwon[5]

本文由 GCTT[6] 原創(chuàng)編譯，Go 中文網(wǎng)[7] 榮譽(yù)推出

推薦閱讀

• 漏桶、令牌桶限流算法的Go語言實現(xiàn)