rustfs/crates/kms/src/cache.rs

// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Caching layer for KMS operations to improve performance

use crate::types::KeyMetadata;
use moka::future::Cache;
use std::time::Duration;

/// KMS cache for storing frequently accessed keys and metadata
pub struct KmsCache {
    key_metadata_cache: Cache<String, KeyMetadata>,
}

impl KmsCache {
    /// Create a new KMS cache with the specified capacity
    ///
    /// # Arguments
    /// * `capacity` - Maximum number of entries in the cache
    ///
    /// # Returns
    /// A new instance of `KmsCache`
    ///
    pub fn new(capacity: u64) -> Self {
        Self {
            key_metadata_cache: Cache::builder()
                .max_capacity(capacity)
                .time_to_live(Duration::from_secs(300)) // 5 minutes default TTL
                .build(),
        }
    }

    /// Get key metadata from cache
    ///
    /// # Arguments
    /// * `key_id` - The ID of the key to retrieve metadata for
    ///
    /// # Returns
    /// An `Option` containing the `KeyMetadata` if found, or `None` if not found
    ///
    pub async fn get_key_metadata(&self, key_id: &str) -> Option<KeyMetadata> {
        self.key_metadata_cache.get(key_id).await
    }

    /// Put key metadata into cache
    ///
    /// # Arguments
    /// * `key_id` - The ID of the key to store metadata for
    /// * `metadata` - The `KeyMetadata` to store in the cache
    ///
    pub async fn put_key_metadata(&mut self, key_id: &str, metadata: &KeyMetadata) {
        self.key_metadata_cache.insert(key_id.to_string(), metadata.clone()).await;
        self.key_metadata_cache.run_pending_tasks().await;
    }

    /// Remove key metadata from cache
    ///
    /// # Arguments
    /// * `key_id` - The ID of the key to remove metadata for
    ///
    pub async fn remove_key_metadata(&mut self, key_id: &str) {
        self.key_metadata_cache.remove(key_id).await;
    }

    /// Clear all cached entries
    pub async fn clear(&mut self) {
        self.key_metadata_cache.invalidate_all();

        // Wait for invalidation to complete
        self.key_metadata_cache.run_pending_tasks().await;
    }

    /// Get cache statistics (hit count, miss count)
    ///
    /// # Returns
    /// A tuple containing total entries and total misses
    ///
    pub fn stats(&self) -> (u64, u64) {
        (
            self.key_metadata_cache.entry_count(),
            0u64, // moka doesn't provide miss count directly
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{KeyState, KeyUsage};
    use jiff::Zoned;
    use std::time::Duration;

    #[derive(Debug, Clone)]
    struct CacheInfo {
        key_metadata_count: u64,
    }

    impl CacheInfo {
        fn total_entries(&self) -> u64 {
            self.key_metadata_count
        }
    }

    impl KmsCache {
        fn with_ttl_for_tests(capacity: u64, metadata_ttl: Duration) -> Self {
            Self {
                key_metadata_cache: Cache::builder().max_capacity(capacity).time_to_live(metadata_ttl).build(),
            }
        }

        fn info_for_tests(&self) -> CacheInfo {
            CacheInfo {
                key_metadata_count: self.key_metadata_cache.entry_count(),
            }
        }

        fn contains_key_metadata_for_tests(&self, key_id: &str) -> bool {
            self.key_metadata_cache.contains_key(key_id)
        }
    }

    #[tokio::test]
    async fn test_cache_operations() {
        let mut cache = KmsCache::new(100);

        // Test key metadata caching
        let metadata = KeyMetadata {
            key_id: "test-key-1".to_string(),
            key_state: KeyState::Enabled,
            key_usage: KeyUsage::EncryptDecrypt,
            description: Some("Test key".to_string()),
            creation_date: Zoned::now(),
            deletion_date: None,
            origin: "KMS".to_string(),
            key_manager: "CUSTOMER".to_string(),
            tags: std::collections::HashMap::new(),
        };

        // Put and get metadata
        cache.put_key_metadata("test-key-1", &metadata).await;
        let retrieved = cache.get_key_metadata("test-key-1").await;
        assert!(retrieved.is_some());
        assert_eq!(retrieved.expect("metadata should be cached").key_id, "test-key-1");

        // Test cache info
        let info = cache.info_for_tests();
        assert_eq!(info.key_metadata_count, 1);
        assert_eq!(info.total_entries(), 1);

        // Test cache clearing
        cache.clear().await;
        let info_after_clear = cache.info_for_tests();
        assert_eq!(info_after_clear.total_entries(), 0);
    }

    #[tokio::test]
    async fn test_cache_with_custom_ttl() {
        let mut cache = KmsCache::with_ttl_for_tests(
            100,
            Duration::from_millis(100), // Short TTL for testing
        );

        let metadata = KeyMetadata {
            key_id: "ttl-test-key".to_string(),
            key_state: KeyState::Enabled,
            key_usage: KeyUsage::EncryptDecrypt,
            description: Some("TTL test key".to_string()),
            creation_date: Zoned::now(),
            deletion_date: None,
            origin: "KMS".to_string(),
            key_manager: "CUSTOMER".to_string(),
            tags: std::collections::HashMap::new(),
        };

        cache.put_key_metadata("ttl-test-key", &metadata).await;

        // Should be present immediately
        assert!(cache.get_key_metadata("ttl-test-key").await.is_some());

        // Wait for TTL to expire
        tokio::time::sleep(Duration::from_millis(150)).await;

        // Should be expired now
        assert!(cache.get_key_metadata("ttl-test-key").await.is_none());
    }

    #[tokio::test]
    async fn test_cache_contains_methods() {
        let mut cache = KmsCache::new(100);

        assert!(!cache.contains_key_metadata_for_tests("nonexistent"));

        let metadata = KeyMetadata {
            key_id: "contains-test".to_string(),
            key_state: KeyState::Enabled,
            key_usage: KeyUsage::EncryptDecrypt,
            description: None,
            creation_date: Zoned::now(),
            deletion_date: None,
            origin: "KMS".to_string(),
            key_manager: "CUSTOMER".to_string(),
            tags: std::collections::HashMap::new(),
        };

        cache.put_key_metadata("contains-test", &metadata).await;

        assert!(cache.contains_key_metadata_for_tests("contains-test"));
    }
}