1use std::collections::HashMap;
27use std::fmt;
28use std::sync::{Arc, RwLock};
29use std::time::Duration;
30
31use chrono::Utc;
32use moka::future::Cache;
33use uuid::Uuid;
34
35const MILLIS_IN_DAY: i64 = 24 * 60 * 60 * 1000;
36
37use super::crypto::{AesGcmCipher, AesKeySize, SecureKey};
38use super::io::EncryptedOutputFile;
39use super::key_metadata::StandardKeyMetadata;
40use super::kms::KeyManagementClient;
41use crate::io::OutputFile;
42use crate::sensitive::SensitiveBytes;
43use crate::spec::{EncryptedKey, FormatVersion, TableMetadataRef};
44use crate::{Error, ErrorKind, Result};
45
46pub const KEK_CREATED_AT_PROPERTY: &str = "KEY_TIMESTAMP";
49
50const DEFAULT_KEK_LIFESPAN_DAYS: i64 = 730;
52
53const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(3600);
55
56#[derive(typed_builder::TypedBuilder)]
60#[builder(mutators(
61 pub fn add_encryption_key(&mut self, key: EncryptedKey) {
63 self.encryption_keys
64 .write()
65 .expect("encryption_keys lock poisoned")
66 .insert(key.key_id().to_string(), key);
67 }
68 pub fn encryption_keys(&mut self, keys: HashMap<String, EncryptedKey>) {
70 self.encryption_keys = RwLock::new(keys);
71 }
72))]
73pub struct EncryptionManager {
74 kms_client: Arc<dyn KeyManagementClient>,
75 #[builder(
76 default = Cache::builder().time_to_live(DEFAULT_CACHE_TTL).build(),
77 setter(skip)
78 )]
79 kek_cache: Cache<String, SensitiveBytes>,
80 #[builder(default = AesKeySize::default())]
82 key_size: AesKeySize,
83 #[builder(setter(into))]
85 table_key_id: String,
86 #[builder(default = RwLock::new(HashMap::new()), via_mutators)]
90 encryption_keys: RwLock<HashMap<String, EncryptedKey>>,
91}
92
93impl fmt::Debug for EncryptionManager {
94 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
95 f.debug_struct("EncryptionManager")
96 .field("key_size", &self.key_size)
97 .field("table_key_id", &self.table_key_id)
98 .finish_non_exhaustive()
99 }
100}
101
102impl EncryptionManager {
103 pub(crate) fn from_table_metadata(
109 kms_client: Option<&Arc<dyn KeyManagementClient>>,
110 metadata: &TableMetadataRef,
111 ) -> Result<Option<Arc<Self>>> {
112 if metadata.format_version() < FormatVersion::V3 {
113 return Ok(None);
114 }
115
116 let table_properties = metadata.table_properties();
117 let encryption_key_id = table_properties.encryption_key_id()?;
118 let Some(table_key_id) = encryption_key_id.as_deref() else {
119 if kms_client.is_some() {
120 tracing::warn!(
121 "KeyManagementClient provided but table does not have encryption.key-id set"
122 );
123 }
124 return Ok(None);
125 };
126
127 let kms_client = kms_client.ok_or_else(|| {
128 Error::new(
129 ErrorKind::PreconditionFailed,
130 "Table has encryption.key-id set but no KeyManagementClient was provided to TableBuilder",
131 )
132 })?;
133
134 let em = EncryptionManager::builder()
135 .kms_client(Arc::clone(kms_client))
136 .table_key_id(table_key_id)
137 .encryption_keys(metadata.encryption_keys.clone())
138 .key_size(table_properties.data_encryption_key_size()?)
139 .build();
140 Ok(Some(Arc::new(em)))
141 }
142
143 pub fn generate_key_metadata(&self) -> StandardKeyMetadata {
146 StandardKeyMetadata::generate(self.key_size)
147 }
148
149 pub fn encrypt(&self, raw_output: OutputFile) -> EncryptedOutputFile {
154 EncryptedOutputFile::new(raw_output, self.generate_key_metadata())
155 }
156
157 pub async fn encrypt_manifest_list_key_metadata(
166 &self,
167 key_metadata: &StandardKeyMetadata,
168 ) -> Result<String> {
169 let kek = match self.find_active_kek()? {
170 Some(existing) => existing,
171 None => self.create_kek().await?,
172 };
173
174 let kek_bytes = self.unwrap_key_encryption_key(&kek).await?;
175
176 let aad = Self::kek_timestamp_aad(&kek)?;
178 let serialized = key_metadata.encode()?;
179 let wrapped_metadata = self.wrap_dek_with_kek(&serialized, &kek_bytes, Some(aad))?;
180
181 let wrapped_key = EncryptedKey::builder()
182 .key_id(Uuid::new_v4().to_string())
183 .encrypted_key_metadata(wrapped_metadata)
184 .encrypted_by_id(kek.key_id())
185 .build();
186
187 let wrapped_key_id = wrapped_key.key_id().to_string();
188 self.insert_encryption_key(wrapped_key);
189 Ok(wrapped_key_id)
190 }
191
192 pub async fn decrypt_manifest_list_key_metadata(
198 &self,
199 encryption_key_id: &str,
200 ) -> Result<StandardKeyMetadata> {
201 let encrypted_key = self
202 .encryption_keys
203 .read()
204 .expect("encryption_keys lock poisoned")
205 .get(encryption_key_id)
206 .cloned()
207 .ok_or_else(|| {
208 Error::new(
209 ErrorKind::DataInvalid,
210 format!("Encryption key '{encryption_key_id}' not found"),
211 )
212 })?;
213
214 let kek_key_id = encrypted_key.encrypted_by_id().ok_or_else(|| {
215 Error::new(
216 ErrorKind::DataInvalid,
217 format!(
218 "EncryptedKey '{}' has no encrypted_by_id",
219 encrypted_key.key_id()
220 ),
221 )
222 })?;
223
224 let bytes = self
225 .decrypt_dek(kek_key_id, encrypted_key.encrypted_key_metadata())
226 .await?;
227
228 StandardKeyMetadata::decode(bytes.as_bytes())
229 }
230
231 pub fn with_encryption_keys<F, R>(&self, f: F) -> R
236 where F: FnOnce(&HashMap<String, EncryptedKey>) -> R {
237 let keys = self
238 .encryption_keys
239 .read()
240 .expect("encryption_keys lock poisoned");
241 f(&keys)
242 }
243
244 fn insert_encryption_key(&self, key: EncryptedKey) {
245 self.encryption_keys
246 .write()
247 .expect("encryption_keys lock poisoned")
248 .insert(key.key_id().to_string(), key);
249 }
250
251 async fn create_kek(&self) -> Result<EncryptedKey> {
254 let (plaintext_kek, wrapped_kek) = if self.kms_client.supports_key_generation() {
255 let result = self.kms_client.generate_key(&self.table_key_id).await?;
256 (result.key().clone(), result.wrapped_key().to_vec())
257 } else {
258 let plaintext_key = SecureKey::generate(self.key_size);
259 let wrapped = self
260 .kms_client
261 .wrap_key(plaintext_key.as_bytes(), &self.table_key_id)
262 .await?;
263
264 (SensitiveBytes::new(plaintext_key.as_bytes()), wrapped)
265 };
266
267 let key_id = Uuid::new_v4().to_string();
268 let now_ms = Utc::now().timestamp_millis();
269
270 let mut properties = HashMap::new();
271 properties.insert(KEK_CREATED_AT_PROPERTY.to_string(), now_ms.to_string());
272
273 self.kek_cache.insert(key_id.clone(), plaintext_kek).await;
274
275 let kek = EncryptedKey::builder()
276 .key_id(key_id)
277 .encrypted_key_metadata(wrapped_kek)
278 .encrypted_by_id(&self.table_key_id)
279 .properties(properties)
280 .build();
281
282 self.insert_encryption_key(kek.clone());
283 Ok(kek)
284 }
285
286 fn is_kek_expired(&self, kek: &EncryptedKey) -> bool {
288 let created_at_ms = match kek
289 .properties()
290 .get(KEK_CREATED_AT_PROPERTY)
291 .and_then(|ts| ts.parse::<i64>().ok())
292 {
293 Some(ts) => ts,
294 None => return true, };
296
297 let now_ms = Utc::now().timestamp_millis();
298 let lifespan_ms = DEFAULT_KEK_LIFESPAN_DAYS * MILLIS_IN_DAY;
299 (now_ms - created_at_ms) >= lifespan_ms
300 }
301
302 fn find_active_kek(&self) -> Result<Option<EncryptedKey>> {
304 let keys = self
305 .encryption_keys
306 .read()
307 .expect("encryption_keys lock poisoned");
308 Ok(keys
309 .values()
310 .filter(|kek| {
311 kek.encrypted_by_id()
312 .map(|id| id == self.table_key_id)
313 .unwrap_or(false)
314 && !self.is_kek_expired(kek)
315 })
316 .max_by_key(|kek| {
317 kek.properties()
318 .get(KEK_CREATED_AT_PROPERTY)
319 .and_then(|ts| ts.parse::<i64>().ok())
320 .unwrap_or(0)
321 })
322 .cloned())
323 }
324
325 async fn unwrap_key_encryption_key(&self, kek: &EncryptedKey) -> Result<SensitiveBytes> {
327 let cache_key = kek.key_id().to_string();
328
329 if let Some(cached) = self.kek_cache.get(&cache_key).await {
330 return Ok(cached);
331 }
332
333 let master_key_id = kek.encrypted_by_id().ok_or_else(|| {
334 Error::new(
335 ErrorKind::DataInvalid,
336 format!("KEK '{}' has no encrypted_by_id", kek.key_id()),
337 )
338 })?;
339
340 let plaintext = self
341 .kms_client
342 .unwrap_key(kek.encrypted_key_metadata(), master_key_id)
343 .await?;
344
345 self.kek_cache.insert(cache_key, plaintext.clone()).await;
346
347 Ok(plaintext)
348 }
349
350 async fn decrypt_dek(&self, kek_key_id: &str, wrapped_dek: &[u8]) -> Result<SensitiveBytes> {
353 let kek = self
354 .encryption_keys
355 .read()
356 .expect("encryption_keys lock poisoned")
357 .get(kek_key_id)
358 .cloned()
359 .ok_or_else(|| {
360 Error::new(
361 ErrorKind::DataInvalid,
362 format!("KEK not found in encryption keys: {kek_key_id}"),
363 )
364 })?;
365
366 let aad = Self::kek_timestamp_aad(&kek)?;
368
369 let kek_bytes = self.unwrap_key_encryption_key(&kek).await?;
370 self.unwrap_dek_with_kek(wrapped_dek, &kek_bytes, Some(aad))
371 .map_err(|e| {
372 Error::new(
373 e.kind(),
374 format!("Failed to unwrap key metadata with KEK '{kek_key_id}'"),
375 )
376 .with_source(e)
377 })
378 }
379
380 fn kek_timestamp_aad(kek: &EncryptedKey) -> Result<&[u8]> {
382 kek.properties()
383 .get(KEK_CREATED_AT_PROPERTY)
384 .map(|ts| ts.as_bytes())
385 .ok_or_else(|| {
386 Error::new(
387 ErrorKind::DataInvalid,
388 format!(
389 "KEK '{}' is missing required '{}' property",
390 kek.key_id(),
391 KEK_CREATED_AT_PROPERTY
392 ),
393 )
394 })
395 }
396
397 fn wrap_dek_with_kek(
399 &self,
400 dek: &[u8],
401 kek: &SensitiveBytes,
402 aad: Option<&[u8]>,
403 ) -> Result<Vec<u8>> {
404 let key = SecureKey::try_from(kek.clone())?;
405 let cipher = AesGcmCipher::new(key);
406 cipher.encrypt(dek, aad)
407 }
408
409 fn unwrap_dek_with_kek(
411 &self,
412 wrapped_dek: &[u8],
413 kek: &SensitiveBytes,
414 aad: Option<&[u8]>,
415 ) -> Result<SensitiveBytes> {
416 let key = SecureKey::try_from(kek.clone())?;
417 let cipher = AesGcmCipher::new(key);
418 cipher.decrypt(wrapped_dek, aad).map(SensitiveBytes::new)
419 }
420}
421
422#[cfg(test)]
423mod tests {
424 use super::*;
425 use crate::encryption::EncryptedInputFile;
426 use crate::encryption::kms::MemoryKeyManagementClient;
427
428 fn create_test_kms() -> Arc<dyn KeyManagementClient> {
429 let kms = MemoryKeyManagementClient::new();
430 kms.add_master_key("master-1").unwrap();
431 Arc::new(kms)
432 }
433
434 fn create_test_manager() -> EncryptionManager {
435 EncryptionManager::builder()
436 .kms_client(create_test_kms())
437 .table_key_id("master-1")
438 .build()
439 }
440
441 #[tokio::test]
442 async fn test_create_kek() {
443 let mgr = create_test_manager();
444 let kek = mgr.create_kek().await.unwrap();
445
446 assert!(!kek.key_id().is_empty());
447 assert!(!kek.encrypted_key_metadata().is_empty());
448 assert_eq!(kek.encrypted_by_id(), Some("master-1"));
449 assert!(kek.properties().contains_key(KEK_CREATED_AT_PROPERTY));
450 }
451
452 fn sample_key_metadata() -> StandardKeyMetadata {
453 StandardKeyMetadata::try_new(b"0123456789abcdef")
454 .unwrap()
455 .with_aad_prefix(b"test-aad-prefix!")
456 }
457
458 #[tokio::test]
459 async fn test_wrap_unwrap_key_metadata_roundtrip() {
460 let mgr = create_test_manager();
461 let plaintext = sample_key_metadata();
462
463 let key_id = mgr
464 .encrypt_manifest_list_key_metadata(&plaintext)
465 .await
466 .unwrap();
467
468 assert_eq!(mgr.with_encryption_keys(|k| k.len()), 2);
470
471 let decrypted = mgr
472 .decrypt_manifest_list_key_metadata(&key_id)
473 .await
474 .unwrap();
475 assert_eq!(decrypted, plaintext);
476 }
477
478 #[tokio::test]
479 async fn test_kek_reuse_when_not_expired() {
480 let mgr = create_test_manager();
481
482 let _id1 = mgr
484 .encrypt_manifest_list_key_metadata(&sample_key_metadata())
485 .await
486 .unwrap();
487 let kek_id = mgr.with_encryption_keys(|keys| {
488 assert_eq!(keys.len(), 2);
489 keys.values()
490 .find(|k| k.encrypted_by_id() == Some("master-1"))
491 .unwrap()
492 .key_id()
493 .to_string()
494 });
495
496 let id2 = mgr
498 .encrypt_manifest_list_key_metadata(&sample_key_metadata())
499 .await
500 .unwrap();
501 let entry2 = mgr.with_encryption_keys(|keys| {
502 assert_eq!(keys.len(), 3);
503 keys.get(&id2).cloned().unwrap()
504 });
505 assert_eq!(entry2.encrypted_by_id(), Some(kek_id.as_str()));
506 }
507
508 #[tokio::test]
509 async fn test_kek_rotation_when_expired() {
510 let kms = create_test_kms();
511
512 let three_years_ago_ms = Utc::now().timestamp_millis() - (3 * 365 * MILLIS_IN_DAY);
514 let mut properties = HashMap::new();
515 properties.insert(
516 KEK_CREATED_AT_PROPERTY.to_string(),
517 three_years_ago_ms.to_string(),
518 );
519
520 let kek_key = SecureKey::generate(AesKeySize::Bits128);
522 let wrapped = kms.wrap_key(kek_key.as_bytes(), "master-1").await.unwrap();
523
524 let old_kek = EncryptedKey::builder()
525 .key_id("expired-kek")
526 .encrypted_key_metadata(wrapped)
527 .encrypted_by_id("master-1")
528 .properties(properties)
529 .build();
530
531 let mgr = EncryptionManager::builder()
533 .kms_client(kms)
534 .table_key_id("master-1")
535 .add_encryption_key(old_kek.clone())
536 .build();
537
538 let new_entry_id = mgr
540 .encrypt_manifest_list_key_metadata(&sample_key_metadata())
541 .await
542 .unwrap();
543 let entry = mgr
544 .with_encryption_keys(|keys| keys.get(&new_entry_id).cloned())
545 .unwrap();
546 let used_kek_id = entry.encrypted_by_id().unwrap();
547 assert_ne!(used_kek_id, old_kek.key_id());
548 }
549
550 #[tokio::test]
551 async fn test_is_kek_expired_no_timestamp() {
552 let mgr = create_test_manager();
553
554 let kek = EncryptedKey::builder()
556 .key_id("no-ts")
557 .encrypted_key_metadata(vec![0u8; 32])
558 .build();
559
560 assert!(mgr.is_kek_expired(&kek));
561 }
562
563 #[tokio::test]
564 async fn test_decrypt_with_unknown_key_id() {
565 let mgr = create_test_manager();
566 let result = mgr.decrypt_manifest_list_key_metadata("nonexistent").await;
567 assert!(result.is_err());
568 }
569
570 #[tokio::test]
571 async fn test_kek_cache_hit() {
572 let mgr = create_test_manager();
573
574 let key_id = mgr
576 .encrypt_manifest_list_key_metadata(&sample_key_metadata())
577 .await
578 .unwrap();
579
580 let _ = mgr
582 .decrypt_manifest_list_key_metadata(&key_id)
583 .await
584 .unwrap();
585 }
586
587 #[tokio::test]
588 async fn test_unwrap_fails_when_kek_missing_timestamp() {
589 let mgr = create_test_manager();
590
591 let entry_id = mgr
593 .encrypt_manifest_list_key_metadata(&sample_key_metadata())
594 .await
595 .unwrap();
596
597 let mut keys = mgr.with_encryption_keys(|k| k.clone());
600 let kek_id = keys
601 .get(&entry_id)
602 .unwrap()
603 .encrypted_by_id()
604 .unwrap()
605 .to_string();
606 let kek = keys.remove(&kek_id).unwrap();
607 let kek_no_ts = EncryptedKey::builder()
608 .key_id(kek.key_id())
609 .encrypted_key_metadata(kek.encrypted_key_metadata())
610 .encrypted_by_id(kek.encrypted_by_id().unwrap())
611 .build();
612 keys.insert(kek_no_ts.key_id().to_string(), kek_no_ts);
613
614 let mgr = EncryptionManager::builder()
615 .kms_client(create_test_kms())
616 .table_key_id("master-1")
617 .encryption_keys(keys)
618 .build();
619
620 let result = mgr.decrypt_manifest_list_key_metadata(&entry_id).await;
621 assert!(result.is_err());
622 let err = result.unwrap_err();
623 assert_eq!(err.kind(), ErrorKind::DataInvalid);
624 assert!(
625 err.to_string().contains(KEK_CREATED_AT_PROPERTY),
626 "error should mention the missing property: {err}"
627 );
628 }
629
630 #[tokio::test]
631 async fn test_unwrap_fails_when_kek_timestamp_tampered() {
632 let mgr = create_test_manager();
633
634 let entry_id = mgr
636 .encrypt_manifest_list_key_metadata(&sample_key_metadata())
637 .await
638 .unwrap();
639
640 let mut keys = mgr.with_encryption_keys(|k| k.clone());
642 let kek_id = keys
643 .get(&entry_id)
644 .unwrap()
645 .encrypted_by_id()
646 .unwrap()
647 .to_string();
648 let kek = keys.remove(&kek_id).unwrap();
649 let mut tampered_properties = kek.properties().clone();
650 tampered_properties.insert(KEK_CREATED_AT_PROPERTY.to_string(), "9999999".to_string());
651 let tampered_kek = EncryptedKey::builder()
652 .key_id(kek.key_id())
653 .encrypted_key_metadata(kek.encrypted_key_metadata())
654 .encrypted_by_id(kek.encrypted_by_id().unwrap())
655 .properties(tampered_properties)
656 .build();
657 keys.insert(tampered_kek.key_id().to_string(), tampered_kek);
658
659 let mgr = EncryptionManager::builder()
660 .kms_client(create_test_kms())
661 .table_key_id("master-1")
662 .encryption_keys(keys)
663 .build();
664
665 let result = mgr.decrypt_manifest_list_key_metadata(&entry_id).await;
667 assert!(
668 result.is_err(),
669 "tampered timestamp should cause decryption failure"
670 );
671 }
672
673 #[tokio::test]
674 async fn test_encrypt_decrypt_roundtrip() {
675 use crate::io::FileIO;
676
677 let io = FileIO::new_with_memory();
678 let path = "memory:///test/encrypt_roundtrip.bin";
679
680 let kms = MemoryKeyManagementClient::new();
681 kms.add_master_key("master-1").unwrap();
682 let mgr = EncryptionManager::builder()
683 .kms_client(Arc::new(kms) as Arc<dyn KeyManagementClient>)
684 .table_key_id("master-1")
685 .build();
686
687 let output = io.new_output(path).unwrap();
688 let encrypted_output = mgr.encrypt(output);
689
690 let plaintext = b"Hello, encrypted Iceberg round-trip!";
691 let serialized_metadata = encrypted_output.key_metadata().encode().unwrap();
692 encrypted_output
693 .write(bytes::Bytes::from(plaintext.to_vec()))
694 .await
695 .unwrap();
696
697 let input = io.new_input(path).unwrap();
698 let parsed_metadata = StandardKeyMetadata::decode(&serialized_metadata).unwrap();
699 let decrypted_file = EncryptedInputFile::new(input, parsed_metadata);
700
701 let content = decrypted_file.read().await.unwrap();
702 assert_eq!(&content[..], plaintext);
703 }
704}