Go 中基於 IP 地址的 HTTP 請求限流
在本教程中，我們將創建一個基於用戶 IP 地址進行速率限制的簡單的中間件。
「乾淨的」HTTP 服務
讓我們從構建一個簡單的 HTTP 服務開始，該服務具有非常簡單的 endpiont。這可能是個非常「重」的 endpoint，因此我們想在這裏添加速率限制。
1package main
2
3import (
4 "log"
5 "net/http"
6)
7
8func main() {
9 mux := http.NewServeMux()
10 mux.HandleFunc("/", okHandler)
11
12 if err := http.ListenAndServe(":8888", mux); err != nil {
13
14 // log.Fatalf("unable to start server: %s", err.Error())
15 }
16}
17
18func okHandler(w http.ResponseWriter, r *http.Request) {
19 // Some very expensive database call
20 w.Write([]byte("alles gut"))
21}
22
23
在 main.go 中，我們啓動服務，該服務監聽 :8888 並擁有一個 endpoint /。
golang.org/x/time/rate
我們將使用 Go 中 x/time/rate 包，該包提供了令牌桶限速算法。rate#Limiter[2] 控制事件發生的頻率。它實現了一個大小爲 b 的「令牌桶」，初始化時是滿的，並且以每秒 r 個令牌的速率重新填充。非正式地，在任意足夠長的時間間隔中，限速器將速率限制在每秒 r 個令牌，最大突發事件爲 b 個。
既然我們想要實現基於 IP 地址的限速器，我們還需要維護一個限速器的 map。
1package main
2
3import (
4 "sync"
5
6 "golang.org/x/time/rate"
7)
8
9// IPRateLimiter .
10type IPRateLimiter struct {
11 ips map[string]*rate.Limiter
12 mu *sync.RWMutex
13 r rate.Limit
14 b int
15}
16
17// NewIPRateLimiter .
18func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
19 i := &IPRateLimiter{
20 ips: make(map[string]*rate.Limiter),
21 mu: &sync.RWMutex{},
22 r: r,
23 b: b,
24 }
25
26 return i
27}
28
29// AddIP creates a new rate limiter and adds it to the ips map,
30// using the IP address as the key
31func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {
32 i.mu.Lock()
33 defer i.mu.Unlock()
34
35 limiter := rate.NewLimiter(i.r, i.b)
36
37 i.ips[ip] = limiter
38
39 return limiter
40}
41
42// GetLimiter returns the rate limiter for the provided IP address if it exists.
43// Otherwise calls AddIP to add IP address to the map
44func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
45 i.mu.Lock()
46 limiter, exists := i.ips[ip]
47
48 if !exists {
49 i.mu.Unlock()
50 return i.AddIP(ip)
51 }
52
53 i.mu.Unlock()
54
55 return limiter
56}
57
58
NewIPRateLimiter 創建了一個 IP 限速器的實例，HTTP 服務需要調用該實例的 GetLimiter 方法去獲取特定 IP 的限速器（從 map 中獲取，或者生成一個新的）。
中間件
讓我們升級我們的 HTTP 服務，將中間件添加到所有的 endpoints 中，因此，如果某 IP 達到了限制速率，服務將會返回 429 Too Many Requests，否則服務將處理該請求。
在 limitMiddleware 函數中，每次中間件收到 HTTP 請求，我們都會調用全侷限速器的 Allow() 方法。如果令牌桶中沒有剩餘的令牌，Allow() 將返回 false，我們返回給用戶 429 Too Many Requests 響應。否則，調用 Allow() 將消耗桶中的一個令牌，我們將控制權傳遞給調用鏈的下一個處理器。
1package main
2
3import (
4 "log"
5 "net/http"
6)
7
8var limiter = NewIPRateLimiter(1, 5)
9
10func main() {
11 mux := http.NewServeMux()
12 mux.HandleFunc("/", okHandler)
13
14 if err := http.ListenAndServe(":8888", limitMiddleware(mux)); err != nil {
15 log.Fatalf("unable to start server: %s", err.Error())
16 }
17}
18
19func limitMiddleware(next http.Handler) http.Handler {
20 return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
21 limiter := limiter.GetLimiter(r.RemoteAddr)
22 if !limiter.Allow() {
23 http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
24 return
25 }
26
27 next.ServeHTTP(w, r)
28 })
29}
30
31func okHandler(w http.ResponseWriter, r *http.Request) {
32 // Some very expensive database call
33 w.Write([]byte("alles gut"))
34}
35
36
構建 & 運行
1go get golang.org/x/time/rate
2go build -o server .
3./server
4
5
測試
有一個非常棒的工具稱作 vegeta，我喜歡在 HTTP 負載測試中使用（它也是用 Go 編寫的）
1brew install vegeta
2
3
1GET http://localhost:8888/
2
3
然後，以每個時間單元 100 個請求的速率攻擊 10 秒。
1vegeta attack -duration=10s -rate=100 -targets=vegeta.conf | vegeta report
2
3
結果，你將看到一些請求返回 200，但大多數返回 429。
via: https://pliutau.com/rate-limit-http-requests/
本文由 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/
本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/aj_hULmeqy3whILr2ytECQ