fn_registry/functions/infra/rate_limit_middleware_test.go

package infra

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestRateLimitMiddleware(t *testing.T) {
	base := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})

	t.Run("permite request bajo el limite", func(t *testing.T) {
		rl := RateLimiterNew(10, 20)
		mw := RateLimitMiddleware(rl)
		handler := mw(base)

		rec := httptest.NewRecorder()
		req := httptest.NewRequest("GET", "/api", nil)
		req.RemoteAddr = "192.0.2.1:1234"
		handler.ServeHTTP(rec, req)

		if rec.Code != http.StatusOK {
			t.Errorf("status=%d, want 200", rec.Code)
		}
		if rec.Header().Get("X-RateLimit-Limit") != "20" {
			t.Errorf("X-RateLimit-Limit=%q, want 20", rec.Header().Get("X-RateLimit-Limit"))
		}
	})

	t.Run("responde 429 cuando se excede el limite", func(t *testing.T) {
		rl := RateLimiterNew(1, 2)
		mw := RateLimitMiddleware(rl)
		handler := mw(base)

		// Consumir burst
		for i := 0; i < 2; i++ {
			rec := httptest.NewRecorder()
			req := httptest.NewRequest("GET", "/api", nil)
			req.RemoteAddr = "192.0.2.2:1234"
			handler.ServeHTTP(rec, req)
		}
		// Tercer request rechazado
		rec := httptest.NewRecorder()
		req := httptest.NewRequest("GET", "/api", nil)
		req.RemoteAddr = "192.0.2.2:1234"
		handler.ServeHTTP(rec, req)

		if rec.Code != http.StatusTooManyRequests {
			t.Errorf("status=%d, want 429", rec.Code)
		}
		if rec.Header().Get("Retry-After") == "" {
			t.Error("Retry-After header no seteado en 429")
		}
	})

	t.Run("IPs distintas tienen buckets independientes", func(t *testing.T) {
		rl := RateLimiterNew(1, 1)
		mw := RateLimitMiddleware(rl)
		handler := mw(base)

		// Agotar IP A
		recA := httptest.NewRecorder()
		reqA := httptest.NewRequest("GET", "/api", nil)
		reqA.RemoteAddr = "10.0.0.1:1000"
		handler.ServeHTTP(recA, reqA)

		recA2 := httptest.NewRecorder()
		reqA2 := httptest.NewRequest("GET", "/api", nil)
		reqA2.RemoteAddr = "10.0.0.1:1000"
		handler.ServeHTTP(recA2, reqA2)
		if recA2.Code != http.StatusTooManyRequests {
			t.Fatal("IP A deberia estar limitada")
		}

		// IP B intacta
		recB := httptest.NewRecorder()
		reqB := httptest.NewRequest("GET", "/api", nil)
		reqB.RemoteAddr = "10.0.0.2:1000"
		handler.ServeHTTP(recB, reqB)
		if recB.Code != http.StatusOK {
			t.Errorf("IP B status=%d, want 200", recB.Code)
		}
	})

	t.Run("X-Forwarded-For tiene prioridad sobre RemoteAddr", func(t *testing.T) {
		rl := RateLimiterNew(1, 1)
		mw := RateLimitMiddleware(rl)
		handler := mw(base)

		// Request 1 con XFF=A, RemoteAddr=R1
		rec1 := httptest.NewRecorder()
		req1 := httptest.NewRequest("GET", "/api", nil)
		req1.Header.Set("X-Forwarded-For", "203.0.113.1")
		req1.RemoteAddr = "192.0.2.99:5000"
		handler.ServeHTTP(rec1, req1)
		if rec1.Code != http.StatusOK {
			t.Fatal("primer request deberia pasar")
		}

		// Request 2 con XFF=A pero RemoteAddr distinto -> mismo bucket A -> rechazado
		rec2 := httptest.NewRecorder()
		req2 := httptest.NewRequest("GET", "/api", nil)
		req2.Header.Set("X-Forwarded-For", "203.0.113.1")
		req2.RemoteAddr = "192.0.2.100:5000"
		handler.ServeHTTP(rec2, req2)
		if rec2.Code != http.StatusTooManyRequests {
			t.Errorf("XFF deberia identificar misma IP, status=%d", rec2.Code)
		}
	})

	t.Run("responde con Retry-After header en 429", func(t *testing.T) {
		rl := RateLimiterNew(1, 1)
		mw := RateLimitMiddleware(rl)
		handler := mw(base)

		// Agotar
		req := httptest.NewRequest("GET", "/api", nil)
		req.RemoteAddr = "172.16.0.1:1000"
		handler.ServeHTTP(httptest.NewRecorder(), req)

		// Rechazado
		rec := httptest.NewRecorder()
		handler.ServeHTTP(rec, req)
		if rec.Code != http.StatusTooManyRequests {
			t.Fatalf("status=%d, want 429", rec.Code)
		}
		if rec.Header().Get("Retry-After") == "" {
			t.Error("Retry-After vacio en 429")
		}
		if rec.Header().Get("Content-Type") != "application/json" {
			t.Errorf("Content-Type=%q, want application/json", rec.Header().Get("Content-Type"))
		}
	})
}

func TestRateLimitClientIP(t *testing.T) {
	t.Run("X-Forwarded-For first hop", func(t *testing.T) {
		req := httptest.NewRequest("GET", "/", nil)
		req.Header.Set("X-Forwarded-For", "203.0.113.1, 198.51.100.1, 10.0.0.1")
		req.RemoteAddr = "192.0.2.1:1234"
		got := rateLimitClientIP(req)
		if got != "203.0.113.1" {
			t.Errorf("got %q, want 203.0.113.1", got)
		}
	})

	t.Run("X-Real-IP fallback", func(t *testing.T) {
		req := httptest.NewRequest("GET", "/", nil)
		req.Header.Set("X-Real-IP", "203.0.113.5")
		req.RemoteAddr = "192.0.2.1:1234"
		got := rateLimitClientIP(req)
		if got != "203.0.113.5" {
			t.Errorf("got %q, want 203.0.113.5", got)
		}
	})

	t.Run("RemoteAddr cuando no hay headers", func(t *testing.T) {
		req := httptest.NewRequest("GET", "/", nil)
		req.RemoteAddr = "192.0.2.1:1234"
		got := rateLimitClientIP(req)
		if got != "192.0.2.1" {
			t.Errorf("got %q, want 192.0.2.1", got)
		}
	})
}