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