1use crate::types::{MemoryError, MemoryResult, MemoryTenantScope};
2use serde::{Deserialize, Serialize};
3use serde_json::{Map, Value};
4use tandem_enterprise_contract::DataClass;
5
6pub const MEMORY_ENVELOPE_METADATA_KEY: &str = "memory_envelope";
7const HOSTED_ENCRYPTION_REQUIRED_ENV: &str = "TANDEM_MEMORY_ENCRYPTION_REQUIRED";
8
9#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
10pub struct MemoryKeyScope {
11 pub org_id: String,
12 pub workspace_id: String,
13 #[serde(default, skip_serializing_if = "Option::is_none")]
14 pub deployment_id: Option<String>,
15 pub data_class: DataClass,
16 #[serde(default, skip_serializing_if = "Option::is_none")]
17 pub source_binding_id: Option<String>,
18}
19
20impl MemoryKeyScope {
21 pub fn new(
22 tenant_scope: &MemoryTenantScope,
23 data_class: DataClass,
24 source_binding_id: Option<String>,
25 ) -> Self {
26 Self {
27 org_id: tenant_scope.org_id.clone(),
28 workspace_id: tenant_scope.workspace_id.clone(),
29 deployment_id: tenant_scope.deployment_id.clone(),
30 data_class,
31 source_binding_id,
32 }
33 }
34
35 pub fn canonical_id(&self) -> String {
36 let deployment = self.deployment_id.as_deref().unwrap_or("default");
37 let class = serde_json::to_value(self.data_class)
38 .ok()
39 .and_then(|value| value.as_str().map(ToOwned::to_owned))
40 .unwrap_or_else(|| "unknown".to_string());
41 match self.source_binding_id.as_deref() {
42 Some(source_binding_id) if !source_binding_id.trim().is_empty() => format!(
43 "tandem/memory/{}/{}/{}/{}/source/{}",
44 self.org_id, self.workspace_id, deployment, class, source_binding_id
45 ),
46 _ => format!(
47 "tandem/memory/{}/{}/{}/{}",
48 self.org_id, self.workspace_id, deployment, class
49 ),
50 }
51 }
52
53 fn validates_against_tenant(&self, tenant_scope: &MemoryTenantScope) -> bool {
54 self.org_id == tenant_scope.org_id
55 && self.workspace_id == tenant_scope.workspace_id
56 && self.deployment_id.as_deref().unwrap_or("")
57 == tenant_scope.deployment_id.as_deref().unwrap_or("")
58 }
59
60 fn validate_partitioned(&self) -> MemoryResult<()> {
61 for (field, value) in [
62 ("org_id", self.org_id.as_str()),
63 ("workspace_id", self.workspace_id.as_str()),
64 ] {
65 if is_wildcard_scope(value) {
66 return Err(MemoryError::InvalidConfig(format!(
67 "memory envelope key scope must not use wildcard `{field}`"
68 )));
69 }
70 }
71 if self
72 .deployment_id
73 .as_deref()
74 .map(is_wildcard_scope)
75 .unwrap_or(false)
76 {
77 return Err(MemoryError::InvalidConfig(
78 "memory envelope key scope must not use wildcard `deployment_id`".to_string(),
79 ));
80 }
81 Ok(())
82 }
83}
84
85#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
86pub struct MemoryEnvelopeMetadata {
87 pub key_scope: MemoryKeyScope,
88 pub kek_id: String,
89 pub kek_version: String,
90 pub wrapped_dek: String,
91 pub algorithm: String,
92 pub encryption_context_hash: String,
93 pub rotation_epoch: u64,
94 pub policy_decision_id: String,
95 pub audit_id: String,
96}
97
98impl MemoryEnvelopeMetadata {
99 pub fn from_metadata(metadata: Option<&Value>) -> MemoryResult<Option<Self>> {
100 let Some(value) = metadata.and_then(|value| value.get(MEMORY_ENVELOPE_METADATA_KEY)) else {
101 return Ok(None);
102 };
103 serde_json::from_value(value.clone())
104 .map(Some)
105 .map_err(MemoryError::from)
106 }
107
108 pub fn attach_to_metadata(&self, metadata: Option<Value>) -> MemoryResult<Value> {
109 let mut object = match metadata {
110 Some(Value::Object(object)) => object,
111 Some(_) => {
112 return Err(MemoryError::InvalidConfig(
113 "memory envelope metadata requires object metadata".to_string(),
114 ));
115 }
116 None => Map::new(),
117 };
118 object.insert(
119 MEMORY_ENVELOPE_METADATA_KEY.to_string(),
120 serde_json::to_value(self)?,
121 );
122 Ok(Value::Object(object))
123 }
124
125 fn validate_required_fields(&self) -> MemoryResult<()> {
126 let required = [
127 ("kek_id", self.kek_id.as_str()),
128 ("kek_version", self.kek_version.as_str()),
129 ("wrapped_dek", self.wrapped_dek.as_str()),
130 ("algorithm", self.algorithm.as_str()),
131 (
132 "encryption_context_hash",
133 self.encryption_context_hash.as_str(),
134 ),
135 ("policy_decision_id", self.policy_decision_id.as_str()),
136 ("audit_id", self.audit_id.as_str()),
137 ];
138 for (field, value) in required {
139 if value.trim().is_empty() {
140 return Err(MemoryError::InvalidConfig(format!(
141 "hosted memory encryption metadata missing `{field}`"
142 )));
143 }
144 }
145 Ok(())
146 }
147}
148
149pub fn hosted_memory_encryption_required() -> bool {
150 std::env::var(HOSTED_ENCRYPTION_REQUIRED_ENV)
151 .map(|value| matches!(value.as_str(), "1" | "true" | "TRUE" | "yes" | "YES"))
152 .unwrap_or(false)
153}
154
155pub fn validate_memory_envelope_for_write(
156 tenant_scope: &MemoryTenantScope,
157 metadata: Option<&Value>,
158) -> MemoryResult<()> {
159 validate_memory_envelope_for_required_write(
160 tenant_scope,
161 metadata,
162 hosted_memory_encryption_required(),
163 )
164}
165
166pub fn validate_memory_envelope_for_required_write(
167 tenant_scope: &MemoryTenantScope,
168 metadata: Option<&Value>,
169 encryption_required: bool,
170) -> MemoryResult<()> {
171 let envelope = MemoryEnvelopeMetadata::from_metadata(metadata)?;
172 let Some(envelope) = envelope else {
173 if encryption_required {
174 return Err(MemoryError::InvalidConfig(
175 "hosted memory encryption requires memory_envelope metadata".to_string(),
176 ));
177 }
178 return Ok(());
179 };
180
181 envelope.validate_required_fields()?;
182 envelope.key_scope.validate_partitioned()?;
183 if !envelope.key_scope.validates_against_tenant(tenant_scope) {
184 return Err(MemoryError::InvalidConfig(
185 "memory envelope key scope does not match tenant scope".to_string(),
186 ));
187 }
188 validate_enterprise_source_binding(metadata, &envelope)
189}
190
191fn is_wildcard_scope(value: &str) -> bool {
192 matches!(
193 value.trim().to_ascii_lowercase().as_str(),
194 "" | "*" | "all" | "global" | "default"
195 )
196}
197
198fn validate_enterprise_source_binding(
199 metadata: Option<&Value>,
200 envelope: &MemoryEnvelopeMetadata,
201) -> MemoryResult<()> {
202 let Some(binding) = metadata.and_then(|value| value.get("enterprise_source_binding")) else {
203 return Ok(());
204 };
205 if let Some(binding_data_class) = binding.get("data_class").and_then(Value::as_str) {
206 let expected = serde_json::to_value(envelope.key_scope.data_class)?
207 .as_str()
208 .unwrap_or_default()
209 .to_string();
210 if binding_data_class != expected {
211 return Err(MemoryError::InvalidConfig(
212 "memory envelope data class does not match enterprise source binding".to_string(),
213 ));
214 }
215 }
216 if let Some(binding_id) = binding.get("binding_id").and_then(Value::as_str) {
217 if envelope.key_scope.source_binding_id.as_deref() != Some(binding_id) {
218 return Err(MemoryError::InvalidConfig(
219 "memory envelope source binding does not match enterprise source binding"
220 .to_string(),
221 ));
222 }
223 }
224 Ok(())
225}
226
227#[cfg(test)]
228mod tests {
229 use super::*;
230
231 fn tenant_scope() -> MemoryTenantScope {
232 MemoryTenantScope {
233 org_id: "acme".to_string(),
234 workspace_id: "finance".to_string(),
235 deployment_id: Some("prod".to_string()),
236 }
237 }
238
239 fn envelope(data_class: DataClass) -> MemoryEnvelopeMetadata {
240 MemoryEnvelopeMetadata {
241 key_scope: MemoryKeyScope::new(
242 &tenant_scope(),
243 data_class,
244 Some("drive-1".to_string()),
245 ),
246 kek_id: "projects/acme/locations/global/keyRings/memory/cryptoKeys/finance".to_string(),
247 kek_version: "1".to_string(),
248 wrapped_dek: "wrapped".to_string(),
249 algorithm: "AES-256-GCM".to_string(),
250 encryption_context_hash: "ctx-hash".to_string(),
251 rotation_epoch: 0,
252 policy_decision_id: "decision-1".to_string(),
253 audit_id: "audit-1".to_string(),
254 }
255 }
256
257 #[test]
258 fn key_scope_canonical_id_includes_tenant_class_and_source() {
259 let scope = MemoryKeyScope::new(
260 &tenant_scope(),
261 DataClass::FinancialRecord,
262 Some("drive-1".to_string()),
263 );
264 assert_eq!(
265 scope.canonical_id(),
266 "tandem/memory/acme/finance/prod/financial_record/source/drive-1"
267 );
268 }
269
270 #[test]
271 fn envelope_round_trips_through_metadata() {
272 let envelope = envelope(DataClass::FinancialRecord);
273 let metadata = envelope
274 .attach_to_metadata(Some(serde_json::json!({"kind": "test"})))
275 .expect("attach metadata");
276 assert_eq!(
277 MemoryEnvelopeMetadata::from_metadata(Some(&metadata))
278 .expect("parse metadata")
279 .as_ref(),
280 Some(&envelope)
281 );
282 }
283
284 #[test]
285 fn validation_rejects_tenant_mismatch() {
286 let mut envelope = envelope(DataClass::FinancialRecord);
287 envelope.key_scope.workspace_id = "hr".to_string();
288 let metadata = envelope.attach_to_metadata(None).expect("metadata");
289
290 let err = validate_memory_envelope_for_write(&tenant_scope(), Some(&metadata))
291 .expect_err("tenant mismatch should fail");
292 assert!(err
293 .to_string()
294 .contains("key scope does not match tenant scope"));
295 }
296
297 #[test]
298 fn validation_rejects_wildcard_key_scope() {
299 let mut envelope = envelope(DataClass::FinancialRecord);
300 envelope.key_scope.org_id = "*".to_string();
301 let metadata = envelope.attach_to_metadata(None).expect("metadata");
302
303 let err = validate_memory_envelope_for_write(&tenant_scope(), Some(&metadata))
304 .expect_err("wildcard key scope should fail");
305 assert!(err.to_string().contains("wildcard `org_id`"));
306 }
307
308 #[test]
309 fn validation_rejects_source_binding_mismatch() {
310 let metadata = envelope(DataClass::FinancialRecord)
311 .attach_to_metadata(Some(serde_json::json!({
312 "enterprise_source_binding": {
313 "binding_id": "other-drive",
314 "data_class": "financial_record"
315 }
316 })))
317 .expect("metadata");
318
319 let err = validate_memory_envelope_for_write(&tenant_scope(), Some(&metadata))
320 .expect_err("source binding mismatch should fail");
321 assert!(err.to_string().contains("source binding does not match"));
322 }
323
324 #[test]
325 fn hosted_required_mode_rejects_missing_envelope() {
326 let err = validate_memory_envelope_for_required_write(&tenant_scope(), None, true)
327 .expect_err("hosted required mode should fail without metadata");
328
329 assert!(err
330 .to_string()
331 .contains("requires memory_envelope metadata"));
332 }
333
334 #[test]
335 fn local_mode_allows_missing_envelope() {
336 validate_memory_envelope_for_required_write(&tenant_scope(), None, false)
337 .expect("local mode should allow missing envelope metadata");
338 }
339}