如果你執行 HTTP 服務，並且希望限制 HTTP 的訪問頻率，那麼你可以藉助一些比較穩定的工具，例如： github.com/didip/tollbooth。不過如果你構建的應用比較簡單，也可以自己來實現。

我們可以使用一個現有的 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) {
 // 某些消耗很高的資料庫請求
 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 建立了一個新的速率限制器，並將其新增到 ips 對映中,// 使用 IP地址作為金鑰
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 返回所提供的IP地址的速率限制器(如果存在的話).
// 否則呼叫 AddIP 將 IP 地址新增到對映中
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 函式中的全域性方法 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("/",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) {
 // 非常重要的資料請求（譯者注：這句話沒理解到位）
 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。

以上就是本文的全部內容，希望對大家的學習有所幫助，也希望大家多多支援我們。