internal/server/auth_limit_test.go

Rama main · réplica de lectura del repositorio original.

package restapi

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
)

func limitedRouter(now *time.Time) *gin.Engine {
	gin.SetMode(gin.TestMode)
	limiter := newAuthFailureLimiter()
	limiter.now = func() time.Time { return *now }
	r := gin.New()
	_ = r.SetTrustedProxies(nil)
	r.Use(apiKeyMiddleware("secret", limiter))
	r.GET("/ping", func(c *gin.Context) { c.Status(http.StatusOK) })
	return r
}

func send(r *gin.Engine, remote, key string) *httptest.ResponseRecorder {
	w := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/ping", nil)
	req.RemoteAddr = remote + ":4000"
	if key != "" {
		req.Header.Set(APIKeyHeader, key)
	}
	r.ServeHTTP(w, req)
	return w
}

func TestAPIKeyMiddlewareLocksOutAClientAfterRepeatedWrongKeys(t *testing.T) {
	now := time.Unix(1_700_000_000, 0)
	r := limitedRouter(&now)

	for i := range maxAuthFailures {
		if w := send(r, "203.0.113.5", "guess"); w.Code != http.StatusUnauthorized {
			t.Fatalf("wrong key %d: got %d, want 401", i+1, w.Code)
		}
	}

	// Locked out: even the right key gets 429 with Retry-After.
	w := send(r, "203.0.113.5", "secret")
	if w.Code != http.StatusTooManyRequests {
		t.Fatalf("right key while locked out: got %d, want 429", w.Code)
	}
	if w.Header().Get("Retry-After") == "" {
		t.Error("429 has no Retry-After header")
	}

	// Another client is not affected.
	if w := send(r, "203.0.113.6", "secret"); w.Code != http.StatusOK {
		t.Errorf("other client: got %d, want 200", w.Code)
	}

	// The lockout ends.
	now = now.Add(authLockout + time.Second)
	if w := send(r, "203.0.113.5", "secret"); w.Code != http.StatusOK {
		t.Errorf("after the lockout: got %d, want 200", w.Code)
	}
}

func TestAPIKeyMiddlewareForgetsOldFailures(t *testing.T) {
	now := time.Unix(1_700_000_000, 0)
	r := limitedRouter(&now)

	for range 3 {
		for range maxAuthFailures - 1 {
			send(r, "203.0.113.5", "guess")
		}
		now = now.Add(authFailureWindow + time.Second)
	}

	if w := send(r, "203.0.113.5", "secret"); w.Code != http.StatusOK {
		t.Errorf("failures spread over several windows locked the client: got %d, want 200", w.Code)
	}
}

func TestAPIKeyMiddlewareMissingKeyDoesNotCountAsAGuess(t *testing.T) {
	now := time.Unix(1_700_000_000, 0)
	r := limitedRouter(&now)

	for i := range maxAuthFailures * 2 {
		if w := send(r, "203.0.113.5", ""); w.Code != http.StatusUnauthorized {
			t.Fatalf("missing key %d: got %d, want 401", i+1, w.Code)
		}
	}
	if w := send(r, "203.0.113.5", "secret"); w.Code != http.StatusOK {
		t.Errorf("right key after missing-key requests: got %d, want 200", w.Code)
	}
}

func TestAuthFailureLimiterTableStaysBounded(t *testing.T) {
	limiter := newAuthFailureLimiter()

	for i := range maxTrackedClients + 500 {
		limiter.fail(fmt.Sprint(i))
	}
	if got := len(limiter.clients); got > maxTrackedClients {
		t.Errorf("tracked clients = %d, want at most %d", got, maxTrackedClients)
	}
}