IPアドレスに基づくGoにおけるHTTP要求の制限

HTTPサーバーを実行している場合は、エンドポイントにRate Level要求を適用する場合は、github.com/didip/tollbooth . しかし、あなたが非常に単純な何かを構築しているなら、それはあなた自身でそれを実装するのは難しいことではない.
既に実験的なパッケージがあるx/time/rate , 我々が使うことができるもの.
このチュートリアルでは、ユーザーの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 シングルエンドポイント/ .

ゴラン.タイム/レート
私たちはx/time/rate トークンバケットレートリミッタアルゴリズムを提供するGOパッケージ.rate#Limiter イベントが頻繁に行われるように制御します.サイズの「トークンバケツ」を実装しますb , 最初はフルでリフィルr 秒あたりのトークン.形式的に、どんな大きな十分な時間間隔ででも、リミッターは毎秒Rトークンへの率を制限します.
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サーバをアップグレードして、すべてのエンドポイントにミドルウェアを加えましょう、したがって、IPが制限に達したなら、それは429あまりに多くの要求に応えます、さもなければ、それは要求を進めます.
にlimitMiddleware 関数は、グローバルリミッターのAllow() ミドルウェアがHTTPリクエストを受け取るたびに.バケツに残っているトークンがなければAllow() を返します.さもなければAllow() バケツから正確に1つのトークンを消費し、チェーン内の次のハンドラに制御を渡します.

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

テスト
HTTPロードテスト用の非常にいいツールがありますvegeta (囲碁).

brew install vegeta

どのような要求を行うのかを簡単に設定できるファイルを作成する必要があります.

GET http://localhost:8888/

そして、時間単位で100リクエストで10秒間攻撃を実行します.

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

その結果、いくつかのリクエストが200を返したことがわかります.