feat(encryption) [13/N] Write encrypted manifest files

crates/iceberg/src/transaction/snapshot.rs

@@ -249,16 +249,25 @@ impl<'a> SnapshotProducer<'a> {
             DataFileFormat::Avro
         );
         let output_file = self.table.file_io().new_output(new_manifest_path)?;
-        let builder = ManifestWriterBuilder::new(
-            output_file,
-            Some(self.snapshot_id),
-            self.table.metadata().current_schema().clone(),
-            self.table
-                .metadata()
-                .default_partition_spec()
-                .as_ref()
-                .clone(),
-        );
+        let partition_spec = self
+            .table
+            .metadata()
+            .default_partition_spec()
+            .as_ref()
+            .clone();
+        let schema = self.table.metadata().current_schema().clone();
+
+        let builder = if let Some(em) = self.table.encryption_manager() {
+            ManifestWriterBuilder::new_from_encrypted(
+                em.encrypt(output_file),
+                Some(self.snapshot_id),
+                schema,
+                partition_spec,
+            )?
+        } else {
+            ManifestWriterBuilder::new(output_file, Some(self.snapshot_id), schema, partition_spec)
+        };
+
         match self.table.metadata().format_version() {
             FormatVersion::V1 => Ok(builder.build_v1()),
             FormatVersion::V2 => match content {
@@ -530,3 +539,179 @@ impl<'a> SnapshotProducer<'a> {
         Ok(ActionCommit::new(updates, requirements))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+    use std::fs::File;
+    use std::io::BufReader;
+    use std::sync::Arc;
+
+    use uuid::Uuid;
+
+    use super::SnapshotProducer;
+    use crate::TableIdent;
+    use crate::encryption::kms::{KeyManagementClient, MemoryKeyManagementClient};
+    use crate::encryption::{SensitiveBytes, StandardKeyMetadata};
+    use crate::io::FileIO;
+    use crate::spec::{
+        DataContentType, DataFileBuilder, DataFileFormat, Manifest, ManifestContentType,
+        ManifestStatus, Struct, TableMetadata,
+    };
+    use crate::table::Table;
+    use crate::test_utils::test_runtime;
+
+    /// Build a table backed by the V3 encryption fixture and an in-memory KMS,
+    /// so it has an [`EncryptionManager`](crate::encryption::EncryptionManager).
+    fn make_encrypted_table() -> Table {
+        let file = File::open(format!(
+            "{}/testdata/table_metadata/{}",
+            env!("CARGO_MANIFEST_DIR"),
+            "TableMetadataV3ValidEncryption.json"
+        ))
+        .unwrap();
+        let reader = BufReader::new(file);
+        let metadata = serde_json::from_reader::<_, TableMetadata>(reader).unwrap();
+
+        let kms: Arc<dyn KeyManagementClient> = {
+            let k = MemoryKeyManagementClient::new();
+            k.add_master_key_bytes(
+                "master-1",
+                SensitiveBytes::new([
+                    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c,
+                    0x0d, 0x0e, 0x0f,
+                ]),
+            )
+            .unwrap();
+            Arc::new(k)
+        };
+
+        Table::builder()
+            .metadata(metadata)
+            .metadata_location("memory:///table/metadata/v1.json")
+            .identifier(TableIdent::from_strs(["ns1", "test1"]).unwrap())
+            .file_io(FileIO::new_with_memory())
+            .kms_client(kms)
+            .runtime(test_runtime())
+            .build()
+            .unwrap()
+    }
+
+    fn data_file() -> crate::spec::DataFile {
+        DataFileBuilder::default()
+            .content(DataContentType::Data)
+            .file_path("memory:///table/data/00000.parquet".to_string())
+            .file_format(DataFileFormat::Parquet)
+            .partition(Struct::empty())
+            .record_count(100)
+            .file_size_in_bytes(4096)
+            .partition_spec_id(0)
+            .build()
+            .unwrap()
+    }
+
+    #[tokio::test]
+    async fn write_encrypted_manifest_roundtrips() {
+        let table = make_encrypted_table();
+        assert!(
+            table.encryption_manager().is_some(),
+            "fixture table should have an EncryptionManager"
+        );
+
+        let mut producer =
+            SnapshotProducer::new(&table, Uuid::new_v4(), HashMap::new(), vec![data_file()]);
+
+        let mut writer = producer
+            .new_manifest_writer(ManifestContentType::Data)
+            .unwrap();
+        writer
+            .add_entry(
+                crate::spec::ManifestEntry::builder()
+                    .status(ManifestStatus::Added)
+                    .data_file(data_file())
+                    .build(),
+            )
+            .unwrap();
+        let manifest_file = writer.write_manifest_file().await.unwrap();
+
+        // The manifest list entry must carry decodable key metadata.
+        let key_metadata_bytes = manifest_file
+            .key_metadata
+            .as_ref()
+            .expect("encrypted manifest must record key metadata");
+        assert!(!key_metadata_bytes.is_empty());
+        StandardKeyMetadata::decode(key_metadata_bytes)
+            .expect("recorded key metadata must decode as StandardKeyMetadata");
+
+        // The bytes on disk must actually be encrypted: a plaintext Avro parse
+        // must fail. (Guards against silently taking the unencrypted branch.)
+        let raw = table
+            .file_io()
+            .new_input(&manifest_file.manifest_path)
+            .unwrap()
+            .read()
+            .await
+            .unwrap();
+        assert!(
+            Manifest::parse_avro(&raw).is_err(),
+            "manifest bytes should not be readable as plaintext Avro"
+        );
+
+        // load_manifest self-decrypts using the recorded key metadata and must
+        // recover the entry we wrote.
+        let manifest = manifest_file.load_manifest(table.file_io()).await.unwrap();
+        assert_eq!(manifest.entries().len(), 1);
+        assert_eq!(
+            manifest.entries()[0].data_file().file_path(),
+            "memory:///table/data/00000.parquet"
+        );
+    }
+
+    #[tokio::test]
+    async fn write_plaintext_manifest_when_no_encryption() {
+        let table = crate::transaction::tests::make_v2_table();
+        assert!(table.encryption_manager().is_none());
+
+        let data_file = DataFileBuilder::default()
+            .content(DataContentType::Data)
+            .file_path("memory:///table/data/00000.parquet".to_string())
+            .file_format(DataFileFormat::Parquet)
+            .partition(Struct::from_iter([Some(crate::spec::Literal::long(0))]))
+            .record_count(100)
+            .file_size_in_bytes(4096)
+            .partition_spec_id(0)
+            .build()
+            .unwrap();
+
+        let mut producer = SnapshotProducer::new(&table, Uuid::new_v4(), HashMap::new(), vec![
+            data_file.clone(),
+        ]);
+
+        let mut writer = producer
+            .new_manifest_writer(ManifestContentType::Data)
+            .unwrap();
+        writer
+            .add_entry(
+                crate::spec::ManifestEntry::builder()
+                    .status(ManifestStatus::Added)
+                    .data_file(data_file)
+                    .build(),
+            )
+            .unwrap();
+        let manifest_file = writer.write_manifest_file().await.unwrap();
+
+        assert!(
+            manifest_file.key_metadata.is_none(),
+            "plaintext manifest must not record key metadata"
+        );
+
+        let raw = table
+            .file_io()
+            .new_input(&manifest_file.manifest_path)
+            .unwrap()
+            .read()
+            .await
+            .unwrap();
+        Manifest::parse_avro(&raw).expect("plaintext manifest bytes must parse as Avro");
+    }
+}