rustberg 0.0.5

//! Idempotency key support for safe retries.
//!
//! This module provides idempotency key handling to enable safe retries of
//! mutation operations (create, update, delete). When a client sends a request
//! with an `Idempotency-Key` header, the server:
//!
//! 1. Checks if a response for that key already exists
//! 2. If yes, returns the cached response
//! 3. If no, processes the request and caches the response
//!
//! # Security
//!
//! The cache has a bounded size (`MAX_CACHE_SIZE`) to prevent memory exhaustion
//! attacks. Uses moka crate for O(1) eviction based on LRU policy.
//!
//! # Usage
//!
//! ```
//! use rustberg::catalog::IdempotencyCache;
//! use std::time::Duration;
//!
//! let cache = IdempotencyCache::new(Duration::from_secs(86400)); // 24h TTL
//! // Use cache.get() and cache.set() in handlers
//! ```
//!
//! # Header Format
//!
//! The `Idempotency-Key` header should contain a unique identifier (preferably UUIDv7).
//! Example: `Idempotency-Key: 01895c3e-8844-7fff-a5cb-7a583a3e51fe`

use axum::body::Bytes;
use axum::http::{HeaderMap, HeaderValue, StatusCode};
use moka::sync::Cache;
use serde::Serialize;
use std::sync::Arc;
use std::time::{Duration, Instant};

/// Header name for idempotency keys.
pub const IDEMPOTENCY_KEY_HEADER: &str = "idempotency-key";

/// Response header indicating an idempotency key was used.
pub const IDEMPOTENCY_KEY_USED_HEADER: &str = "idempotency-key-used";

/// Default TTL for idempotency keys (24 hours).
pub const DEFAULT_TTL: Duration = Duration::from_secs(86400);

/// Maximum length for idempotency keys.
const MAX_KEY_LENGTH: usize = 256;

/// Maximum number of entries in the idempotency cache.
/// SEC-026: Bounded cache size prevents memory exhaustion attacks.
const MAX_CACHE_SIZE: usize = 100_000;

// ============================================================================
// Idempotency Key
// ============================================================================

/// Represents a validated idempotency key.
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct IdempotencyKey {
    key: String,
    /// Method + path combination to scope the key
    scope: String,
}

impl IdempotencyKey {
    /// Creates a new idempotency key with scope.
    pub fn new(key: impl Into<String>, method: &str, path: &str) -> Option<Self> {
        let key = key.into();

        // Validate key length
        if key.is_empty() || key.len() > MAX_KEY_LENGTH {
            return None;
        }

        // Validate key characters (alphanumeric, hyphens, underscores)
        if !key
            .chars()
            .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
        {
            return None;
        }

        Some(Self {
            key,
            scope: format!("{}:{}", method, path),
        })
    }

    /// Extracts an idempotency key from request headers.
    pub fn from_headers(headers: &HeaderMap, method: &str, path: &str) -> Option<Self> {
        headers
            .get(IDEMPOTENCY_KEY_HEADER)
            .and_then(|v| v.to_str().ok())
            .and_then(|key| Self::new(key, method, path))
    }

    /// Returns the raw key value.
    pub fn value(&self) -> &str {
        &self.key
    }
}

// ============================================================================
// Cached Response
// ============================================================================

/// A cached response for idempotent requests.
#[derive(Debug, Clone)]
pub struct CachedResponse {
    /// HTTP status code.
    pub status: StatusCode,
    /// Response body.
    pub body: Bytes,
    /// Content-Type header.
    pub content_type: Option<String>,
    /// When this response was cached.
    pub cached_at: Instant,
}

impl CachedResponse {
    /// Creates a new cached response.
    pub fn new(status: StatusCode, body: Bytes, content_type: Option<String>) -> Self {
        Self {
            status,
            body,
            content_type,
            cached_at: Instant::now(),
        }
    }

    /// Creates a cached response from JSON.
    ///
    /// Returns `None` if the value cannot be serialized (logged as warning).
    pub fn from_json<T: Serialize>(status: StatusCode, value: &T) -> Option<Self> {
        match serde_json::to_vec(value) {
            Ok(body) => Some(Self {
                status,
                body: Bytes::from(body),
                content_type: Some("application/json".to_string()),
                cached_at: Instant::now(),
            }),
            Err(e) => {
                tracing::warn!(error = %e, "Failed to serialize response for idempotency cache");
                None
            }
        }
    }

    /// Checks if this response has expired.
    pub fn is_expired(&self, ttl: Duration) -> bool {
        self.cached_at.elapsed() > ttl
    }

    /// Converts to an Axum response.
    pub fn into_axum_response(self) -> axum::response::Response {
        use axum::http::header::CONTENT_TYPE;
        use axum::response::IntoResponse;

        let mut response = (self.status, self.body).into_response();

        if let Some(content_type) = self.content_type {
            if let Ok(value) = HeaderValue::from_str(&content_type) {
                response.headers_mut().insert(CONTENT_TYPE, value);
            }
        }

        // Add header indicating this was a cached response
        response.headers_mut().insert(
            IDEMPOTENCY_KEY_USED_HEADER,
            HeaderValue::from_static("true"),
        );

        response
    }
}

// ============================================================================
// Idempotency Cache
// ============================================================================

use crate::catalog::idempotency_store::{IdempotencyEntry, IdempotencyStore};

/// Thread-safe cache for idempotent responses.
///
/// Uses moka crate for O(1) eviction based on LRU policy and TTL expiration.
/// Supports optional persistent backing store for durability across restarts.
/// When a persistent store is configured:
/// - `set` writes to both in-memory cache and persistent store (async background)
/// - `bootstrap_from_store` loads persisted entries into memory at startup
#[derive(Clone)]
pub struct IdempotencyCache {
    /// Moka cache with automatic TTL-based eviction and bounded capacity.
    cache: Cache<IdempotencyKey, CachedResponse>,
    /// Time-to-live for cached responses.
    ttl: Duration,
    /// Optional persistent backing store for durability.
    persistent_store: Option<Arc<dyn IdempotencyStore>>,
}

impl IdempotencyCache {
    /// Creates a new idempotency cache with the specified TTL.
    ///
    /// Uses moka for O(1) eviction based on LRU policy.
    pub fn new(ttl: Duration) -> Self {
        Self {
            cache: Cache::builder()
                .max_capacity(MAX_CACHE_SIZE as u64)
                .time_to_live(ttl)
                .build(),
            ttl,
            persistent_store: None,
        }
    }

    /// Creates a new idempotency cache with persistent backing store.
    pub fn with_persistent_store(ttl: Duration, store: Arc<dyn IdempotencyStore>) -> Self {
        Self {
            cache: Cache::builder()
                .max_capacity(MAX_CACHE_SIZE as u64)
                .time_to_live(ttl)
                .build(),
            ttl,
            persistent_store: Some(store),
        }
    }

    /// Creates a new idempotency cache with default TTL (24 hours).
    pub fn default_cache() -> Self {
        Self::new(DEFAULT_TTL)
    }

    /// Bootstraps the in-memory cache from the persistent store.
    ///
    /// Call this at startup to load persisted idempotency keys.
    /// Also cleans up expired entries from the persistent store.
    pub async fn bootstrap_from_store(&self) -> crate::error::Result<usize> {
        let store = match &self.persistent_store {
            Some(s) => s,
            None => return Ok(0),
        };

        // First, clean up expired entries
        let _ = store.cleanup_expired().await?;

        // Count how many we have (entries are valid since we just cleaned up)
        let count = store.count().await?;

        tracing::info!(
            entries = count,
            "Bootstrapped idempotency cache from persistent store"
        );

        Ok(count)
    }

    /// Gets a cached response for the given key.
    ///
    /// Returns `Some(response)` if found and not expired, `None` otherwise.
    /// Expiration is handled automatically by moka's TTL feature.
    /// Note: Only checks in-memory cache for performance. Use `bootstrap_from_store`
    /// at startup to populate from persistent store.
    pub fn get(&self, key: &IdempotencyKey) -> Option<CachedResponse> {
        self.cache.get(key)
    }

    /// Checks the persistent store for a key (async version of get).
    ///
    /// Returns `Some(CachedResponse)` if found in persistent store.
    pub async fn get_from_persistent(&self, key: &IdempotencyKey) -> Option<CachedResponse> {
        let store = self.persistent_store.as_ref()?;

        match store.get(&key.scope, key.value()).await {
            Ok(Some(entry)) => {
                // Convert IdempotencyEntry to CachedResponse and cache it
                let response = CachedResponse::new(
                    StatusCode::from_u16(entry.status_code).unwrap_or(StatusCode::OK),
                    Bytes::from(entry.response_body.clone()),
                    entry.content_type.clone(),
                );

                // Populate in-memory cache for future sync lookups
                self.cache.insert(key.clone(), response.clone());

                Some(response)
            }
            _ => None,
        }
    }

    /// Stores a response for the given key.
    ///
    /// Moka handles eviction automatically via LRU + TTL, so no manual eviction needed.
    /// If a persistent store is configured, also writes there asynchronously.
    pub fn set(&self, key: IdempotencyKey, response: CachedResponse) {
        // Write to persistent store in background (non-blocking)
        if let Some(store) = &self.persistent_store {
            let store = store.clone();
            let key_value = key.value().to_string();
            let scope = key.scope.clone();
            let status_code = response.status.as_u16();
            let response_body = response.body.to_vec();
            let content_type = response.content_type.clone();
            let ttl = self.ttl;

            tokio::spawn(async move {
                let entry = IdempotencyEntry::new(
                    key_value,
                    scope,
                    status_code,
                    response_body,
                    content_type,
                    ttl,
                );

                if let Err(e) = store.set(entry).await {
                    tracing::warn!(error = %e, "Failed to persist idempotency entry");
                }
            });
        }

        // Write to moka cache (O(1) operation with automatic eviction)
        self.cache.insert(key, response);
    }

    /// Removes a cached response.
    pub fn remove(&self, key: &IdempotencyKey) {
        self.cache.invalidate(key);
    }

    /// Checks if a key is already being processed.
    ///
    /// This can be used for in-flight request detection.
    pub fn contains(&self, key: &IdempotencyKey) -> bool {
        self.cache.contains_key(key)
    }

    /// Begins processing a request with an idempotency key.
    ///
    /// Returns `Ok(guard)` if the key is not already cached, inserting a
    /// "processing" placeholder. Call `guard.complete(response)` when the
    /// request finishes successfully. If the guard is dropped without
    /// `complete()`, the placeholder is removed so the client can retry.
    ///
    /// Returns `Err(cached_response)` if the key already has a cached result.
    pub fn try_begin(&self, key: IdempotencyKey) -> Result<IdempotencyGuard<'_>, CachedResponse> {
        // Check if already cached
        if let Some(existing) = self.get(&key) {
            return Err(existing);
        }

        // Insert a processing placeholder so concurrent requests see it
        let placeholder = CachedResponse::new(
            StatusCode::ACCEPTED,
            Bytes::from_static(b""),
            Some("application/json".to_string()),
        );
        self.cache.insert(key.clone(), placeholder);

        Ok(IdempotencyGuard::new(self, key))
    }

    /// Cleans up expired entries.
    ///
    /// With moka, this triggers eager eviction of expired entries.
    /// Normally moka handles this automatically, but this can be called
    /// to force immediate cleanup.
    pub fn cleanup(&self) {
        // Moka handles TTL eviction automatically, but we can run pending tasks
        self.cache.run_pending_tasks();
    }

    /// Returns the number of cached entries.
    pub fn len(&self) -> usize {
        self.cache.entry_count() as usize
    }

    /// Returns true if the cache is empty.
    pub fn is_empty(&self) -> bool {
        self.cache.entry_count() == 0
    }

    /// Returns the configured TTL.
    pub fn ttl(&self) -> Duration {
        self.ttl
    }
}

impl Default for IdempotencyCache {
    fn default() -> Self {
        Self::default_cache()
    }
}

// ============================================================================
// Idempotency Guard
// ============================================================================

/// Guard for idempotency key processing.
///
/// This is returned when a request with an idempotency key is being processed.
/// When dropped without `complete()` being called, the entry is removed to allow retries.
pub struct IdempotencyGuard<'a> {
    cache: &'a IdempotencyCache,
    key: IdempotencyKey,
    completed: bool,
}

impl<'a> IdempotencyGuard<'a> {
    /// Creates a new guard.
    fn new(cache: &'a IdempotencyCache, key: IdempotencyKey) -> Self {
        Self {
            cache,
            key,
            completed: false,
        }
    }

    /// Marks the operation as complete and stores the response.
    pub fn complete(mut self, response: CachedResponse) {
        self.completed = true;
        self.cache.set(self.key.clone(), response);
    }
}

impl<'a> Drop for IdempotencyGuard<'a> {
    fn drop(&mut self) {
        if !self.completed {
            // Request failed or was cancelled, allow retry
            self.cache.remove(&self.key);
        }
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;

    #[test]
    fn test_idempotency_key_validation() {
        // Valid keys
        assert!(IdempotencyKey::new("abc123", "POST", "/v1/tables").is_some());
        assert!(IdempotencyKey::new("uuid-with-dashes", "POST", "/v1/tables").is_some());
        assert!(IdempotencyKey::new("key_with_underscores", "POST", "/v1/tables").is_some());

        // Invalid keys
        assert!(IdempotencyKey::new("", "POST", "/v1/tables").is_none()); // Empty
        assert!(IdempotencyKey::new("key with spaces", "POST", "/v1/tables").is_none()); // Spaces
        assert!(IdempotencyKey::new("key@symbol", "POST", "/v1/tables").is_none()); // Invalid char

        // Too long
        let long_key = "a".repeat(MAX_KEY_LENGTH + 1);
        assert!(IdempotencyKey::new(&long_key, "POST", "/v1/tables").is_none());
    }

    #[test]
    fn test_idempotency_key_scoping() {
        let key1 = IdempotencyKey::new("same-key", "POST", "/v1/tables").unwrap();
        let key2 = IdempotencyKey::new("same-key", "DELETE", "/v1/tables").unwrap();
        let key3 = IdempotencyKey::new("same-key", "POST", "/v1/namespaces").unwrap();

        // Same key value but different scopes should be different
        assert_ne!(key1, key2);
        assert_ne!(key1, key3);
        assert_ne!(key2, key3);
    }

    #[test]
    fn test_idempotency_key_from_headers() {
        let mut headers = HeaderMap::new();
        headers.insert(
            IDEMPOTENCY_KEY_HEADER,
            HeaderValue::from_static("test-key-123"),
        );

        let key = IdempotencyKey::from_headers(&headers, "POST", "/v1/tables").unwrap();
        assert_eq!(key.value(), "test-key-123");
    }

    #[test]
    fn test_cached_response_expiry() {
        let response = CachedResponse::new(
            StatusCode::OK,
            Bytes::from("test"),
            Some("application/json".to_string()),
        );

        // Not expired immediately
        assert!(!response.is_expired(Duration::from_secs(60)));

        // Would be expired if TTL was 0
        assert!(response.is_expired(Duration::from_nanos(1)));
    }

    #[test]
    fn test_idempotency_cache_basic() {
        let cache = IdempotencyCache::new(Duration::from_secs(60));
        let key = IdempotencyKey::new("test-key", "POST", "/v1/tables").unwrap();

        // Initially empty
        assert!(cache.get(&key).is_none());

        // Store response
        let response = CachedResponse::new(
            StatusCode::CREATED,
            Bytes::from(r#"{"result": "ok"}"#),
            Some("application/json".to_string()),
        );
        cache.set(key.clone(), response);

        // Should be retrievable
        let cached = cache.get(&key).unwrap();
        assert_eq!(cached.status, StatusCode::CREATED);
    }

    #[test]
    fn test_idempotency_cache_expiry() {
        let cache = IdempotencyCache::new(Duration::from_millis(10));
        let key = IdempotencyKey::new("test-key", "POST", "/v1/tables").unwrap();

        let response = CachedResponse::new(StatusCode::OK, Bytes::from("test"), None);
        cache.set(key.clone(), response);

        // Available immediately
        assert!(cache.get(&key).is_some());

        // Wait for expiry
        thread::sleep(Duration::from_millis(20));

        // Should be expired now
        assert!(cache.get(&key).is_none());
    }

    #[test]
    fn test_idempotency_cache_cleanup() {
        // Use longer TTL so entries don't expire during setup
        let cache = IdempotencyCache::new(Duration::from_millis(100));

        // Add some entries
        for i in 0..5 {
            let key = IdempotencyKey::new(format!("key-{}", i), "POST", "/v1/tables").unwrap();
            let response = CachedResponse::new(StatusCode::OK, Bytes::from("test"), None);
            cache.set(key, response);
        }

        // Force moka to process pending tasks
        cache.cleanup();

        // All entries should still be present (not expired yet)
        assert!(
            cache.len() >= 4,
            "Expected at least 4 entries, got {}",
            cache.len()
        );

        // Wait for expiry
        thread::sleep(Duration::from_millis(150));

        // Cleanup should remove all expired entries
        cache.cleanup();
        assert_eq!(cache.len(), 0);
    }

    #[test]
    fn test_cached_response_from_json() {
        #[derive(Serialize)]
        struct TestResponse {
            message: String,
        }

        let value = TestResponse {
            message: "success".to_string(),
        };

        let response = CachedResponse::from_json(StatusCode::CREATED, &value).unwrap();
        assert_eq!(response.status, StatusCode::CREATED);
        assert_eq!(response.content_type, Some("application/json".to_string()));
        assert!(std::str::from_utf8(&response.body)
            .unwrap()
            .contains("success"));
    }

    #[test]
    fn test_idempotency_cache_bounded_size() {
        // SEC-026: Test that cache evicts entries when at capacity
        // Note: moka handles capacity automatically with LRU eviction
        let cache = IdempotencyCache::new(Duration::from_secs(3600));

        // Add entries
        let test_size = 1000;
        for i in 0..test_size {
            let key = IdempotencyKey::new(format!("key-{}", i), "POST", "/v1/tables").unwrap();
            let response = CachedResponse::new(StatusCode::OK, Bytes::from("test"), None);
            cache.set(key, response);
        }

        // Force moka to process pending tasks for accurate count
        cache.cleanup();

        // All entries should be present (moka handles eviction lazily)
        // The actual count may be slightly less due to moka's internal eviction
        assert!(cache.len() <= test_size);

        // Verify we can still get entries
        let key = IdempotencyKey::new("key-0", "POST", "/v1/tables").unwrap();
        assert!(cache.get(&key).is_some());
    }

    #[test]
    fn test_idempotency_guard_complete() {
        let cache = IdempotencyCache::new(Duration::from_secs(3600));
        let key = IdempotencyKey::new("guard-ok", "POST", "/v1/tables").unwrap();

        // Start processing
        let guard = cache.try_begin(key.clone()).expect("should acquire guard");

        // Complete with real response
        let response = CachedResponse::new(StatusCode::OK, Bytes::from("done"), None);
        guard.complete(response);

        // Cache now has the real response
        let cached = cache.get(&key).expect("entry should exist after complete");
        assert_eq!(cached.status, StatusCode::OK);
        assert_eq!(&cached.body[..], b"done");
    }

    #[test]
    fn test_idempotency_guard_drop_without_complete() {
        let cache = IdempotencyCache::new(Duration::from_secs(3600));
        let key = IdempotencyKey::new("guard-drop", "POST", "/v1/tables").unwrap();

        {
            // Start processing but drop the guard without completing
            let _guard = cache.try_begin(key.clone()).expect("should acquire guard");
            // guard dropped here
        }

        // Entry should be removed so the client can retry
        assert!(
            cache.get(&key).is_none(),
            "entry should be removed on guard drop"
        );
    }

    #[test]
    fn test_try_begin_returns_cached_response() {
        let cache = IdempotencyCache::new(Duration::from_secs(3600));
        let key = IdempotencyKey::new("try-begin", "POST", "/v1/tables").unwrap();

        // Pre-populate cache
        let response = CachedResponse::new(StatusCode::CREATED, Bytes::from("existing"), None);
        cache.set(key.clone(), response);

        // try_begin should return the cached response
        let result = cache.try_begin(key);
        assert!(result.is_err(), "should return Err with cached response");
        let cached = match result {
            Err(resp) => resp,
            Ok(_) => panic!("expected Err with cached response"),
        };
        assert_eq!(cached.status, StatusCode::CREATED);
    }
}