notifier/internal/auth/keystore_test.go

package auth

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/igodwin/notifier/internal/logging"
)

// fakeKeyDB is an in-memory keyDatabase for exercising HybridKeyStore.
type fakeKeyDB struct {
	byHash   map[string]*APIKey
	saveErr  error
	saveCnt  int
	closeCnt int
}

func newFakeKeyDB() *fakeKeyDB {
	return &fakeKeyDB{byHash: make(map[string]*APIKey)}
}

func (f *fakeKeyDB) SaveKey(_ context.Context, key *APIKey, _ string) error {
	f.saveCnt++
	if f.saveErr != nil {
		return f.saveErr
	}
	stored := *key
	stored.Key = ""
	f.byHash[key.KeyHash] = &stored
	return nil
}

func (f *fakeKeyDB) GetKeyByHash(_ context.Context, keyHash string) (*APIKey, error) {
	key, ok := f.byHash[keyHash]
	if !ok {
		return nil, ErrKeyNotFound
	}
	return key, nil
}

func (f *fakeKeyDB) GetKeyByName(_ context.Context, name string) (*APIKey, error) {
	for _, key := range f.byHash {
		if key.Name == name {
			return key, nil
		}
	}
	return nil, ErrKeyNotFound
}

func (f *fakeKeyDB) ListKeys(_ context.Context, clientID string) ([]*APIKey, error) {
	var keys []*APIKey
	for _, key := range f.byHash {
		if key.ClientID == clientID && key.IsActive {
			keys = append(keys, key)
		}
	}
	return keys, nil
}

func (f *fakeKeyDB) DeactivateKeyByHash(_ context.Context, keyHash string, _ string) error {
	key, ok := f.byHash[keyHash]
	if !ok {
		return ErrKeyNotFound
	}
	key.IsActive = false
	return nil
}

func (f *fakeKeyDB) UpdateLastUsed(_ context.Context, _ string) error { return nil }

func (f *fakeKeyDB) LoadAllKeys(_ context.Context) ([]*APIKey, error) {
	var keys []*APIKey
	for _, key := range f.byHash {
		if key.IsActive {
			keys = append(keys, key)
		}
	}
	return keys, nil
}

func (f *fakeKeyDB) GetAuditLog(_ context.Context, _ string, _ int) ([]map[string]interface{}, error) {
	return nil, nil
}

func (f *fakeKeyDB) GetAuditLogByName(_ context.Context, _ string, _ int) ([]map[string]interface{}, error) {
	return nil, nil
}

func (f *fakeKeyDB) Close() error {
	f.closeCnt++
	return nil
}

func newHybridWithFake(db keyDatabase) (*HybridKeyStore, *APIKeyStore) {
	cache := NewAPIKeyStore()
	h := &HybridKeyStore{cache: cache, db: db}
	return h, cache
}

func TestAPIKeyStoreDoesNotRetainRawKey(t *testing.T) {
	store := NewAPIKeyStore()
	apiKey, err := store.CreateKey("client-a", []string{"notify-email"}, 10, nil)
	if err != nil {
		t.Fatalf("CreateKey: %v", err)
	}
	if apiKey.Key == "" {
		t.Fatal("creation response must include the raw key")
	}

	// The store must be indexed by digest, not raw key, and stored copies
	// must not carry the raw secret.
	store.mu.RLock()
	defer store.mu.RUnlock()
	if _, ok := store.keys[apiKey.Key]; ok {
		t.Error("store is keyed by raw key; expected digest")
	}
	stored, ok := store.keys[apiKey.KeyHash]
	if !ok {
		t.Fatal("store missing entry under key digest")
	}
	if stored.Key != "" {
		t.Error("stored record retains raw key")
	}
	if stored.KeyPreview == "" {
		t.Error("stored record missing preview")
	}
}

func TestValidateKeyByRawValue(t *testing.T) {
	store := NewAPIKeyStore()
	apiKey, err := store.CreateKey("client-a", []string{"admin"}, 0, nil)
	if err != nil {
		t.Fatalf("CreateKey: %v", err)
	}

	got, err := store.ValidateKey(apiKey.Key)
	if err != nil {
		t.Fatalf("ValidateKey with raw key: %v", err)
	}
	if got.ClientID != "client-a" {
		t.Errorf("ClientID = %q, want client-a", got.ClientID)
	}

	if _, err := store.ValidateKey("nk_bogus"); !errors.Is(err, ErrInvalidKey) {
		t.Errorf("bogus key error = %v, want ErrInvalidKey", err)
	}

	if err := store.DeactivateKey(apiKey.Key); err != nil {
		t.Fatalf("DeactivateKey: %v", err)
	}
	if _, err := store.ValidateKey(apiKey.Key); !errors.Is(err, ErrKeyInactive) {
		t.Errorf("inactive key error = %v, want ErrKeyInactive", err)
	}
}

func TestValidateKeyExpiry(t *testing.T) {
	store := NewAPIKeyStore()
	expires := -time.Minute // already expired
	apiKey, err := store.CreateKey("client-a", []string{"admin"}, 0, &expires)
	if err != nil {
		t.Fatalf("CreateKey: %v", err)
	}
	if _, err := store.ValidateKey(apiKey.Key); !errors.Is(err, ErrKeyExpired) {
		t.Errorf("expired key error = %v, want ErrKeyExpired", err)
	}
}

func TestHybridCreateKeyIsWriteThrough(t *testing.T) {
	db := newFakeKeyDB()
	db.saveErr = errors.New("db down")
	h, cache := newHybridWithFake(db)

	_, err := h.CreateKey(context.Background(), "client-a", []string{"admin"}, 0, nil, "tester")
	if err == nil {
		t.Fatal("expected error when DB write fails")
	}

	// The failed key must not be usable from the cache.
	cache.mu.RLock()
	n := len(cache.keys)
	cache.mu.RUnlock()
	if n != 0 {
		t.Errorf("cache has %d key(s) after failed DB write, want 0", n)
	}
}

func TestHybridCreateKeySucceedsAndCaches(t *testing.T) {
	db := newFakeKeyDB()
	h, _ := newHybridWithFake(db)

	apiKey, err := h.CreateKey(context.Background(), "client-a", []string{"admin"}, 0, nil, "tester")
	if err != nil {
		t.Fatalf("CreateKey: %v", err)
	}
	if db.saveCnt != 1 {
		t.Errorf("saveCnt = %d, want 1", db.saveCnt)
	}
	if _, err := h.ValidateKey(context.Background(), apiKey.Key); err != nil {
		t.Errorf("ValidateKey after create: %v", err)
	}
}

func TestHybridValidateKeyFallsBackToDatabase(t *testing.T) {
	db := newFakeKeyDB()
	h, cache := newHybridWithFake(db)

	// Simulate a key created before a restart: present in DB, absent in cache.
	apiKey, err := generateAPIKey("client-a", []string{"admin"}, 5, nil)
	if err != nil {
		t.Fatalf("generateAPIKey: %v", err)
	}
	if err := db.SaveKey(context.Background(), apiKey, "tester"); err != nil {
		t.Fatalf("SaveKey: %v", err)
	}

	got, err := h.ValidateKey(context.Background(), apiKey.Key)
	if err != nil {
		t.Fatalf("ValidateKey via DB fallback: %v", err)
	}
	if got.ClientID != "client-a" {
		t.Errorf("ClientID = %q, want client-a", got.ClientID)
	}

	// The fallback must repopulate the cache (including a rate limiter).
	if _, err := cache.ValidateKey(apiKey.Key); err != nil {
		t.Errorf("cache not repopulated after fallback: %v", err)
	}
	if ok, err := h.CheckRateLimit(apiKey.Key); err != nil || !ok {
		t.Errorf("CheckRateLimit after fallback: ok=%v err=%v", ok, err)
	}
}

func TestHybridInitializeFromDatabase(t *testing.T) {
	db := newFakeKeyDB()
	h, cache := newHybridWithFake(db)

	for i := 0; i < 3; i++ {
		apiKey, err := generateAPIKey("client-a", []string{"admin"}, 0, nil)
		if err != nil {
			t.Fatalf("generateAPIKey: %v", err)
		}
		if err := db.SaveKey(context.Background(), apiKey, "tester"); err != nil {
			t.Fatalf("SaveKey: %v", err)
		}
	}

	loaded, err := h.InitializeFromDatabase(context.Background())
	if err != nil {
		t.Fatalf("InitializeFromDatabase: %v", err)
	}
	if loaded != 3 {
		t.Errorf("loaded = %d, want 3", loaded)
	}
	cache.mu.RLock()
	n := len(cache.keys)
	cache.mu.RUnlock()
	if n != 3 {
		t.Errorf("cache has %d key(s), want 3", n)
	}
}

func TestHybridWithoutDatabase(t *testing.T) {
	h := NewHybridKeyStore(NewAPIKeyStore(), nil)

	if h.HasDatabase() {
		t.Fatal("HasDatabase() = true with nil db")
	}

	// All of these must work (or fail cleanly) without panicking.
	apiKey, err := h.CreateKey(context.Background(), "client-a", []string{"admin"}, 0, nil, "tester")
	if err != nil {
		t.Fatalf("CreateKey without db: %v", err)
	}
	if _, err := h.ValidateKey(context.Background(), apiKey.Key); err != nil {
		t.Errorf("ValidateKey without db: %v", err)
	}
	if keys, err := h.ListKeys(context.Background(), "client-a"); err != nil || len(keys) != 1 {
		t.Errorf("ListKeys without db: keys=%d err=%v", len(keys), err)
	}
	if err := h.UpdateLastUsed(context.Background(), apiKey.Key); err != nil {
		t.Errorf("UpdateLastUsed without db: %v", err)
	}
	if err := h.DeactivateKeyByName(context.Background(), apiKey.Name, "tester"); err != nil {
		t.Errorf("DeactivateKeyByName without db: %v", err)
	}
	if _, err := h.ValidateKey(context.Background(), apiKey.Key); err == nil {
		t.Error("key still valid after revocation")
	}
	if _, err := h.GetAuditLogByName(context.Background(), apiKey.Name, 10); err == nil {
		t.Error("expected audit log error without db")
	}
	if err := h.Close(); err != nil {
		t.Errorf("Close without db: %v", err)
	}
}

func TestRegisterAdminKeyInMemoryHashesKey(t *testing.T) {
	store := NewAPIKeyStore()
	logger, _ := logging.NewFromConfig("error", "stdout")

	raw := "nk_test_admin_key_value"
	apiKey, err := RegisterAdminKeyInMemory(store, raw, logger)
	if err != nil {
		t.Fatalf("RegisterAdminKeyInMemory: %v", err)
	}
	if apiKey.KeyHash != HashKey(raw) {
		t.Error("registered key missing correct digest")
	}
	if _, err := store.ValidateKey(raw); err != nil {
		t.Errorf("ValidateKey after admin registration: %v", err)
	}

	store.mu.RLock()
	defer store.mu.RUnlock()
	if _, ok := store.keys[raw]; ok {
		t.Error("admin key stored under raw value; expected digest")
	}
}