Day 059 — 2026-10-02 | ミドルウェア設計パターン: ルートごとに上限が違うレート制限——トークンバケットを60行で書き、守りの「どこ」に置くかを10個のテストで固める

📋 コピー// server.go（守りの組み立てのうち、レート制限をどこに置くかを決める）
package main

import (
	"fmt"
	"net/http"
)

type Access int

const (
	accessUnset Access = iota // ゼロ値は「決めていない」。登録時に panic
	Public
	User
	Admin
)

type Route struct {
	Pattern string
	Access  Access
	Limit   Limit // Day 057 の宣言表に足した列。ゼロ値は「決めていない」
	Handler http.HandlerFunc
}

func routeTable() []Route {
	return []Route{
		{"GET /api/v1/catalog", Public, Limit{PerSecond: 5, Burst: 10}, catalog},
		{"GET /api/v1/orders/{id}", User, Limit{PerSecond: 20, Burst: 40}, getOrder},
		{"POST /api/v1/orders", User, Limit{PerSecond: 1, Burst: 5}, createOrder},
		{"GET /admin/users", Admin, Limit{PerSecond: 2, Burst: 4}, listUsers},
	}
}

func newServer(cfg Config, l *Limiter) http.Handler { return newServerFrom(cfg, routeTable(), l) }

func newServerFrom(cfg Config, routes []Route, l *Limiter) http.Handler {
	mux := http.NewServeMux()
	for _, rt := range routes {
		var guards []Middleware
		switch rt.Access {
		case Public:
		case User:
			guards = append(guards, Authenticate(cfg))
		case Admin:
			guards = append(guards, Authenticate(cfg), RequireAdmin)
		default:
			panic(fmt.Sprintf("route %q: Access が決まっていない", rt.Pattern))
		}
		// TODO: l.Middleware(rt.Pattern, rt.Limit) を guards のどこに入れるか決める
		mux.Handle(rt.Pattern, Chain(rt.Handler, guards...))
	}
	return mux
}

📋 コピー// ratelimit.go（ここを実装する）
package main

import (
	"math"
	"net/http"
	"time"
)

// Limit は1ルート・1主体あたりの上限。PerSecond は定常の許可レート、Burst はバケツの大きさ。
// ゼロ値は「決めていない」として登録時に panic させること（R8）。
type Limit struct {
	PerSecond float64
	Burst     int
}

// Unlimited は「無制限にする」と明示する値。
var Unlimited = Limit{PerSecond: math.Inf(1), Burst: math.MaxInt32}

type Limiter struct {
	// TODO: 状態を持つ。時計は now から取る（テストが差し替える）
}

func NewLimiter(now func() time.Time) *Limiter { panic("TODO") }

// size は保持しているバケツの数（R10 のテストが使う）。
func (l *Limiter) size() int { panic("TODO") }

// Middleware は pattern・lim 用のミドルウェアを返す。上限超過なら 429 と Retry-After を返す。
func (l *Limiter) Middleware(pattern string, lim Limit) Middleware { panic("TODO") }

var _ http.Handler // import を残すための行。実装したら消してよい

📋 コピーpackage main

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type fakeClock struct {
	mu sync.Mutex
	t  time.Time
}

func (c *fakeClock) Now() time.Time          { c.mu.Lock(); defer c.mu.Unlock(); return c.t }
func (c *fakeClock) Advance(d time.Duration) { c.mu.Lock(); c.t = c.t.Add(d); c.mu.Unlock() }

var testCfg = Config{Tokens: map[string]Principal{
	"tok-alice": {"alice", "user"},
	"tok-bob":   {"bob", "user"},
	"tok-root":  {"root", "admin"},
}}

func setup(routes []Route) (http.Handler, *fakeClock, *Limiter) {
	clk := &fakeClock{t: time.Date(2026, 10, 2, 0, 0, 0, 0, time.UTC)}
	l := NewLimiter(clk.Now)
	if routes == nil {
		routes = routeTable()
	}
	return newServerFrom(testCfg, routes, l), clk, l
}

func do(h http.Handler, method, path, token, remote string) *httptest.ResponseRecorder {
	req := httptest.NewRequest(method, path, nil)
	if token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	if remote != "" {
		req.RemoteAddr = remote
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	return rec
}

func TestBurstThen429WithRetryAfter(t *testing.T) { // R1, R2
	h, _, _ := setup(nil)
	for i := 0; i < 5; i++ { // POST /orders は Burst 5
		if c := do(h, "POST", "/api/v1/orders", "tok-alice", "").Code; c != 201 {
			t.Fatalf("%d 回目: %d", i+1, c)
		}
	}
	rec := do(h, "POST", "/api/v1/orders", "tok-alice", "")
	if rec.Code != 429 || rec.Header().Get("Retry-After") != "1" {
		t.Fatalf("code=%d Retry-After=%q", rec.Code, rec.Header().Get("Retry-After"))
	}
}

func TestRefillAfterWaiting(t *testing.T) { // R2: 言われた秒数だけ待てば通る
	h, clk, _ := setup(nil)
	for i := 0; i < 5; i++ {
		do(h, "POST", "/api/v1/orders", "tok-alice", "")
	}
	rec := do(h, "POST", "/api/v1/orders", "tok-alice", "")
	clk.Advance(1 * time.Second) // Retry-After: 1
	if c := do(h, "POST", "/api/v1/orders", "tok-alice", "").Code; c != 201 {
		t.Fatalf("Retry-After=%s 待ったのに %d", rec.Header().Get("Retry-After"), c)
	}
	if c := do(h, "POST", "/api/v1/orders", "tok-alice", "").Code; c != 429 {
		t.Fatalf("補充は1つ分だけのはずが %d", c)
	}
}

func TestRetryAfterIsCeiling(t *testing.T) { // R2: 0.4/s なら 2.5 秒 → 3
	routes := []Route{{"POST /x", User, Limit{PerSecond: 0.4, Burst: 1}, createOrder}}
	h, clk, _ := setup(routes)
	do(h, "POST", "/x", "tok-alice", "")
	rec := do(h, "POST", "/x", "tok-alice", "")
	if got := rec.Header().Get("Retry-After"); got != "3" {
		t.Fatalf("Retry-After=%q, want 3", got)
	}
	clk.Advance(2500 * time.Millisecond)
	if c := do(h, "POST", "/x", "tok-alice", "").Code; c != 201 {
		t.Fatalf("2.5 秒後に %d", c)
	}
}

func TestPerPrincipal(t *testing.T) { // R3
	h, _, _ := setup(nil)
	for i := 0; i < 6; i++ {
		do(h, "POST", "/api/v1/orders", "tok-alice", "")
	}
	if c := do(h, "POST", "/api/v1/orders", "tok-bob", "").Code; c != 201 {
		t.Fatalf("bob が alice の巻き添え: %d", c)
	}
}

func TestPerRoute(t *testing.T) { // R4
	h, _, _ := setup(nil)
	for i := 0; i < 6; i++ {
		do(h, "POST", "/api/v1/orders", "tok-alice", "")
	}
	if c := do(h, "GET", "/api/v1/orders/7", "tok-alice", "").Code; c != 200 {
		t.Fatalf("別ルートが巻き添え: %d", c)
	}
}

func TestAnonymousByIPAndIgnoresXFF(t *testing.T) { // R5
	h, _, _ := setup(nil)
	for i := 0; i < 10; i++ { // catalog は Burst 10
		req := httptest.NewRequest("GET", "/api/v1/catalog", nil)
		req.RemoteAddr = "203.0.113.9:5000" // 接続ごとにポートは変わるが IP は同じ
		req.Header.Set("X-Forwarded-For", fmt.Sprintf("198.51.100.%d", i))
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, req)
		if rec.Code != 200 {
			t.Fatalf("%d 回目: %d", i+1, rec.Code)
		}
	}
	if c := do(h, "GET", "/api/v1/catalog", "", "203.0.113.9:5001").Code; c != 429 {
		t.Fatalf("同じ IP・別ポートが通った: %d", c)
	}
	if c := do(h, "GET", "/api/v1/catalog", "", "203.0.113.10:5000").Code; c != 200 {
		t.Fatalf("別 IP が巻き添え: %d", c)
	}
}

func TestRejectedDoesNotConsumeBudget(t *testing.T) { // R6: 401・403 は予算を使わない
	h, _, _ := setup(nil)
	for i := 0; i < 100; i++ {
		do(h, "POST", "/api/v1/orders", "wrong", "")
		do(h, "GET", "/admin/users", "tok-alice", "") // 403
	}
	for i := 0; i < 5; i++ {
		if c := do(h, "POST", "/api/v1/orders", "tok-alice", "").Code; c != 201 {
			t.Fatalf("alice の %d 回目: %d", i+1, c)
		}
	}
	for i := 0; i < 4; i++ {
		if c := do(h, "GET", "/admin/users", "tok-root", "").Code; c != 200 {
			t.Fatalf("root の %d 回目: %d", i+1, c)
		}
	}
}

func TestLimitedRequestNeverReachesHandler(t *testing.T) { // R7
	var n atomic.Int32
	routes := []Route{{"POST /x", User, Limit{PerSecond: 1, Burst: 2}, func(w http.ResponseWriter, r *http.Request) { n.Add(1) }}}
	h, _, _ := setup(routes)
	for i := 0; i < 10; i++ {
		do(h, "POST", "/x", "tok-alice", "")
	}
	if n.Load() != 2 {
		t.Fatalf("ハンドラが %d 回呼ばれた（2 のはず）", n.Load())
	}
}

func TestUndecidedLimitPanics(t *testing.T) { // R8
	defer func() {
		if recover() == nil {
			t.Fatal("Limit のゼロ値が登録できてしまった")
		}
	}()
	setup([]Route{{"GET /x", Public, Limit{}, catalog}})
}

func TestUnlimitedIsExplicit(t *testing.T) { // R8
	h, _, _ := setup([]Route{{"GET /x", Public, Unlimited, catalog}})
	for i := 0; i < 1000; i++ {
		if c := do(h, "GET", "/x", "", "").Code; c != 200 {
			t.Fatalf("%d 回目: %d", i+1, c)
		}
	}
}

func TestConcurrentExactlyBurst(t *testing.T) { // R9: -race で通り、許可数がちょうど Burst
	routes := []Route{{"GET /x", Public, Limit{PerSecond: 0.001, Burst: 50}, catalog}}
	h, _, _ := setup(routes)
	var ok atomic.Int32
	var wg sync.WaitGroup
	for i := 0; i < 400; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if do(h, "GET", "/x", "", "198.51.100.1:1").Code == 200 {
				ok.Add(1)
			}
		}()
	}
	wg.Wait()
	if ok.Load() != 50 {
		t.Fatalf("許可 %d 件（50 のはず）", ok.Load())
	}
}

func TestIdleBucketsAreSwept(t *testing.T) { // R10: 来なくなった主体のバケツは消える
	h, clk, l := setup(nil)
	for i := 0; i < 1000; i++ {
		do(h, "GET", "/api/v1/catalog", "", fmt.Sprintf("10.0.%d.%d:1", i/256, i%256))
	}
	if l.size() != 1000 {
		t.Fatalf("size=%d", l.size())
	}
	clk.Advance(10 * time.Second) // catalog は 5/s なので 2 秒で満タンに戻る
	clk.Advance(time.Minute)      // sweepEvery を越える
	do(h, "GET", "/api/v1/catalog", "", "192.0.2.1:1")
	if l.size() != 1 {
		t.Fatalf("掃除後 size=%d（1 のはず）", l.size())
	}
}

func TestClockGoingBackwardsDoesNotMintTokens(t *testing.T) {
	routes := []Route{{"GET /x", Public, Limit{PerSecond: 1, Burst: 1}, catalog}}
	h, clk, _ := setup(routes)
	do(h, "GET", "/x", "", "")
	clk.Advance(-time.Hour)
	if c := do(h, "GET", "/x", "", "").Code; c != 429 {
		t.Fatalf("時計が戻ったら通った: %d", c)
	}
}