silk/server-go/internal/middleware/state_test.go

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)
	}
}