oikos/internal/httpapi/ratelimit_test.go

package httpapi

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"
)

// okHandler is a sentinel upstream that records it was reached.
func okHandler(t *testing.T, reached *bool) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		*reached = true
		w.WriteHeader(http.StatusOK)
	})
}

func TestRateLimiterDisabledWhenRPSZero(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	rl := newRateLimiterFromConfig(ctx, 0, 0)
	if rl.rps != 0 {
		t.Fatalf("rps should be 0 when disabled, got %v", rl.rps)
	}
	// Disabled limiter is a pass-through: requests always reach upstream.
	reached := false
	h := rl.middleware(okHandler(t, &reached))
	for i := 0; i < 50; i++ {
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/entities", nil))
		if rec.Code != http.StatusOK {
			t.Fatalf("disabled limiter request %d: want 200, got %d", i, rec.Code)
		}
	}
	if !reached {
		t.Fatal("disabled limiter never reached upstream")
	}
}

func TestRateLimiterThrottlesAfterBurst(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	rl := newRateLimiter(ctx, 1, 3) // 1 rps, burst 3
	reached := false
	h := rl.middleware(okHandler(t, &reached))

	var last429, okCount int
	for i := 0; i < 6; i++ {
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/entities", nil))
		switch rec.Code {
		case http.StatusOK:
			okCount++
		case http.StatusTooManyRequests:
			last429 = i
		}
	}
	if okCount < 1 {
		t.Fatal("expected at least one request through within burst")
	}
	if last429 == 0 {
		t.Fatal("expected at least one 429 once burst exhausted")
	}
	if !reached {
		t.Fatal("upstream never reached")
	}
}

func TestRateLimiterExemptsHealthz(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	rl := newRateLimiter(ctx, 1, 1) // tiny burst
	reached := false
	h := rl.middleware(okHandler(t, &reached))
	// /healthz must never be throttled, even under a flood.
	for i := 0; i < 20; i++ {
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/healthz", nil))
		if rec.Code != http.StatusOK {
			t.Fatalf("healthz request %d throttled: want 200, got %d", i, rec.Code)
		}
	}
	if !reached {
		t.Fatal("healthz never reached upstream")
	}
}

func TestClientIPTakesRightmostXFF(t *testing.T) {
	// The reverse proxy appends the real client as the LAST hop; earlier hops
	// are spoofable and must be ignored.
	cases := []struct {
		name   string
		xff    string
		remote string
		wantIP string
	}{
		{"single xff", "203.0.113.7", "10.0.0.1:4000", "203.0.113.7"},
		{"multi hop takes rightmost", "spoofed-attacker, 203.0.113.7", "10.0.0.1:4000", "203.0.113.7"},
		{"no xff falls back to remote", "", "198.51.100.2:4000", "198.51.100.2"},
		{"blank xff falls back to remote", "  ", "198.51.100.2:4000", "198.51.100.2"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, "/", nil)
			req.RemoteAddr = tc.remote
			if strings.TrimSpace(tc.xff) != "" {
				req.Header.Set("X-Forwarded-For", tc.xff)
			}
			if got := clientIP(req); got != tc.wantIP {
				t.Fatalf("clientIP: want %q, got %q", tc.wantIP, got)
			}
		})
	}
}

func TestRateLimiterConcurrency(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	rl := newRateLimiter(ctx, 1000, 100) // generous; ensures no deadlock/panic under contention
	h := rl.middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))

	var wg sync.WaitGroup
	for i := 0; i < 20; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			rec := httptest.NewRecorder()
			h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/entities", nil))
		}()
	}
	wg.Wait()
}

func TestNewRateLimiterFromConfigBurstDefault(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()                             // release the sweep goroutine + ticker
	rl := newRateLimiterFromConfig(ctx, 10, 0) // burst unset → defaults to 2x
	if rl.burst != 20 {
		t.Fatalf("default burst should be 2x rps (20), got %d", rl.burst)
	}
}

// TestSweepStopsOnContextCancel verifies the ticker is released when the
// server context is cancelled (no process-lifetime goroutine/ticker leak).
func TestSweepStopsOnContextCancel(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	rl := newRateLimiter(ctx, 1, 1)
	cancel()
	// Give the sweeper a moment to observe cancellation. It must return
	// without blocking; the deferred ticker.Stop() fires on return.
	time.Sleep(20 * time.Millisecond)
	// Limiter remains usable for the brief test lifetime.
	if !rl.get("10.0.0.1").Allow() {
		t.Fatal("limiter should still allow within burst after sweep stops")
	}
}