hammerwork 2.5.0

//! Core encryption engine implementation for Hammerwork.
//!
//! This module provides the actual cryptographic operations for encrypting and
//! decrypting job payloads using various algorithms and key management strategies.

use super::envelope::{self, DbStore, EnvelopeKey, KmsKeyWrapper, PoolStore, WrappedKeyStore};
#[cfg(feature = "encryption")]
use super::key_manager::KeyManagerBackend;
use super::{
    EncryptedPayload, EncryptionAlgorithm, EncryptionConfig, EncryptionError, EncryptionMetadata,
    EncryptionStats, KeySource, RetentionPolicy, SecretBytes, kms,
};
use serde_json::Value;
use sqlx::Pool;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use tracing::{info, warn};

#[cfg(feature = "encryption")]
use {
    aes_gcm::{
        Aes256Gcm, Key, KeyInit, Nonce,
        aead::{Aead, Payload},
    },
    base64::Engine,
    chacha20poly1305::{ChaCha20Poly1305, Key as ChaChaKey, Nonce as ChaChaNonce},
    rand::{RngCore, rngs::OsRng},
    sha2::{Digest, Sha256},
};

/// Core encryption engine for job payload encryption and decryption.
///
/// The engine manages encryption keys, provides encryption/decryption operations,
/// handles PII field detection and protection, and maintains statistics about
/// encryption operations.
///
/// # Thread Safety
///
/// The engine is designed to be used from multiple threads safely. Internal
/// state is protected with appropriate synchronization primitives.
///
/// The engine does not depend on a database backend: it works with PostgreSQL, MySQL or
/// no database at all. Only [`EncryptionEngine::new_with_pool`] and
/// [`EncryptionEngine::rotate_kms_key`] take a connection pool (PostgreSQL or MySQL), to
/// store KMS-wrapped data keys.
///
/// To encrypt job payloads at rest, give the engine to a queue with
/// [`JobQueue::with_encryption`](crate::JobQueue::with_encryption).
///
/// `aws://` and `gcp://` key sources use envelope encryption: the data key is generated by
/// the KMS once, stored KMS-encrypted in the database, and decrypted by the KMS on every
/// later start. Create engines for those sources with [`EncryptionEngine::new_with_pool`].
///
/// # Examples
///
/// ```rust,no_run
/// # #[cfg(feature = "encryption")]
/// # {
/// use hammerwork::encryption::{EncryptionAlgorithm, EncryptionConfig, EncryptionEngine, KeySource};
/// use serde_json::json;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
///     // HAMMERWORK_ENCRYPTION_KEY holds a base64-encoded 32-byte key
///     let config = EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
///         .with_key_source(KeySource::Environment("HAMMERWORK_ENCRYPTION_KEY".to_string()));
///     let engine = EncryptionEngine::new(config).await?;
///
///     let payload = json!({
///         "user_id": "123",
///         "credit_card": "4111-1111-1111-1111",
///         "amount": 99.99
///     });
///
///     let pii_fields = vec!["credit_card"];
///     let encrypted = engine.encrypt_payload(&payload, &pii_fields).await?;
///     let decrypted = engine.decrypt_payload(&encrypted).await?;
///
///     assert_eq!(payload, decrypted);
///     Ok(())
/// }
/// # }
/// ```
pub struct EncryptionEngine {
    config: EncryptionConfig,
    /// Keys by key id. Key material is shared (never copied for an operation) and wiped
    /// from memory when the engine drops it.
    keys: Arc<Mutex<HashMap<String, KeyVersions>>>,
    /// KMS that wraps the data keys and the database that stores them, for `aws://` and
    /// `gcp://` key sources
    kms: Option<KmsBacking>,
    stats: Arc<Mutex<EncryptionStats>>,
}

/// Key material shared without copying, wiped from memory when the last reference goes.
type SharedKey = Arc<SecretBytes>;

/// The keys an engine holds for one key id.
struct KeyVersions {
    /// The key used to encrypt under this key id (and tried first to decrypt)
    active: SharedKey,
    /// Stored version of `active`, for KMS-wrapped keys (0 for other keys)
    active_version: u32,
    /// Earlier versions of a KMS-wrapped key, newest first, used to decrypt payloads
    /// written before a rotation
    previous: Vec<(u32, SharedKey)>,
}

impl KeyVersions {
    /// A key that is not KMS-wrapped (it has no versions).
    fn single(key: SecretBytes) -> Self {
        Self {
            active: Arc::new(key),
            active_version: 0,
            previous: Vec::new(),
        }
    }

    /// KMS-wrapped key versions loaded from the store; `None` if there are none.
    fn from_envelope_keys(keys: Vec<EnvelopeKey>) -> Option<Self> {
        let mut keys = keys.into_iter();
        let first = keys.next()?;
        let mut versions = Self {
            active: Arc::new(first.key),
            active_version: first.version,
            previous: Vec::new(),
        };
        versions.add(keys);
        Some(versions)
    }

    /// Versions of a KMS-wrapped key held.
    fn versions(&self) -> Vec<u32> {
        std::iter::once(self.active_version)
            .chain(self.previous.iter().map(|(version, _)| *version))
            .collect()
    }

    /// Add versions loaded from the store. A newer active version becomes the key used
    /// for encryption, and the one it replaces is kept for decryption.
    fn add(&mut self, keys: impl IntoIterator<Item = EnvelopeKey>) {
        for key in keys {
            if self.versions().contains(&key.version) {
                continue;
            }
            let version = key.version;
            let shared = Arc::new(key.key);
            if key.active && version > self.active_version {
                let old = std::mem::replace(&mut self.active, shared);
                self.previous.push((self.active_version, old));
                self.active_version = version;
            } else {
                self.previous.push((version, shared));
            }
        }
        self.previous
            .sort_by_key(|(version, _)| std::cmp::Reverse(*version));
    }

    /// Every key to try for decryption: the active one first, then earlier versions.
    fn candidates(&self) -> Vec<SharedKey> {
        std::iter::once(self.active.clone())
            .chain(self.previous.iter().map(|(_, key)| key.clone()))
            .collect()
    }
}

/// Where an engine with an `aws://` or `gcp://` key source gets its data keys.
struct KmsBacking {
    kms: Arc<dyn KmsKeyWrapper>,
    /// The database that stores the KMS-wrapped keys, to reload versions rotated by
    /// other processes
    store: Arc<dyn WrappedKeyStore>,
}

fn key_lock_error() -> EncryptionError {
    EncryptionError::KeyManagement("Failed to acquire key lock".to_string())
}

impl std::fmt::Debug for EncryptionEngine {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // Never print key material
        f.debug_struct("EncryptionEngine")
            .field("algorithm", &self.config.algorithm)
            .field("key_id", &self.key_id())
            .finish_non_exhaustive()
    }
}

impl EncryptionEngine {
    /// Creates an encryption engine whose key may be persisted in the database.
    ///
    /// For `aws://` and `gcp://` key sources the data key uses envelope encryption: on
    /// first use the KMS generates it and only its KMS-encrypted form is stored in
    /// `hammerwork_kms_data_keys` (migration 017); later engines, in any process, read
    /// that blob and decrypt it with the KMS, so they all use the same key. Engines that
    /// start at the same time converge on one stored key. Earlier versions (see
    /// [`EncryptionEngine::rotate_kms_key`]) are loaded too, so payloads encrypted before
    /// a rotation can still be decrypted.
    ///
    /// The engine keeps a handle on `pool`: when it meets a payload encrypted with a key
    /// version it does not hold (another process rotated the key with
    /// [`EncryptionEngine::rotate_kms_key`]), or with another key id stored in the same
    /// table, it reloads the stored versions and retries before failing. From then on it
    /// also encrypts with the newest active version. See also
    /// [`EncryptionEngine::reload_kms_keys`].
    ///
    /// Other key sources behave exactly as with [`EncryptionEngine::new`].
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # #[cfg(all(feature = "encryption", feature = "postgres"))]
    /// # {
    /// use hammerwork::encryption::{EncryptionAlgorithm, EncryptionConfig, EncryptionEngine, KeySource};
    ///
    /// # async fn example(pool: sqlx::PgPool) -> Result<(), hammerwork::encryption::EncryptionError> {
    /// let config = EncryptionConfig::new(EncryptionAlgorithm::AES256GCM).with_key_source(
    ///     KeySource::External("aws://alias/hammerwork?region=us-east-1".to_string()),
    /// );
    /// let engine = EncryptionEngine::new_with_pool(config, &pool).await?;
    /// # Ok(())
    /// # }
    /// # }
    /// ```
    pub async fn new_with_pool<DB: KeyManagerBackend>(
        config: EncryptionConfig,
        pool: &Pool<DB>,
    ) -> Result<Self, EncryptionError> {
        let kms = match envelope::parse_kms_source(&config.key_source)? {
            Some(source) => source.connect().await?,
            None => return Self::from_config(config).await,
        };
        Self::with_kms(config, pool, kms).await
    }

    /// Create an engine whose data key is wrapped by `kms` and stored in `pool`.
    async fn with_kms<DB: KeyManagerBackend>(
        config: EncryptionConfig,
        pool: &Pool<DB>,
        kms: Arc<dyn KmsKeyWrapper>,
    ) -> Result<Self, EncryptionError> {
        Self::with_kms_store(config, Arc::new(PoolStore(pool.clone())), kms).await
    }

    /// Create an engine whose data key is wrapped by `kms` and stored in `store`.
    async fn with_kms_store(
        config: EncryptionConfig,
        store: Arc<dyn WrappedKeyStore>,
        kms: Arc<dyn KmsKeyWrapper>,
    ) -> Result<Self, EncryptionError> {
        let key_id = config
            .key_id
            .clone()
            .unwrap_or_else(|| "default".to_string());
        let key_name = kms_key_name(&key_id);
        let size = config.key_size_bytes();

        let active =
            envelope::load_or_create(store.as_ref(), kms.as_ref(), &key_name, size).await?;
        let active_version = active.version;
        let mut versions = KeyVersions {
            active: Arc::new(active.key),
            active_version,
            previous: Vec::new(),
        };
        if active_version > 1 {
            versions.add(
                envelope::load_new(store.as_ref(), kms.as_ref(), &key_name, size, |v| {
                    v == active_version
                })
                .await?,
            );
        }

        Ok(Self {
            config,
            keys: Arc::new(Mutex::new(HashMap::from([(key_id, versions)]))),
            kms: Some(KmsBacking { kms, store }),
            stats: Arc::new(Mutex::new(EncryptionStats::new())),
        })
    }

    /// Rotates a KMS-wrapped data key (`aws://` and `gcp://` key sources).
    ///
    /// The KMS generates a new data key; its KMS-encrypted form becomes the active version
    /// in `hammerwork_kms_data_keys` and the previous version is retired but kept. This
    /// engine encrypts with the new key from now on and can still decrypt payloads
    /// encrypted with earlier versions.
    ///
    /// Other engines (other workers and processes) load the new version from the database
    /// the first time they meet a payload encrypted with it, and from then on encrypt
    /// with it too; [`EncryptionEngine::reload_kms_keys`] loads it right away.
    ///
    /// Returns the new version number. Fails with
    /// [`EncryptionError::InvalidConfiguration`] for engines not created by
    /// [`EncryptionEngine::new_with_pool`] with an `aws://` or `gcp://` source.
    pub async fn rotate_kms_key<DB: KeyManagerBackend>(
        &self,
        pool: &Pool<DB>,
    ) -> Result<u32, EncryptionError> {
        self.rotate_kms_key_in(&DbStore(pool)).await
    }

    /// [`EncryptionEngine::rotate_kms_key`] with the keys stored in `store`.
    async fn rotate_kms_key_in(&self, store: &dyn WrappedKeyStore) -> Result<u32, EncryptionError> {
        let kms = self.kms.as_ref().ok_or_else(|| {
            EncryptionError::InvalidConfiguration(
                "rotate_kms_key needs an engine created by new_with_pool with an aws:// or \
                 gcp:// key source"
                    .to_string(),
            )
        })?;
        let key_id = self
            .config
            .key_id
            .clone()
            .unwrap_or_else(|| "default".to_string());
        let rotated = envelope::rotate(
            store,
            kms.kms.as_ref(),
            &kms_key_name(&key_id),
            self.config.key_size_bytes(),
        )
        .await?;
        let version = rotated.version;

        let mut keys = self.keys.lock().map_err(|_| key_lock_error())?;
        match keys.get_mut(&key_id) {
            Some(versions) => versions.add([rotated]),
            None => {
                if let Some(versions) = KeyVersions::from_envelope_keys(vec![rotated]) {
                    keys.insert(key_id, versions);
                }
            }
        }
        Ok(version)
    }

    /// Loads the versions of this engine's KMS-wrapped data key that other processes
    /// stored since it started (see [`EncryptionEngine::rotate_kms_key`]).
    ///
    /// A newer active version becomes the key this engine encrypts with; earlier ones
    /// stay available for decryption. Returns whether a version was added. Engines reload
    /// on their own when a payload cannot be decrypted, so this is only needed to switch
    /// encryption to a rotated key without waiting for such a payload.
    ///
    /// Only engines created by [`EncryptionEngine::new_with_pool`] with an `aws://` or
    /// `gcp://` key source have stored versions; for other engines this does nothing and
    /// returns `false`.
    pub async fn reload_kms_keys(&self) -> Result<bool, EncryptionError> {
        self.reload_kms_key(self.key_id()).await
    }

    /// Load the stored versions of `key_id` that this engine does not hold yet. Returns
    /// whether a version was added.
    async fn reload_kms_key(&self, key_id: &str) -> Result<bool, EncryptionError> {
        let Some(backing) = &self.kms else {
            return Ok(false);
        };
        // The lock is not held across the queries below.
        let known: Vec<u32> = self
            .keys
            .lock()
            .map_err(|_| key_lock_error())?
            .get(key_id)
            .map(KeyVersions::versions)
            .unwrap_or_default();
        let loaded = envelope::load_new(
            backing.store.as_ref(),
            backing.kms.as_ref(),
            &kms_key_name(key_id),
            self.config.key_size_bytes(),
            |version| known.contains(&version),
        )
        .await?;
        if loaded.is_empty() {
            return Ok(false);
        }
        info!(
            "Loaded {} new version(s) of KMS-wrapped data key {}",
            loaded.len(),
            key_id
        );

        let mut keys = self.keys.lock().map_err(|_| key_lock_error())?;
        match keys.get_mut(key_id) {
            Some(versions) => versions.add(loaded),
            None => {
                if let Some(versions) = KeyVersions::from_envelope_keys(loaded) {
                    keys.insert(key_id.to_string(), versions);
                }
            }
        }
        Ok(true)
    }
}

/// Name of an engine's KMS-wrapped data key in `hammerwork_kms_data_keys`.
fn kms_key_name(key_id: &str) -> String {
    format!("engine/{}", key_id)
}

// Constructors and key loading.
impl EncryptionEngine {
    /// Create an engine with the key named by `config.key_source`.
    async fn from_config(
        #[cfg_attr(not(feature = "encryption"), allow(unused_variables))] config: EncryptionConfig,
    ) -> Result<Self, EncryptionError> {
        #[cfg(not(feature = "encryption"))]
        {
            return Err(EncryptionError::InvalidConfiguration(
                "Encryption feature is not enabled. Enable the 'encryption' feature flag."
                    .to_string(),
            ));
        }

        #[cfg(feature = "encryption")]
        {
            let mut keys = HashMap::new();
            let key_id = config
                .key_id
                .clone()
                .unwrap_or_else(|| "default".to_string());

            // Load the encryption key
            let key = Self::load_key(&config.key_source, config.key_size_bytes()).await?;
            keys.insert(key_id, KeyVersions::single(key));

            Ok(Self {
                config,
                keys: Arc::new(Mutex::new(keys)),
                kms: None,
                stats: Arc::new(Mutex::new(EncryptionStats::new())),
            })
        }
    }

    /// The key for [`KeySource::Generated`]`(location)`.
    ///
    /// Reuses the key already stored at `location` (a `file://` URL or plain path, or an
    /// `env://` variable), so an engine restarted with the same configuration can still
    /// decrypt what it encrypted. Only when nothing is stored there is a new random key
    /// generated and stored. `stdout://` always generates a key and prints it (for
    /// initial setup). A stored value that is not a valid key is an error; it is never
    /// overwritten.
    #[cfg(feature = "encryption")]
    fn load_or_generate_key(
        location: &str,
        expected_size: usize,
    ) -> Result<SecretBytes, EncryptionError> {
        use base64::Engine;
        use std::io::Write;

        let decode = |stored: &str, place: &str| -> Result<SecretBytes, EncryptionError> {
            let key = SecretBytes::new(
                base64::engine::general_purpose::STANDARD
                    .decode(stored.trim())
                    .map_err(|e| {
                        EncryptionError::KeyManagement(format!(
                            "Invalid base64 key stored in {}: {}",
                            place, e
                        ))
                    })?,
            );
            if key.len() != expected_size {
                return Err(EncryptionError::KeyManagement(format!(
                    "Key size mismatch in {}: expected {} bytes, got {}",
                    place,
                    expected_size,
                    key.len()
                )));
            }
            Ok(key)
        };
        let storage_error = |e: std::io::Error| {
            EncryptionError::KeyManagement(format!(
                "Failed to store generated key at {}: {}",
                location, e
            ))
        };

        let mut key = SecretBytes::new(vec![0u8; expected_size]);
        OsRng.fill_bytes(&mut key);
        let encoded =
            zeroize::Zeroizing::new(base64::engine::general_purpose::STANDARD.encode(&key));

        if let Some(var_name) = location.strip_prefix("env://") {
            if let Ok(stored) = std::env::var(var_name) {
                return decode(&zeroize::Zeroizing::new(stored), var_name);
            }
            // SAFETY: as for any `set_var`, the caller must not read or write the
            // environment from other threads meanwhile.
            unsafe {
                std::env::set_var(var_name, encoded.as_str());
            }
            info!("Stored generated key in environment variable: {}", var_name);
            return Ok(key);
        }
        if location.starts_with("stdout://") {
            println!("Generated encryption key: {}", encoded.as_str());
            println!("Store this key securely and set it as an environment variable.");
            return Ok(key);
        }

        let file_path = location.strip_prefix("file://").unwrap_or(location);
        if let Some(parent) = std::path::Path::new(file_path).parent()
            && !parent.as_os_str().is_empty()
        {
            std::fs::create_dir_all(parent).map_err(storage_error)?;
        }
        let mut options = std::fs::OpenOptions::new();
        // `create_new`: never replace a key another process stored meanwhile.
        options.write(true).create_new(true);
        #[cfg(unix)]
        {
            use std::os::unix::fs::OpenOptionsExt;
            options.mode(0o600);
        }
        match options.open(file_path) {
            Ok(mut file) => {
                file.write_all(encoded.as_bytes())
                    .and_then(|()| file.flush())
                    .map_err(storage_error)?;
                info!("Stored generated key in file: {}", file_path);
                Ok(key)
            }
            Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
                let stored =
                    zeroize::Zeroizing::new(std::fs::read_to_string(file_path).map_err(|e| {
                        EncryptionError::KeyManagement(format!(
                            "Failed to read stored key {}: {}",
                            file_path, e
                        ))
                    })?);
                info!("Loaded previously generated key from file: {}", file_path);
                decode(&stored, file_path)
            }
            Err(e) => Err(storage_error(e)),
        }
    }

    #[cfg(feature = "encryption")]
    async fn load_key(
        key_source: &KeySource,
        expected_size: usize,
    ) -> Result<SecretBytes, EncryptionError> {
        match key_source {
            KeySource::Static(key_str) => {
                let key_bytes = SecretBytes::new(
                    base64::engine::general_purpose::STANDARD
                        .decode(key_str)
                        .map_err(|e| {
                            EncryptionError::KeyManagement(format!("Invalid base64 key: {}", e))
                        })?,
                );

                if key_bytes.len() != expected_size {
                    return Err(EncryptionError::KeyManagement(format!(
                        "Key size mismatch: expected {} bytes, got {}",
                        expected_size,
                        key_bytes.len()
                    )));
                }

                Ok(key_bytes)
            }
            KeySource::Environment(var_name) => {
                let key_str = zeroize::Zeroizing::new(std::env::var(var_name).map_err(|_| {
                    EncryptionError::KeyManagement(format!(
                        "Environment variable {} not found",
                        var_name
                    ))
                })?);

                let key_bytes = SecretBytes::new(
                    base64::engine::general_purpose::STANDARD
                        .decode(key_str.as_str())
                        .map_err(|e| {
                            EncryptionError::KeyManagement(format!(
                                "Invalid base64 key in {}: {}",
                                var_name, e
                            ))
                        })?,
                );

                if key_bytes.len() != expected_size {
                    return Err(EncryptionError::KeyManagement(format!(
                        "Key size mismatch in {}: expected {} bytes, got {}",
                        var_name,
                        expected_size,
                        key_bytes.len()
                    )));
                }

                Ok(key_bytes)
            }
            KeySource::External(service_config) => {
                load_external_key(service_config, expected_size).await
            }
            KeySource::Generated(location) => Self::load_or_generate_key(location, expected_size),
        }
    }
}

impl EncryptionEngine {
    /// Creates a new encryption engine with the specified configuration.
    ///
    /// This will attempt to load the encryption key according to the
    /// configuration's key source. If the key cannot be loaded, an error
    /// will be returned.
    ///
    /// # Arguments
    ///
    /// * `config` - The encryption configuration to use
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The encryption key cannot be loaded from the specified source
    /// - The encryption feature is not enabled
    /// - The configuration is invalid
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "encryption")]
    /// # {
    /// use hammerwork::encryption::{EncryptionEngine, EncryptionConfig, EncryptionAlgorithm, KeySource};
    ///
    /// # async fn example() -> Result<(), hammerwork::encryption::EncryptionError> {
    /// // With environment variable key (base64-encoded 32-byte key)
    /// let config = EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
    ///     .with_key_source(KeySource::Environment("MY_ENCRYPTION_KEY".to_string()));
    /// let engine = EncryptionEngine::new(config).await?;
    ///
    /// // With static key (for testing only)
    /// let config = EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
    ///     .with_key_source(KeySource::Static(
    ///         "QUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUE=".to_string(),
    ///     ));
    /// let engine = EncryptionEngine::new(config).await?;
    /// # Ok(())
    /// # }
    /// # }
    /// ```
    pub async fn new(config: EncryptionConfig) -> Result<Self, EncryptionError> {
        Self::from_config(config).await
    }

    /// Adds a decrypt-only key, for payloads encrypted under another key id.
    ///
    /// The engine always encrypts with its configured key (`config.key_id`, or
    /// `"default"`). Each encrypted payload records the id of the key it was encrypted
    /// with; after switching to a new key id, add the previous key here so payloads
    /// written before the switch can still be decrypted.
    ///
    /// Fails if the key cannot be loaded or has the wrong size. Adding the engine's own
    /// key id is an error.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "encryption")]
    /// # {
    /// use hammerwork::encryption::{EncryptionAlgorithm, EncryptionConfig, EncryptionEngine, KeySource};
    ///
    /// # async fn example() -> Result<(), hammerwork::encryption::EncryptionError> {
    /// let config = EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
    ///     .with_key_id("key-2026-10")
    ///     .with_key_source(KeySource::Environment("HAMMERWORK_KEY_2026_10".to_string()));
    /// let engine = EncryptionEngine::new(config)
    ///     .await?
    ///     .with_decryption_key(
    ///         "key-2026-04",
    ///         &KeySource::Environment("HAMMERWORK_KEY_2026_04".to_string()),
    ///     )
    ///     .await?;
    /// # Ok(())
    /// # }
    /// # }
    /// ```
    pub async fn with_decryption_key(
        self,
        key_id: impl Into<String>,
        key_source: &KeySource,
    ) -> Result<Self, EncryptionError> {
        let key_id = key_id.into();
        if key_id == self.key_id() {
            return Err(EncryptionError::InvalidConfiguration(format!(
                "Key id '{}' is the engine's encryption key",
                key_id
            )));
        }
        let key = Self::load_key(key_source, self.config.key_size_bytes()).await?;
        self.keys
            .lock()
            .map_err(|_| key_lock_error())?
            .insert(key_id, KeyVersions::single(key));
        Ok(self)
    }

    /// The engine's configuration.
    pub fn config(&self) -> &EncryptionConfig {
        &self.config
    }

    /// Id of the key the engine encrypts with (`config.key_id`, or `"default"`).
    pub fn key_id(&self) -> &str {
        self.config.key_id.as_deref().unwrap_or("default")
    }

    /// The algorithm the engine encrypts with.
    pub fn algorithm(&self) -> &EncryptionAlgorithm {
        &self.config.algorithm
    }

    /// The key used for encryption (shared, not copied).
    fn active_key(&self) -> Result<SharedKey, EncryptionError> {
        let key_id = self.key_id();
        let keys = self.keys.lock().map_err(|_| key_lock_error())?;
        keys.get(key_id)
            .map(|versions| versions.active.clone())
            .ok_or_else(|| EncryptionError::KeyManagement(format!("Key not found: {}", key_id)))
    }

    /// Encrypts a job payload with optional PII field protection.
    ///
    /// This method encrypts the entire payload and optionally applies special
    /// handling to fields containing personally identifiable information (PII).
    /// PII fields are tracked in the encryption metadata for compliance purposes.
    ///
    /// # Arguments
    ///
    /// * `payload` - The JSON payload to encrypt
    /// * `pii_fields` - List of field names that contain PII
    ///
    /// # Returns
    ///
    /// An encrypted payload container with all necessary metadata for decryption.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "encryption")]
    /// # {
    /// use serde_json::json;
    ///
    /// # async fn example(engine: hammerwork::encryption::EncryptionEngine) -> Result<(), Box<dyn std::error::Error>> {
    /// let payload = json!({
    ///     "user_id": "123",
    ///     "email": "user@example.com",
    ///     "ssn": "123-45-6789",
    ///     "amount": 99.99
    /// });
    ///
    /// let pii_fields = vec!["email", "ssn"];
    /// let encrypted = engine.encrypt_payload(&payload, &pii_fields).await?;
    /// # Ok(())
    /// # }
    /// # }
    /// ```
    pub async fn encrypt_payload(
        &self,
        payload: &Value,
        pii_fields: &[impl AsRef<str>],
    ) -> Result<EncryptedPayload, EncryptionError> {
        self.encrypt_payload_with_retention(payload, pii_fields, RetentionPolicy::UseDefault)
            .await
    }

    /// Encrypts a job payload with a custom retention policy.
    ///
    /// Similar to `encrypt_payload` but allows specifying a custom retention
    /// policy for this specific payload.
    ///
    /// # Arguments
    ///
    /// * `payload` - The JSON payload to encrypt
    /// * `pii_fields` - List of field names that contain PII
    /// * `retention_policy` - Custom retention policy for this payload
    pub async fn encrypt_payload_with_retention(
        &self,
        payload: &Value,
        pii_fields: &[impl AsRef<str>],
        retention_policy: RetentionPolicy,
    ) -> Result<EncryptedPayload, EncryptionError> {
        self.encrypt_payload_with_associated_data(payload, pii_fields, retention_policy, &[])
            .await
    }

    /// Encrypts a payload and authenticates `associated_data` with it.
    ///
    /// The associated data is not stored: [`EncryptionEngine::decrypt_payload_with_associated_data`]
    /// must be given the same bytes, or decryption fails. Use it to bind the ciphertext
    /// to its context (for job payloads, [`JobQueue`](crate::JobQueue) binds it to the
    /// job id, queue name, key id, algorithm and PII fields; see
    /// [`job_payload`](super::job_payload)), so a ciphertext copied into another context
    /// cannot be decrypted there. Empty associated data is the same as none
    /// ([`EncryptionEngine::encrypt_payload_with_retention`]).
    pub async fn encrypt_payload_with_associated_data(
        &self,
        payload: &Value,
        pii_fields: &[impl AsRef<str>],
        retention_policy: RetentionPolicy,
        associated_data: &[u8],
    ) -> Result<EncryptedPayload, EncryptionError> {
        let result = self.encrypt_value(payload, pii_fields, retention_policy, associated_data);
        if result.is_err()
            && let Ok(mut stats) = self.stats.lock()
        {
            stats.record_encryption_error();
        }
        result
    }

    fn encrypt_value(
        &self,
        payload: &Value,
        pii_fields: &[impl AsRef<str>],
        retention_policy: RetentionPolicy,
        associated_data: &[u8],
    ) -> Result<EncryptedPayload, EncryptionError> {
        let start_time = Instant::now();
        let key = self.active_key()?;

        let payload_bytes = serde_json::to_vec(payload).map_err(EncryptionError::Serialization)?;

        // Keyed integrity hash of the plaintext. It is stored next to the ciphertext, so it
        // must not be a plain digest: a SHA-256 of a low-entropy value (an SSN, a card
        // number) can be reversed by brute force.
        let payload_hash = keyed_payload_hash(&key, &payload_bytes)?;

        let data_to_encrypt = if self.config.compression_enabled {
            self.compress_data(&payload_bytes)?
        } else {
            payload_bytes
        };

        let (ciphertext, nonce, tag) =
            self.encrypt_data(&key, &data_to_encrypt, associated_data)?;

        let pii_field_names: Vec<String> =
            pii_fields.iter().map(|f| f.as_ref().to_string()).collect();

        let mut metadata = EncryptionMetadata::new(
            &self.config,
            pii_field_names.clone(),
            retention_policy,
            payload_hash,
        );
        metadata.key_id = self.key_id().to_string();
        metadata.delete_at = metadata.retention_policy.calculate_deletion_time(
            metadata.encrypted_at,
            None,
            self.config.default_retention,
        );

        let duration_ms = start_time.elapsed().as_millis() as f64;
        if let Ok(mut stats) = self.stats.lock() {
            stats.record_encryption(&self.config.algorithm, data_to_encrypt.len(), duration_ms);
            stats.record_pii_encryption(pii_field_names.len());
        }

        Ok(EncryptedPayload::new(ciphertext, nonce, tag, metadata))
    }

    /// Decrypts an encrypted payload back to its original JSON form.
    ///
    /// This method handles key lookup, algorithm selection, decompression,
    /// and integrity verification automatically based on the metadata
    /// contained in the encrypted payload.
    ///
    /// # Arguments
    ///
    /// * `encrypted_payload` - The encrypted payload to decrypt
    ///
    /// # Returns
    ///
    /// The original JSON payload.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The decryption key is not available
    /// - The ciphertext has been corrupted
    /// - The algorithm is not supported
    /// - The integrity check fails
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "encryption")]
    /// # {
    /// # async fn example(engine: hammerwork::encryption::EncryptionEngine, encrypted_payload: hammerwork::encryption::EncryptedPayload) -> Result<(), Box<dyn std::error::Error>> {
    /// let decrypted_payload = engine.decrypt_payload(&encrypted_payload).await?;
    /// println!("Decrypted: {}", decrypted_payload);
    /// # Ok(())
    /// # }
    /// # }
    /// ```
    pub async fn decrypt_payload(
        &self,
        encrypted_payload: &EncryptedPayload,
    ) -> Result<Value, EncryptionError> {
        self.decrypt_payload_with_associated_data(encrypted_payload, &[])
            .await
    }

    /// Decrypts a payload encrypted by
    /// [`EncryptionEngine::encrypt_payload_with_associated_data`].
    ///
    /// Fails with [`EncryptionError::DecryptionFailed`] unless `associated_data` is
    /// exactly the associated data the payload was encrypted with.
    pub async fn decrypt_payload_with_associated_data(
        &self,
        encrypted_payload: &EncryptedPayload,
        associated_data: &[u8],
    ) -> Result<Value, EncryptionError> {
        let mut result = self.decrypt_value(encrypted_payload, associated_data);
        // The payload may use a key version (or key id) that another process stored after
        // this engine loaded its keys, e.g. after a rotation: reload and try again.
        if result.is_err() && self.kms.is_some() {
            match self
                .reload_kms_key(&encrypted_payload.metadata.key_id)
                .await
            {
                Ok(true) => result = self.decrypt_value(encrypted_payload, associated_data),
                Ok(false) => {}
                Err(e) => warn!(
                    "Could not reload KMS-wrapped data key {} after a decryption failure: {}",
                    encrypted_payload.metadata.key_id, e
                ),
            }
        }
        if result.is_err()
            && let Ok(mut stats) = self.stats.lock()
        {
            stats.record_decryption_error();
        }
        result
    }

    fn decrypt_value(
        &self,
        encrypted_payload: &EncryptedPayload,
        associated_data: &[u8],
    ) -> Result<Value, EncryptionError> {
        let start_time = Instant::now();

        // The key named by the payload, then earlier versions of a rotated KMS-wrapped key
        let key_id = &encrypted_payload.metadata.key_id;
        let candidates = self
            .keys
            .lock()
            .map_err(|_| key_lock_error())?
            .get(key_id)
            .map(KeyVersions::candidates)
            .ok_or_else(|| {
                EncryptionError::KeyManagement(format!(
                    "Key not found: {} (the payload was encrypted with a key this engine \
                     does not have; see EncryptionEngine::with_decryption_key)",
                    key_id
                ))
            })?;

        let ciphertext = encrypted_payload.decode_ciphertext()?;
        let nonce = encrypted_payload.decode_nonce()?;
        let tag = encrypted_payload.decode_tag()?;
        let algorithm = &encrypted_payload.metadata.algorithm;

        let mut first_error = None;
        let mut decrypted = None;
        for key in &candidates {
            match self.decrypt_data(&ciphertext, &nonce, &tag, key, algorithm, associated_data) {
                Ok(data) => {
                    decrypted = Some((key, data));
                    break;
                }
                Err(err) => {
                    first_error.get_or_insert(err);
                }
            }
        }
        let (key, decrypted_data) = match decrypted {
            Some(found) => found,
            None => {
                return Err(first_error.unwrap_or_else(|| {
                    EncryptionError::DecryptionFailed("no key could decrypt the payload".into())
                }));
            }
        };

        let payload_bytes = if encrypted_payload.metadata.compressed {
            self.decompress_data(&decrypted_data)?
        } else {
            decrypted_data
        };

        // Verify integrity. Payloads written before keyed hashes carry a plain SHA-256.
        let expected = &encrypted_payload.metadata.payload_hash;
        let calculated = if expected.starts_with(KEYED_HASH_PREFIX) {
            keyed_payload_hash(key, &payload_bytes)?
        } else {
            self.calculate_payload_hash(&payload_bytes)
        };
        if calculated != *expected {
            return Err(EncryptionError::DecryptionFailed(
                "Payload integrity check failed".to_string(),
            ));
        }

        let payload: Value =
            serde_json::from_slice(&payload_bytes).map_err(EncryptionError::Serialization)?;

        let duration_ms = start_time.elapsed().as_millis() as f64;
        if let Ok(mut stats) = self.stats.lock() {
            stats.record_decryption(payload_bytes.len(), duration_ms);
        }

        Ok(payload)
    }

    /// Identifies PII fields in a JSON payload based on field names and patterns.
    ///
    /// This method scans the payload for common PII field names and patterns
    /// to automatically identify fields that should be encrypted.
    ///
    /// # Arguments
    ///
    /// * `payload` - The JSON payload to scan
    ///
    /// # Returns
    ///
    /// A list of field names that likely contain PII.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # #[cfg(feature = "encryption")]
    /// # {
    /// use serde_json::json;
    ///
    /// # fn example(engine: &hammerwork::encryption::EncryptionEngine) {
    /// let payload = json!({
    ///     "user_id": "123",
    ///     "email": "user@example.com",
    ///     "credit_card_number": "4111-1111-1111-1111",
    ///     "amount": 99.99
    /// });
    ///
    /// let pii_fields = engine.identify_pii_fields(&payload);
    /// // Returns: ["email", "credit_card_number"]
    /// # }
    /// # }
    /// ```
    pub fn identify_pii_fields(&self, payload: &Value) -> Vec<String> {
        let mut pii_fields = Vec::new();

        // Common PII field name patterns
        let pii_patterns = [
            "ssn",
            "social_security",
            "social_security_number",
            "credit_card",
            "credit_card_number",
            "card_number",
            "cc_number",
            "email",
            "email_address",
            "e_mail",
            "phone",
            "phone_number",
            "telephone",
            "mobile",
            "passport",
            "passport_number",
            "driver_license",
            "drivers_license",
            "license_number",
            "bank_account",
            "account_number",
            "routing_number",
            "password",
            "secret",
            "private_key",
            "address",
            "street_address",
            "home_address",
            "date_of_birth",
            "birth_date",
            "dob",
            "tax_id",
            "taxpayer_id",
            "ein",
        ];

        self.scan_object_for_pii(payload, &pii_patterns, &mut pii_fields, "");
        pii_fields
    }

    /// Gets the current encryption statistics.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # #[cfg(feature = "encryption")]
    /// # {
    /// # fn example(engine: &hammerwork::encryption::EncryptionEngine) {
    /// let stats = engine.get_stats();
    /// println!("Jobs encrypted: {}", stats.jobs_encrypted);
    /// println!("Success rate: {:.2}%", stats.encryption_success_rate());
    /// # }
    /// # }
    /// ```
    pub fn get_stats(&self) -> EncryptionStats {
        self.stats
            .lock()
            .map(|stats| stats.clone())
            .unwrap_or_else(|_| {
                // If lock is poisoned, return default stats
                EncryptionStats::new()
            })
    }

    /// Cleans up expired encrypted data based on retention policies.
    ///
    /// This method should be called periodically to remove encrypted
    /// payloads that have exceeded their retention period.
    ///
    /// # Arguments
    ///
    /// * `encrypted_payloads` - List of encrypted payloads to check
    ///
    /// # Returns
    ///
    /// The number of payloads that should be deleted.
    pub fn cleanup_expired_data(&self, encrypted_payloads: &[EncryptedPayload]) -> usize {
        let expired_count = encrypted_payloads
            .iter()
            .filter(|payload| payload.should_delete_now())
            .count();

        if expired_count > 0
            && let Ok(mut stats) = self.stats.lock()
        {
            stats.record_retention_cleanup(expired_count as u64);
        }

        expired_count
    }

    // Private helper methods

    #[allow(clippy::type_complexity)]
    fn encrypt_data(
        &self,
        key: &[u8],
        data: &[u8],
        aad: &[u8],
    ) -> Result<(Vec<u8>, Vec<u8>, Vec<u8>), EncryptionError> {
        let payload = Payload { msg: data, aad };
        check_key_size(key)?;
        match self.config.algorithm {
            EncryptionAlgorithm::AES256GCM => {
                let cipher_key = Key::<Aes256Gcm>::from_slice(key);
                let cipher = Aes256Gcm::new(cipher_key);

                // Generate random nonce
                let mut nonce_bytes = vec![0u8; 12];
                OsRng.fill_bytes(&mut nonce_bytes);
                let nonce = Nonce::from_slice(&nonce_bytes);

                // Encrypt the data
                let ciphertext_with_tag = cipher.encrypt(nonce, payload).map_err(|e| {
                    EncryptionError::EncryptionFailed(format!("AES encryption failed: {}", e))
                })?;

                // Split ciphertext and tag
                let tag_start = ciphertext_with_tag.len() - 16;
                let ciphertext = ciphertext_with_tag[..tag_start].to_vec();
                let tag = ciphertext_with_tag[tag_start..].to_vec();

                Ok((ciphertext, nonce_bytes, tag))
            }
            EncryptionAlgorithm::ChaCha20Poly1305 => {
                let cipher_key = ChaChaKey::from_slice(key);
                let cipher = ChaCha20Poly1305::new(cipher_key);

                // Generate random nonce
                let mut nonce_bytes = vec![0u8; 12];
                OsRng.fill_bytes(&mut nonce_bytes);
                let nonce = ChaChaNonce::from_slice(&nonce_bytes);

                // Encrypt the data
                let ciphertext_with_tag = cipher.encrypt(nonce, payload).map_err(|e| {
                    EncryptionError::EncryptionFailed(format!("ChaCha20 encryption failed: {}", e))
                })?;

                // Split ciphertext and tag
                let tag_start = ciphertext_with_tag.len() - 16;
                let ciphertext = ciphertext_with_tag[..tag_start].to_vec();
                let tag = ciphertext_with_tag[tag_start..].to_vec();

                Ok((ciphertext, nonce_bytes, tag))
            }
        }
    }

    #[cfg(feature = "encryption")]
    fn decrypt_data(
        &self,
        ciphertext: &[u8],
        nonce: &[u8],
        tag: &[u8],
        key: &[u8],
        algorithm: &EncryptionAlgorithm,
        aad: &[u8],
    ) -> Result<Vec<u8>, EncryptionError> {
        // The nonce and tag come from storage: reject wrong sizes instead of panicking
        // in `from_slice`.
        check_key_size(key)?;
        if nonce.len() != NONCE_SIZE || tag.len() != TAG_SIZE {
            return Err(EncryptionError::DecryptionFailed(format!(
                "Invalid nonce or tag size ({} and {} bytes, expected {} and {})",
                nonce.len(),
                tag.len(),
                NONCE_SIZE,
                TAG_SIZE
            )));
        }

        // Reconstruct the ciphertext with tag
        let mut ciphertext_with_tag = ciphertext.to_vec();
        ciphertext_with_tag.extend_from_slice(tag);

        match algorithm {
            EncryptionAlgorithm::AES256GCM => {
                let cipher_key = Key::<Aes256Gcm>::from_slice(key);
                let cipher = Aes256Gcm::new(cipher_key);
                let nonce = Nonce::from_slice(nonce);

                cipher
                    .decrypt(
                        nonce,
                        Payload {
                            msg: ciphertext_with_tag.as_slice(),
                            aad,
                        },
                    )
                    .map_err(|e| {
                        EncryptionError::DecryptionFailed(format!("AES decryption failed: {}", e))
                    })
            }
            EncryptionAlgorithm::ChaCha20Poly1305 => {
                let cipher_key = ChaChaKey::from_slice(key);
                let cipher = ChaCha20Poly1305::new(cipher_key);
                let nonce = ChaChaNonce::from_slice(nonce);

                cipher
                    .decrypt(
                        nonce,
                        Payload {
                            msg: ciphertext_with_tag.as_slice(),
                            aad,
                        },
                    )
                    .map_err(|e| {
                        EncryptionError::DecryptionFailed(format!(
                            "ChaCha20 decryption failed: {}",
                            e
                        ))
                    })
            }
        }
    }

    #[cfg(feature = "encryption")]
    fn compress_data(&self, data: &[u8]) -> Result<Vec<u8>, EncryptionError> {
        use flate2::{Compression, write::GzEncoder};
        use std::io::Write;

        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
        encoder
            .write_all(data)
            .map_err(|e| EncryptionError::FieldProcessing(format!("Compression failed: {}", e)))?;
        encoder.finish().map_err(|e| {
            EncryptionError::FieldProcessing(format!("Compression finalization failed: {}", e))
        })
    }

    #[cfg(feature = "encryption")]
    fn decompress_data(&self, data: &[u8]) -> Result<Vec<u8>, EncryptionError> {
        use flate2::read::GzDecoder;
        use std::io::Read;

        let mut decoder = GzDecoder::new(data);
        let mut decompressed = Vec::new();
        decoder.read_to_end(&mut decompressed).map_err(|e| {
            EncryptionError::FieldProcessing(format!("Decompression failed: {}", e))
        })?;
        Ok(decompressed)
    }

    #[cfg(feature = "encryption")]
    fn calculate_payload_hash(&self, data: &[u8]) -> String {
        let mut hasher = Sha256::new();
        hasher.update(data);
        hex::encode(hasher.finalize())
    }

    #[allow(clippy::only_used_in_recursion)]
    fn scan_object_for_pii(
        &self,
        value: &Value,
        patterns: &[&str],
        pii_fields: &mut Vec<String>,
        prefix: &str,
    ) {
        match value {
            Value::Object(map) => {
                for (key, val) in map {
                    let field_name = if prefix.is_empty() {
                        key.clone()
                    } else {
                        format!("{}.{}", prefix, key)
                    };

                    // Check if this key matches any PII patterns
                    let key_lower = key.to_lowercase();
                    if patterns.iter().any(|pattern| {
                        key_lower.contains(pattern)
                            || key_lower
                                .replace('_', "")
                                .contains(&pattern.replace('_', ""))
                    }) {
                        pii_fields.push(field_name.clone());
                    }

                    // Recursively scan nested objects
                    if matches!(val, Value::Object(_) | Value::Array(_)) {
                        self.scan_object_for_pii(val, patterns, pii_fields, &field_name);
                    }
                }
            }
            Value::Array(arr) => {
                for (i, val) in arr.iter().enumerate() {
                    let field_name = format!("{}[{}]", prefix, i);
                    if matches!(val, Value::Object(_) | Value::Array(_)) {
                        self.scan_object_for_pii(val, patterns, pii_fields, &field_name);
                    }
                }
            }
            _ => {} // Primitive values don't contain nested fields
        }
    }
}

/// Nonce size of both supported AEAD algorithms.
const NONCE_SIZE: usize = 12;
/// Authentication tag size of both supported AEAD algorithms.
const TAG_SIZE: usize = 16;
/// Key size of both supported AEAD algorithms.
const KEY_SIZE: usize = 32;
/// Prefix of keyed payload hashes (`EncryptionMetadata::payload_hash`).
const KEYED_HASH_PREFIX: &str = "hmac-sha256:";

fn check_key_size(key: &[u8]) -> Result<(), EncryptionError> {
    if key.len() == KEY_SIZE {
        Ok(())
    } else {
        Err(EncryptionError::KeyManagement(format!(
            "Invalid key size: expected {} bytes, got {}",
            KEY_SIZE,
            key.len()
        )))
    }
}

/// HMAC-SHA256 of `data` under a key derived from the data key, as
/// `hmac-sha256:<hex>`.
fn keyed_payload_hash(key: &[u8], data: &[u8]) -> Result<String, EncryptionError> {
    use hmac::{Hmac, Mac};

    // Derive a separate MAC key so the data key itself is never used outside the AEAD.
    let mac_key = Sha256::new()
        .chain_update(b"hammerwork payload hash v1")
        .chain_update(key)
        .finalize();
    let mut mac = <Hmac<Sha256> as Mac>::new_from_slice(&mac_key)
        .map_err(|e| EncryptionError::Cryptographic(format!("HMAC key error: {}", e)))?;
    mac.update(data);
    Ok(format!(
        "{}{}",
        KEYED_HASH_PREFIX,
        hex::encode(mac.finalize().into_bytes())
    ))
}

// External key loaders. They fail closed: an unavailable KMS, missing credentials or a
// KMS source whose cargo feature is not enabled is an error, never a fallback key.

/// Load a key from the external key management service named by `service_config`.
async fn load_external_key(
    service_config: &str,
    expected_size: usize,
) -> Result<SecretBytes, EncryptionError> {
    if service_config.starts_with("aws://") || service_config.starts_with("gcp://") {
        // Validates the configuration and the cargo feature first
        envelope::parse_kms_source(&KeySource::External(service_config.to_string()))?;
        Err(EncryptionError::InvalidConfiguration(format!(
            "{} keys are generated once by the KMS and stored KMS-encrypted in the database; \
             create the engine with EncryptionEngine::new_with_pool",
            &service_config[..service_config.find("://").unwrap_or(0) + 3]
        )))
    } else if service_config.starts_with("vault://") {
        load_from_vault(service_config, expected_size).await
    } else if service_config.starts_with("azure://") {
        load_from_azure_kv(service_config, expected_size).await
    } else {
        Err(EncryptionError::KeyManagement(format!(
            "Unknown external key management service: {}",
            service_config
        )))
    }
}

/// Load a key from a HashiCorp Vault KV v2 secret (`vault://<mount>/<path>?addr=<url>`).
///
/// The secret's string `key` field must be a base64-encoded key of exactly
/// `expected_size` bytes; anything else is an [`EncryptionError::InvalidConfiguration`].
async fn load_from_vault(
    service_config: &str,
    expected_size: usize,
) -> Result<SecretBytes, EncryptionError> {
    let (secret_path, params) = kms::parse_service_config(service_config, "vault://");
    let vault_addr = kms::vault_address(&params)?;
    let (mount, secret) = kms::vault_mount_and_path(secret_path)?;

    info!(
        "Loading key from HashiCorp Vault: path={}, addr={}",
        secret_path, vault_addr
    );

    #[cfg(not(feature = "vault-kms"))]
    {
        let _ = (mount, secret, expected_size);
        Err(kms::kms_feature_disabled("vault://", "vault-kms"))
    }

    #[cfg(feature = "vault-kms")]
    {
        let key_str = kms::vault_read_key_field(&vault_addr, mount, secret).await?;
        let key = kms::exact_vault_key(&key_str, expected_size)?;
        info!("Successfully loaded encryption key from HashiCorp Vault");
        Ok(key)
    }
}

/// Load a key from Azure Key Vault (`azure://<vault-host>/keys/<key-name>`, key name
/// defaults to `encryption-key`). The key material must be exactly `expected_size`
/// bytes; anything else is an [`EncryptionError::InvalidConfiguration`].
async fn load_from_azure_kv(
    service_config: &str,
    expected_size: usize,
) -> Result<SecretBytes, EncryptionError> {
    let (vault_url, key_name) = kms::azure_vault_and_key(service_config, "encryption-key")?;

    info!(
        "Loading key from Azure Key Vault: vault={}, key={}",
        vault_url, key_name
    );

    #[cfg(not(feature = "azure-kv"))]
    {
        let _ = (vault_url, key_name, expected_size);
        Err(kms::kms_feature_disabled("azure://", "azure-kv"))
    }

    #[cfg(feature = "azure-kv")]
    {
        let decoded_key = super::azure::fetch_key_material(&vault_url, &key_name)
            .await
            .map_err(|e| {
                EncryptionError::KeyManagement(format!(
                    "Failed to load key from Azure Key Vault: {}",
                    e
                ))
            })?;

        let key = kms::exact_size(
            decoded_key,
            expected_size,
            &format!("Azure Key Vault key {} in {}", key_name, vault_url),
        )?;
        info!("Successfully loaded encryption key from Azure Key Vault");
        Ok(key)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    use std::time::Duration;

    #[cfg(feature = "encryption")]
    #[tokio::test]
    async fn test_encryption_engine_creation() {
        // Set up a test key in environment
        unsafe {
            std::env::set_var(
                "TEST_ENCRYPTION_KEY",
                "dGVzdGtleTE5ODc2NTQzMjEwOTg3NjU0MzIxMHRlc3Q=",
            ); // 32 bytes base64
        }

        let config = EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
            .with_key_source(KeySource::Environment("TEST_ENCRYPTION_KEY".to_string()));

        let engine = EncryptionEngine::new(config).await;
        assert!(engine.is_ok());
    }

    #[cfg(feature = "encryption")]
    #[tokio::test]
    async fn test_payload_encryption_decryption() {
        unsafe {
            std::env::set_var(
                "TEST_ENCRYPTION_KEY",
                "dGVzdGtleTE5ODc2NTQzMjEwOTg3NjU0MzIxMHRlc3Q=",
            );
        }

        let config = EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
            .with_key_source(KeySource::Environment("TEST_ENCRYPTION_KEY".to_string()));

        let engine = EncryptionEngine::new(config).await.unwrap();

        let payload = json!({
            "user_id": "123",
            "email": "test@example.com",
            "amount": 99.99
        });

        let pii_fields = vec!["email"];
        let encrypted = engine.encrypt_payload(&payload, &pii_fields).await.unwrap();
        let decrypted = engine.decrypt_payload(&encrypted).await.unwrap();

        assert_eq!(payload, decrypted);
        assert!(encrypted.contains_pii_field("email"));
        assert!(!encrypted.contains_pii_field("user_id"));
    }

    #[tokio::test]
    async fn test_pii_field_identification() {
        // It used to build an engine with the default key source (an unset environment
        // variable) and silently return when that failed, so it never ran.
        let engine = EncryptionEngine::new(static_engine("k", KEY_A))
            .await
            .unwrap();

        let payload = json!({
            "user_id": "123",
            "email": "test@example.com",
            "credit_card_number": "4111-1111-1111-1111",
            "ssn": "123-45-6789",
            "amount": 99.99,
            "metadata": {
                "phone_number": "+1-555-123-4567"
            }
        });

        let pii_fields = engine.identify_pii_fields(&payload);

        assert!(pii_fields.contains(&"email".to_string()));
        assert!(pii_fields.contains(&"credit_card_number".to_string()));
        assert!(pii_fields.contains(&"ssn".to_string()));
        assert!(pii_fields.contains(&"metadata.phone_number".to_string()));
        assert!(!pii_fields.contains(&"user_id".to_string()));
        assert!(!pii_fields.contains(&"amount".to_string()));
    }

    #[cfg(feature = "encryption")]
    #[tokio::test]
    async fn test_encryption_with_compression() {
        unsafe {
            std::env::set_var(
                "TEST_ENCRYPTION_KEY",
                "dGVzdGtleTE5ODc2NTQzMjEwOTg3NjU0MzIxMHRlc3Q=",
            );
        }

        let config = EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
            .with_key_source(KeySource::Environment("TEST_ENCRYPTION_KEY".to_string()))
            .with_compression_enabled(true);

        let engine = EncryptionEngine::new(config).await.unwrap();

        let payload = json!({
            "data": "a".repeat(1000), // Large repeating data that compresses well
            "user_id": "123"
        });

        let encrypted = engine
            .encrypt_payload(&payload, &Vec::<String>::new())
            .await
            .unwrap();
        let decrypted = engine.decrypt_payload(&encrypted).await.unwrap();

        assert_eq!(payload, decrypted);
        assert!(encrypted.metadata.compressed);
    }

    #[cfg(feature = "encryption")]
    #[tokio::test]
    async fn test_encryption_stats() {
        unsafe {
            std::env::set_var(
                "TEST_ENCRYPTION_KEY",
                "dGVzdGtleTE5ODc2NTQzMjEwOTg3NjU0MzIxMHRlc3Q=",
            );
        }

        let config = EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
            .with_key_source(KeySource::Environment("TEST_ENCRYPTION_KEY".to_string()));

        let engine = EncryptionEngine::new(config).await.unwrap();

        let payload = json!({"test": "data"});

        let initial_stats = engine.get_stats();
        assert_eq!(initial_stats.jobs_encrypted, 0);

        let _encrypted = engine.encrypt_payload(&payload, &["test"]).await.unwrap();

        let stats = engine.get_stats();
        assert_eq!(stats.jobs_encrypted, 1);
        assert_eq!(stats.pii_fields_encrypted, 1);
        assert!(stats.avg_encryption_time_ms >= 0.0);
    }

    #[tokio::test]
    async fn test_retention_policy_cleanup() {
        let engine = EncryptionEngine::new(static_engine("k", KEY_A))
            .await
            .unwrap();

        // Create a mock encrypted payload with immediate deletion policy
        let metadata = EncryptionMetadata::new(
            &EncryptionConfig::default(),
            vec![],
            RetentionPolicy::DeleteAfter(Duration::from_secs(1)),
            "test_hash".to_string(),
        );

        let payload = EncryptedPayload::new(
            b"test".to_vec(),
            b"nonce".to_vec(),
            b"tag".to_vec(),
            metadata,
        );

        // Should not be expired immediately
        assert!(!payload.should_delete_now());

        // Wait for expiration (in a real test, we'd mock the time)
        tokio::time::sleep(Duration::from_millis(1100)).await;

        let expired_count = engine.cleanup_expired_data(&[payload]);
        assert_eq!(expired_count, 1);
    }

    fn static_engine(key_id: &str, key: &str) -> EncryptionConfig {
        EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
            .with_key_id(key_id)
            .with_key_source(KeySource::Static(key.to_string()))
    }

    const KEY_A: &str = "dGVzdGtleTE5ODc2NTQzMjEwOTg3NjU0MzIxMHRlc3Q=";
    const KEY_B: &str = "QUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUE=";

    #[tokio::test]
    async fn test_decryption_key_for_previous_key_id() {
        let old = EncryptionEngine::new(static_engine("old", KEY_A))
            .await
            .unwrap();
        let payload = json!({"ssn": "123-45-6789"});
        let encrypted = old.encrypt_payload(&payload, &["ssn"]).await.unwrap();
        assert_eq!(encrypted.metadata.key_id, "old");

        let new = EncryptionEngine::new(static_engine("new", KEY_B))
            .await
            .unwrap();
        assert!(matches!(
            new.decrypt_payload(&encrypted).await,
            Err(EncryptionError::KeyManagement(_))
        ));

        let new = new
            .with_decryption_key("old", &KeySource::Static(KEY_A.to_string()))
            .await
            .unwrap();
        assert_eq!(new.decrypt_payload(&encrypted).await.unwrap(), payload);
        // New payloads still use the engine's own key
        let fresh = new.encrypt_payload(&payload, &["ssn"]).await.unwrap();
        assert_eq!(fresh.metadata.key_id, "new");
        assert!(old.decrypt_payload(&fresh).await.is_err());

        // The engine's own key id cannot be replaced
        assert!(matches!(
            new.with_decryption_key("new", &KeySource::Static(KEY_A.to_string()))
                .await,
            Err(EncryptionError::InvalidConfiguration(_))
        ));
    }

    #[tokio::test]
    async fn test_payload_hash_is_keyed() {
        let engine = EncryptionEngine::new(static_engine("k", KEY_A))
            .await
            .unwrap();
        let payload = json!({"ssn": "123-45-6789"});
        let encrypted = engine.encrypt_payload(&payload, &["ssn"]).await.unwrap();

        // Not the plain SHA-256 of the plaintext, which could be brute-forced
        let plain = engine.calculate_payload_hash(&serde_json::to_vec(&payload).unwrap());
        assert!(encrypted.metadata.payload_hash.starts_with("hmac-sha256:"));
        assert!(!encrypted.metadata.payload_hash.contains(&plain));

        // A different key gives a different hash
        let other = EncryptionEngine::new(static_engine("k", KEY_B))
            .await
            .unwrap();
        let other_encrypted = other.encrypt_payload(&payload, &["ssn"]).await.unwrap();
        assert_ne!(
            encrypted.metadata.payload_hash,
            other_encrypted.metadata.payload_hash
        );

        // Tampering with the hash is detected
        let mut tampered = encrypted.clone();
        tampered.metadata.payload_hash = format!("hmac-sha256:{}", "0".repeat(64));
        assert!(matches!(
            engine.decrypt_payload(&tampered).await,
            Err(EncryptionError::DecryptionFailed(_))
        ));
    }

    #[tokio::test]
    async fn test_legacy_unkeyed_hash_still_decrypts() {
        let engine = EncryptionEngine::new(static_engine("k", KEY_A))
            .await
            .unwrap();
        let payload = json!({"a": 1});
        let mut encrypted = engine.encrypt_payload(&payload, &["a"]).await.unwrap();
        encrypted.metadata.payload_hash =
            engine.calculate_payload_hash(&serde_json::to_vec(&payload).unwrap());
        assert_eq!(engine.decrypt_payload(&encrypted).await.unwrap(), payload);
    }

    #[tokio::test]
    async fn test_bad_nonce_or_tag_is_an_error_not_a_panic() {
        use base64::Engine as _;
        let engine = EncryptionEngine::new(static_engine("k", KEY_A))
            .await
            .unwrap();
        let encrypted = engine
            .encrypt_payload(&json!({"a": 1}), &["a"])
            .await
            .unwrap();

        let mut short_nonce = encrypted.clone();
        short_nonce.nonce = base64::engine::general_purpose::STANDARD.encode([0u8; 4]);
        assert!(engine.decrypt_payload(&short_nonce).await.is_err());

        let mut short_tag = encrypted.clone();
        short_tag.tag = base64::engine::general_purpose::STANDARD.encode([0u8; 3]);
        assert!(engine.decrypt_payload(&short_tag).await.is_err());
        assert_eq!(engine.get_stats().decryption_errors, 2);
    }

    #[tokio::test]
    async fn test_debug_does_not_print_keys() {
        let engine = EncryptionEngine::new(static_engine("k", KEY_A))
            .await
            .unwrap();
        let debug = format!("{:?}", engine);
        assert!(debug.contains("AES256GCM"));
        assert!(!debug.contains(KEY_A));
        assert!(!debug.contains("116, 101, 115, 116")); // "test" as key bytes
    }

    // ---- Fail-closed external key loading (#16); no cloud credentials needed ----

    #[tokio::test]
    async fn test_external_key_invalid_config_fails() {
        let cases = [
            "aws://?region=us-east-1",
            "vault://key-only?addr=http://127.0.0.1:1",
            "gcp://not-a-resource",
            "azure://",
        ];
        for source in cases {
            let err = load_external_key(source, 32).await.unwrap_err();
            assert!(
                matches!(err, EncryptionError::InvalidConfiguration(_)),
                "{source}: {err}"
            );
        }
        let err = load_external_key("ftp://key", 32).await.unwrap_err();
        assert!(matches!(err, EncryptionError::KeyManagement(_)), "{err}");
    }

    #[tokio::test]
    async fn test_external_key_without_feature_fails() {
        #[allow(unused_mut)]
        let mut cases: Vec<(&str, &str)> = Vec::new();
        #[cfg(not(feature = "aws-kms"))]
        cases.push(("aws://alias/key?region=us-east-1", "aws-kms"));
        #[cfg(not(feature = "vault-kms"))]
        cases.push(("vault://secret/key?addr=http://127.0.0.1:1", "vault-kms"));
        #[cfg(not(feature = "gcp-kms"))]
        cases.push((
            "gcp://projects/p/locations/global/keyRings/r/cryptoKeys/k",
            "gcp-kms",
        ));
        #[cfg(not(feature = "azure-kv"))]
        cases.push(("azure://my-vault.vault.azure.net/keys/k", "azure-kv"));

        for (source, feature) in cases {
            let err = load_external_key(source, 32).await.unwrap_err();
            assert!(
                matches!(err, EncryptionError::InvalidConfiguration(_)),
                "{source}: {err}"
            );
            assert!(err.to_string().contains(feature), "{source}: {err}");
        }
    }

    /// Without a database there is nowhere to keep the KMS-wrapped key, so `new` refuses
    /// `aws://` and `gcp://` sources instead of generating a key per process (#26).
    #[tokio::test]
    async fn test_kms_source_without_pool_fails() {
        #[allow(unused_mut)]
        let mut cases: Vec<&str> = Vec::new();
        #[cfg(feature = "aws-kms")]
        cases.push("aws://alias/key?region=us-east-1");
        #[cfg(feature = "gcp-kms")]
        cases.push("gcp://projects/p/locations/global/keyRings/r/cryptoKeys/k");
        for source in cases {
            let err = load_external_key(source, 32).await.unwrap_err();
            assert!(
                matches!(err, EncryptionError::InvalidConfiguration(_)),
                "{source}: {err}"
            );
            assert!(err.to_string().contains("new_with_pool"), "{source}: {err}");
        }
    }

    #[cfg(feature = "vault-kms")]
    #[tokio::test]
    async fn test_vault_key_unavailable_fails() {
        let result = load_from_vault("vault://secret/hammerwork?addr=http://127.0.0.1:1", 32).await;
        assert!(matches!(result, Err(EncryptionError::KeyManagement(_))));
    }

    // ---- KMS-wrapped data keys (#26), with a local mock KMS ----

    use crate::encryption::envelope::test_support::MockKms;

    fn same_kms(kms: &MockKms) -> Arc<dyn KmsKeyWrapper> {
        Arc::new(MockKms::new(kms.kms_key_id(), [0x5a; 32]))
    }

    fn active_key(engine: &EncryptionEngine, key_id: &str) -> Vec<u8> {
        engine.keys.lock().unwrap()[key_id].active.to_vec()
    }

    fn previous_keys(engine: &EncryptionEngine, key_id: &str) -> Vec<Vec<u8>> {
        engine
            .keys
            .lock()
            .unwrap()
            .get(key_id)
            .map(|versions| versions.previous.iter().map(|(_, k)| k.to_vec()).collect())
            .unwrap_or_default()
    }

    /// Two engines (two workers, or two processes) share one key store. After one of them
    /// rotates the data key, the other decrypts payloads encrypted with the new version
    /// by reloading the stored versions, and then encrypts with it too (#64 H4).
    #[tokio::test]
    async fn test_engines_sharing_a_store_survive_a_rotation() {
        use crate::encryption::envelope::test_support::MemoryStore;

        let kms = MockKms::unique();
        let store: Arc<dyn WrappedKeyStore> = Arc::new(MemoryStore::default());
        let config = || EncryptionConfig::new(EncryptionAlgorithm::AES256GCM);
        let payload = json!({"ssn": "123-45-6789"});

        let rotating = EncryptionEngine::with_kms_store(config(), store.clone(), same_kms(&kms))
            .await
            .unwrap();
        let other = EncryptionEngine::with_kms_store(config(), store.clone(), same_kms(&kms))
            .await
            .unwrap();
        let v1 = active_key(&other, "default");
        // Nothing new to load yet.
        assert!(!other.reload_kms_keys().await.unwrap());

        assert_eq!(rotating.rotate_kms_key_in(store.as_ref()).await.unwrap(), 2);
        let after_rotation = rotating.encrypt_payload(&payload, &["ssn"]).await.unwrap();

        // The other engine only held v1: it reloads v2 and decrypts.
        assert_eq!(
            other.decrypt_payload(&after_rotation).await.unwrap(),
            payload
        );
        let v2 = active_key(&rotating, "default");
        assert_ne!(v1, v2);
        assert_eq!(active_key(&other, "default"), v2, "now encrypts with v2");
        assert_eq!(previous_keys(&other, "default"), vec![v1.clone()]);
        assert_eq!(other.get_stats().decryption_errors, 0);

        // What it encrypts now is readable by an engine that only knows v2 and v1.
        let from_other = other.encrypt_payload(&payload, &["ssn"]).await.unwrap();
        assert_eq!(
            rotating.decrypt_payload(&from_other).await.unwrap(),
            payload
        );

        // A third rotation is picked up by an explicit reload.
        rotating.rotate_kms_key_in(store.as_ref()).await.unwrap();
        assert!(other.reload_kms_keys().await.unwrap());
        assert_eq!(
            active_key(&other, "default"),
            active_key(&rotating, "default")
        );
        assert_eq!(previous_keys(&other, "default").len(), 2);

        // A payload under another key id stored in the same table is looked up by the
        // key id it records.
        let other_id = EncryptionEngine::with_kms_store(
            config().with_key_id("tenant-b"),
            store.clone(),
            same_kms(&kms),
        )
        .await
        .unwrap();
        let tenant_b = other_id.encrypt_payload(&payload, &["ssn"]).await.unwrap();
        assert_eq!(other.decrypt_payload(&tenant_b).await.unwrap(), payload);

        // A tampered payload still fails, after one reload attempt that finds nothing.
        let mut tampered = rotating.encrypt_payload(&payload, &["ssn"]).await.unwrap();
        tampered.metadata.payload_hash = "hmac-sha256:00".to_string();
        assert!(other.decrypt_payload(&tampered).await.is_err());

        // Engines without a KMS have nothing to reload.
        let static_engine = EncryptionEngine::new(static_engine("k", KEY_A))
            .await
            .unwrap();
        assert!(!static_engine.reload_kms_keys().await.unwrap());
    }

    /// Engines in different processes (or after a restart) load the same stored data key,
    /// and keep earlier versions after a rotation.
    async fn engine_kms_roundtrip<DB: KeyManagerBackend>(pool: Pool<DB>) {
        let kms = MockKms::unique();
        let config = || EncryptionConfig::new(EncryptionAlgorithm::AES256GCM);

        let first = EncryptionEngine::with_kms(config(), &pool, same_kms(&kms))
            .await
            .unwrap();
        let v1 = active_key(&first, "default");
        assert_eq!(v1.len(), 32);

        // Restart: the same key is recovered from the stored KMS-wrapped blob
        let second = EncryptionEngine::with_kms(config(), &pool, same_kms(&kms))
            .await
            .unwrap();
        assert_eq!(active_key(&second, "default"), v1);
        assert!(previous_keys(&second, "default").is_empty());

        assert_eq!(second.rotate_kms_key(&pool).await.unwrap(), 2);
        let v2 = active_key(&second, "default");
        assert_ne!(v1, v2);
        assert_eq!(previous_keys(&second, "default"), vec![v1.clone()]);

        let third = EncryptionEngine::with_kms(config(), &pool, same_kms(&kms))
            .await
            .unwrap();
        assert_eq!(active_key(&third, "default"), v2);
        assert_eq!(previous_keys(&third, "default"), vec![v1.clone()]);

        // Another key id is a separate key
        let other =
            EncryptionEngine::with_kms(config().with_key_id("other"), &pool, same_kms(&kms))
                .await
                .unwrap();
        assert_ne!(active_key(&other, "other"), v1);
        assert_ne!(active_key(&other, "other"), v2);

        let stored = DB::load_kms_wrapped_keys(&pool, "engine/default", "aws", kms.kms_key_id())
            .await
            .unwrap();
        assert_eq!(stored.len(), 2);

        // rotate_kms_key needs a KMS source
        let static_engine = EncryptionEngine::new_with_pool(
            config().with_key_source(KeySource::Static(
                "dGVzdGtleTE5ODc2NTQzMjEwOTg3NjU0MzIxMHRlc3Q=".to_string(),
            )),
            &pool,
        )
        .await
        .unwrap();
        assert!(matches!(
            static_engine.rotate_kms_key(&pool).await,
            Err(EncryptionError::InvalidConfiguration(_))
        ));
    }

    /// Engines starting at the same time converge on one stored data key.
    async fn engine_kms_concurrent_first_load<DB: KeyManagerBackend>(pool: Pool<DB>) {
        let kms = MockKms::unique();
        let tasks: Vec<_> = (0..6)
            .map(|_| {
                let pool = pool.clone();
                let kms = same_kms(&kms);
                tokio::spawn(async move {
                    let config = EncryptionConfig::new(EncryptionAlgorithm::AES256GCM);
                    let engine = EncryptionEngine::with_kms(config, &pool, kms)
                        .await
                        .unwrap();
                    active_key(&engine, "default")
                })
            })
            .collect();
        let mut keys = Vec::new();
        for task in tasks {
            keys.push(task.await.unwrap());
        }
        assert!(keys.windows(2).all(|w| w[0] == w[1]));
        let stored = DB::load_kms_wrapped_keys(&pool, "engine/default", "aws", kms.kms_key_id())
            .await
            .unwrap();
        assert_eq!(stored.len(), 1);
    }

    /// Two processes share the KMS-wrapped key in the database. One rotates the data key
    /// and enqueues; a worker of the other, started before the rotation, still runs the
    /// job with its plaintext instead of failing it until it goes Dead (#64 H4).
    async fn worker_survives_a_rotation_by_another_process<DB>(pool: Pool<DB>)
    where
        DB: KeyManagerBackend + Send + Sync + 'static,
        crate::JobQueue<DB>: crate::queue::DatabaseQueue<Database = DB> + Send + Sync + 'static,
    {
        use crate::queue::DatabaseQueue;
        use crate::{Job, JobQueue, JobStatus, Worker};

        let kms = MockKms::unique();
        let config = || EncryptionConfig::new(EncryptionAlgorithm::AES256GCM);
        let producer_engine = EncryptionEngine::with_kms(config(), &pool, same_kms(&kms))
            .await
            .unwrap();
        let worker_engine = EncryptionEngine::with_kms(config(), &pool, same_kms(&kms))
            .await
            .unwrap();
        let queue_name = format!("kms_rotation_{}", uuid::Uuid::new_v4().simple());

        // The producer rotates after the worker loaded version 1
        producer_engine.rotate_kms_key(&pool).await.unwrap();
        let producer = JobQueue::new(pool.clone())
            .with_encryption(producer_engine)
            .with_encrypted_queues([queue_name.clone()]);
        let payload = json!({"ssn": "123-45-6789"});
        let id = producer
            .enqueue(Job::new(queue_name.clone(), payload.clone()).with_max_attempts(1))
            .await
            .unwrap();

        let consumer = Arc::new(JobQueue::new(pool.clone()).with_encryption(worker_engine));
        let seen = Arc::new(Mutex::new(Vec::new()));
        let handler_seen = Arc::clone(&seen);
        let handler: crate::worker::JobHandler = Arc::new(move |job: Job| {
            let seen = Arc::clone(&handler_seen);
            Box::pin(async move {
                seen.lock().unwrap().push(job.payload.clone());
                Ok(())
            })
        });
        let worker = Worker::new(Arc::clone(&consumer), queue_name, handler)
            .with_poll_interval(Duration::from_millis(50));
        let (shutdown_tx, shutdown_rx) = tokio::sync::mpsc::channel(1);
        let task = tokio::spawn(async move { worker.run(shutdown_rx).await });

        let deadline = std::time::Instant::now() + Duration::from_secs(15);
        let job = loop {
            let job = consumer.get_job(id).await.unwrap().unwrap();
            if !matches!(job.status, JobStatus::Pending | JobStatus::Running)
                || std::time::Instant::now() > deadline
            {
                break job;
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
        };
        shutdown_tx.send(()).await.unwrap();
        task.await.unwrap().unwrap();

        assert_eq!(job.status, JobStatus::Completed, "{:?}", job.error_message);
        assert_eq!(*seen.lock().unwrap(), vec![payload]);
        consumer.delete_job(id).await.unwrap();
    }

    #[cfg(feature = "postgres")]
    #[tokio::test]
    #[ignore = "requires PostgreSQL: DATABASE_URL"]
    async fn test_worker_survives_a_rotation_by_another_process_postgres() {
        worker_survives_a_rotation_by_another_process(postgres_pool().await).await;
    }

    #[cfg(feature = "mysql")]
    #[tokio::test]
    #[ignore = "requires MySQL: MYSQL_DATABASE_URL"]
    async fn test_worker_survives_a_rotation_by_another_process_mysql() {
        worker_survives_a_rotation_by_another_process(mysql_pool().await).await;
    }

    /// Payloads encrypted before a restart and before a rotation can be decrypted.
    async fn engine_kms_payload_roundtrip<DB: KeyManagerBackend>(pool: Pool<DB>) {
        let kms = MockKms::unique();
        let config = || EncryptionConfig::new(EncryptionAlgorithm::AES256GCM);
        let payload = json!({"ssn": "123-45-6789", "amount": 10});

        let first = EncryptionEngine::with_kms(config(), &pool, same_kms(&kms))
            .await
            .unwrap();
        let before_rotation = first.encrypt_payload(&payload, &["ssn"]).await.unwrap();

        let second = EncryptionEngine::with_kms(config(), &pool, same_kms(&kms))
            .await
            .unwrap();
        assert_eq!(
            second.decrypt_payload(&before_rotation).await.unwrap(),
            payload
        );
        second.rotate_kms_key(&pool).await.unwrap();
        let after_rotation = second.encrypt_payload(&payload, &["ssn"]).await.unwrap();
        assert_eq!(
            second.decrypt_payload(&before_rotation).await.unwrap(),
            payload
        );

        let third = EncryptionEngine::with_kms(config(), &pool, same_kms(&kms))
            .await
            .unwrap();
        assert_eq!(
            third.decrypt_payload(&before_rotation).await.unwrap(),
            payload
        );
        assert_eq!(
            third.decrypt_payload(&after_rotation).await.unwrap(),
            payload
        );
        // The engine that did not rotate loads the new version from the database (#64 H4)
        assert_eq!(
            first.decrypt_payload(&after_rotation).await.unwrap(),
            payload
        );
        assert_eq!(
            active_key(&first, "default"),
            active_key(&second, "default")
        );
    }

    #[cfg(feature = "postgres")]
    async fn postgres_pool() -> sqlx::PgPool {
        use crate::migrations::{MigrationManager, postgres::PostgresMigrationRunner};
        let url = std::env::var("DATABASE_URL")
            .expect("DATABASE_URL must point at a PostgreSQL test database");
        let pool = sqlx::PgPool::connect(&url).await.unwrap();
        MigrationManager::new(Box::new(PostgresMigrationRunner::new(pool.clone())))
            .run_migrations()
            .await
            .unwrap();
        pool
    }

    #[cfg(feature = "mysql")]
    async fn mysql_pool() -> sqlx::MySqlPool {
        use crate::migrations::{MigrationManager, mysql::MySqlMigrationRunner};
        let url = std::env::var("MYSQL_DATABASE_URL")
            .expect("MYSQL_DATABASE_URL must point at a MySQL test database");
        let pool = sqlx::MySqlPool::connect(&url).await.unwrap();
        MigrationManager::new(Box::new(MySqlMigrationRunner::new(pool.clone())))
            .run_migrations()
            .await
            .unwrap();
        pool
    }

    #[cfg(feature = "postgres")]
    #[tokio::test]
    #[ignore = "requires PostgreSQL: DATABASE_URL"]
    async fn test_engine_kms_payload_roundtrip_postgres() {
        engine_kms_payload_roundtrip(postgres_pool().await).await;
    }

    #[cfg(feature = "mysql")]
    #[tokio::test]
    #[ignore = "requires MySQL: MYSQL_DATABASE_URL"]
    async fn test_engine_kms_payload_roundtrip_mysql() {
        engine_kms_payload_roundtrip(mysql_pool().await).await;
    }

    #[cfg(feature = "postgres")]
    #[tokio::test]
    #[ignore = "requires PostgreSQL: DATABASE_URL"]
    async fn test_engine_kms_roundtrip_postgres() {
        engine_kms_roundtrip(postgres_pool().await).await;
    }

    #[cfg(feature = "mysql")]
    #[tokio::test]
    #[ignore = "requires MySQL: MYSQL_DATABASE_URL"]
    async fn test_engine_kms_roundtrip_mysql() {
        engine_kms_roundtrip(mysql_pool().await).await;
    }

    #[cfg(feature = "postgres")]
    #[tokio::test]
    #[ignore = "requires PostgreSQL: DATABASE_URL"]
    async fn test_engine_kms_concurrent_first_load_postgres() {
        engine_kms_concurrent_first_load(postgres_pool().await).await;
    }

    #[cfg(feature = "mysql")]
    #[tokio::test]
    #[ignore = "requires MySQL: MYSQL_DATABASE_URL"]
    async fn test_engine_kms_concurrent_first_load_mysql() {
        engine_kms_concurrent_first_load(mysql_pool().await).await;
    }

    // ---- Key sources ----

    fn generated(location: &str) -> EncryptionConfig {
        EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
            .with_key_source(KeySource::Generated(location.to_string()))
    }

    #[tokio::test]
    async fn test_generated_key_is_stored_and_reused() {
        use base64::Engine as _;

        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("keys/engine.key");
        let location = format!("file://{}", path.display());

        let first = EncryptionEngine::new(generated(&location)).await.unwrap();
        let stored = std::fs::read_to_string(&path).unwrap();
        assert_eq!(
            base64::engine::general_purpose::STANDARD
                .decode(&stored)
                .unwrap()
                .len(),
            32
        );
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let mode = std::fs::metadata(&path).unwrap().permissions().mode();
            assert_eq!(mode & 0o777, 0o600, "the key file is private");
        }
        let encrypted = first
            .encrypt_payload(&json!({"secret": 1}), &["secret"])
            .await
            .unwrap();

        // A restart with the same configuration reuses the stored key: what was
        // encrypted before can still be decrypted, and the file is unchanged.
        let restarted = EncryptionEngine::new(generated(&location)).await.unwrap();
        assert_eq!(
            restarted.decrypt_payload(&encrypted).await.unwrap(),
            json!({"secret": 1})
        );
        assert_eq!(std::fs::read_to_string(&path).unwrap(), stored);

        // A plain path works the same way.
        let plain = EncryptionEngine::new(generated(path.to_str().unwrap()))
            .await
            .unwrap();
        assert!(plain.decrypt_payload(&encrypted).await.is_ok());
    }

    #[tokio::test]
    async fn test_generated_key_file_that_is_not_a_key_is_an_error() {
        let dir = tempfile::tempdir().unwrap();
        for (contents, expected) in [
            ("not base64!", "Invalid base64 key"),
            ("c2hvcnQ=", "expected 32 bytes, got 5"),
        ] {
            let path = dir.path().join(format!("{}.key", contents.len()));
            std::fs::write(&path, contents).unwrap();
            let err = EncryptionEngine::new(generated(path.to_str().unwrap()))
                .await
                .unwrap_err();
            assert!(err.to_string().contains(expected), "{err}");
            assert_eq!(
                std::fs::read_to_string(&path).unwrap(),
                contents,
                "an invalid stored key is never overwritten"
            );
        }
    }

    #[tokio::test]
    async fn test_generated_key_storage_failure_is_an_error() {
        let dir = tempfile::tempdir().unwrap();
        // The parent "directory" is a file, so the key cannot be stored.
        let blocker = dir.path().join("blocker");
        std::fs::write(&blocker, "x").unwrap();
        let err = EncryptionEngine::new(generated(&format!(
            "file://{}/engine.key",
            blocker.display()
        )))
        .await
        .unwrap_err();
        assert!(
            err.to_string().contains("Failed to store generated key"),
            "{err}"
        );
    }

    #[tokio::test]
    async fn test_generated_key_in_environment_and_stdout() {
        let var = format!(
            "HAMMERWORK_TEST_GENERATED_{}",
            uuid::Uuid::new_v4().simple()
        );
        let location = format!("env://{var}");
        let first = EncryptionEngine::new(generated(&location)).await.unwrap();
        assert!(std::env::var(&var).is_ok(), "the key is stored in {var}");
        let encrypted = first
            .encrypt_payload(&json!({"a": 1}), &[] as &[&str])
            .await
            .unwrap();
        let second = EncryptionEngine::new(generated(&location)).await.unwrap();
        assert!(second.decrypt_payload(&encrypted).await.is_ok());

        // stdout:// prints a new key every time.
        let printed = EncryptionEngine::new(generated("stdout://")).await.unwrap();
        assert!(printed.decrypt_payload(&encrypted).await.is_err());
    }

    #[tokio::test]
    async fn test_static_and_environment_key_errors() {
        let err = EncryptionEngine::new(static_engine("k", "not base64!"))
            .await
            .unwrap_err();
        assert!(err.to_string().contains("Invalid base64 key"), "{err}");
        let err = EncryptionEngine::new(static_engine("k", "c2hvcnQ="))
            .await
            .unwrap_err();
        assert!(
            err.to_string().contains("expected 32 bytes, got 5"),
            "{err}"
        );

        let missing = format!("HAMMERWORK_TEST_MISSING_{}", uuid::Uuid::new_v4().simple());
        let err = EncryptionEngine::new(
            EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
                .with_key_source(KeySource::Environment(missing.clone())),
        )
        .await
        .unwrap_err();
        assert!(err.to_string().contains(&missing), "{err}");

        for (value, expected) in [
            ("%%%", "Invalid base64 key in"),
            ("c2hvcnQ=", "Key size mismatch in"),
        ] {
            let var = format!("HAMMERWORK_TEST_BAD_{}", uuid::Uuid::new_v4().simple());
            unsafe {
                std::env::set_var(&var, value);
            }
            let err = EncryptionEngine::new(
                EncryptionConfig::new(EncryptionAlgorithm::AES256GCM)
                    .with_key_source(KeySource::Environment(var)),
            )
            .await
            .unwrap_err();
            assert!(err.to_string().contains(expected), "{err}");
        }
    }

    // ---- Algorithms, tampering, compression ----

    #[tokio::test]
    async fn test_chacha20_roundtrip_and_tampering() {
        use base64::Engine as _;

        let config = EncryptionConfig::new(EncryptionAlgorithm::ChaCha20Poly1305)
            .with_key_id("chacha")
            .with_key_source(KeySource::Static(KEY_B.to_string()));
        let engine = EncryptionEngine::new(config).await.unwrap();
        assert_eq!(engine.algorithm(), &EncryptionAlgorithm::ChaCha20Poly1305);
        let payload = json!({"card": "4111", "amount": 5});
        let encrypted = engine.encrypt_payload(&payload, &["card"]).await.unwrap();
        assert_eq!(encrypted.decode_nonce().unwrap().len(), 12);
        assert_eq!(encrypted.decode_tag().unwrap().len(), 16);
        assert_eq!(engine.decrypt_payload(&encrypted).await.unwrap(), payload);

        let mut tampered = encrypted.clone();
        let mut ciphertext = encrypted.decode_ciphertext().unwrap();
        ciphertext[0] ^= 0xff;
        tampered.ciphertext = base64::engine::general_purpose::STANDARD.encode(ciphertext);
        let err = engine.decrypt_payload(&tampered).await.unwrap_err();
        assert!(
            err.to_string().contains("ChaCha20 decryption failed"),
            "{err}"
        );

        // Decryption follows the algorithm recorded with the payload, so an engine
        // configured for AES with the same key can still read it; another key cannot.
        let aes = EncryptionEngine::new(static_engine("chacha", KEY_B))
            .await
            .unwrap();
        assert_eq!(aes.decrypt_payload(&encrypted).await.unwrap(), payload);
        let other_key = EncryptionEngine::new(
            EncryptionConfig::new(EncryptionAlgorithm::ChaCha20Poly1305)
                .with_key_id("chacha")
                .with_key_source(KeySource::Static(KEY_A.to_string())),
        )
        .await
        .unwrap();
        assert!(other_key.decrypt_payload(&encrypted).await.is_err());
    }

    #[tokio::test]
    async fn test_associated_data_must_match() {
        let engine = EncryptionEngine::new(static_engine("k", KEY_A))
            .await
            .unwrap();
        let payload = json!({"a": 1});
        let encrypted = engine
            .encrypt_payload_with_associated_data(
                &payload,
                &[] as &[&str],
                RetentionPolicy::UseDefault,
                b"job-1",
            )
            .await
            .unwrap();
        assert_eq!(
            engine
                .decrypt_payload_with_associated_data(&encrypted, b"job-1")
                .await
                .unwrap(),
            payload
        );
        assert!(
            engine
                .decrypt_payload_with_associated_data(&encrypted, b"job-2")
                .await
                .is_err()
        );
    }

    #[tokio::test]
    async fn test_compression_roundtrip_and_corruption() {
        let engine =
            EncryptionEngine::new(static_engine("k", KEY_A).with_compression_enabled(true))
                .await
                .unwrap();
        let payload = json!({"text": "a".repeat(10_000)});
        let encrypted = engine
            .encrypt_payload(&payload, &[] as &[&str])
            .await
            .unwrap();
        assert!(encrypted.metadata.compressed);
        assert!(
            encrypted.decode_ciphertext().unwrap().len() < 1_000,
            "compressible payloads shrink"
        );
        assert_eq!(engine.decrypt_payload(&encrypted).await.unwrap(), payload);

        // Valid ciphertext of something that is not gzip data.
        let plain = EncryptionEngine::new(static_engine("k", KEY_A))
            .await
            .unwrap();
        let mut not_gzip = plain
            .encrypt_payload(&payload, &[] as &[&str])
            .await
            .unwrap();
        not_gzip.metadata.compressed = true;
        let err = engine.decrypt_payload(&not_gzip).await.unwrap_err();
        assert!(err.to_string().contains("Decompression failed"), "{err}");
    }

    // ---- PII detection, retention, stats ----

    #[tokio::test]
    async fn test_identify_pii_fields_nested_and_in_arrays() {
        let engine = EncryptionEngine::new(static_engine("k", KEY_A))
            .await
            .unwrap();
        let payload = json!({
            "user_id": 1,
            "Email": "a@example.com",
            "creditCardNumber": "4111",
            "billing": {"home_address": "1 Main St", "zip_code": "12345"},
            "contacts": [{"phone": "555"}, "not an object", {"name": "x"}],
            "items": [1, 2, 3]
        });
        let mut fields = engine.identify_pii_fields(&payload);
        fields.sort();
        assert_eq!(
            fields,
            vec![
                "Email",
                "billing.home_address",
                "contacts[0].phone",
                "creditCardNumber",
            ]
        );
        assert!(engine.identify_pii_fields(&json!([1, 2])).is_empty());
        assert!(engine.identify_pii_fields(&json!("ssn")).is_empty());
    }

    #[tokio::test]
    async fn test_cleanup_expired_data_counts_and_records() {
        let engine = EncryptionEngine::new(static_engine("k", KEY_A))
            .await
            .unwrap();
        let expired = engine
            .encrypt_payload_with_retention(
                &json!({"a": 1}),
                &[] as &[&str],
                RetentionPolicy::DeleteAt(chrono::Utc::now() - chrono::Duration::hours(1)),
            )
            .await
            .unwrap();
        let kept = engine
            .encrypt_payload_with_retention(
                &json!({"a": 1}),
                &[] as &[&str],
                RetentionPolicy::KeepIndefinitely,
            )
            .await
            .unwrap();
        assert_eq!(
            engine.cleanup_expired_data(&[expired.clone(), kept.clone()]),
            1
        );
        assert_eq!(engine.cleanup_expired_data(&[kept]), 0);
        assert_eq!(engine.get_stats().retention_cleanups, 1);
    }

    #[test]
    fn test_vault_key_must_be_exactly_the_key_size() {
        use base64::Engine as _;
        let b64 = |bytes: &[u8]| base64::engine::general_purpose::STANDARD.encode(bytes);

        // The right size is used as is.
        let exact: Vec<u8> = (0..32).collect();
        assert_eq!(
            kms::exact_vault_key(&b64(&exact), 32).unwrap().as_slice(),
            exact.as_slice()
        );
        // Longer or shorter material is rejected, not truncated or padded.
        for wrong in [(0..40).collect::<Vec<u8>>(), b"short".to_vec(), vec![7]] {
            let err = kms::exact_vault_key(&b64(&wrong), 32).unwrap_err();
            assert!(
                matches!(err, EncryptionError::InvalidConfiguration(_)),
                "{err}"
            );
            assert!(
                err.to_string().contains("expected exactly 32 bytes"),
                "{err}"
            );
        }
        // A passphrase (not base64) is rejected, not hashed into a key.
        let err = kms::exact_vault_key("a passphrase!", 32).unwrap_err();
        assert!(
            matches!(err, EncryptionError::InvalidConfiguration(_)),
            "{err}"
        );
        assert!(err.to_string().contains("not base64"), "{err}");
    }

    #[test]
    fn test_azure_key_material_must_be_exactly_the_key_size() {
        let exact: Vec<u8> = (0..32).collect();
        assert_eq!(
            kms::exact_size(SecretBytes::new(exact.clone()), 32, "Azure key")
                .unwrap()
                .as_slice(),
            exact.as_slice()
        );
        for wrong in [(0..40).collect::<Vec<u8>>(), b"0123456789".to_vec()] {
            let err = kms::exact_size(SecretBytes::new(wrong), 32, "Azure key").unwrap_err();
            assert!(
                matches!(err, EncryptionError::InvalidConfiguration(_)),
                "{err}"
            );
            assert!(
                err.to_string()
                    .starts_with("Invalid configuration: Azure key"),
                "{err}"
            );
        }
    }

    #[cfg(feature = "azure-kv")]
    #[tokio::test]
    async fn test_azure_key_with_invalid_vault_fails() {
        let err = load_from_azure_kv("azure://", 32).await.unwrap_err();
        assert!(
            matches!(err, EncryptionError::InvalidConfiguration(_)),
            "{err}"
        );
        // A vault host that is not a valid URL fails before any request is made.
        let err = load_from_azure_kv("azure://bad host/keys/k", 32)
            .await
            .unwrap_err();
        assert!(
            err.to_string()
                .contains("Failed to load key from Azure Key Vault"),
            "{err}"
        );
    }
}