trelent-hyok 0.1.12

//! DEK persistence and storage management.
//!
//! This module provides a flexible system for storing and retrieving Data Encryption Keys (DEKs)
//! across different storage backends. It includes:
//!
//! - Cloud provider integrations (AWS Secrets Manager, Azure Key Vault)
//! - In-memory storage for testing
//! - Custom storage backend support
//! - Optional caching layer
//!
//! The persistence system ensures DEKs are encrypted with a Customer Managed Key (CMK)
//! before storage and decrypted upon retrieval.
//!
use std::sync::Arc;

use async_trait::async_trait;
use futures_util::FutureExt;

mod mem_persistence;
#[cfg(feature = "debug")]
pub use mem_persistence::*;
mod azure_persistence;
#[cfg(feature = "azure")]
pub use azure_persistence::*;
mod aws_persistence;
mod custom_persistence;
pub use custom_persistence::*;

#[cfg(feature = "aws")]
pub use aws_persistence::*;

use crate::cache::DEKCache;
use crate::cmk::CMK;
use crate::error::generator::PersistError;

/// A service that handles the persistence of Data Encryption Keys (DEKs),
/// with optional caching support.
///
/// This service provides:
/// - Transparent encryption/decryption of DEKs using a CMK
/// - Configurable storage backend
/// - Optional caching for improved performance
/// - Context-based key management
///
/// # Example
/// ```no_run
/// use std::sync::Arc;
/// use hyokashi::{CMK, DEKPersister, DEKPersistService};
///
/// // Create a service with AWS CMK and persister
/// let cmk = Arc::new(CMK::AWS(/* ... */));
/// let persister = Arc::new(DEKPersister::AWS(/* ... */));
/// let service = DEKPersistService::new(persister, cmk, None);
///
/// // Store and retrieve a DEK
///     let dek = vec![1, 2, 3];
///     let stored = service.save_dek("my-key".to_string(), dek).await?;
///     let retrieved = service.fetch_dek("my-key".to_string()).await?;
///
/// ```
pub struct DEKPersistService {
    cmk: Arc<CMK>,
    persister: Arc<DEKPersister>,
    cache: Option<Arc<DEKCache>>,
}

impl DEKPersistService {
    /// Creates a new `DEKPersistService` with the specified components.
    ///
    /// # Arguments
    ///
    /// * `persister` - The storage backend for encrypted DEKs
    /// * `cmk` - The Customer Managed Key used to encrypt/decrypt DEKs
    /// * `cache` - Optional cache for faster DEK retrieval
    pub fn new(persister: Arc<DEKPersister>, cmk: Arc<CMK>, cache: Option<Arc<DEKCache>>) -> Self {
        Self {
            persister,
            cmk,
            cache,
        }
    }

    /// Saves a DEK by first encrypting it with the CMK and then storing it
    /// using the configured persister. If caching is enabled, the unencrypted
    /// DEK is also cached.
    ///
    /// # Arguments
    ///
    /// * `context` - Identifier for the DEK
    /// * `dek` - The key bytes to store
    ///
    /// # Errors
    ///
    /// Returns a `PersistError` if:
    /// - CMK encryption fails
    /// - Persistence fails
    pub async fn save_dek(&self, context: String, dek: Vec<u8>) -> Result<Vec<u8>, PersistError> {
        let encrypted_dek: Vec<u8> = self.cmk
            .encrypt(dek.clone()).await
            .map_err(|e| PersistError::CMKError(e))?;
        let value = self.persister.persist(&encrypted_dek, context.clone()).await?;
        if let Some(cache) = &self.cache {
            let _ = cache.set(context, dek);
        }
        Ok(value)
    }

    /// Retrieves a DEK by first checking the cache (if enabled), then falling back
    /// to fetching from the persister and decrypting with the CMK.
    ///
    /// # Arguments
    ///
    /// * `context` - Identifier for the DEK to retrieve
    ///
    /// # Errors
    ///
    /// Returns a `PersistError` if:
    /// - The DEK cannot be found
    /// - CMK decryption fails
    pub async fn fetch_dek(&self, context: String) -> Result<Vec<u8>, PersistError> {
        let cached_value = if let Some(cache) = &self.cache { cache.get(&context) } else { None };
        let dek = match cached_value {
            Some(value) => { value }
            None => {
                let dek = self.persister.fetch(context.clone()).await?;
                let dek = self.cmk.decrypt(dek).await.map_err(|e| PersistError::CMKError(e))?;
                if let Some(cache) = &self.cache {
                    let _ = cache.set(context, dek.clone());
                }
                dek
            }
        };

        Ok(dek)
    }
}

/// A trait defining the interface for DEK storage backends.
///
/// This trait allows implementing custom storage solutions while maintaining
/// a consistent interface across the library. Implementors must provide
/// methods for storing and retrieving DEKs using a context string as the identifier.
///
/// # Safety
///
/// Implementations should ensure:
/// - Secure storage of sensitive key material
/// - Proper error handling for all operations
/// - Thread-safety for concurrent access
/// - Proper handling of context strings
///
/// # Example
/// ```no_run
/// use async_trait::async_trait;
/// use hyokashi::{DEKPersisterTrait, PersistError};
///
/// struct MyStorage {
///     // Storage implementation details...
/// }
///
/// #[async_trait]
/// impl DEKPersisterTrait for MyStorage {
///     async fn persist(&self, dek: &Vec<u8>, context: String) -> Result<Vec<u8>, PersistError> {
///         // Implement secure storage...
///         # Ok(vec![])
///     }
///
///     async fn fetch(&self, context: String) -> Result<Vec<u8>, PersistError> {
///         // Implement secure retrieval...
///         # Ok(vec![])
///     }
/// }
/// ```
#[async_trait]
pub trait DEKPersisterTrait {
    /// Retrieves a DEK using the provided context.
    ///
    /// # Errors
    ///
    /// Returns a `PersistError` if retrieval fails
    async fn fetch(&self, context: String) -> Result<Vec<u8>, PersistError>;

    /// Persists a DEK with the provided context.
    ///
    /// # Errors
    ///
    /// Returns a `PersistError` if storage fails
    async fn persist(&self, dek: &Vec<u8>, context: String) -> Result<Vec<u8>, PersistError>;
}

/// Available persistence backends for DEK storage.
///
/// This enum provides a unified interface to different storage solutions:
///
/// - `AWS`: AWS Secrets Manager integration (requires "aws" feature)
/// - `Azure`: Azure Key Vault integration (requires "azure" feature)
/// - `Mem`: In-memory storage for testing (requires "debug" feature)
/// - `Custom`: User-defined storage implementation
///
/// # Feature Flags
///
/// Different backends can be enabled via feature flags:
/// - `aws`: Enables AWS Secrets Manager support
/// - `azure`: Enables Azure Key Vault support
/// - `debug`: Enables in-memory storage
///
/// # Example
/// ```no_run
/// use hyokashi::DEKPersister;
///
/// // Create an AWS Secrets Manager persister
/// #[cfg(feature = "aws")]
/// let persister = DEKPersister::AWS(AWSPersister::new(/* ... */));
///
/// // Create a custom persister
/// let persister = DEKPersister::Custom(CustomPersister::new(/* ... */));
/// ```
pub enum DEKPersister {
    #[cfg(feature = "debug")]
    /// In-memory storage for testing
    Mem(MemPersister),
    #[cfg(feature = "azure")]
    /// Azure Key Vault storage
    Azure(AzurePersister),
    #[cfg(feature = "aws")]
    /// AWS Secrets Manager storage
    AWS(AWSPersister),
    /// Custom storage implementation
    Custom(CustomPersister),
}

impl DEKPersister {
    /// Retrieves a DEK using the configured backend.
    ///
    /// # Errors
    ///
    /// Returns a `PersistError` if retrieval fails
    pub async fn fetch(&self, context: String) -> Result<Vec<u8>, PersistError> {
        (match self {
            #[cfg(feature = "debug")]
            DEKPersister::Mem(persister) => persister.fetch(context).boxed(),
            #[cfg(feature = "azure")]
            DEKPersister::Azure(persister) => persister.fetch(context).boxed(),
            #[cfg(feature = "aws")]
            DEKPersister::AWS(persister) => persister.fetch(context).boxed(),
            DEKPersister::Custom(persister) => persister.fetch(context).boxed(),
        }).await
    }

    /// Persists a DEK using the configured backend.
    ///
    /// # Errors
    ///
    /// Returns a `PersistError` if storage fails
    pub async fn persist(&self, dek: &Vec<u8>, context: String) -> Result<Vec<u8>, PersistError> {
        (match self {
            #[cfg(feature = "debug")]
            DEKPersister::Mem(persister) => persister.persist(dek, context).boxed(),
            #[cfg(feature = "azure")]
            DEKPersister::Azure(persister) => persister.persist(dek, context).boxed(),
            #[cfg(feature = "aws")]
            DEKPersister::AWS(persister) => persister.persist(dek, context).boxed(),
            DEKPersister::Custom(persister) => persister.persist(dek, context).boxed(),
        }).await
    }
}