summercms/surf/limiter_test.go

package surf

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"git.golem15.com/golem15/summercms/backpack"
	"git.golem15.com/golem15/summercms/bouncer"
	"git.golem15.com/golem15/summercms/pact"
	"git.golem15.com/golem15/summercms/party"
)

func TestFixedWindowLimiterMemoryStoreWindow(t *testing.T) {
	s := NewMemoryStore(0)
	decay := 80 * time.Millisecond

	if c := s.Hit("k", decay); c != 1 {
		t.Fatalf("hit 1 = %d", c)
	}
	if c := s.Hit("k", decay); c != 2 {
		t.Fatalf("hit 2 = %d", c)
	}
	if s.TooManyAttempts("k", 3) {
		t.Fatal("TooManyAttempts at count==max-1")
	}
	if c := s.Hit("k", decay); c != 3 {
		t.Fatalf("hit 3 = %d", c)
	}
	if !s.TooManyAttempts("k", 3) {
		t.Fatal("TooManyAttempts at count==max")
	}

	time.Sleep(decay + 20*time.Millisecond)
	if s.TooManyAttempts("k", 3) {
		t.Fatal("TooManyAttempts after window elapsed")
	}
	if c := s.Hit("k", decay); c != 1 {
		t.Fatalf("fresh window hit = %d", c)
	}
}

func TestFixedWindowLimiterMemoryStoreFirstHitWins(t *testing.T) {
	s := NewMemoryStore(0)
	decay := 200 * time.Millisecond
	s.Hit("k", decay)
	time.Sleep(120 * time.Millisecond)
	s.Hit("k", decay) // must not extend the original window
	time.Sleep(100 * time.Millisecond)
	if s.TooManyAttempts("k", 1) {
		t.Fatal("window was extended by a later hit")
	}
}

func TestFixedWindowLimiterSuccessHeaders(t *testing.T) {
	lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
	h := lim.Middleware("3,1")(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d", rec.Code)
	}
	if rec.Header().Get("X-RateLimit-Limit") != "3" {
		t.Fatalf("limit = %q", rec.Header().Get("X-RateLimit-Limit"))
	}
	if rec.Header().Get("X-RateLimit-Remaining") != "2" {
		t.Fatalf("remaining = %q", rec.Header().Get("X-RateLimit-Remaining"))
	}
	if rec.Header().Get("Retry-After") != "" || rec.Header().Get("X-RateLimit-Reset") != "" {
		t.Fatal("Retry-After / X-RateLimit-Reset on success")
	}
}

func TestFixedWindowLimiterTooManyAttemptsHeaders(t *testing.T) {
	lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
	h := lim.Middleware("1,1")(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	req := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
	ok := httptest.NewRecorder()
	h.ServeHTTP(ok, req)
	if ok.Code != http.StatusOK {
		t.Fatalf("first status = %d", ok.Code)
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusTooManyRequests {
		t.Fatalf("status = %d", rec.Code)
	}
	if rec.Header().Get("Retry-After") == "" {
		t.Fatal("missing Retry-After")
	}
	if rec.Header().Get("X-RateLimit-Reset") == "" {
		t.Fatal("missing X-RateLimit-Reset")
	}
	if rec.Header().Get("X-RateLimit-Limit") != "1" {
		t.Fatalf("limit = %q", rec.Header().Get("X-RateLimit-Limit"))
	}
	if rec.Header().Get("X-RateLimit-Remaining") != "0" {
		t.Fatalf("remaining = %q", rec.Header().Get("X-RateLimit-Remaining"))
	}
	if rec.Body.String() != `{"message":"Too Many Attempts."}` {
		t.Fatalf("body = %q", rec.Body.String())
	}
	var payload map[string]any
	if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
		t.Fatal(err)
	}
	if payload["message"] != "Too Many Attempts." {
		t.Fatalf("payload = %v", payload)
	}
}

func TestFixedWindowLimiterStackedBuckets(t *testing.T) {
	okHandler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	t.Run("exhaust A", func(t *testing.T) {
		lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
		if err := lim.RegisterBucket("demo", "bucket-a", Bucket{
			Max: 1, Decay: time.Minute,
			Key: func(*http.Request) string { return "a" },
		}); err != nil {
			t.Fatal(err)
		}
		if err := lim.RegisterBucket("demo", "bucket-b", Bucket{
			Max: 100, Decay: time.Minute,
			Key: func(*http.Request) string { return "b" },
		}); err != nil {
			t.Fatal(err)
		}
		h := stackThrottle(lim, okHandler, "bucket-a", "bucket-b")
		req := httptest.NewRequest(http.MethodGet, "/x", nil)
		first := httptest.NewRecorder()
		h.ServeHTTP(first, req)
		if first.Code != http.StatusOK {
			t.Fatalf("first = %d", first.Code)
		}
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, req)
		if rec.Code != http.StatusTooManyRequests {
			t.Fatalf("exhausted A status = %d", rec.Code)
		}
	})

	t.Run("exhaust B", func(t *testing.T) {
		lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
		if err := lim.RegisterBucket("demo", "bucket-a", Bucket{
			Max: 100, Decay: time.Minute,
			Key: func(*http.Request) string { return "a-fresh" },
		}); err != nil {
			t.Fatal(err)
		}
		if err := lim.RegisterBucket("demo", "bucket-b", Bucket{
			Max: 1, Decay: time.Minute,
			Key: func(*http.Request) string { return "b-only" },
		}); err != nil {
			t.Fatal(err)
		}
		h := stackThrottle(lim, okHandler, "bucket-a", "bucket-b")
		req := httptest.NewRequest(http.MethodGet, "/x", nil)
		first := httptest.NewRecorder()
		h.ServeHTTP(first, req)
		if first.Code != http.StatusOK {
			t.Fatalf("first = %d", first.Code)
		}
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, req)
		if rec.Code != http.StatusTooManyRequests {
			t.Fatalf("exhausted B status = %d", rec.Code)
		}
	})
}

func TestFixedWindowLimiterInlineThrottleKeys(t *testing.T) {
	okHandler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	t.Run("principals differ", func(t *testing.T) {
		lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
		h := lim.Middleware("1,1")(okHandler)
		reqA := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
		reqA.RemoteAddr = "192.0.2.1:1"
		reqA = reqA.WithContext(bouncer.WithUser(reqA.Context(), &bouncer.Principal{ID: 1}))
		reqB := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
		reqB.RemoteAddr = "192.0.2.1:1"
		reqB = reqB.WithContext(bouncer.WithUser(reqB.Context(), &bouncer.Principal{ID: 2}))

		a1 := httptest.NewRecorder()
		h.ServeHTTP(a1, reqA)
		a2 := httptest.NewRecorder()
		h.ServeHTTP(a2, reqA)
		if a2.Code != http.StatusTooManyRequests {
			t.Fatalf("user 1 second status = %d", a2.Code)
		}
		b1 := httptest.NewRecorder()
		h.ServeHTTP(b1, reqB)
		if b1.Code != http.StatusOK {
			t.Fatalf("user 2 should have a distinct key, status = %d", b1.Code)
		}
	})

	t.Run("anonymous same IP+Host", func(t *testing.T) {
		lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
		h := lim.Middleware("1,1")(okHandler)
		req1 := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
		req1.RemoteAddr = "192.0.2.1:1"
		req2 := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
		req2.RemoteAddr = "192.0.2.1:9"
		first := httptest.NewRecorder()
		h.ServeHTTP(first, req1)
		if first.Code != http.StatusOK {
			t.Fatalf("first = %d", first.Code)
		}
		second := httptest.NewRecorder()
		h.ServeHTTP(second, req2)
		if second.Code != http.StatusTooManyRequests {
			t.Fatalf("anonymous same IP+Host should share key, status = %d", second.Code)
		}
	})
}

func TestFixedWindowLimiterUnknownBucketFailsAssemble(t *testing.T) {
	p := assemblePlugin{id: "golem15.demo", use: []string{"throttle:missing"}}
	_, err := Assemble(backpack.New(nil), []party.Plugin{p})
	if err == nil || !strings.Contains(err.Error(), "missing") {
		t.Fatalf("want unknown throttle in error, got %v", err)
	}
}

func TestFixedWindowLimiterDuplicateBucket(t *testing.T) {
	lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
	b := Bucket{Max: 1, Decay: time.Minute, Key: func(*http.Request) string { return "k" }}
	if err := lim.RegisterBucket("one", "shared", b); err != nil {
		t.Fatal(err)
	}
	err := lim.RegisterBucket("two", "shared", b)
	if err == nil || !strings.Contains(err.Error(), "shared") || !strings.Contains(err.Error(), "one") {
		t.Fatalf("got %v", err)
	}
}

func stackThrottle(lim *FixedWindowLimiter, next http.Handler, names ...string) http.Handler {
	h := next
	for i := len(names) - 1; i >= 0; i-- {
		h = lim.Middleware(names[i])(h)
	}
	return h
}

var _ pact.Router = (*Router)(nil)