Golang 基於IP地址的HTTP限速請求
使用 golang.org/x/time
在本教程中，我們將基於用戶的IP地址創建一個簡單的中間件來進行速率限制。
純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，它提供令牌桶速率限制器算法。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 Server並將中間件添加到所有端點，因此，如果IP達到限制，它將響應429 Too Many Requests，否則，它將繼續該請求。
在該limitMiddleware函數中，Allow()每次收到HTTP請求時，我們都調用全侷限制器的方法。如果存儲桶中沒有剩餘令牌，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
測試
ab -n 1000 -c 100 http://localhost:8888/
其中－n表示請求數，－c表示併發數