【Go言語】レートリミッターを実装する

type Config struct {
    // Next defines a function to skip this middleware when returned true.
    //
    // Optional. Default: nil
    Next func(c *fiber.Ctx) bool
 
    // Max number of recent connections during `Expiration` seconds before sending a 429 response
    //
    // Default: 5
    Max int
 
    // KeyGenerator allows you to generate custom keys, by default c.IP() is used
    //
    // Default: func(c *fiber.Ctx) string {
    //   return c.IP()
    // }
    KeyGenerator func(*fiber.Ctx) string
 
    // Expiration is the time on how long to keep records of requests in memory
    //
    // Default: 1 * time.Minute
    Expiration time.Duration
 
    // LimitReached is called when a request hits the limit
    //
    // Default: func(c *fiber.Ctx) error {
    //   return c.SendStatus(fiber.StatusTooManyRequests)
    // }
    LimitReached fiber.Handler
 
    // When set to true, requests with StatusCode >= 400 won't be counted.
    //
    // Default: false
    SkipFailedRequests bool
 
    // When set to true, requests with StatusCode < 400 won't be counted.
    //
    // Default: false
    SkipSuccessfulRequests bool
 
    // Store is used to store the state of the middleware
    //
    // Default: an in memory store for this process only
    Storage fiber.Storage
 
    // LimiterMiddleware is the struct that implements a limiter middleware.
    //
    // Default: a new Fixed Window Rate Limiter
    LimiterMiddleware LimiterHandler
}

package limiter
 
import (
    "strconv"
    "sync"
    "sync/atomic"
 
    "github.com/gofiber/fiber/v2"
    "github.com/gofiber/fiber/v2/utils"
)
 
type FixedWindow struct{}
 
// 固定ウィンドウ方式のレートリミッターを実装するミドルウェアの生成
func (FixedWindow) New(cfg Config) fiber.Handler {
    var (
        // 同時アクセスを管理するためのミューテックス
        mux        = &sync.RWMutex{}
        // 最大リクエスト数を文字列に変換（レスポンスヘッダー用）
        max        = strconv.Itoa(cfg.Max)
        // 有効期限（秒単位）を計算
        expiration = uint64(cfg.Expiration.Seconds())
    )
 
    // ストレージ操作を簡略化するためのマネージャーを生成
    manager := newManager(cfg.Storage)
 
    // タイムスタンプを毎秒更新する（グローバルタイムスタンプを使用）
    utils.StartTimeStampUpdater()
 
    // ミドルウェアとして使用される関数を返す
    return func(c *fiber.Ctx) error {
        // cfg.Next が設定されており、その条件を満たす場合は次の処理を実行
        if cfg.Next != nil && cfg.Next(c) {
            return c.Next()
        }
 
        // リクエストから一意のキーを生成（通常はクライアントのIPアドレス）
        key := cfg.KeyGenerator(c)
 
        // ミューテックスをロックして並列処理を防ぐ
        mux.Lock()
 
        // 該当するキーのエントリを取得
        e := manager.get(key)
 
        // 現在のタイムスタンプを取得
        ts := uint64(atomic.LoadUint32(&utils.Timestamp))
 
        // エントリが存在しない場合、または有効期限が切れている場合は初期化
        if e.exp == 0 {
            // 新しい有効期限を設定
            e.exp = ts + expiration
        } else if ts >= e.exp {
            // 有効期限切れの場合、カウンターをリセットし、新しい有効期限を設定
            e.currHits = 0
            e.exp = ts + expiration
        }
 
        // 現在のリクエスト数をインクリメント
        e.currHits++
 
        // ウィンドウリセットまでの残り時間を計算
        resetInSec := e.exp - ts
 
        // 残りの許可されたリクエスト数を計算
        remaining := cfg.Max - e.currHits
 
        // エントリをストレージに更新
        manager.set(key, e, cfg.Expiration)
 
        // ミューテックスを解除
        mux.Unlock()
 
        // 残りのリクエスト数が0を下回る場合、制限に達したと判断
        if remaining < 0 {
            // Retry-After ヘッダーを設定し、制限解除までの時間を通知
            c.Set(fiber.HeaderRetryAfter, strconv.FormatUint(resetInSec, 10))
 
            // 制限に達した場合のハンドラーを呼び出す
            return cfg.LimitReached(c)
        }
 
        // 制限を超えていない場合は、次のミドルウェアまたはハンドラーを実行
        err := c.Next()
 
        // 成功/失敗に応じたリクエストのカウント調整
        if (cfg.SkipSuccessfulRequests && c.Response().StatusCode() < fiber.StatusBadRequest) ||
            (cfg.SkipFailedRequests && c.Response().StatusCode() >= fiber.StatusBadRequest) {
            // ミューテックスを再びロックしてカウンターを調整
            mux.Lock()
            e = manager.get(key)
            e.currHits-- // カウンターをデクリメント
            remaining++
            manager.set(key, e, cfg.Expiration)
            mux.Unlock()
        }
 
        // レートリミット関連のヘッダーをレスポンスに設定
        c.Set(xRateLimitLimit, max)
        c.Set(xRateLimitRemaining, strconv.Itoa(remaining))
        c.Set(xRateLimitReset, strconv.FormatUint(resetInSec, 10))
 
        return err
    }
}

package limiter
 
import (
    "strconv"
    "sync"
    "sync/atomic"
 
    "github.com/gofiber/fiber/v2"
    "github.com/gofiber/fiber/v2/utils"
)
 
type FixedWindow struct{}
 
// New は固定ウィンドウ方式のレートリミッターを実装する
func (FixedWindow) New(cfg Config) fiber.Handler {
    // ミューテックスと設定値を初期化
    mux := &sync.RWMutex{}
    max := strconv.Itoa(cfg.Max)
    expiration := uint64(cfg.Expiration.Seconds())
 
    // ストレージ管理用のマネージャーを作成
    manager := newManager(cfg.Storage)
 
    // タイムスタンプの更新をスタート
    utils.StartTimeStampUpdater()
 
    // レートリミッターの処理を持つ関数を宣言
    handler := func(c *fiber.Ctx) error {
        // 次のミドルウェアをスキップする条件を確認
        if cfg.Next != nil && cfg.Next(c) {
            return c.Next()
        }
 
        // リクエストから一意のキーを生成
        key := cfg.KeyGenerator(c)
 
        // ミューテックスをロック
        mux.Lock()
 
        // キーに対応するエントリを取得
        e := manager.get(key)
 
        // 現在のタイムスタンプを取得
        ts := uint64(atomic.LoadUint32(&utils.Timestamp))
 
        // 新しいエントリを初期化、または期限切れのエントリをリセット
        if e.exp == 0 || ts >= e.exp {
            e.currHits = 0
            e.exp = ts + expiration
        }
 
        // リクエストカウントをインクリメント
        e.currHits++
 
        // 残りのリクエスト数とリセットまでの時間を計算
        resetInSec := e.exp - ts
        remaining := cfg.Max - e.currHits
 
        // エントリを更新
        manager.set(key, e, cfg.Expiration)
 
        // ミューテックスを解除
        mux.Unlock()
 
        // 許可されたリクエスト数を超えた場合の処理
        if remaining < 0 {
            c.Set(fiber.HeaderRetryAfter, strconv.FormatUint(resetInSec, 10))
            return cfg.LimitReached(c)
        }
 
        // 次のハンドラーに処理を渡す
        err := c.Next()
 
        // 成功または失敗リクエストをスキップする設定を確認し、カウントを調整
        if (cfg.SkipSuccessfulRequests && c.Response().StatusCode() < fiber.StatusBadRequest) ||
            (cfg.SkipFailedRequests && c.Response().StatusCode() >= fiber.StatusBadRequest) {
            mux.Lock()
            e = manager.get(key)
            e.currHits-- // カウントをデクリメント
            remaining++
            manager.set(key, e, cfg.Expiration)
            mux.Unlock()
        }
 
        // レートリミット情報をレスポンスヘッダーに追加
        c.Set(xRateLimitLimit, max)
        c.Set(xRateLimitRemaining, strconv.Itoa(remaining))
        c.Set(xRateLimitReset, strconv.FormatUint(resetInSec, 10))
 
        return err
    }
 
    return handler
}