tandem-memory 0.6.9

use crate::types::{MemoryError, MemoryResult, MemoryTenantScope};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use tandem_enterprise_contract::DataClass;

pub const MEMORY_ENVELOPE_METADATA_KEY: &str = "memory_envelope";
pub const MEMORY_ENVELOPE_ALGORITHM: &str = "AES-256-GCM";
const HOSTED_ENCRYPTION_REQUIRED_ENV: &str = "TANDEM_MEMORY_ENCRYPTION_REQUIRED";

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MemoryKeyScope {
    pub org_id: String,
    pub workspace_id: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub deployment_id: Option<String>,
    /// Department (`owner_org_unit_id`) that owns the wrapped DEK (TAN-662).
    /// Makes department a **cryptographic** key dimension, not just an
    /// access-control one: each `(tenant × department × data_class × source)`
    /// gets a distinct DEK, so a leaked department key cannot decrypt another
    /// department's ciphertext in the same tenant + data class. `None` =
    /// tenant-wide (no department), mirroring the `owner_org_unit_id` row column.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub org_unit: Option<String>,
    /// Optional private owner. Hosted search payloads bind their DEK authority
    /// to this subject so a peer principal cannot unwrap owner-private memory.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub owner_subject: Option<String>,
    pub data_class: DataClass,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub source_binding_id: Option<String>,
}

impl MemoryKeyScope {
    pub fn new(
        tenant_scope: &MemoryTenantScope,
        data_class: DataClass,
        source_binding_id: Option<String>,
    ) -> Self {
        Self {
            org_id: tenant_scope.org_id.clone(),
            workspace_id: tenant_scope.workspace_id.clone(),
            deployment_id: tenant_scope.deployment_id.clone(),
            org_unit: None,
            owner_subject: None,
            data_class,
            source_binding_id,
        }
    }

    /// Bind this key scope to a department (TAN-662). A trimmed-empty value is
    /// treated as no department.
    pub fn with_org_unit(mut self, org_unit: Option<String>) -> Self {
        self.org_unit = org_unit
            .map(|unit| unit.trim().to_string())
            .filter(|unit| !unit.is_empty());
        self
    }

    pub fn with_owner_subject(mut self, owner_subject: Option<String>) -> Self {
        self.owner_subject = owner_subject
            .map(|subject| subject.trim().to_string())
            .filter(|subject| !subject.is_empty());
        self
    }

    pub fn canonical_id(&self) -> String {
        let deployment = self.deployment_id.as_deref().unwrap_or("default");
        let class = serde_json::to_value(self.data_class)
            .ok()
            .and_then(|value| value.as_str().map(ToOwned::to_owned))
            .unwrap_or_else(|| "unknown".to_string());
        // The department segment must be structurally distinct from the
        // tenant-wide form so `dept/x` can never collide with a data-class or
        // source segment, keeping per-department DEKs unambiguous. The org-unit
        // value is caller-derived (`{taxonomy_id}/{unit_id}`) and legitimately
        // contains `/`, so it is percent-encoded to prevent delimiter injection:
        // without it, `org_unit="a/source/b"` (no source) would collide with
        // `org_unit="a", source="b"` and share a DEK across departments.
        let dept = match self.org_unit.as_deref() {
            Some(org_unit) if !org_unit.trim().is_empty() => {
                format!("/dept/{}", encode_scope_segment(org_unit))
            }
            _ => String::new(),
        };
        let subject = match self.owner_subject.as_deref() {
            Some(owner_subject) => {
                format!("/subject/{}", encode_scope_segment(owner_subject))
            }
            None => String::new(),
        };
        match self.source_binding_id.as_deref() {
            Some(source_binding_id) if !source_binding_id.trim().is_empty() => format!(
                "tandem/memory/{}/{}/{}/{}{}{}/source/{}",
                self.org_id, self.workspace_id, deployment, class, dept, subject, source_binding_id
            ),
            _ => format!(
                "tandem/memory/{}/{}/{}/{}{}{}",
                self.org_id, self.workspace_id, deployment, class, dept, subject
            ),
        }
    }

    fn validates_against_tenant(&self, tenant_scope: &MemoryTenantScope) -> bool {
        self.org_id == tenant_scope.org_id
            && self.workspace_id == tenant_scope.workspace_id
            && self.deployment_id.as_deref().unwrap_or("")
                == tenant_scope.deployment_id.as_deref().unwrap_or("")
    }

    /// Validate that this scope is safe to seal a DEK under: no wildcard tenant,
    /// deployment, department, or subject segment (any of which would collapse
    /// per-scope DEKs into a shared key). Called before wrapping.
    pub fn validate_for_envelope(&self) -> MemoryResult<()> {
        self.validate_partitioned()
    }

    fn validate_partitioned(&self) -> MemoryResult<()> {
        for (field, value) in [
            ("org_id", self.org_id.as_str()),
            ("workspace_id", self.workspace_id.as_str()),
        ] {
            if is_wildcard_scope(value) {
                return Err(MemoryError::InvalidConfig(format!(
                    "memory envelope key scope must not use wildcard `{field}`"
                )));
            }
        }
        if self
            .deployment_id
            .as_deref()
            .map(is_wildcard_scope)
            .unwrap_or(false)
        {
            return Err(MemoryError::InvalidConfig(
                "memory envelope key scope must not use wildcard `deployment_id`".to_string(),
            ));
        }
        // A wildcard department would collapse per-department DEKs back into one
        // shared key, defeating at-rest department separation (TAN-662).
        if self
            .org_unit
            .as_deref()
            .map(is_wildcard_scope)
            .unwrap_or(false)
        {
            return Err(MemoryError::InvalidConfig(
                "memory envelope key scope must not use wildcard `org_unit`".to_string(),
            ));
        }
        if self
            .owner_subject
            .as_deref()
            .map(is_wildcard_scope)
            .unwrap_or(false)
        {
            return Err(MemoryError::InvalidConfig(
                "memory envelope key scope must not use a wildcard owner subject".to_string(),
            ));
        }
        Ok(())
    }
}

/// Caller-supplied authority that a hosted decrypt is expected to satisfy.
///
/// This value must come from trusted request/store context, never from the
/// untrusted envelope being decrypted. It binds tenant, department, data class,
/// source, policy decision, and audit evidence into one exact authorization
/// contract, including an optional private owner subject.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MemoryEnvelopeAuthority {
    pub key_scope: MemoryKeyScope,
    pub policy_decision_id: String,
    pub audit_id: String,
}

impl MemoryEnvelopeAuthority {
    pub fn new(
        key_scope: MemoryKeyScope,
        policy_decision_id: impl Into<String>,
        audit_id: impl Into<String>,
    ) -> Self {
        Self {
            key_scope,
            policy_decision_id: policy_decision_id.into(),
            audit_id: audit_id.into(),
        }
    }

    pub fn validate(&self) -> MemoryResult<()> {
        self.key_scope.validate_for_envelope()?;
        for (field, value) in [
            ("policy_decision_id", self.policy_decision_id.as_str()),
            ("audit_id", self.audit_id.as_str()),
        ] {
            if is_wildcard_scope(value) {
                return Err(MemoryError::InvalidConfig(format!(
                    "memory envelope authority must use an explicit `{field}`"
                )));
            }
        }
        Ok(())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MemoryEnvelopeMetadata {
    pub key_scope: MemoryKeyScope,
    pub kek_id: String,
    pub kek_version: String,
    pub wrapped_dek: String,
    pub algorithm: String,
    pub encryption_context_hash: String,
    pub rotation_epoch: u64,
    pub policy_decision_id: String,
    pub audit_id: String,
}

impl MemoryEnvelopeMetadata {
    pub fn from_metadata(metadata: Option<&Value>) -> MemoryResult<Option<Self>> {
        let Some(value) = metadata.and_then(|value| value.get(MEMORY_ENVELOPE_METADATA_KEY)) else {
            return Ok(None);
        };
        serde_json::from_value(value.clone())
            .map(Some)
            .map_err(MemoryError::from)
    }

    pub fn attach_to_metadata(&self, metadata: Option<Value>) -> MemoryResult<Value> {
        let mut object = match metadata {
            Some(Value::Object(object)) => object,
            Some(_) => {
                return Err(MemoryError::InvalidConfig(
                    "memory envelope metadata requires object metadata".to_string(),
                ));
            }
            None => Map::new(),
        };
        object.insert(
            MEMORY_ENVELOPE_METADATA_KEY.to_string(),
            serde_json::to_value(self)?,
        );
        Ok(Value::Object(object))
    }

    /// Validate persisted envelope metadata before a caller attempts a scoped
    /// decrypt. File-backed stores persist this structure outside the ciphertext,
    /// so malformed or wildcard scope metadata must fail before reaching KMS.
    pub fn validate_for_storage(&self) -> MemoryResult<()> {
        self.validate_required_fields()?;
        self.key_scope.validate_partitioned()
    }

    pub fn authority(&self) -> MemoryEnvelopeAuthority {
        MemoryEnvelopeAuthority::new(
            self.key_scope.clone(),
            self.policy_decision_id.clone(),
            self.audit_id.clone(),
        )
    }

    /// Validate every persisted authority anchor against trusted caller context
    /// and against the KMS AAD hash. This must run before consulting the DEK
    /// cache, otherwise edited metadata could reuse a cached key.
    pub fn validate_cryptographic_binding(
        &self,
        expected: &MemoryEnvelopeAuthority,
    ) -> MemoryResult<()> {
        self.validate_for_storage()?;
        expected.validate()?;
        if self.key_scope != expected.key_scope {
            return Err(MemoryError::InvalidConfig(
                "memory envelope key scope does not match trusted expected scope".to_string(),
            ));
        }
        if self.policy_decision_id != expected.policy_decision_id {
            return Err(MemoryError::InvalidConfig(
                "memory envelope policy decision does not match trusted authority".to_string(),
            ));
        }
        if self.audit_id != expected.audit_id {
            return Err(MemoryError::InvalidConfig(
                "memory envelope audit id does not match trusted authority".to_string(),
            ));
        }
        if self.algorithm != MEMORY_ENVELOPE_ALGORITHM {
            return Err(MemoryError::InvalidConfig(format!(
                "unsupported memory envelope algorithm `{}`",
                self.algorithm
            )));
        }
        let expected_hash = memory_encryption_context_hash(
            &self.key_scope,
            &self.kek_id,
            &self.kek_version,
            &self.algorithm,
            self.rotation_epoch,
            &self.policy_decision_id,
            &self.audit_id,
        )?;
        if self.encryption_context_hash != expected_hash {
            return Err(MemoryError::InvalidConfig(
                "memory envelope authority context hash does not match persisted metadata"
                    .to_string(),
            ));
        }
        Ok(())
    }

    fn validate_required_fields(&self) -> MemoryResult<()> {
        let required = [
            ("kek_id", self.kek_id.as_str()),
            ("kek_version", self.kek_version.as_str()),
            ("wrapped_dek", self.wrapped_dek.as_str()),
            ("algorithm", self.algorithm.as_str()),
            (
                "encryption_context_hash",
                self.encryption_context_hash.as_str(),
            ),
            ("policy_decision_id", self.policy_decision_id.as_str()),
            ("audit_id", self.audit_id.as_str()),
        ];
        for (field, value) in required {
            if value.trim().is_empty() {
                return Err(MemoryError::InvalidConfig(format!(
                    "hosted memory encryption metadata missing `{field}`"
                )));
            }
        }
        Ok(())
    }
}

/// Canonical KMS additional-authenticated-data hash for every authority-bearing
/// envelope field. Length-prefixing avoids delimiter ambiguity.
pub fn memory_encryption_context_hash(
    key_scope: &MemoryKeyScope,
    kek_id: &str,
    kek_version: &str,
    algorithm: &str,
    rotation_epoch: u64,
    policy_decision_id: &str,
    audit_id: &str,
) -> MemoryResult<String> {
    let mut hasher = Sha256::new();
    hasher.update(b"tandem-memory-envelope-authority-v2");
    update_hash_field(&mut hasher, &serde_json::to_vec(key_scope)?);
    update_hash_field(&mut hasher, kek_id.as_bytes());
    update_hash_field(&mut hasher, kek_version.as_bytes());
    update_hash_field(&mut hasher, algorithm.as_bytes());
    update_hash_field(&mut hasher, &rotation_epoch.to_le_bytes());
    update_hash_field(&mut hasher, policy_decision_id.as_bytes());
    update_hash_field(&mut hasher, audit_id.as_bytes());
    Ok(format!("{:x}", hasher.finalize()))
}

fn update_hash_field(hasher: &mut Sha256, value: &[u8]) {
    hasher.update((value.len() as u64).to_le_bytes());
    hasher.update(value);
}

pub fn hosted_memory_encryption_required() -> bool {
    std::env::var(HOSTED_ENCRYPTION_REQUIRED_ENV)
        .map(|value| matches!(value.as_str(), "1" | "true" | "TRUE" | "yes" | "YES"))
        .unwrap_or(false)
}

pub fn validate_memory_envelope_for_write(
    tenant_scope: &MemoryTenantScope,
    metadata: Option<&Value>,
) -> MemoryResult<()> {
    validate_memory_envelope_for_required_write(
        tenant_scope,
        metadata,
        hosted_memory_encryption_required(),
    )
}

pub fn validate_memory_envelope_for_required_write(
    tenant_scope: &MemoryTenantScope,
    metadata: Option<&Value>,
    encryption_required: bool,
) -> MemoryResult<()> {
    let envelope = MemoryEnvelopeMetadata::from_metadata(metadata)?;
    let Some(envelope) = envelope else {
        if encryption_required {
            return Err(MemoryError::InvalidConfig(
                "hosted memory encryption requires memory_envelope metadata".to_string(),
            ));
        }
        return Ok(());
    };

    envelope.validate_for_storage()?;
    if !envelope.key_scope.validates_against_tenant(tenant_scope) {
        return Err(MemoryError::InvalidConfig(
            "memory envelope key scope does not match tenant scope".to_string(),
        ));
    }
    // The key scope's department must match the row's owner_org_unit_id (TAN-662),
    // so a row can never be sealed under another department's DEK. Both `None`
    // (tenant-wide) is the matching case for undepartmented rows.
    let row_org_unit = crate::types::owner_org_unit_id_from_metadata(metadata);
    if envelope.key_scope.org_unit != row_org_unit {
        return Err(MemoryError::InvalidConfig(
            "memory envelope key scope org_unit does not match row owner_org_unit_id".to_string(),
        ));
    }
    validate_enterprise_source_binding(metadata, &envelope)
}

/// Percent-encode `%` and `/` so a caller-derived scope segment cannot inject a
/// structural delimiter (`/dept/`, `/source/`) into a `canonical_id` and collide
/// with a different scope (TAN-662). `%` is encoded first to keep the mapping
/// unambiguous (so a literal `%2F` cannot be confused with an encoded `/`).
fn encode_scope_segment(value: &str) -> String {
    value.replace('%', "%25").replace('/', "%2F")
}

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

fn validate_enterprise_source_binding(
    metadata: Option<&Value>,
    envelope: &MemoryEnvelopeMetadata,
) -> MemoryResult<()> {
    let Some(binding) = metadata.and_then(|value| value.get("enterprise_source_binding")) else {
        return Ok(());
    };
    if let Some(binding_data_class) = binding.get("data_class").and_then(Value::as_str) {
        let expected = serde_json::to_value(envelope.key_scope.data_class)?
            .as_str()
            .unwrap_or_default()
            .to_string();
        if binding_data_class != expected {
            return Err(MemoryError::InvalidConfig(
                "memory envelope data class does not match enterprise source binding".to_string(),
            ));
        }
    }
    if let Some(binding_id) = binding.get("binding_id").and_then(Value::as_str) {
        if envelope.key_scope.source_binding_id.as_deref() != Some(binding_id) {
            return Err(MemoryError::InvalidConfig(
                "memory envelope source binding does not match enterprise source binding"
                    .to_string(),
            ));
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn tenant_scope() -> MemoryTenantScope {
        MemoryTenantScope {
            org_id: "acme".to_string(),
            workspace_id: "finance".to_string(),
            deployment_id: Some("prod".to_string()),
        }
    }

    fn envelope(data_class: DataClass) -> MemoryEnvelopeMetadata {
        MemoryEnvelopeMetadata {
            key_scope: MemoryKeyScope::new(
                &tenant_scope(),
                data_class,
                Some("drive-1".to_string()),
            ),
            kek_id: "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance".to_string(),
            kek_version: "1".to_string(),
            wrapped_dek: "wrapped".to_string(),
            algorithm: "AES-256-GCM".to_string(),
            encryption_context_hash: "ctx-hash".to_string(),
            rotation_epoch: 0,
            policy_decision_id: "decision-1".to_string(),
            audit_id: "audit-1".to_string(),
        }
    }

    #[test]
    fn key_scope_canonical_id_includes_tenant_class_and_source() {
        let scope = MemoryKeyScope::new(
            &tenant_scope(),
            DataClass::FinancialRecord,
            Some("drive-1".to_string()),
        );
        assert_eq!(
            scope.canonical_id(),
            "tandem/memory/acme/finance/prod/financial_record/source/drive-1"
        );
    }

    #[test]
    fn key_scope_canonical_id_is_distinct_per_department() {
        // TAN-662: each department gets its own DEK scope, so canonical ids must
        // differ across departments in the same tenant + data class + source.
        let base = MemoryKeyScope::new(&tenant_scope(), DataClass::Internal, None);
        let sales = base
            .clone()
            .with_org_unit(Some("department/sales".to_string()));
        let engineering = base
            .clone()
            .with_org_unit(Some("department/engineering".to_string()));

        // The org-unit segment is percent-encoded so its internal `/` cannot be
        // confused with a structural delimiter.
        assert_eq!(
            sales.canonical_id(),
            "tandem/memory/acme/finance/prod/internal/dept/department%2Fsales"
        );
        assert_ne!(sales.canonical_id(), engineering.canonical_id());
        // The tenant-wide (no-department) scope is distinct from any department.
        assert_ne!(base.canonical_id(), sales.canonical_id());

        // With a source binding the department segment precedes the source.
        let sales_sourced = MemoryKeyScope::new(
            &tenant_scope(),
            DataClass::Internal,
            Some("drive-1".to_string()),
        )
        .with_org_unit(Some("department/sales".to_string()));
        assert_eq!(
            sales_sourced.canonical_id(),
            "tandem/memory/acme/finance/prod/internal/dept/department%2Fsales/source/drive-1"
        );
    }

    #[test]
    fn key_scope_canonical_id_resists_delimiter_injection() {
        // TAN-662 (review, P1): a department id that embeds the reserved
        // `/source/` delimiter must not collide with a distinct department+source
        // scope, and a literal `%2F` must not collide with an encoded `/`.
        let injected = MemoryKeyScope::new(&tenant_scope(), DataClass::Internal, None)
            .with_org_unit(Some("department/sales/source/drive-1".to_string()));
        let genuine = MemoryKeyScope::new(
            &tenant_scope(),
            DataClass::Internal,
            Some("drive-1".to_string()),
        )
        .with_org_unit(Some("department/sales".to_string()));
        assert_ne!(injected.canonical_id(), genuine.canonical_id());

        let literal_percent = MemoryKeyScope::new(&tenant_scope(), DataClass::Internal, None)
            .with_org_unit(Some("department%2Fsales".to_string()));
        let real_slash = MemoryKeyScope::new(&tenant_scope(), DataClass::Internal, None)
            .with_org_unit(Some("department/sales".to_string()));
        assert_ne!(literal_percent.canonical_id(), real_slash.canonical_id());
    }

    #[test]
    fn validation_accepts_matching_department_and_rejects_mismatch() {
        // Envelope key scope bound to `department/sales` and a row stamped the
        // same department validates…
        let mut envelope = envelope(DataClass::Internal);
        envelope.key_scope.source_binding_id = None;
        envelope.key_scope = envelope
            .key_scope
            .with_org_unit(Some("department/sales".to_string()));
        let metadata = envelope
            .attach_to_metadata(Some(serde_json::json!({
                "owner_org_unit_id": "department/sales"
            })))
            .expect("metadata");
        validate_memory_envelope_for_write(&tenant_scope(), Some(&metadata))
            .expect("matching department should validate");

        // …but a row owned by a different department is rejected: it must never be
        // sealed under another department's DEK.
        let mismatched = envelope
            .attach_to_metadata(Some(serde_json::json!({
                "owner_org_unit_id": "department/engineering"
            })))
            .expect("metadata");
        let err = validate_memory_envelope_for_write(&tenant_scope(), Some(&mismatched))
            .expect_err("department mismatch should fail");
        assert!(err
            .to_string()
            .contains("org_unit does not match row owner_org_unit_id"));
    }

    #[test]
    fn validation_rejects_wildcard_org_unit() {
        let mut envelope = envelope(DataClass::Internal);
        envelope.key_scope.source_binding_id = None;
        envelope.key_scope.org_unit = Some("*".to_string());
        let metadata = envelope.attach_to_metadata(None).expect("metadata");

        let err = validate_memory_envelope_for_write(&tenant_scope(), Some(&metadata))
            .expect_err("wildcard org_unit should fail");
        assert!(err.to_string().contains("wildcard `org_unit`"));
    }

    #[test]
    fn envelope_round_trips_through_metadata() {
        let envelope = envelope(DataClass::FinancialRecord);
        let metadata = envelope
            .attach_to_metadata(Some(serde_json::json!({"kind": "test"})))
            .expect("attach metadata");
        assert_eq!(
            MemoryEnvelopeMetadata::from_metadata(Some(&metadata))
                .expect("parse metadata")
                .as_ref(),
            Some(&envelope)
        );
    }

    #[test]
    fn validation_rejects_tenant_mismatch() {
        let mut envelope = envelope(DataClass::FinancialRecord);
        envelope.key_scope.workspace_id = "hr".to_string();
        let metadata = envelope.attach_to_metadata(None).expect("metadata");

        let err = validate_memory_envelope_for_write(&tenant_scope(), Some(&metadata))
            .expect_err("tenant mismatch should fail");
        assert!(err
            .to_string()
            .contains("key scope does not match tenant scope"));
    }

    #[test]
    fn validation_rejects_wildcard_key_scope() {
        let mut envelope = envelope(DataClass::FinancialRecord);
        envelope.key_scope.org_id = "*".to_string();
        let metadata = envelope.attach_to_metadata(None).expect("metadata");

        let err = validate_memory_envelope_for_write(&tenant_scope(), Some(&metadata))
            .expect_err("wildcard key scope should fail");
        assert!(err.to_string().contains("wildcard `org_id`"));
    }

    #[test]
    fn validation_rejects_source_binding_mismatch() {
        let metadata = envelope(DataClass::FinancialRecord)
            .attach_to_metadata(Some(serde_json::json!({
                "enterprise_source_binding": {
                    "binding_id": "other-drive",
                    "data_class": "financial_record"
                }
            })))
            .expect("metadata");

        let err = validate_memory_envelope_for_write(&tenant_scope(), Some(&metadata))
            .expect_err("source binding mismatch should fail");
        assert!(err.to_string().contains("source binding does not match"));
    }

    #[test]
    fn hosted_required_mode_rejects_missing_envelope() {
        let err = validate_memory_envelope_for_required_write(&tenant_scope(), None, true)
            .expect_err("hosted required mode should fail without metadata");

        assert!(err
            .to_string()
            .contains("requires memory_envelope metadata"));
    }

    #[test]
    fn local_mode_allows_missing_envelope() {
        validate_memory_envelope_for_required_write(&tenant_scope(), None, false)
            .expect("local mode should allow missing envelope metadata");
    }
}