tandem-memory 0.6.9

use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::process::{Child, Command, ExitStatus, Stdio};
use std::thread;
use std::time::{Duration, Instant};

use base64::Engine;
use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::decrypt_broker::{
    MemoryDecryptBrokerConfig, MemoryDecryptPurpose, MemoryDekUnwrapProvider,
    MemoryDekUnwrapProviderBox, MemoryDekUnwrapTicket, MemoryDekWrapProvider,
    MemoryDekWrapProviderBox, MemoryDekWrapRequest, MemorySecretFamily,
};
use crate::types::{MemoryError, MemoryResult};

pub const GOOGLE_CLOUD_KMS_PROVIDER_ID: &str = "google_cloud_kms";
const GOOGLE_KMS_DECRYPT_COMMAND_ENV: &str = "TANDEM_MEMORY_GOOGLE_KMS_DECRYPT_COMMAND";
const GOOGLE_KMS_ENCRYPT_COMMAND_ENV: &str = "TANDEM_MEMORY_GOOGLE_KMS_ENCRYPT_COMMAND";
const MEMORY_DEK_LEN: usize = 32;
const GOOGLE_KMS_COMMAND_TIMEOUT: Duration = Duration::from_secs(10);
const GOOGLE_KMS_COMMAND_POLL_INTERVAL: Duration = Duration::from_millis(10);

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct GoogleCloudKmsDecryptRequest {
    pub crypto_key_id: String,
    pub kek_version: String,
    pub ciphertext: Vec<u8>,
    pub additional_authenticated_data: Vec<u8>,
    pub runtime_principal_id: String,
    pub principal_id: String,
    pub purpose: MemoryDecryptPurpose,
    pub key_scope_id: String,
    pub audit_id: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct GoogleCloudKmsCommandRequest {
    crypto_key_id: String,
    kek_version: String,
    crypto_key_version_id: String,
    ciphertext_base64: String,
    additional_authenticated_data_base64: String,
    runtime_principal_id: String,
    principal_id: String,
    purpose: MemoryDecryptPurpose,
    key_scope_id: String,
    audit_id: String,
}

pub trait GoogleCloudKmsDecryptClient {
    fn decrypt(&self, request: &GoogleCloudKmsDecryptRequest) -> MemoryResult<Vec<u8>>;
}

#[derive(Debug, Clone)]
pub struct GoogleCloudKmsDekUnwrapProvider<C> {
    client: C,
    runtime_principal_id: String,
}

impl<C> GoogleCloudKmsDekUnwrapProvider<C> {
    pub fn new(client: C, runtime_principal_id: impl Into<String>) -> MemoryResult<Self> {
        let runtime_principal_id = runtime_principal_id.into();
        if is_wildcard_or_blank(&runtime_principal_id) {
            return Err(MemoryError::InvalidConfig(
                "google cloud kms memory provider requires a scoped runtime principal".to_string(),
            ));
        }
        Ok(Self {
            client,
            runtime_principal_id,
        })
    }
}

impl<C> MemoryDekUnwrapProvider for GoogleCloudKmsDekUnwrapProvider<C>
where
    C: GoogleCloudKmsDecryptClient + Send + Sync,
{
    fn provider_id(&self) -> &str {
        GOOGLE_CLOUD_KMS_PROVIDER_ID
    }

    fn secret_family(&self) -> MemorySecretFamily {
        MemorySecretFamily::MemoryEnvelope
    }

    fn unwrap_dek(&self, ticket: &MemoryDekUnwrapTicket) -> MemoryResult<Vec<u8>> {
        if !provider_is_google_cloud_kms(&ticket.provider) {
            return Err(MemoryError::InvalidConfig(format!(
                "google cloud kms provider cannot unwrap provider `{}`",
                ticket.provider
            )));
        }
        if ticket.runtime_principal_id != self.runtime_principal_id {
            return Err(MemoryError::InvalidConfig(
                "memory KMS runtime principal does not match configured provider principal"
                    .to_string(),
            ));
        }
        validate_google_cloud_kms_key_id(&ticket.kek_id)?;
        validate_google_cloud_kms_version(&ticket.kek_id, &ticket.kek_version)?;
        let ciphertext = decode_wrapped_dek(&ticket.wrapped_dek)?;
        let plaintext = self.client.decrypt(&GoogleCloudKmsDecryptRequest {
            crypto_key_id: ticket.kek_id.clone(),
            kek_version: ticket.kek_version.clone(),
            ciphertext,
            additional_authenticated_data: ticket.encryption_context_hash.as_bytes().to_vec(),
            runtime_principal_id: ticket.runtime_principal_id.clone(),
            principal_id: ticket.principal_id.clone(),
            purpose: ticket.purpose,
            key_scope_id: ticket.key_scope_id.clone(),
            audit_id: ticket.audit_id.clone(),
        })?;
        if plaintext.len() != MEMORY_DEK_LEN {
            return Err(MemoryError::InvalidConfig(format!(
                "google cloud kms returned {}-byte memory DEK; expected {MEMORY_DEK_LEN}",
                plaintext.len()
            )));
        }
        Ok(plaintext)
    }
}

#[derive(Debug, Clone)]
pub struct GoogleCloudKmsExternalCommandClient {
    command_path: PathBuf,
    timeout: Duration,
}

impl GoogleCloudKmsExternalCommandClient {
    pub fn new(command_path: impl Into<PathBuf>) -> MemoryResult<Self> {
        let command_path = command_path.into();
        if command_path.as_os_str().is_empty() {
            return Err(MemoryError::InvalidConfig(
                "google cloud kms decrypt command path is required".to_string(),
            ));
        }
        Ok(Self {
            command_path,
            timeout: GOOGLE_KMS_COMMAND_TIMEOUT,
        })
    }

    pub fn from_env() -> MemoryResult<Option<Self>> {
        let Ok(raw) = std::env::var(GOOGLE_KMS_DECRYPT_COMMAND_ENV) else {
            return Ok(None);
        };
        let trimmed = raw.trim();
        if trimmed.is_empty() {
            return Ok(None);
        }
        Self::new(trimmed).map(Some)
    }

    pub fn command_path(&self) -> &Path {
        &self.command_path
    }
}

impl GoogleCloudKmsDecryptClient for GoogleCloudKmsExternalCommandClient {
    fn decrypt(&self, request: &GoogleCloudKmsDecryptRequest) -> MemoryResult<Vec<u8>> {
        let crypto_key_version_id =
            validate_google_cloud_kms_version(&request.crypto_key_id, &request.kek_version)?;
        let command_request = GoogleCloudKmsCommandRequest {
            crypto_key_id: request.crypto_key_id.clone(),
            kek_version: request.kek_version.clone(),
            crypto_key_version_id,
            ciphertext_base64: base64::engine::general_purpose::STANDARD
                .encode(&request.ciphertext),
            additional_authenticated_data_base64: base64::engine::general_purpose::STANDARD
                .encode(&request.additional_authenticated_data),
            runtime_principal_id: request.runtime_principal_id.clone(),
            principal_id: request.principal_id.clone(),
            purpose: request.purpose,
            key_scope_id: request.key_scope_id.clone(),
            audit_id: request.audit_id.clone(),
        };
        let input = serde_json::to_vec(&command_request)?;
        let child = Command::new(&self.command_path)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .map_err(|err| {
                MemoryError::InvalidConfig(format!(
                    "failed to spawn google cloud kms decrypt command `{}`: {err}",
                    self.command_path.display()
                ))
            })?;
        let output = wait_for_kms_command(child, input, "decrypt", self.timeout)?;
        if !output.status.success() {
            return Err(MemoryError::InvalidConfig(format!(
                "google cloud kms decrypt command exited with status {}",
                output.status
            )));
        }
        decode_plaintext_dek_output(&output.stdout, &request.crypto_key_id, &request.kek_version)
    }
}

#[derive(Debug, Clone, Copy)]
pub struct UnconfiguredGoogleCloudKmsClient;

impl GoogleCloudKmsDecryptClient for UnconfiguredGoogleCloudKmsClient {
    fn decrypt(&self, _request: &GoogleCloudKmsDecryptRequest) -> MemoryResult<Vec<u8>> {
        Err(MemoryError::InvalidConfig(
            "google cloud kms memory decrypt provider is configured, but no decrypt client command is available"
                .to_string(),
        ))
    }
}

pub fn provider_is_google_cloud_kms(provider: &str) -> bool {
    matches!(
        normalize_provider_id(provider).as_str(),
        GOOGLE_CLOUD_KMS_PROVIDER_ID | "google_kms" | "gcp_kms"
    )
}

pub fn memory_dek_unwrap_provider_from_config(
    config: &MemoryDecryptBrokerConfig,
) -> MemoryResult<Option<MemoryDekUnwrapProviderBox>> {
    if !config.hosted_required && provider_is_local(&config.provider) {
        return Ok(None);
    }
    config.validate()?;
    if provider_is_google_cloud_kms(&config.provider) {
        if let Some(client) = GoogleCloudKmsExternalCommandClient::from_env()? {
            return Ok(Some(Box::new(GoogleCloudKmsDekUnwrapProvider::new(
                client,
                config.runtime_principal_id.clone(),
            )?)));
        }
        return Ok(Some(Box::new(GoogleCloudKmsDekUnwrapProvider::new(
            UnconfiguredGoogleCloudKmsClient,
            config.runtime_principal_id.clone(),
        )?)));
    }
    Err(MemoryError::InvalidConfig(format!(
        "unsupported hosted memory decrypt provider `{}`",
        config.provider
    )))
}

/// A KMS `encrypt` request that wraps a fresh memory DEK under a KEK. Mirror of
/// [`GoogleCloudKmsDecryptRequest`] for the write path.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct GoogleCloudKmsEncryptRequest {
    pub crypto_key_id: String,
    pub kek_version: String,
    pub plaintext: Vec<u8>,
    pub additional_authenticated_data: Vec<u8>,
    pub runtime_principal_id: String,
    pub key_scope_id: String,
    pub audit_id: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct GoogleCloudKmsCommandEncryptRequest {
    crypto_key_id: String,
    kek_version: String,
    crypto_key_version_id: String,
    plaintext_base64: String,
    additional_authenticated_data_base64: String,
    runtime_principal_id: String,
    key_scope_id: String,
    audit_id: String,
}

pub trait GoogleCloudKmsEncryptClient {
    fn encrypt(&self, request: &GoogleCloudKmsEncryptRequest) -> MemoryResult<Vec<u8>>;
}

#[derive(Debug, Clone)]
pub struct GoogleCloudKmsDekWrapProvider<C> {
    client: C,
    runtime_principal_id: String,
}

impl<C> GoogleCloudKmsDekWrapProvider<C> {
    pub fn new(client: C, runtime_principal_id: impl Into<String>) -> MemoryResult<Self> {
        let runtime_principal_id = runtime_principal_id.into();
        if is_wildcard_or_blank(&runtime_principal_id) {
            return Err(MemoryError::InvalidConfig(
                "google cloud kms memory provider requires a scoped runtime principal".to_string(),
            ));
        }
        Ok(Self {
            client,
            runtime_principal_id,
        })
    }
}

impl<C> MemoryDekWrapProvider for GoogleCloudKmsDekWrapProvider<C>
where
    C: GoogleCloudKmsEncryptClient + Send + Sync,
{
    fn provider_id(&self) -> &str {
        GOOGLE_CLOUD_KMS_PROVIDER_ID
    }

    fn secret_family(&self) -> MemorySecretFamily {
        MemorySecretFamily::MemoryEnvelope
    }

    fn wrap_dek(&self, request: &MemoryDekWrapRequest) -> MemoryResult<Vec<u8>> {
        if !provider_is_google_cloud_kms(&request.provider) {
            return Err(MemoryError::InvalidConfig(format!(
                "google cloud kms provider cannot wrap provider `{}`",
                request.provider
            )));
        }
        if request.runtime_principal_id != self.runtime_principal_id {
            return Err(MemoryError::InvalidConfig(
                "memory KMS runtime principal does not match configured provider principal"
                    .to_string(),
            ));
        }
        validate_google_cloud_kms_key_id(&request.kek_id)?;
        validate_google_cloud_kms_version(&request.kek_id, &request.kek_version)?;
        if request.plaintext_dek.len() != MEMORY_DEK_LEN {
            return Err(MemoryError::InvalidConfig(format!(
                "memory DEK to wrap must be {MEMORY_DEK_LEN} bytes; got {}",
                request.plaintext_dek.len()
            )));
        }
        self.client.encrypt(&GoogleCloudKmsEncryptRequest {
            crypto_key_id: request.kek_id.clone(),
            kek_version: request.kek_version.clone(),
            plaintext: request.plaintext_dek.clone(),
            additional_authenticated_data: request.encryption_context_hash.as_bytes().to_vec(),
            runtime_principal_id: request.runtime_principal_id.clone(),
            key_scope_id: request.key_scope_id.clone(),
            audit_id: request.audit_id.clone(),
        })
    }
}

#[derive(Debug, Clone)]
pub struct GoogleCloudKmsExternalEncryptCommandClient {
    command_path: PathBuf,
    timeout: Duration,
}

impl GoogleCloudKmsExternalEncryptCommandClient {
    pub fn new(command_path: impl Into<PathBuf>) -> MemoryResult<Self> {
        let command_path = command_path.into();
        if command_path.as_os_str().is_empty() {
            return Err(MemoryError::InvalidConfig(
                "google cloud kms encrypt command path is required".to_string(),
            ));
        }
        Ok(Self {
            command_path,
            timeout: GOOGLE_KMS_COMMAND_TIMEOUT,
        })
    }

    pub fn from_env() -> MemoryResult<Option<Self>> {
        let Ok(raw) = std::env::var(GOOGLE_KMS_ENCRYPT_COMMAND_ENV) else {
            return Ok(None);
        };
        let trimmed = raw.trim();
        if trimmed.is_empty() {
            return Ok(None);
        }
        Self::new(trimmed).map(Some)
    }

    pub fn command_path(&self) -> &Path {
        &self.command_path
    }
}

impl GoogleCloudKmsEncryptClient for GoogleCloudKmsExternalEncryptCommandClient {
    fn encrypt(&self, request: &GoogleCloudKmsEncryptRequest) -> MemoryResult<Vec<u8>> {
        let crypto_key_version_id =
            validate_google_cloud_kms_version(&request.crypto_key_id, &request.kek_version)?;
        let command_request = GoogleCloudKmsCommandEncryptRequest {
            crypto_key_id: request.crypto_key_id.clone(),
            kek_version: request.kek_version.clone(),
            crypto_key_version_id,
            plaintext_base64: base64::engine::general_purpose::STANDARD.encode(&request.plaintext),
            additional_authenticated_data_base64: base64::engine::general_purpose::STANDARD
                .encode(&request.additional_authenticated_data),
            runtime_principal_id: request.runtime_principal_id.clone(),
            key_scope_id: request.key_scope_id.clone(),
            audit_id: request.audit_id.clone(),
        };
        let input = serde_json::to_vec(&command_request)?;
        let child = Command::new(&self.command_path)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .map_err(|err| {
                MemoryError::InvalidConfig(format!(
                    "failed to spawn google cloud kms encrypt command `{}`: {err}",
                    self.command_path.display()
                ))
            })?;
        let output = wait_for_kms_command(child, input, "encrypt", self.timeout)?;
        if !output.status.success() {
            return Err(MemoryError::InvalidConfig(format!(
                "google cloud kms encrypt command exited with status {}",
                output.status
            )));
        }
        decode_wrapped_dek_output(&output.stdout, &request.crypto_key_id, &request.kek_version)
    }
}

#[derive(Debug, Clone, Copy)]
pub struct UnconfiguredGoogleCloudKmsEncryptClient;

impl GoogleCloudKmsEncryptClient for UnconfiguredGoogleCloudKmsEncryptClient {
    fn encrypt(&self, _request: &GoogleCloudKmsEncryptRequest) -> MemoryResult<Vec<u8>> {
        Err(MemoryError::InvalidConfig(
            "google cloud kms memory encrypt provider is configured, but no encrypt client command is available"
                .to_string(),
        ))
    }
}

pub fn memory_dek_wrap_provider_from_config(
    config: &MemoryDecryptBrokerConfig,
) -> MemoryResult<Option<MemoryDekWrapProviderBox>> {
    if !config.hosted_required && provider_is_local(&config.provider) {
        return Ok(None);
    }
    config.validate()?;
    if provider_is_google_cloud_kms(&config.provider) {
        if let Some(client) = GoogleCloudKmsExternalEncryptCommandClient::from_env()? {
            return Ok(Some(Box::new(GoogleCloudKmsDekWrapProvider::new(
                client,
                config.runtime_principal_id.clone(),
            )?)));
        }
        return Ok(Some(Box::new(GoogleCloudKmsDekWrapProvider::new(
            UnconfiguredGoogleCloudKmsEncryptClient,
            config.runtime_principal_id.clone(),
        )?)));
    }
    Err(MemoryError::InvalidConfig(format!(
        "unsupported hosted memory encrypt provider `{}`",
        config.provider
    )))
}

struct KmsCommandOutput {
    status: ExitStatus,
    stdout: Vec<u8>,
}

fn wait_for_kms_command(
    mut child: Child,
    input: Vec<u8>,
    operation: &str,
    timeout: Duration,
) -> MemoryResult<KmsCommandOutput> {
    let mut stdin = child.stdin.take().ok_or_else(|| {
        MemoryError::InvalidConfig(format!(
            "google cloud kms {operation} command did not expose stdin"
        ))
    })?;
    let mut stdout = child.stdout.take().ok_or_else(|| {
        MemoryError::InvalidConfig(format!(
            "google cloud kms {operation} command did not expose stdout"
        ))
    })?;
    let mut stderr = child.stderr.take().ok_or_else(|| {
        MemoryError::InvalidConfig(format!(
            "google cloud kms {operation} command did not expose stderr"
        ))
    })?;

    let stdin_writer = thread::spawn(move || {
        stdin.write_all(&input)?;
        drop(stdin);
        std::io::Result::Ok(())
    });
    let stdout_reader = thread::spawn(move || {
        let mut bytes = Vec::new();
        stdout.read_to_end(&mut bytes)?;
        std::io::Result::Ok(bytes)
    });
    let stderr_reader = thread::spawn(move || {
        let mut bytes = Vec::new();
        stderr.read_to_end(&mut bytes)?;
        std::io::Result::Ok(bytes)
    });

    let started = Instant::now();
    let status = loop {
        if let Some(status) = child.try_wait().map_err(|err| {
            MemoryError::InvalidConfig(format!(
                "failed to wait for google cloud kms {operation} command: {err}"
            ))
        })? {
            break status;
        }
        let elapsed = started.elapsed();
        if elapsed >= timeout {
            let termination_error = match child.kill() {
                Ok(()) => child.wait().err().map(|err| err.to_string()),
                Err(kill_error) => match child.try_wait() {
                    Ok(Some(_)) => None,
                    Ok(None) => Some(format!("failed to terminate child: {kill_error}")),
                    Err(wait_error) => Some(format!(
                        "failed to terminate child: {kill_error}; failed to inspect child: {wait_error}"
                    )),
                },
            };
            let suffix = termination_error
                .map(|error| format!("; {error}"))
                .unwrap_or_default();
            return Err(MemoryError::InvalidConfig(format!(
                "google cloud kms {operation} command timed out after {}ms and was terminated{suffix}",
                timeout.as_millis()
            )));
        }
        thread::sleep(GOOGLE_KMS_COMMAND_POLL_INTERVAL.min(timeout.saturating_sub(elapsed)));
    };

    while !stdin_writer.is_finished()
        || !stdout_reader.is_finished()
        || !stderr_reader.is_finished()
    {
        let elapsed = started.elapsed();
        if elapsed >= timeout {
            return Err(MemoryError::InvalidConfig(format!(
                "google cloud kms {operation} command I/O timed out after {}ms",
                timeout.as_millis()
            )));
        }
        thread::sleep(GOOGLE_KMS_COMMAND_POLL_INTERVAL.min(timeout.saturating_sub(elapsed)));
    }

    join_kms_io(stdin_writer, operation, "write stdin")?;
    let stdout = join_kms_io(stdout_reader, operation, "read stdout")?;
    let _stderr = join_kms_io(stderr_reader, operation, "read stderr")?;
    Ok(KmsCommandOutput { status, stdout })
}

fn join_kms_io<T>(
    handle: thread::JoinHandle<std::io::Result<T>>,
    operation: &str,
    action: &str,
) -> MemoryResult<T> {
    handle
        .join()
        .map_err(|_| {
            MemoryError::InvalidConfig(format!(
                "google cloud kms {operation} command I/O worker panicked while attempting to {action}"
            ))
        })?
        .map_err(|err| {
            MemoryError::InvalidConfig(format!(
                "failed to {action} for google cloud kms {operation} command: {err}"
            ))
        })
}

fn decode_wrapped_dek_output(
    output: &[u8],
    crypto_key_id: &str,
    kek_version: &str,
) -> MemoryResult<Vec<u8>> {
    let trimmed = String::from_utf8_lossy(output).trim().to_string();
    if trimmed.is_empty() {
        return Err(MemoryError::InvalidConfig(
            "google cloud kms encrypt command returned empty output".to_string(),
        ));
    }
    let value = serde_json::from_str::<Value>(&trimmed).map_err(|_| {
        MemoryError::InvalidConfig(
            "google cloud kms encrypt command must return version-attested JSON".to_string(),
        )
    })?;
    validate_attested_kek_version(&value, crypto_key_id, kek_version, "encrypt")?;
    for key in [
        "wrapped_dek",
        "wrapped_dek_base64",
        "ciphertext",
        "ciphertext_base64",
    ] {
        if let Some(encoded) = value.get(key).and_then(Value::as_str) {
            return decode_dek_bytes(encoded);
        }
    }
    Err(MemoryError::InvalidConfig(
        "google cloud kms encrypt command JSON must include wrapped_dek_base64".to_string(),
    ))
}

fn decode_plaintext_dek_output(
    output: &[u8],
    crypto_key_id: &str,
    kek_version: &str,
) -> MemoryResult<Vec<u8>> {
    let trimmed = String::from_utf8_lossy(output).trim().to_string();
    if trimmed.is_empty() {
        return Err(MemoryError::InvalidConfig(
            "google cloud kms decrypt command returned empty output".to_string(),
        ));
    }
    let value = serde_json::from_str::<Value>(&trimmed).map_err(|_| {
        MemoryError::InvalidConfig(
            "google cloud kms decrypt command must return version-attested JSON".to_string(),
        )
    })?;
    validate_attested_kek_version(&value, crypto_key_id, kek_version, "decrypt")?;
    for key in ["plaintext", "plaintext_base64", "dek", "dek_base64"] {
        if let Some(encoded) = value.get(key).and_then(Value::as_str) {
            return decode_dek_bytes(encoded);
        }
    }
    Err(MemoryError::InvalidConfig(
        "google cloud kms decrypt command JSON must include plaintext_base64".to_string(),
    ))
}

fn validate_attested_kek_version(
    value: &Value,
    crypto_key_id: &str,
    configured_version: &str,
    operation: &str,
) -> MemoryResult<()> {
    let versioned_resource = validate_google_cloud_kms_version(crypto_key_id, configured_version)?;
    let version = versioned_resource
        .rsplit('/')
        .next()
        .expect("validated version resource has a final segment");
    let attestations = [
        "kek_version",
        "crypto_key_version",
        "crypto_key_version_id",
        "key_version",
    ]
    .into_iter()
    .filter_map(|key| value.get(key).and_then(Value::as_str))
    .collect::<Vec<_>>();
    if attestations.is_empty() {
        return Err(MemoryError::InvalidConfig(format!(
            "google cloud kms {operation} command JSON must attest kek_version"
        )));
    }
    for attested in attestations {
        let attested = attested.trim();
        if attested != configured_version.trim()
            && attested != version
            && attested != versioned_resource
        {
            return Err(MemoryError::InvalidConfig(format!(
                "google cloud kms {operation} command attested key version `{attested}` but configured version is `{}`",
                configured_version.trim()
            )));
        }
    }
    Ok(())
}

fn decode_wrapped_dek(raw: &str) -> MemoryResult<Vec<u8>> {
    decode_dek_bytes(raw).map_err(|_| {
        MemoryError::InvalidConfig(
            "memory envelope wrapped_dek must be base64/base64url or hex encoded".to_string(),
        )
    })
}

fn decode_dek_bytes(raw: &str) -> MemoryResult<Vec<u8>> {
    let trimmed = raw.trim();
    if let Some(decoded) = decode_hex(trimmed) {
        return Ok(decoded);
    }
    base64::engine::general_purpose::STANDARD
        .decode(trimmed)
        .or_else(|_| base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(trimmed))
        .or_else(|_| base64::engine::general_purpose::URL_SAFE.decode(trimmed))
        .map_err(|_| MemoryError::InvalidConfig("memory DEK encoding is invalid".to_string()))
}

fn validate_google_cloud_kms_key_id(kek_id: &str) -> MemoryResult<()> {
    let value = kek_id.trim();
    if is_wildcard_or_blank(value) {
        return Err(MemoryError::InvalidConfig(
            "google cloud kms key id must be explicit".to_string(),
        ));
    }
    let parts: Vec<&str> = value.split('/').collect();
    if parts.len() != 8
        || parts[0] != "projects"
        || parts[2] != "locations"
        || parts[4] != "keyRings"
        || parts[6] != "cryptoKeys"
        || parts.iter().any(|part| part.trim().is_empty())
    {
        return Err(MemoryError::InvalidConfig(
            "google cloud kms key id must be a full projects/*/locations/*/keyRings/*/cryptoKeys/* resource".to_string(),
        ));
    }
    Ok(())
}

fn validate_google_cloud_kms_version(
    kek_id: &str,
    configured_version: &str,
) -> MemoryResult<String> {
    validate_google_cloud_kms_key_id(kek_id)?;
    let value = configured_version.trim();
    if is_wildcard_or_blank(value) {
        return Err(MemoryError::InvalidConfig(
            "google cloud kms requires an explicit key version".to_string(),
        ));
    }
    let prefix = format!("{}/cryptoKeyVersions/", kek_id.trim());
    if value.contains('/') {
        let Some(version) = value.strip_prefix(&prefix) else {
            return Err(MemoryError::InvalidConfig(
                "google cloud kms key version resource must belong to the configured key"
                    .to_string(),
            ));
        };
        if is_wildcard_or_blank(version) || version.contains('/') {
            return Err(MemoryError::InvalidConfig(
                "google cloud kms key version resource must be explicit".to_string(),
            ));
        }
        return Ok(value.to_string());
    }
    Ok(format!("{prefix}{value}"))
}

fn provider_is_local(provider: &str) -> bool {
    let normalized = normalize_provider_id(provider);
    normalized.is_empty() || normalized == "disabled" || normalized.starts_with("local")
}

fn normalize_provider_id(value: &str) -> String {
    value.trim().to_ascii_lowercase().replace(['-', '.'], "_")
}

fn is_wildcard_or_blank(value: &str) -> bool {
    matches!(
        value.trim().to_ascii_lowercase().as_str(),
        "" | "*" | "all" | "global" | "default"
    )
}

fn decode_hex(raw: &str) -> Option<Vec<u8>> {
    let raw = raw.trim();
    if raw.is_empty() || !raw.len().is_multiple_of(2) {
        return None;
    }
    let mut out = Vec::with_capacity(raw.len() / 2);
    let bytes = raw.as_bytes();
    let mut index = 0;
    while index < bytes.len() {
        let hi = (bytes[index] as char).to_digit(16)?;
        let lo = (bytes[index + 1] as char).to_digit(16)?;
        out.push(((hi << 4) | lo) as u8);
        index += 2;
    }
    Some(out)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::decrypt_broker::MemoryDekUnwrapTicket;
    #[cfg(unix)]
    use std::os::unix::fs::PermissionsExt;

    #[derive(Debug, Clone)]
    struct FixtureGoogleKmsClient;

    impl GoogleCloudKmsDecryptClient for FixtureGoogleKmsClient {
        fn decrypt(&self, request: &GoogleCloudKmsDecryptRequest) -> MemoryResult<Vec<u8>> {
            assert_eq!(
                request.crypto_key_id,
                "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance"
            );
            assert_eq!(request.kek_version, "1");
            assert_eq!(request.runtime_principal_id, "runtime-memory-decryptor");
            assert_eq!(
                request.key_scope_id,
                "tandem/memory/acme/finance/prod/internal"
            );
            assert_eq!(request.additional_authenticated_data, b"ctx-hash");
            Ok(vec![7u8; MEMORY_DEK_LEN])
        }
    }

    #[derive(Debug, Clone)]
    struct HexWrappedDekGoogleKmsClient;

    impl GoogleCloudKmsDecryptClient for HexWrappedDekGoogleKmsClient {
        fn decrypt(&self, request: &GoogleCloudKmsDecryptRequest) -> MemoryResult<Vec<u8>> {
            assert_eq!(request.ciphertext, vec![0xaau8; MEMORY_DEK_LEN]);
            Ok(vec![8u8; MEMORY_DEK_LEN])
        }
    }

    fn ticket() -> MemoryDekUnwrapTicket {
        MemoryDekUnwrapTicket {
            provider: "google_cloud_kms".to_string(),
            runtime_principal_id: "runtime-memory-decryptor".to_string(),
            principal_id: "kb-mcp-retrieval-gateway".to_string(),
            purpose: MemoryDecryptPurpose::RetrievalGateway,
            key_scope_id: "tandem/memory/acme/finance/prod/internal".to_string(),
            kek_id: "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance".to_string(),
            kek_version: "1".to_string(),
            wrapped_dek: base64::engine::general_purpose::STANDARD.encode(b"wrapped-dek"),
            algorithm: "AES-256-GCM".to_string(),
            encryption_context_hash: "ctx-hash".to_string(),
            policy_decision_id: "decision-1".to_string(),
            audit_id: "audit-1".to_string(),
            key_lifecycle_decision: None,
        }
    }

    #[test]
    fn google_kms_provider_unwraps_authorized_ticket_through_client() {
        let provider = GoogleCloudKmsDekUnwrapProvider::new(
            FixtureGoogleKmsClient,
            "runtime-memory-decryptor",
        )
        .expect("provider");
        assert_eq!(provider.provider_id(), GOOGLE_CLOUD_KMS_PROVIDER_ID);
        assert_eq!(provider.secret_family(), MemorySecretFamily::MemoryEnvelope);
        assert_eq!(
            provider.unwrap_dek(&ticket()).unwrap(),
            vec![7u8; MEMORY_DEK_LEN]
        );
    }

    #[test]
    fn google_kms_provider_decodes_ambiguous_hex_wrapped_dek_as_hex() {
        let provider = GoogleCloudKmsDekUnwrapProvider::new(
            HexWrappedDekGoogleKmsClient,
            "runtime-memory-decryptor",
        )
        .expect("provider");
        let mut ticket = ticket();
        ticket.wrapped_dek = "aa".repeat(MEMORY_DEK_LEN);
        assert_eq!(
            provider.unwrap_dek(&ticket).unwrap(),
            vec![8u8; MEMORY_DEK_LEN]
        );
    }

    #[test]
    fn google_kms_provider_rejects_wrong_runtime_principal() {
        let provider = GoogleCloudKmsDekUnwrapProvider::new(
            FixtureGoogleKmsClient,
            "runtime-memory-decryptor",
        )
        .expect("provider");
        let mut ticket = ticket();
        ticket.runtime_principal_id = "broad-runtime".to_string();
        let err = provider
            .unwrap_dek(&ticket)
            .expect_err("principal mismatch");
        assert!(err.to_string().contains("runtime principal"));
    }

    #[test]
    fn google_kms_provider_rejects_versioned_key_resource() {
        let provider = GoogleCloudKmsDekUnwrapProvider::new(
            FixtureGoogleKmsClient,
            "runtime-memory-decryptor",
        )
        .expect("provider");
        let mut ticket = ticket();
        ticket.kek_id.push_str("/cryptoKeyVersions/1");
        let err = provider.unwrap_dek(&ticket).expect_err("versioned key");
        assert!(err
            .to_string()
            .contains("full projects/*/locations/*/keyRings/*/cryptoKeys/* resource"));
    }

    #[test]
    fn google_kms_provider_rejects_malformed_key_resource() {
        let provider = GoogleCloudKmsDekUnwrapProvider::new(
            FixtureGoogleKmsClient,
            "runtime-memory-decryptor",
        )
        .expect("provider");
        let mut ticket = ticket();
        ticket.kek_id = "projects/acme/cryptoKeys/finance/locations/global".to_string();
        let err = provider.unwrap_dek(&ticket).expect_err("malformed key");
        assert!(err
            .to_string()
            .contains("full projects/*/locations/*/keyRings/*/cryptoKeys/* resource"));
    }

    #[test]
    fn provider_factory_keeps_local_mode_disabled() {
        let config = MemoryDecryptBrokerConfig::local_disabled();
        assert!(memory_dek_unwrap_provider_from_config(&config)
            .unwrap()
            .is_none());
    }

    #[test]
    fn provider_factory_instantiates_google_provider_from_hosted_config() {
        let config =
            MemoryDecryptBrokerConfig::hosted("google-cloud-kms", "runtime-memory-decryptor")
                .expect("hosted config");
        let provider = memory_dek_unwrap_provider_from_config(&config)
            .expect("factory")
            .expect("provider");
        assert_eq!(provider.provider_id(), GOOGLE_CLOUD_KMS_PROVIDER_ID);
        assert_eq!(provider.secret_family(), MemorySecretFamily::MemoryEnvelope);
    }

    #[test]
    fn command_output_accepts_json_plaintext_base64() {
        let encoded = base64::engine::general_purpose::STANDARD.encode([9u8; MEMORY_DEK_LEN]);
        let out = format!(r#"{{"plaintext_base64":"{encoded}","kek_version":"1"}}"#);
        assert_eq!(
            decode_plaintext_dek_output(
                out.as_bytes(),
                "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance",
                "1",
            )
            .unwrap(),
            vec![9u8; MEMORY_DEK_LEN]
        );
    }

    #[test]
    fn command_output_decodes_ambiguous_hex_plaintext_as_hex() {
        let encoded = "aa".repeat(MEMORY_DEK_LEN);
        let out = format!(
            r#"{{"plaintext":"{encoded}","crypto_key_version_id":"projects/acme/locations/global/keyRings/memory/cryptoKeys/finance/cryptoKeyVersions/1"}}"#
        );
        assert_eq!(
            decode_plaintext_dek_output(
                out.as_bytes(),
                "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance",
                "1",
            )
            .unwrap(),
            vec![0xaau8; MEMORY_DEK_LEN]
        );
    }

    #[test]
    fn command_output_rejects_unattested_legacy_bytes() {
        let out = "aa".repeat(MEMORY_DEK_LEN);
        let error = decode_plaintext_dek_output(
            out.as_bytes(),
            "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance",
            "1",
        )
        .expect_err("unattested output must fail closed");
        assert!(error.to_string().contains("version-attested JSON"));
    }

    #[cfg(unix)]
    fn write_executable_script(path: &Path, body: &str) {
        std::fs::write(path, body).expect("write KMS fixture");
        let mut permissions = std::fs::metadata(path)
            .expect("fixture metadata")
            .permissions();
        permissions.set_mode(0o700);
        std::fs::set_permissions(path, permissions).expect("make fixture executable");
    }

    #[cfg(unix)]
    #[test]
    fn external_decrypt_rejects_rotation_mismatch_and_passes_configured_version() {
        let dir = tempfile::tempdir().expect("tempdir");
        let script_path = dir.path().join("decrypt-kms.sh");
        let request_path = dir.path().join("decrypt-request.json");
        let plaintext = base64::engine::general_purpose::STANDARD.encode([7u8; MEMORY_DEK_LEN]);
        write_executable_script(
            &script_path,
            &format!(
                "#!/bin/sh\ncat > '{}'\nprintf '%s\\n' '{}'\n",
                request_path.display(),
                format!(r#"{{"plaintext_base64":"{plaintext}","kek_version":"2"}}"#)
            ),
        );
        let client = GoogleCloudKmsExternalCommandClient::new(&script_path).expect("client");
        let request = GoogleCloudKmsDecryptRequest {
            crypto_key_id: "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance"
                .to_string(),
            kek_version: "1".to_string(),
            ciphertext: vec![1u8; MEMORY_DEK_LEN],
            additional_authenticated_data: b"ctx-hash".to_vec(),
            runtime_principal_id: "runtime-memory-decryptor".to_string(),
            principal_id: "kb-mcp-retrieval-gateway".to_string(),
            purpose: MemoryDecryptPurpose::RetrievalGateway,
            key_scope_id: "tandem/memory/acme/finance/prod/internal".to_string(),
            audit_id: "audit-rotation".to_string(),
        };

        let error = client
            .decrypt(&request)
            .expect_err("rotated decrypt version must not be accepted");
        assert!(error.to_string().contains("attested key version `2`"));
        let command_request: Value = serde_json::from_slice(
            &std::fs::read(&request_path).expect("captured decrypt request"),
        )
        .expect("decrypt request JSON");
        assert_eq!(command_request["kek_version"], "1");
        assert_eq!(
            command_request["crypto_key_version_id"],
            "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance/cryptoKeyVersions/1"
        );
    }

    #[cfg(unix)]
    #[test]
    fn external_encrypt_rejects_rotation_mismatch_and_passes_configured_version() {
        let dir = tempfile::tempdir().expect("tempdir");
        let script_path = dir.path().join("encrypt-kms.sh");
        let request_path = dir.path().join("encrypt-request.json");
        let wrapped = base64::engine::general_purpose::STANDARD.encode([8u8; MEMORY_DEK_LEN]);
        write_executable_script(
            &script_path,
            &format!(
                "#!/bin/sh\ncat > '{}'\nprintf '%s\\n' '{}'\n",
                request_path.display(),
                format!(r#"{{"wrapped_dek_base64":"{wrapped}","kek_version":"2"}}"#)
            ),
        );
        let client = GoogleCloudKmsExternalEncryptCommandClient::new(&script_path).expect("client");
        let request = GoogleCloudKmsEncryptRequest {
            crypto_key_id: "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance"
                .to_string(),
            kek_version: "1".to_string(),
            plaintext: vec![1u8; MEMORY_DEK_LEN],
            additional_authenticated_data: b"ctx-hash".to_vec(),
            runtime_principal_id: "runtime-memory-encryptor".to_string(),
            key_scope_id: "tandem/memory/acme/finance/prod/internal".to_string(),
            audit_id: "audit-rotation".to_string(),
        };

        let error = client
            .encrypt(&request)
            .expect_err("rotated encrypt version must not be accepted");
        assert!(error.to_string().contains("attested key version `2`"));
        let command_request: Value = serde_json::from_slice(
            &std::fs::read(&request_path).expect("captured encrypt request"),
        )
        .expect("encrypt request JSON");
        assert_eq!(command_request["kek_version"], "1");
        assert_eq!(
            command_request["crypto_key_version_id"],
            "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance/cryptoKeyVersions/1"
        );
    }

    #[cfg(unix)]
    #[test]
    fn external_kms_command_times_out_and_terminates_child() {
        let dir = tempfile::tempdir().expect("tempdir");
        let script_path = dir.path().join("slow-kms.sh");
        let pid_path = dir.path().join("slow-kms.pid");
        std::fs::write(
            &script_path,
            format!(
                "#!/bin/sh\necho $$ > '{}'\nexec sleep 30\n",
                pid_path.display()
            ),
        )
        .expect("write KMS fixture");
        let mut permissions = std::fs::metadata(&script_path)
            .expect("fixture metadata")
            .permissions();
        permissions.set_mode(0o700);
        std::fs::set_permissions(&script_path, permissions).expect("make fixture executable");

        let client = GoogleCloudKmsExternalCommandClient {
            command_path: script_path,
            timeout: Duration::from_millis(100),
        };
        let request = GoogleCloudKmsDecryptRequest {
            crypto_key_id: "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance"
                .to_string(),
            kek_version: "1".to_string(),
            ciphertext: vec![1u8; MEMORY_DEK_LEN],
            additional_authenticated_data: b"ctx-hash".to_vec(),
            runtime_principal_id: "runtime-memory-decryptor".to_string(),
            principal_id: "kb-mcp-retrieval-gateway".to_string(),
            purpose: MemoryDecryptPurpose::RetrievalGateway,
            key_scope_id: "tandem/memory/acme/finance/prod/internal".to_string(),
            audit_id: "audit-timeout".to_string(),
        };

        let started = Instant::now();
        let error = client.decrypt(&request).expect_err("command must time out");
        assert!(
            started.elapsed() < Duration::from_secs(2),
            "timeout was not bounded"
        );
        assert!(error.to_string().contains("timed out after 100ms"));

        let pid = std::fs::read_to_string(&pid_path).expect("fixture pid");
        let still_running = Command::new("kill")
            .args(["-0", pid.trim()])
            .stderr(Stdio::null())
            .status()
            .expect("inspect fixture process")
            .success();
        assert!(!still_running, "timed-out KMS child was not terminated");
    }
}