notifier/internal/auth/keystore_hybrid.go

package auth

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"
)

// keyDatabase is the persistence surface HybridKeyStore needs. *KeyStoreDB
// implements it; tests substitute a fake.
type keyDatabase interface {
	SaveKey(ctx context.Context, key *APIKey, createdBy string) error
	GetKeyByHash(ctx context.Context, keyHash string) (*APIKey, error)
	GetKeyByName(ctx context.Context, name string) (*APIKey, error)
	ListKeys(ctx context.Context, clientID string) ([]*APIKey, error)
	DeactivateKeyByHash(ctx context.Context, keyHash string, deactivatedBy string) error
	UpdateLastUsed(ctx context.Context, keyHash string) error
	LoadAllKeys(ctx context.Context) ([]*APIKey, error)
	GetAuditLog(ctx context.Context, keyHash string, limit int) ([]map[string]interface{}, error)
	GetAuditLogByName(ctx context.Context, name string, limit int) ([]map[string]interface{}, error)
	Close() error
}

// HybridKeyStore combines an in-memory cache with an optional persistent
// database backend. Write-through strategy: writes go to the DB first and the
// cache is only updated after the DB write succeeds. Without a database the
// store degrades to in-memory-only operation.
type HybridKeyStore struct {
	cache *APIKeyStore
	db    keyDatabase // nil when no database is configured
	mu    sync.RWMutex
}

// NewHybridKeyStore creates a new hybrid key store. db may be nil, in which
// case all operations are served from the in-memory cache only.
func NewHybridKeyStore(cache *APIKeyStore, db *KeyStoreDB) *HybridKeyStore {
	h := &HybridKeyStore{cache: cache}
	if db != nil {
		h.db = db
	}
	return h
}

// HasDatabase reports whether a persistent backend is configured.
func (h *HybridKeyStore) HasDatabase() bool {
	return h.db != nil
}

// InitializeFromDatabase loads all active keys from the database into the
// cache. Call once at startup so previously issued keys survive restarts.
// Returns the number of keys loaded.
func (h *HybridKeyStore) InitializeFromDatabase(ctx context.Context) (int, error) {
	if h.db == nil {
		return 0, nil
	}

	keys, err := h.db.LoadAllKeys(ctx)
	if err != nil {
		return 0, fmt.Errorf("failed to load keys from database: %w", err)
	}

	for _, key := range keys {
		h.cache.RegisterKey(key)
	}

	return len(keys), nil
}

// CreateKey generates a new API key, persists it (when a database is
// configured), and only then registers it in the cache. If the database write
// fails the key is not usable anywhere.
func (h *HybridKeyStore) CreateKey(ctx context.Context, clientID string, roles []string, rateLimit int, expiresIn *time.Duration, createdBy string) (*APIKey, error) {
	h.mu.Lock()
	defer h.mu.Unlock()

	apiKey, err := generateAPIKey(clientID, roles, rateLimit, expiresIn)
	if err != nil {
		return nil, err
	}

	if h.db != nil {
		if err := h.db.SaveKey(ctx, apiKey, createdBy); err != nil {
			return nil, err
		}
	}

	h.cache.RegisterKey(apiKey)
	return apiKey, nil
}

// EnsurePersisted upserts an externally created key (e.g. a bootstrap admin
// key loaded from a Kubernetes secret) into the database, if one is configured.
func (h *HybridKeyStore) EnsurePersisted(ctx context.Context, apiKey *APIKey, createdBy string) error {
	if h.db == nil {
		return nil
	}
	return h.db.SaveKey(ctx, apiKey, createdBy)
}

// ValidateKey checks if a raw key is valid. The cache is consulted first; on
// a miss the database is checked and the cache repopulated, so keys created
// by another instance (or before a restart) still authenticate.
func (h *HybridKeyStore) ValidateKey(ctx context.Context, keyStr string) (*APIKey, error) {
	if key, err := h.cache.ValidateKey(keyStr); err == nil {
		return key, nil
	} else if !errors.Is(err, ErrInvalidKey) {
		// Present in cache but inactive/expired — no point hitting the DB.
		return nil, err
	}

	if h.db == nil {
		return nil, ErrKeyNotFound
	}

	key, err := h.db.GetKeyByHash(ctx, HashKey(keyStr))
	if err != nil {
		return nil, err
	}

	h.cache.RegisterKey(key)
	return h.cache.ValidateKey(keyStr)
}

// ListKeys returns all active keys for a client.
func (h *HybridKeyStore) ListKeys(ctx context.Context, clientID string) ([]*APIKey, error) {
	if h.db == nil {
		return h.cache.ListKeys(clientID), nil
	}
	return h.db.ListKeys(ctx, clientID)
}

// DeactivateKeyByName deactivates a key identified by name in both the
// database and the cache.
func (h *HybridKeyStore) DeactivateKeyByName(ctx context.Context, name string, deactivatedBy string) error {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.db == nil {
		return h.deactivateCachedByName(name)
	}

	key, err := h.db.GetKeyByName(ctx, name)
	if err != nil {
		return err
	}

	if err := h.db.DeactivateKeyByHash(ctx, key.KeyHash, deactivatedBy); err != nil {
		return err
	}

	h.cache.RemoveKeyByHash(key.KeyHash)
	return nil
}

// deactivateCachedByName handles revocation when running without a database.
func (h *HybridKeyStore) deactivateCachedByName(name string) error {
	h.cache.mu.Lock()
	defer h.cache.mu.Unlock()

	for hash, key := range h.cache.keys {
		if key.Name == name {
			delete(h.cache.keys, hash)
			delete(h.cache.rateLimits, hash)
			return nil
		}
	}
	return ErrKeyNotFound
}

// UpdateLastUsed updates the last used timestamp in the database.
// Cache is not updated to avoid contention.
func (h *HybridKeyStore) UpdateLastUsed(ctx context.Context, keyStr string) error {
	if h.db == nil {
		return h.cache.UpdateLastUsed(keyStr)
	}
	return h.db.UpdateLastUsed(ctx, HashKey(keyStr))
}

// CheckRateLimit checks if a key has exceeded its rate limit
func (h *HybridKeyStore) CheckRateLimit(keyStr string) (bool, error) {
	return h.cache.CheckRateLimit(keyStr)
}

// GetAuditLog retrieves audit log for a raw key
func (h *HybridKeyStore) GetAuditLog(ctx context.Context, keyStr string, limit int) ([]map[string]interface{}, error) {
	if h.db == nil {
		return nil, fmt.Errorf("audit log requires a database backend")
	}
	return h.db.GetAuditLog(ctx, HashKey(keyStr), limit)
}

// GetAuditLogByName retrieves audit log for a key identified by name
func (h *HybridKeyStore) GetAuditLogByName(ctx context.Context, name string, limit int) ([]map[string]interface{}, error) {
	if h.db == nil {
		return nil, fmt.Errorf("audit log requires a database backend")
	}
	return h.db.GetAuditLogByName(ctx, name, limit)
}

// Close closes the database connection, if any.
func (h *HybridKeyStore) Close() error {
	if h.db == nil {
		return nil
	}
	return h.db.Close()
}

// SyncCache performs a full cache refresh from the database. Useful for
// multi-instance deployments where keys may be created or revoked elsewhere.
func (h *HybridKeyStore) SyncCache(ctx context.Context) error {
	if h.db == nil {
		return nil
	}

	h.mu.Lock()
	defer h.mu.Unlock()

	keys, err := h.db.LoadAllKeys(ctx)
	if err != nil {
		return fmt.Errorf("failed to sync cache: %w", err)
	}

	h.cache.mu.Lock()
	h.cache.keys = make(map[string]*APIKey)
	h.cache.rateLimits = make(map[string]*RateLimiter)
	h.cache.mu.Unlock()

	for _, key := range keys {
		h.cache.RegisterKey(key)
	}

	return nil
}