scratchbox/internal/http/middleware_test.go

package httpapi

import (
	"encoding/json"
	"io"
	"log"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"strings"
	"testing"
	"time"
)

func TestAllowlistMiddlewareDeniesAndAllowsByIP(t *testing.T) {
	t.Parallel()

	logger := log.New(io.Discard, "", 0)
	allowed := []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}
	mw := AllowlistMiddleware(allowed, false, nil, logger)

	next := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	})
	handler := mw(next)

	reqAllowed := httptest.NewRequest(http.MethodPost, "/api/scratch", nil)
	reqAllowed.RemoteAddr = "10.1.2.3:1234"
	recAllowed := httptest.NewRecorder()
	handler.ServeHTTP(recAllowed, reqAllowed)
	if recAllowed.Code != http.StatusNoContent {
		t.Fatalf("allowed request status = %d, want %d", recAllowed.Code, http.StatusNoContent)
	}

	reqDenied := httptest.NewRequest(http.MethodPost, "/api/scratch", nil)
	reqDenied.RemoteAddr = "192.168.10.5:1234"
	recDenied := httptest.NewRecorder()
	handler.ServeHTTP(recDenied, reqDenied)
	if recDenied.Code != http.StatusForbidden {
		t.Fatalf("denied request status = %d, want %d", recDenied.Code, http.StatusForbidden)
	}
}

func TestAllowlistMiddlewareHonorsProxyHeaderWhenTrusted(t *testing.T) {
	t.Parallel()

	logger := log.New(io.Discard, "", 0)
	allowed := []netip.Prefix{netip.MustParsePrefix("203.0.113.5/32")}
	trustedProxies := []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}
	mw := AllowlistMiddleware(allowed, true, trustedProxies, logger)
	handler := mw(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	}))

	req := httptest.NewRequest(http.MethodPost, "/api/scratch", nil)
	req.RemoteAddr = "10.0.0.10:9999"
	req.Header.Set("X-Forwarded-For", "203.0.113.5, 198.51.100.2")
	rec := httptest.NewRecorder()

	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusNoContent {
		t.Fatalf("status = %d, want %d", rec.Code, http.StatusNoContent)
	}
}

func TestRateLimitMiddlewareKeysByClientIP(t *testing.T) {
	t.Parallel()

	limiter := NewRateLimiter(60, 1)
	handler := RateLimitMiddleware(limiter, false, nil)(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	}))

	req1 := httptest.NewRequest(http.MethodPost, "/api/scratch", nil)
	req1.RemoteAddr = "198.51.100.7:1111"
	rec1 := httptest.NewRecorder()
	handler.ServeHTTP(rec1, req1)
	if rec1.Code != http.StatusNoContent {
		t.Fatalf("first request status = %d, want %d", rec1.Code, http.StatusNoContent)
	}

	req2 := httptest.NewRequest(http.MethodPost, "/api/scratch", nil)
	req2.RemoteAddr = "198.51.100.7:2222"
	rec2 := httptest.NewRecorder()
	handler.ServeHTTP(rec2, req2)
	if rec2.Code != http.StatusTooManyRequests {
		t.Fatalf("second same-ip request status = %d, want %d", rec2.Code, http.StatusTooManyRequests)
	}

	req3 := httptest.NewRequest(http.MethodPost, "/api/scratch", nil)
	req3.RemoteAddr = "198.51.100.8:3333"
	rec3 := httptest.NewRecorder()
	handler.ServeHTTP(rec3, req3)
	if rec3.Code != http.StatusNoContent {
		t.Fatalf("different-ip request status = %d, want %d", rec3.Code, http.StatusNoContent)
	}
}

func TestAllowlistMiddlewareNoRestrictionsReturnsNext(t *testing.T) {
	t.Parallel()

	handler := AllowlistMiddleware(nil, false, nil, log.New(io.Discard, "", 0))(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusAccepted)
	}))
	req := httptest.NewRequest(http.MethodPost, "/api/scratch", nil)
	req.RemoteAddr = "not-a-valid-remote-addr"
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusAccepted {
		t.Fatalf("status = %d, want %d", rec.Code, http.StatusAccepted)
	}
}

func TestAllowlistMiddlewareRejectsUnknownClientIP(t *testing.T) {
	t.Parallel()

	logger := log.New(io.Discard, "", 0)
	allowed := []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}
	handler := AllowlistMiddleware(allowed, false, nil, logger)(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	}))
	req := httptest.NewRequest(http.MethodPost, "/api/scratch", nil)
	req.RemoteAddr = "garbage"
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusForbidden {
		t.Fatalf("status = %d, want %d", rec.Code, http.StatusForbidden)
	}
}

func TestRateLimitMiddlewareUnknownIP(t *testing.T) {
	t.Parallel()

	limiter := NewRateLimiter(60, 1)
	handler := RateLimitMiddleware(limiter, false, nil)(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	}))
	req := httptest.NewRequest(http.MethodPost, "/api/scratch", nil)
	req.RemoteAddr = "bad-addr"
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusTooManyRequests {
		t.Fatalf("status = %d, want %d", rec.Code, http.StatusTooManyRequests)
	}
}

func TestRateLimitMiddleware429PayloadAndRetryHeader(t *testing.T) {
	t.Parallel()

	limiter := NewRateLimiter(1, 1)
	handler := RateLimitMiddleware(limiter, false, nil)(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	}))
	req1 := httptest.NewRequest(http.MethodPost, "/api/scratch", nil)
	req1.RemoteAddr = "198.18.0.1:1111"
	handler.ServeHTTP(httptest.NewRecorder(), req1)

	req2 := httptest.NewRequest(http.MethodPost, "/api/scratch", nil)
	req2.RemoteAddr = "198.18.0.1:1112"
	rec2 := httptest.NewRecorder()
	handler.ServeHTTP(rec2, req2)
	if rec2.Code != http.StatusTooManyRequests {
		t.Fatalf("status = %d, want %d", rec2.Code, http.StatusTooManyRequests)
	}
	if got := rec2.Header().Get("Retry-After"); got != "60" {
		t.Fatalf("Retry-After = %q, want %q", got, "60")
	}
	if ct := rec2.Header().Get("Content-Type"); !strings.Contains(ct, "application/json") {
		t.Fatalf("Content-Type = %q, want application/json", ct)
	}
	var payload map[string]string
	if err := json.Unmarshal(rec2.Body.Bytes(), &payload); err != nil {
		t.Fatalf("json unmarshal error = %v", err)
	}
	if payload["error"] == "" {
		t.Fatalf("expected error payload")
	}
}

func TestRateLimiterPrunesStaleVisitors(t *testing.T) {
	t.Parallel()

	limiter := NewRateLimiter(60, 1)
	limiter.ttl = time.Nanosecond
	limiter.visitors["old"] = &visitor{lastSeen: time.Now().Add(-time.Hour)}
	if !limiter.Allow("new") {
		t.Fatalf("Allow(new) = false, want true")
	}
	if _, ok := limiter.visitors["old"]; ok {
		t.Fatalf("stale visitor should be removed")
	}
}

func TestExtractClientIPFallbacks(t *testing.T) {
	t.Parallel()

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.RemoteAddr = "203.0.113.9"
	ip, ok := extractClientIP(req, false, nil)
	if !ok || ip.String() != "203.0.113.9" {
		t.Fatalf("extract direct ip = (%v,%v), want (203.0.113.9,true)", ip, ok)
	}

	req2 := httptest.NewRequest(http.MethodGet, "/", nil)
	req2.RemoteAddr = "192.0.2.50:8080"
	req2.Header.Set("X-Forwarded-For", "bad,198.51.100.1")
	req2.Header.Set("X-Real-IP", "198.51.100.5")
	ip2, ok2 := extractClientIP(req2, true, []netip.Prefix{netip.MustParsePrefix("192.0.2.0/24")})
	if !ok2 || ip2.String() != "198.51.100.5" {
		t.Fatalf("extract x-real-ip = (%v,%v), want (198.51.100.5,true)", ip2, ok2)
	}

	req3 := httptest.NewRequest(http.MethodGet, "/", nil)
	req3.RemoteAddr = "still-not-an-ip"
	if _, ok3 := extractClientIP(req3, false, nil); ok3 {
		t.Fatalf("extractClientIP() ok = true, want false")
	}
}

func TestExtractClientIPIgnoresHeadersFromUntrustedProxy(t *testing.T) {
	t.Parallel()

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.RemoteAddr = "203.0.113.10:8080"
	req.Header.Set("X-Forwarded-For", "198.51.100.5")
	ip, ok := extractClientIP(req, true, []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")})
	if !ok || ip.String() != "203.0.113.10" {
		t.Fatalf("extractClientIP() = (%v,%v), want (203.0.113.10,true)", ip, ok)
	}
}