package middleware import ( "context" "net/http" "net/http/httptest" "sync" "testing" "time" "github.com/gin-gonic/gin" ) type memoryState struct { mu sync.Mutex revoked map[string]time.Time failures map[string]int locked map[string]time.Time } func newMemoryState() *memoryState { return &memoryState{ revoked: map[string]time.Time{}, failures: map[string]int{}, locked: map[string]time.Time{}, } } func (m *memoryState) RevokeToken(_ context.Context, id string, exp time.Time) error { m.mu.Lock() defer m.mu.Unlock() m.revoked[id] = exp return nil } func (m *memoryState) IsTokenRevoked(_ context.Context, id string) (bool, error) { m.mu.Lock() defer m.mu.Unlock() exp, ok := m.revoked[id] if !ok { return false, nil } if time.Now().After(exp) { delete(m.revoked, id) return false, nil } return true, nil } func (m *memoryState) CheckLoginLock(_ context.Context, key string) (bool, time.Duration, error) { m.mu.Lock() defer m.mu.Unlock() until, ok := m.locked[key] if !ok { return false, 0, nil } if time.Now().After(until) { delete(m.locked, key) return false, 0, nil } return true, time.Until(until), nil } func (m *memoryState) RecordLoginFailure(_ context.Context, key string) error { m.mu.Lock() defer m.mu.Unlock() m.failures[key]++ if m.failures[key] >= maxFailures { m.locked[key] = time.Now().Add(lockDuration) } return nil } func (m *memoryState) RecordLoginSuccess(_ context.Context, key string) error { m.mu.Lock() defer m.mu.Unlock() delete(m.failures, key) delete(m.locked, key) return nil } func withMemoryState(t *testing.T) *memoryState { t.Helper() store := newMemoryState() InitState(store, "test") t.Cleanup(func() { InitState(nil, "test") }) return store } func TestRevokeTokenIsCrossInstanceState(t *testing.T) { withMemoryState(t) claims := &JWTClaims{} claims.ID = "token-1" if err := RevokeToken(claims, "token", time.Now().Add(time.Hour)); err != nil { t.Fatalf("RevokeToken failed: %v", err) } revoked, err := IsRevoked(claims) if err != nil { t.Fatalf("IsRevoked failed: %v", err) } if !revoked { t.Fatal("token should be revoked") } } func TestRevokeTokenUnavailableFailsClosed(t *testing.T) { InitState(nil, "production") t.Cleanup(func() { InitState(nil, "test") }) claims := &JWTClaims{} claims.ID = "x" if err := RevokeToken(claims, "token", time.Now().Add(time.Hour)); err == nil { t.Fatal("state unavailable should fail closed") } } func TestLoginLockStateSemantics(t *testing.T) { store := withMemoryState(t) for i := 0; i < maxFailures; i++ { if err := store.RecordLoginFailure(context.Background(), "1.1.1.1|user"); err != nil { t.Fatalf("record failure failed: %v", err) } } locked, _, err := store.CheckLoginLock(context.Background(), "1.1.1.1|user") if err != nil { t.Fatalf("check lock failed: %v", err) } if !locked { t.Fatal("expected login lock after max failures") } if err := store.RecordLoginSuccess(context.Background(), "1.1.1.1|user"); err != nil { t.Fatalf("record success failed: %v", err) } locked, _, _ = store.CheckLoginLock(context.Background(), "1.1.1.1|user") if locked { t.Fatal("login lock should be cleared after success") } } func TestCheckLoginLockUnavailableReturns503(t *testing.T) { InitState(nil, "test") t.Cleanup(func() { InitState(nil, "test") }) gin.SetMode(gin.TestMode) rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = httptest.NewRequest(http.MethodPost, "/auth/login", nil) if !CheckLoginLock(c, "user") { t.Fatal("state unavailable should abort login") } if rec.Code != http.StatusServiceUnavailable { t.Fatalf("status = %d, want 503", rec.Code) } }