hammerwork 2.3.0

//! Envelope encryption for keys that come from AWS KMS (`aws://`) and GCP KMS (`gcp://`).
//!
//! A KMS never hands out the same plaintext key twice, so a key source backed by one
//! cannot simply ask the KMS for a key on every load: each process would get a different
//! key. Instead, on first use a data key is generated with the KMS and only its
//! KMS-encrypted form (AWS `CiphertextBlob`, GCP `Encrypt` ciphertext) is stored in
//! `hammerwork_kms_data_keys`. Later loads read the stored blob and call KMS `Decrypt`, so
//! every process and every restart recovers the same plaintext key. Plaintext key
//! material is never written to the database.
//!
//! Stored keys are scoped by `(key_name, kms_provider, kms_key_id)` and versioned:
//! [`rotate`] adds a new active version and retires the previous one, which can still be
//! loaded (and decrypted by the KMS) to read data written under it.
//!
//! The KMS-encrypted blob is bound to `key_name` (AWS encryption context, GCP additional
//! authenticated data), so a blob stored for one purpose cannot be decrypted as another.

use super::key_manager::{KeyManagerBackend, KmsWrappedKey};
use super::{EncryptionError, KeySource, SecretBytes, kms};
use chrono::Utc;
use sqlx::Pool;
use std::sync::Arc;
use uuid::Uuid;

/// A data key generated by a KMS: the plaintext and its KMS-encrypted form.
pub(crate) struct DataKey {
    pub plaintext: SecretBytes,
    pub ciphertext: Vec<u8>,
}

/// The KMS operations envelope encryption needs. Implemented for AWS KMS and GCP KMS, and
/// by a local mock in tests.
#[async_trait::async_trait]
pub(crate) trait KmsKeyWrapper: Send + Sync {
    /// `aws` or `gcp` (stored in `kms_provider`).
    fn provider(&self) -> &'static str;

    /// The KMS key that wraps data keys (stored in `kms_key_id`).
    fn kms_key_id(&self) -> &str;

    /// Generate a `size`-byte data key and return it with its KMS-encrypted form, bound
    /// to `context`.
    async fn generate_data_key(
        &self,
        size: usize,
        context: &str,
    ) -> Result<DataKey, EncryptionError>;

    /// Decrypt a data key produced by [`KmsKeyWrapper::generate_data_key`] with the same
    /// `context`.
    async fn decrypt(
        &self,
        ciphertext: &[u8],
        context: &str,
    ) -> Result<SecretBytes, EncryptionError>;
}

/// A parsed `aws://` or `gcp://` key source.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum KmsSource {
    Aws {
        key_id: String,
        region: String,
        endpoint: Option<String>,
    },
    Gcp {
        resource: String,
        location: String,
    },
}

/// Parse a key source that needs envelope encryption.
///
/// Returns `Ok(None)` for sources that do not (environment, static, generated, Vault,
/// Azure). Fails for an invalid `aws://` / `gcp://` configuration and for one whose cargo
/// feature (`aws-kms`, `gcp-kms`) is not enabled. Does not contact the KMS.
pub(crate) fn parse_kms_source(source: &KeySource) -> Result<Option<KmsSource>, EncryptionError> {
    let KeySource::External(service_config) = source else {
        return Ok(None);
    };
    if service_config.starts_with("aws://") {
        let (key_id, region, endpoint) = kms::aws_key_config(service_config)?;
        if !cfg!(feature = "aws-kms") {
            return Err(kms::kms_feature_disabled("aws://", "aws-kms"));
        }
        Ok(Some(KmsSource::Aws {
            key_id: key_id.to_string(),
            region: region.to_string(),
            endpoint: endpoint.map(str::to_string),
        }))
    } else if service_config.starts_with("gcp://") {
        let (resource, _) = kms::parse_service_config(service_config, "gcp://");
        let location = kms::gcp_location(resource)?;
        if !cfg!(feature = "gcp-kms") {
            return Err(kms::kms_feature_disabled("gcp://", "gcp-kms"));
        }
        Ok(Some(KmsSource::Gcp {
            resource: resource.to_string(),
            location,
        }))
    } else {
        Ok(None)
    }
}

impl KmsSource {
    /// Create a client for this KMS.
    pub(crate) async fn connect(&self) -> Result<Arc<dyn KmsKeyWrapper>, EncryptionError> {
        match self {
            KmsSource::Aws {
                key_id,
                region,
                endpoint,
            } => {
                #[cfg(feature = "aws-kms")]
                {
                    Ok(Arc::new(
                        aws::AwsKms::connect(key_id, region, endpoint.as_deref()).await,
                    ))
                }
                #[cfg(not(feature = "aws-kms"))]
                {
                    let _ = (key_id, region, endpoint);
                    Err(kms::kms_feature_disabled("aws://", "aws-kms"))
                }
            }
            KmsSource::Gcp { resource, location } => {
                #[cfg(feature = "gcp-kms")]
                {
                    Ok(Arc::new(gcp::GcpKms::connect(resource, location).await?))
                }
                #[cfg(not(feature = "gcp-kms"))]
                {
                    let _ = (resource, location);
                    Err(kms::kms_feature_disabled("gcp://", "gcp-kms"))
                }
            }
        }
    }
}

/// Where KMS-wrapped keys are stored.
#[async_trait::async_trait]
pub(crate) trait WrappedKeyStore: Send + Sync {
    /// Insert `key` unless a key with the same scope and version already exists.
    async fn insert_if_absent(&self, key: &KmsWrappedKey) -> Result<(), EncryptionError>;

    /// Retire the active versions in `key`'s scope and insert `key`, atomically. Fails if
    /// the version already exists (a concurrent rotation won).
    async fn insert_rotated(&self, key: &KmsWrappedKey) -> Result<(), EncryptionError>;

    /// Every version in a scope, newest first.
    async fn load_versions(
        &self,
        key_name: &str,
        kms_provider: &str,
        kms_key_id: &str,
    ) -> Result<Vec<KmsWrappedKey>, EncryptionError>;
}

/// [`WrappedKeyStore`] backed by `hammerwork_kms_data_keys`.
pub(crate) struct DbStore<'a, DB: sqlx::Database>(pub &'a Pool<DB>);

#[async_trait::async_trait]
impl<DB: KeyManagerBackend> WrappedKeyStore for DbStore<'_, DB> {
    async fn insert_if_absent(&self, key: &KmsWrappedKey) -> Result<(), EncryptionError> {
        DB::insert_kms_wrapped_key_if_absent(self.0, key).await
    }

    async fn insert_rotated(&self, key: &KmsWrappedKey) -> Result<(), EncryptionError> {
        DB::insert_rotated_kms_wrapped_key(self.0, key).await
    }

    async fn load_versions(
        &self,
        key_name: &str,
        kms_provider: &str,
        kms_key_id: &str,
    ) -> Result<Vec<KmsWrappedKey>, EncryptionError> {
        DB::load_kms_wrapped_keys(self.0, key_name, kms_provider, kms_key_id).await
    }
}

/// [`WrappedKeyStore`] backed by `hammerwork_kms_data_keys`, owning its pool, for an
/// engine that reloads keys later (after another process rotated them).
pub(crate) struct PoolStore<DB: sqlx::Database>(pub Pool<DB>);

#[async_trait::async_trait]
impl<DB: KeyManagerBackend> WrappedKeyStore for PoolStore<DB> {
    async fn insert_if_absent(&self, key: &KmsWrappedKey) -> Result<(), EncryptionError> {
        DbStore(&self.0).insert_if_absent(key).await
    }

    async fn insert_rotated(&self, key: &KmsWrappedKey) -> Result<(), EncryptionError> {
        DbStore(&self.0).insert_rotated(key).await
    }

    async fn load_versions(
        &self,
        key_name: &str,
        kms_provider: &str,
        kms_key_id: &str,
    ) -> Result<Vec<KmsWrappedKey>, EncryptionError> {
        DbStore(&self.0)
            .load_versions(key_name, kms_provider, kms_key_id)
            .await
    }
}

/// A plaintext key recovered from the KMS, with its stored version.
pub(crate) struct EnvelopeKey {
    pub version: u32,
    pub key: SecretBytes,
    /// Whether this is the active version (the one new data is encrypted with)
    pub active: bool,
}

fn new_record(
    kms: &dyn KmsKeyWrapper,
    key_name: &str,
    version: u32,
    size: usize,
    wrapped_key: Vec<u8>,
) -> Result<KmsWrappedKey, EncryptionError> {
    Ok(KmsWrappedKey {
        id: Uuid::new_v4(),
        key_name: key_name.to_string(),
        kms_provider: kms.provider().to_string(),
        kms_key_id: kms.kms_key_id().to_string(),
        version,
        key_size: u32::try_from(size).map_err(|_| {
            EncryptionError::KeyManagement(format!("Unsupported key size {}", size))
        })?,
        wrapped_key,
        active: true,
        created_at: Utc::now(),
        retired_at: None,
    })
}

/// Decrypt a stored key with the KMS and check its size.
async fn unwrap(
    kms: &dyn KmsKeyWrapper,
    record: &KmsWrappedKey,
    size: usize,
) -> Result<EnvelopeKey, EncryptionError> {
    if record.key_size as usize != size {
        return Err(EncryptionError::KeyManagement(format!(
            "Stored KMS data key {} version {} is {} bytes, expected {}",
            record.key_name, record.version, record.key_size, size
        )));
    }
    let key = kms
        .decrypt(&record.wrapped_key, &record.key_name)
        .await
        .map_err(|e| {
            EncryptionError::KeyManagement(format!(
                "Failed to decrypt stored data key {} version {} with {} KMS key {}: {}",
                record.key_name, record.version, record.kms_provider, record.kms_key_id, e
            ))
        })?;
    if key.len() != size {
        return Err(EncryptionError::KeyManagement(format!(
            "KMS returned a {}-byte key for {} version {}, expected {}",
            key.len(),
            record.key_name,
            record.version,
            size
        )));
    }
    Ok(EnvelopeKey {
        version: record.version,
        key,
        active: record.active,
    })
}

async fn versions(
    store: &dyn WrappedKeyStore,
    kms: &dyn KmsKeyWrapper,
    key_name: &str,
) -> Result<Vec<KmsWrappedKey>, EncryptionError> {
    store
        .load_versions(key_name, kms.provider(), kms.kms_key_id())
        .await
}

fn active(versions: &[KmsWrappedKey]) -> Option<&KmsWrappedKey> {
    versions.iter().find(|key| key.active)
}

/// Load the active key named `key_name`, generating and storing it on first use.
///
/// When several processes start at once, each may generate a data key, but only the first
/// insert is kept (version 1 is unique per scope); every process then reads back and
/// decrypts the stored key, so all of them converge on the same key.
pub(crate) async fn load_or_create(
    store: &dyn WrappedKeyStore,
    kms: &dyn KmsKeyWrapper,
    key_name: &str,
    size: usize,
) -> Result<EnvelopeKey, EncryptionError> {
    let existing = versions(store, kms, key_name).await?;
    if let Some(record) = active(&existing) {
        return unwrap(kms, record, size).await;
    }
    if !existing.is_empty() {
        return Err(EncryptionError::KeyManagement(format!(
            "No active KMS data key for {} ({} key {}); refusing to generate a new one \
             while retired versions exist",
            key_name,
            kms.provider(),
            kms.kms_key_id()
        )));
    }

    let data_key = kms.generate_data_key(size, key_name).await?;
    if data_key.plaintext.len() != size {
        return Err(EncryptionError::KeyManagement(format!(
            "KMS generated a {}-byte key, expected {}",
            data_key.plaintext.len(),
            size
        )));
    }
    let record = new_record(kms, key_name, 1, size, data_key.ciphertext)?;
    store.insert_if_absent(&record).await?;

    let stored = versions(store, kms, key_name).await?;
    let winner = active(&stored).ok_or_else(|| {
        EncryptionError::KeyManagement(format!(
            "KMS data key for {} was not found after storing it",
            key_name
        ))
    })?;
    if winner.id == record.id {
        tracing::info!(
            "Generated and stored a new {} KMS data key for {}",
            kms.provider(),
            key_name
        );
        Ok(EnvelopeKey {
            version: record.version,
            key: data_key.plaintext,
            active: true,
        })
    } else {
        // Another process stored its key first: use that one.
        unwrap(kms, winner, size).await
    }
}

/// Generate a new version of `key_name`, make it the active one and retire the previous
/// version (which stays stored and decryptable). Returns the new key.
pub(crate) async fn rotate(
    store: &dyn WrappedKeyStore,
    kms: &dyn KmsKeyWrapper,
    key_name: &str,
    size: usize,
) -> Result<EnvelopeKey, EncryptionError> {
    let existing = versions(store, kms, key_name).await?;
    let next_version = existing.iter().map(|key| key.version).max().unwrap_or(0) + 1;
    let data_key = kms.generate_data_key(size, key_name).await?;
    if data_key.plaintext.len() != size {
        return Err(EncryptionError::KeyManagement(format!(
            "KMS generated a {}-byte key, expected {}",
            data_key.plaintext.len(),
            size
        )));
    }
    let record = new_record(kms, key_name, next_version, size, data_key.ciphertext)?;
    store.insert_rotated(&record).await?;
    tracing::info!(
        "Rotated {} KMS data key {} to version {}",
        kms.provider(),
        key_name,
        next_version
    );
    Ok(EnvelopeKey {
        version: next_version,
        key: data_key.plaintext,
        active: true,
    })
}

/// Decrypt every stored version of `key_name`, newest first.
pub(crate) async fn load_all(
    store: &dyn WrappedKeyStore,
    kms: &dyn KmsKeyWrapper,
    key_name: &str,
    size: usize,
) -> Result<Vec<EnvelopeKey>, EncryptionError> {
    load_new(store, kms, key_name, size, |_| false).await
}

/// Decrypt the stored versions of `key_name` for which `known` is false, newest first.
///
/// Reading the stored versions is one database query; the KMS is only called for
/// versions that are not known yet.
pub(crate) async fn load_new(
    store: &dyn WrappedKeyStore,
    kms: &dyn KmsKeyWrapper,
    key_name: &str,
    size: usize,
    known: impl Fn(u32) -> bool,
) -> Result<Vec<EnvelopeKey>, EncryptionError> {
    let mut keys = Vec::new();
    for record in versions(store, kms, key_name).await? {
        if !known(record.version) {
            keys.push(unwrap(kms, &record, size).await?);
        }
    }
    Ok(keys)
}

#[cfg(feature = "aws-kms")]
mod aws {
    use super::{DataKey, KmsKeyWrapper, SecretBytes};
    use crate::encryption::EncryptionError;
    use aws_sdk_kms::{
        Client,
        error::DisplayErrorContext,
        primitives::Blob,
        types::{DataKeySpec, EncryptionAlgorithmSpec},
    };

    /// Encryption context key binding a wrapped data key to what it is for.
    const CONTEXT_KEY: &str = "hammerwork:key-name";

    pub(super) struct AwsKms {
        client: Client,
        key_id: String,
    }

    impl AwsKms {
        pub(super) async fn connect(key_id: &str, region: &str, endpoint: Option<&str>) -> Self {
            let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest())
                .region(aws_config::Region::new(region.to_string()));
            if let Some(endpoint) = endpoint {
                loader = loader.endpoint_url(endpoint);
            }
            let config = loader.load().await;
            Self {
                client: Client::new(&config),
                key_id: key_id.to_string(),
            }
        }

        /// A client for a local fake KMS at `endpoint`, with static credentials.
        #[cfg(test)]
        pub(super) fn for_endpoint(key_id: &str, endpoint: &str) -> Self {
            let config = aws_sdk_kms::Config::builder()
                .behavior_version(aws_sdk_kms::config::BehaviorVersion::latest())
                .region(aws_sdk_kms::config::Region::new("us-east-1"))
                .endpoint_url(endpoint)
                .credentials_provider(aws_sdk_kms::config::Credentials::new(
                    "AKIDTEST", "secret", None, None, "test",
                ))
                .build();
            Self {
                client: Client::from_conf(config),
                key_id: key_id.to_string(),
            }
        }
    }

    #[async_trait::async_trait]
    impl KmsKeyWrapper for AwsKms {
        fn provider(&self) -> &'static str {
            "aws"
        }

        fn kms_key_id(&self) -> &str {
            &self.key_id
        }

        async fn generate_data_key(
            &self,
            size: usize,
            context: &str,
        ) -> Result<DataKey, EncryptionError> {
            let key_spec = match size {
                32 => DataKeySpec::Aes256,
                16 => DataKeySpec::Aes128,
                _ => {
                    return Err(EncryptionError::KeyManagement(format!(
                        "Unsupported key size for AWS KMS: {} bytes",
                        size
                    )));
                }
            };
            let response = self
                .client
                .generate_data_key()
                .key_id(&self.key_id)
                .key_spec(key_spec)
                .encryption_context(CONTEXT_KEY, context)
                .send()
                .await
                .map_err(|e| {
                    EncryptionError::KeyManagement(format!(
                        "AWS KMS GenerateDataKey with key {} failed: {}",
                        self.key_id,
                        DisplayErrorContext(&e)
                    ))
                })?;
            let plaintext = response.plaintext.ok_or_else(|| {
                EncryptionError::KeyManagement(
                    "AWS KMS GenerateDataKey response has no plaintext key".to_string(),
                )
            })?;
            let ciphertext = response.ciphertext_blob.ok_or_else(|| {
                EncryptionError::KeyManagement(
                    "AWS KMS GenerateDataKey response has no CiphertextBlob".to_string(),
                )
            })?;
            Ok(DataKey {
                plaintext: SecretBytes::new(plaintext.into_inner()),
                ciphertext: ciphertext.into_inner(),
            })
        }

        async fn decrypt(
            &self,
            ciphertext: &[u8],
            context: &str,
        ) -> Result<SecretBytes, EncryptionError> {
            let response = self
                .client
                .decrypt()
                .key_id(&self.key_id)
                .ciphertext_blob(Blob::new(ciphertext))
                .encryption_algorithm(EncryptionAlgorithmSpec::SymmetricDefault)
                .encryption_context(CONTEXT_KEY, context)
                .send()
                .await
                .map_err(|e| {
                    EncryptionError::KeyManagement(format!(
                        "AWS KMS Decrypt with key {} failed: {}",
                        self.key_id,
                        DisplayErrorContext(&e)
                    ))
                })?;
            response
                .plaintext
                .map(|plaintext| SecretBytes::new(plaintext.into_inner()))
                .ok_or_else(|| {
                    EncryptionError::KeyManagement(
                        "AWS KMS Decrypt response has no plaintext".to_string(),
                    )
                })
        }
    }
}

#[cfg(feature = "gcp-kms")]
mod gcp {
    use super::{DataKey, KmsKeyWrapper, SecretBytes};
    use crate::encryption::EncryptionError;
    use google_cloud_kms::client::{Client, ClientConfig};
    use google_cloud_kms::grpc::kms::v1::{
        DecryptRequest, EncryptRequest, GenerateRandomBytesRequest,
    };

    pub(super) struct GcpKms {
        client: Client,
        resource: String,
        location: String,
    }

    impl GcpKms {
        pub(super) async fn connect(
            resource: &str,
            location: &str,
        ) -> Result<Self, EncryptionError> {
            let config = ClientConfig::default().with_auth().await.map_err(|e| {
                EncryptionError::KeyManagement(format!("Failed to configure GCP KMS client: {}", e))
            })?;
            crate::ensure_rustls_crypto_provider();
            let client = Client::new(config).await.map_err(|e| {
                EncryptionError::KeyManagement(format!("Failed to create GCP KMS client: {}", e))
            })?;
            Ok(Self {
                client,
                resource: resource.to_string(),
                location: location.to_string(),
            })
        }

        /// A client for a local fake Cloud KMS at `endpoint`, without authentication.
        #[cfg(test)]
        pub(super) async fn for_endpoint(
            resource: &str,
            location: &str,
            endpoint: &str,
        ) -> Result<Self, EncryptionError> {
            /// Hands out a fixed bearer token.
            #[derive(Debug)]
            struct StaticToken;
            #[async_trait::async_trait]
            impl token_source::TokenSource for StaticToken {
                async fn token(&self) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
                    Ok("Bearer test-token".to_string())
                }
            }
            impl token_source::TokenSourceProvider for StaticToken {
                fn token_source(&self) -> std::sync::Arc<dyn token_source::TokenSource> {
                    std::sync::Arc::new(StaticToken)
                }
            }

            let config = ClientConfig {
                endpoint: endpoint.to_string(),
                token_source_provider: Box::new(StaticToken),
                ..ClientConfig::default()
            };
            crate::ensure_rustls_crypto_provider();
            let client = Client::new(config).await.map_err(|e| {
                EncryptionError::KeyManagement(format!("Failed to create GCP KMS client: {}", e))
            })?;
            Ok(Self {
                client,
                resource: resource.to_string(),
                location: location.to_string(),
            })
        }
    }

    #[async_trait::async_trait]
    impl KmsKeyWrapper for GcpKms {
        fn provider(&self) -> &'static str {
            "gcp"
        }

        fn kms_key_id(&self) -> &str {
            &self.resource
        }

        async fn generate_data_key(
            &self,
            size: usize,
            context: &str,
        ) -> Result<DataKey, EncryptionError> {
            let length_bytes = i32::try_from(size).map_err(|_| {
                EncryptionError::KeyManagement(format!(
                    "Unsupported key size for GCP KMS: {}",
                    size
                ))
            })?;
            let plaintext = SecretBytes::new(
                self.client
                    .generate_random_bytes(
                        GenerateRandomBytesRequest {
                            location: self.location.clone(),
                            length_bytes,
                            protection_level: 1, // SOFTWARE
                        },
                        None,
                    )
                    .await
                    .map_err(|e| {
                        EncryptionError::KeyManagement(format!(
                            "GCP KMS GenerateRandomBytes failed: {}",
                            e
                        ))
                    })?
                    .data,
            );
            let ciphertext = self
                .client
                .encrypt(
                    EncryptRequest {
                        name: self.resource.clone(),
                        plaintext: plaintext.to_vec(),
                        additional_authenticated_data: context.as_bytes().to_vec(),
                        ..Default::default()
                    },
                    None,
                )
                .await
                .map_err(|e| {
                    EncryptionError::KeyManagement(format!(
                        "GCP KMS Encrypt with key {} failed: {}",
                        self.resource, e
                    ))
                })?
                .ciphertext;
            Ok(DataKey {
                plaintext,
                ciphertext,
            })
        }

        async fn decrypt(
            &self,
            ciphertext: &[u8],
            context: &str,
        ) -> Result<SecretBytes, EncryptionError> {
            Ok(SecretBytes::new(
                self.client
                    .decrypt(
                        DecryptRequest {
                            name: self.resource.clone(),
                            ciphertext: ciphertext.to_vec(),
                            additional_authenticated_data: context.as_bytes().to_vec(),
                            ..Default::default()
                        },
                        None,
                    )
                    .await
                    .map_err(|e| {
                        EncryptionError::KeyManagement(format!(
                            "GCP KMS Decrypt with key {} failed: {}",
                            self.resource, e
                        ))
                    })?
                    .plaintext,
            ))
        }
    }
}

/// A local stand-in for a KMS and an in-memory key store, for tests.
#[cfg(test)]
pub(crate) mod test_support {
    use super::*;
    use aes_gcm::{
        Aes256Gcm, KeyInit, Nonce,
        aead::{Aead, Payload},
    };
    use rand::{RngCore, rngs::OsRng};
    use std::sync::Mutex;
    use std::sync::atomic::{AtomicUsize, Ordering};

    /// Wraps data keys with a fixed AES-256-GCM key, like a KMS key that never leaves
    /// the service. Two `MockKms` with the same `kms_key_id` and secret behave like the
    /// same KMS key seen from two processes.
    pub(crate) struct MockKms {
        kms_key_id: String,
        secret: [u8; 32],
        pub generated: AtomicUsize,
        pub decrypted: AtomicUsize,
    }

    impl MockKms {
        pub(crate) fn new(kms_key_id: &str, secret: [u8; 32]) -> Self {
            Self {
                kms_key_id: kms_key_id.to_string(),
                secret,
                generated: AtomicUsize::new(0),
                decrypted: AtomicUsize::new(0),
            }
        }

        /// A mock KMS key with a unique id, so tests never share stored keys.
        pub(crate) fn unique() -> Self {
            Self::new(&format!("mock-key-{}", Uuid::new_v4()), [0x5a; 32])
        }
    }

    #[async_trait::async_trait]
    impl KmsKeyWrapper for MockKms {
        fn provider(&self) -> &'static str {
            "aws"
        }

        fn kms_key_id(&self) -> &str {
            &self.kms_key_id
        }

        async fn generate_data_key(
            &self,
            size: usize,
            context: &str,
        ) -> Result<DataKey, EncryptionError> {
            self.generated.fetch_add(1, Ordering::SeqCst);
            let mut plaintext = SecretBytes::new(vec![0u8; size]);
            OsRng.fill_bytes(&mut plaintext);
            let cipher = Aes256Gcm::new_from_slice(&self.secret).unwrap();
            let mut nonce = [0u8; 12];
            OsRng.fill_bytes(&mut nonce);
            let mut ciphertext = nonce.to_vec();
            ciphertext.extend(
                cipher
                    .encrypt(
                        Nonce::from_slice(&nonce),
                        Payload {
                            msg: &plaintext,
                            aad: context.as_bytes(),
                        },
                    )
                    .unwrap(),
            );
            Ok(DataKey {
                plaintext,
                ciphertext,
            })
        }

        async fn decrypt(
            &self,
            ciphertext: &[u8],
            context: &str,
        ) -> Result<SecretBytes, EncryptionError> {
            self.decrypted.fetch_add(1, Ordering::SeqCst);
            if ciphertext.len() < 12 {
                return Err(EncryptionError::DecryptionFailed("short".to_string()));
            }
            let (nonce, body) = ciphertext.split_at(12);
            Aes256Gcm::new_from_slice(&self.secret)
                .unwrap()
                .decrypt(
                    Nonce::from_slice(nonce),
                    Payload {
                        msg: body,
                        aad: context.as_bytes(),
                    },
                )
                .map(SecretBytes::new)
                .map_err(|_| EncryptionError::DecryptionFailed("mock KMS decrypt".to_string()))
        }
    }

    /// In-memory [`WrappedKeyStore`] with the same uniqueness rule as the tables.
    #[derive(Default)]
    pub(crate) struct MemoryStore {
        pub keys: Mutex<Vec<KmsWrappedKey>>,
    }

    fn same_scope(a: &KmsWrappedKey, b: &KmsWrappedKey) -> bool {
        a.key_name == b.key_name && a.kms_provider == b.kms_provider && a.kms_key_id == b.kms_key_id
    }

    #[async_trait::async_trait]
    impl WrappedKeyStore for MemoryStore {
        async fn insert_if_absent(&self, key: &KmsWrappedKey) -> Result<(), EncryptionError> {
            let mut keys = self.keys.lock().unwrap();
            if !keys
                .iter()
                .any(|k| same_scope(k, key) && k.version == key.version)
            {
                keys.push(key.clone());
            }
            Ok(())
        }

        async fn insert_rotated(&self, key: &KmsWrappedKey) -> Result<(), EncryptionError> {
            let mut keys = self.keys.lock().unwrap();
            if keys
                .iter()
                .any(|k| same_scope(k, key) && k.version == key.version)
            {
                return Err(EncryptionError::DatabaseError("duplicate".to_string()));
            }
            for k in keys.iter_mut().filter(|k| same_scope(k, key)) {
                k.active = false;
                k.retired_at = Some(Utc::now());
            }
            keys.push(key.clone());
            Ok(())
        }

        async fn load_versions(
            &self,
            key_name: &str,
            kms_provider: &str,
            kms_key_id: &str,
        ) -> Result<Vec<KmsWrappedKey>, EncryptionError> {
            let mut found: Vec<_> = self
                .keys
                .lock()
                .unwrap()
                .iter()
                .filter(|k| {
                    k.key_name == key_name
                        && k.kms_provider == kms_provider
                        && k.kms_key_id == kms_key_id
                })
                .cloned()
                .collect();
            found.sort_by_key(|k| std::cmp::Reverse(k.version));
            Ok(found)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::test_support::{MemoryStore, MockKms};
    use super::*;
    use std::sync::atomic::Ordering;

    #[tokio::test]
    async fn test_load_twice_returns_same_key_and_stores_only_ciphertext() {
        let store = MemoryStore::default();
        let kms = MockKms::unique();

        let first = load_or_create(&store, &kms, "key-manager/master", 32)
            .await
            .unwrap();
        assert_eq!(first.version, 1);
        assert_eq!(first.key.len(), 32);

        // A "restarted process": a new client for the same KMS key
        let restarted = MockKms::new(kms.kms_key_id(), [0x5a; 32]);
        let second = load_or_create(&store, &restarted, "key-manager/master", 32)
            .await
            .unwrap();
        assert_eq!(first.key, second.key);
        assert_eq!(restarted.generated.load(Ordering::SeqCst), 0);
        assert_eq!(restarted.decrypted.load(Ordering::SeqCst), 1);

        let keys = store.keys.lock().unwrap();
        assert_eq!(keys.len(), 1);
        assert!(
            !keys[0]
                .wrapped_key
                .windows(32)
                .any(|w| w == first.key.as_slice())
        );
        assert_eq!(keys[0].kms_key_id, kms.kms_key_id());
        assert_eq!(keys[0].kms_provider, "aws");
    }

    #[tokio::test]
    async fn test_key_names_are_separate_and_bound_to_ciphertext() {
        let store = MemoryStore::default();
        let kms = MockKms::unique();
        let a = load_or_create(&store, &kms, "engine/a", 32).await.unwrap();
        let b = load_or_create(&store, &kms, "engine/b", 32).await.unwrap();
        assert_ne!(a.key, b.key);

        // A blob moved to another key name does not decrypt (KMS context mismatch)
        let mut keys = store.keys.lock().unwrap().clone();
        keys[0].key_name = "engine/b".to_string();
        assert!(unwrap(&kms, &keys[0], 32).await.is_err());
    }

    #[tokio::test]
    async fn test_concurrent_first_load_converges() {
        let store = Arc::new(MemoryStore::default());
        let kms_key_id = format!("mock-key-{}", Uuid::new_v4());
        let tasks: Vec<_> = (0..8)
            .map(|_| {
                let store = store.clone();
                let kms_key_id = kms_key_id.clone();
                tokio::spawn(async move {
                    let kms = MockKms::new(&kms_key_id, [0x5a; 32]);
                    load_or_create(store.as_ref(), &kms, "key-manager/master", 32)
                        .await
                        .unwrap()
                        .key
                })
            })
            .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]));
        assert_eq!(store.keys.lock().unwrap().len(), 1);
    }

    #[tokio::test]
    async fn test_rotation_keeps_old_versions_decryptable() {
        let store = MemoryStore::default();
        let kms = MockKms::unique();
        let v1 = load_or_create(&store, &kms, "engine/default", 32)
            .await
            .unwrap();
        let v2 = rotate(&store, &kms, "engine/default", 32).await.unwrap();
        assert_eq!(v2.version, 2);
        assert_ne!(v1.key, v2.key);

        let active = load_or_create(&store, &kms, "engine/default", 32)
            .await
            .unwrap();
        assert_eq!((active.version, &active.key), (2, &v2.key));

        let all = load_all(&store, &kms, "engine/default", 32).await.unwrap();
        let all: Vec<_> = all.into_iter().map(|k| (k.version, k.key)).collect();
        assert_eq!(all, vec![(2, v2.key), (1, v1.key)]);
    }

    #[tokio::test]
    async fn test_wrong_kms_key_or_size_fails_closed() {
        let store = MemoryStore::default();
        let kms = MockKms::unique();
        load_or_create(&store, &kms, "key-manager/master", 32)
            .await
            .unwrap();

        // Same KMS key id, different key material: decrypt fails, no new key is generated
        let wrong = MockKms::new(kms.kms_key_id(), [0x11; 32]);
        assert!(
            load_or_create(&store, &wrong, "key-manager/master", 32)
                .await
                .is_err()
        );
        assert_eq!(wrong.generated.load(Ordering::SeqCst), 0);

        assert!(
            load_or_create(&store, &kms, "key-manager/master", 16)
                .await
                .is_err()
        );
    }

    /// Round trip against a LocalStack (or other AWS-compatible) KMS through the
    /// `&endpoint=` parameter: create a KMS key, load the data key twice, rotate.
    /// Returns early when `AWS_ENDPOINT_URL` is not set. Needs AWS credentials in the
    /// environment (any values work with LocalStack).
    #[cfg(feature = "aws-kms")]
    #[tokio::test]
    #[ignore = "requires LocalStack: AWS_ENDPOINT_URL"]
    async fn test_aws_kms_localstack_roundtrip() {
        let Ok(endpoint) = std::env::var("AWS_ENDPOINT_URL") else {
            eprintln!("AWS_ENDPOINT_URL is not set; skipping the LocalStack KMS test");
            return;
        };
        let region = "us-east-1";
        let config = aws_config::defaults(aws_config::BehaviorVersion::latest())
            .region(aws_config::Region::new(region))
            .endpoint_url(&endpoint)
            .load()
            .await;
        let key_id = aws_sdk_kms::Client::new(&config)
            .create_key()
            .send()
            .await
            .expect("failed to create a KMS key")
            .key_metadata
            .expect("no key metadata")
            .key_id;

        let source = KeySource::External(format!(
            "aws://{}?region={}&endpoint={}",
            key_id, region, endpoint
        ));
        let store = MemoryStore::default();
        let first = parse_kms_source(&source)
            .unwrap()
            .unwrap()
            .connect()
            .await
            .unwrap();
        let v1 = load_or_create(&store, first.as_ref(), "key-manager/master", 32)
            .await
            .unwrap();

        let second = parse_kms_source(&source)
            .unwrap()
            .unwrap()
            .connect()
            .await
            .unwrap();
        let again = load_or_create(&store, second.as_ref(), "key-manager/master", 32)
            .await
            .unwrap();
        assert_eq!(v1.key, again.key);

        let v2 = rotate(&store, second.as_ref(), "key-manager/master", 32)
            .await
            .unwrap();
        let all = load_all(&store, first.as_ref(), "key-manager/master", 32)
            .await
            .unwrap();
        assert_eq!(all.len(), 2);
        assert_eq!(all[0].key, v2.key);
        assert_eq!(all[1].key, v1.key);
    }

    #[test]
    fn test_parse_kms_source() {
        assert_eq!(
            parse_kms_source(&KeySource::Environment("X".to_string())).unwrap(),
            None
        );
        assert_eq!(
            parse_kms_source(&KeySource::External("vault://a/b".to_string())).unwrap(),
            None
        );
        assert!(matches!(
            parse_kms_source(&KeySource::External("aws://?region=x".to_string())),
            Err(EncryptionError::InvalidConfiguration(_))
        ));
        assert!(matches!(
            parse_kms_source(&KeySource::External("gcp://my-key".to_string())),
            Err(EncryptionError::InvalidConfiguration(_))
        ));

        let aws = parse_kms_source(&KeySource::External(
            "aws://alias/k?region=eu-west-1&endpoint=http://localhost:4566".to_string(),
        ));
        let gcp = parse_kms_source(&KeySource::External(
            "gcp://projects/p/locations/us/keyRings/r/cryptoKeys/k".to_string(),
        ));
        #[cfg(feature = "aws-kms")]
        assert_eq!(
            aws.unwrap(),
            Some(KmsSource::Aws {
                key_id: "alias/k".to_string(),
                region: "eu-west-1".to_string(),
                endpoint: Some("http://localhost:4566".to_string()),
            })
        );
        #[cfg(not(feature = "aws-kms"))]
        assert!(aws.unwrap_err().to_string().contains("aws-kms"));
        #[cfg(feature = "gcp-kms")]
        assert_eq!(
            gcp.unwrap(),
            Some(KmsSource::Gcp {
                resource: "projects/p/locations/us/keyRings/r/cryptoKeys/k".to_string(),
                location: "projects/p/locations/us".to_string(),
            })
        );
        #[cfg(not(feature = "gcp-kms"))]
        assert!(gcp.unwrap_err().to_string().contains("gcp-kms"));
    }

    /// A fake AWS KMS (JSON 1.1 protocol) for `aws::AwsKms`. "Wraps" a data key as
    /// `context || 0 || key`, and refuses to unwrap it under another encryption context.
    /// Key id `missing-plaintext` returns responses without the key material.
    #[cfg(feature = "aws-kms")]
    async fn fake_aws_kms() -> String {
        use base64::Engine as _;
        use tokio::io::{AsyncReadExt, AsyncWriteExt};

        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = format!("http://{}", listener.local_addr().unwrap());
        tokio::spawn(async move {
            while let Ok((mut socket, _)) = listener.accept().await {
                tokio::spawn(async move {
                    let b64 = base64::engine::general_purpose::STANDARD;
                    loop {
                        // Read one request (headers, then content-length bytes).
                        let mut buffer = Vec::new();
                        let mut chunk = [0u8; 8192];
                        let (head, body) = loop {
                            if let Some(end) = buffer.windows(4).position(|w| w == b"\r\n\r\n") {
                                let head = String::from_utf8_lossy(&buffer[..end]).to_lowercase();
                                let length: usize = head
                                    .lines()
                                    .find_map(|l| l.strip_prefix("content-length:"))
                                    .and_then(|v| v.trim().parse().ok())
                                    .unwrap_or(0);
                                if buffer.len() >= end + 4 + length {
                                    break (head, buffer[end + 4..end + 4 + length].to_vec());
                                }
                            }
                            match socket.read(&mut chunk).await {
                                Ok(0) | Err(_) => return,
                                Ok(n) => buffer.extend_from_slice(&chunk[..n]),
                            }
                        };
                        let request: serde_json::Value =
                            serde_json::from_slice(&body).unwrap_or_default();
                        let target = head
                            .lines()
                            .find_map(|l| l.strip_prefix("x-amz-target:"))
                            .unwrap_or("")
                            .trim()
                            .to_string();
                        let context = request["EncryptionContext"]["hammerwork:key-name"]
                            .as_str()
                            .unwrap_or("")
                            .to_string();
                        let missing = request["KeyId"] == "missing-plaintext";
                        let (status, reply) = match target.as_str() {
                            "trentservice.generatedatakey" => {
                                let size = if request["KeySpec"] == "AES_128" {
                                    16
                                } else {
                                    32
                                };
                                let key = vec![7u8; size];
                                let mut wrapped = context.into_bytes();
                                wrapped.push(0);
                                wrapped.extend_from_slice(&key);
                                if missing {
                                    (200, serde_json::json!({ "KeyId": "k" }))
                                } else {
                                    (
                                        200,
                                        serde_json::json!({
                                            "KeyId": request["KeyId"],
                                            "Plaintext": b64.encode(&key),
                                            "CiphertextBlob": b64.encode(&wrapped),
                                        }),
                                    )
                                }
                            }
                            "trentservice.decrypt" => {
                                let blob = b64
                                    .decode(request["CiphertextBlob"].as_str().unwrap_or(""))
                                    .unwrap_or_default();
                                let split = blob.iter().position(|b| *b == 0);
                                match split {
                                    _ if missing => (200, serde_json::json!({ "KeyId": "k" })),
                                    Some(at) if blob[..at] == *context.as_bytes() => (
                                        200,
                                        serde_json::json!({
                                            "KeyId": request["KeyId"],
                                            "Plaintext": b64.encode(&blob[at + 1..]),
                                        }),
                                    ),
                                    _ => (
                                        400,
                                        serde_json::json!({
                                            "__type": "InvalidCiphertextException",
                                            "message": "context mismatch"
                                        }),
                                    ),
                                }
                            }
                            _ => (
                                400,
                                serde_json::json!({
                                    "__type": "UnsupportedOperationException",
                                    "message": target
                                }),
                            ),
                        };
                        let reply = reply.to_string();
                        let response = format!(
                            "HTTP/1.1 {status} X\r\ncontent-type: application/x-amz-json-1.1\r\ncontent-length: {}\r\n\r\n{reply}",
                            reply.len()
                        );
                        if socket.write_all(response.as_bytes()).await.is_err() {
                            return;
                        }
                    }
                });
            }
        });
        url
    }

    #[cfg(feature = "aws-kms")]
    #[tokio::test]
    async fn test_aws_kms_wraps_and_unwraps_with_context() {
        let endpoint = fake_aws_kms().await;
        let kms = aws::AwsKms::for_endpoint("alias/test", &endpoint);
        assert_eq!(kms.provider(), "aws");
        assert_eq!(kms.kms_key_id(), "alias/test");

        let data_key = kms.generate_data_key(32, "engine/k1").await.unwrap();
        assert_eq!(data_key.plaintext.as_slice(), [7u8; 32].as_slice());
        assert_ne!(data_key.ciphertext, data_key.plaintext.to_vec());
        assert_eq!(
            kms.decrypt(&data_key.ciphertext, "engine/k1")
                .await
                .unwrap(),
            data_key.plaintext
        );
        assert_eq!(
            kms.generate_data_key(16, "engine/k1")
                .await
                .unwrap()
                .plaintext
                .len(),
            16
        );

        // The encryption context binds the wrapped key to its name.
        let err = kms
            .decrypt(&data_key.ciphertext, "engine/other")
            .await
            .unwrap_err();
        assert!(matches!(err, EncryptionError::KeyManagement(_)));
        assert!(
            err.to_string()
                .contains("AWS KMS Decrypt with key alias/test failed"),
            "{err}"
        );

        // Through the envelope functions, with a store.
        let store = test_support::MemoryStore::default();
        let created = load_or_create(&store, &kms, "engine/k1", 32).await.unwrap();
        let again = load_or_create(&store, &kms, "engine/k1", 32).await.unwrap();
        assert_eq!(created.key, again.key);
    }

    #[cfg(feature = "aws-kms")]
    #[tokio::test]
    async fn test_aws_kms_errors() {
        let endpoint = fake_aws_kms().await;
        let kms = aws::AwsKms::for_endpoint("alias/test", &endpoint);
        let err = kms
            .generate_data_key(24, "c")
            .await
            .err()
            .expect("an error");
        assert!(
            err.to_string()
                .contains("Unsupported key size for AWS KMS: 24 bytes"),
            "{err}"
        );

        let missing = aws::AwsKms::for_endpoint("missing-plaintext", &endpoint);
        let err = missing
            .generate_data_key(32, "c")
            .await
            .err()
            .expect("an error");
        assert!(err.to_string().contains("has no plaintext key"), "{err}");
        let err = missing.decrypt(b"x", "c").await.unwrap_err();
        assert!(err.to_string().contains("has no plaintext"), "{err}");

        // Nothing listening: a request error, not a panic.
        let unreachable = aws::AwsKms::for_endpoint("alias/test", "http://127.0.0.1:1");
        let err = unreachable
            .generate_data_key(32, "c")
            .await
            .err()
            .expect("an error");
        assert!(
            err.to_string()
                .contains("AWS KMS GenerateDataKey with key alias/test failed"),
            "{err}"
        );
    }

    #[cfg(feature = "aws-kms")]
    #[tokio::test]
    async fn test_aws_kms_source_connects_with_region_and_endpoint() {
        let endpoint = fake_aws_kms().await;
        let source = KeySource::External(format!(
            "aws://alias/x?region=eu-west-1&endpoint={endpoint}"
        ));
        let kms = parse_kms_source(&source)
            .unwrap()
            .unwrap()
            .connect()
            .await
            .unwrap();
        assert_eq!(kms.provider(), "aws");
        assert_eq!(kms.kms_key_id(), "alias/x");
    }

    /// Creating a Cloud KMS client used to panic: the Google Cloud clients need a
    /// process-default rustls crypto provider, and with both of rustls' backends in the
    /// dependency tree none could be chosen. A connection failure is now an error.
    #[cfg(feature = "gcp-kms")]
    #[tokio::test]
    async fn test_gcp_kms_client_failure_is_an_error_not_a_panic() {
        let result = gcp::GcpKms::for_endpoint(
            "projects/p/locations/l/keyRings/r/cryptoKeys/k",
            "projects/p/locations/l",
            "http://127.0.0.1:1",
        )
        .await;
        let err = result.err().expect("nothing listens on port 1");
        assert!(matches!(err, EncryptionError::KeyManagement(_)), "{err}");
        assert!(
            err.to_string().contains("Failed to create GCP KMS client"),
            "{err}"
        );
        assert!(rustls::crypto::CryptoProvider::get_default().is_some());
    }
}