diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index c2013294e3..fb6d001063 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1228,12 +1228,14 @@ pub struct iceberg::puffin::PuffinReader impl iceberg::puffin::PuffinReader pub async fn iceberg::puffin::PuffinReader::blob(&self, blob_metadata: &iceberg::puffin::BlobMetadata) -> iceberg::Result pub async fn iceberg::puffin::PuffinReader::file_metadata(&self) -> iceberg::Result<&iceberg::puffin::FileMetadata> -pub fn iceberg::puffin::PuffinReader::new(input_file: iceberg::io::InputFile) -> Self +pub async fn iceberg::puffin::PuffinReader::new(input_file: iceberg::io::InputFile) -> iceberg::Result +pub async fn iceberg::puffin::PuffinReader::new_from_encrypted(encrypted_input: iceberg::encryption::EncryptedInputFile) -> iceberg::Result pub struct iceberg::puffin::PuffinWriter impl iceberg::puffin::PuffinWriter pub async fn iceberg::puffin::PuffinWriter::add(&mut self, blob: iceberg::puffin::Blob, compression_codec: iceberg::compression::CompressionCodec) -> iceberg::Result<()> pub async fn iceberg::puffin::PuffinWriter::close(self) -> iceberg::Result<()> pub async fn iceberg::puffin::PuffinWriter::new(output_file: &iceberg::io::OutputFile, properties: std::collections::hash::map::HashMap, compress_footer: bool) -> iceberg::Result +pub async fn iceberg::puffin::PuffinWriter::new_from_encrypted(encrypted_output: &iceberg::encryption::EncryptedOutputFile, properties: std::collections::hash::map::HashMap, compress_footer: bool) -> iceberg::Result pub const iceberg::puffin::APACHE_DATASKETCHES_THETA_V1: &str pub const iceberg::puffin::CREATED_BY_PROPERTY: &str pub const iceberg::puffin::DELETION_VECTOR_V1: &str diff --git a/crates/iceberg/src/puffin/metadata.rs b/crates/iceberg/src/puffin/metadata.rs index 10cfd265da..3c1d25b1d9 100644 --- a/crates/iceberg/src/puffin/metadata.rs +++ b/crates/iceberg/src/puffin/metadata.rs @@ -21,7 +21,7 @@ use bytes::Bytes; use serde::{Deserialize, Serialize}; use crate::compression::CompressionCodec; -use crate::io::{FileRead, InputFile}; +use crate::io::FileRead; use crate::{Error, ErrorKind, Result}; /// Human-readable identification of the application writing the file, along with its version. @@ -277,21 +277,14 @@ impl FileMetadata { } /// Returns the file metadata about a Puffin file - pub(crate) async fn read(input_file: &InputFile) -> Result { - let file_read = input_file.reader().await?; - + pub(crate) async fn read(file_read: &dyn FileRead, file_length: u64) -> Result { let first_four_bytes = file_read.read(0..FileMetadata::MAGIC_LENGTH.into()).await?; FileMetadata::check_magic(&first_four_bytes)?; - let input_file_length = input_file.metadata().await?.size; let footer_payload_length = - FileMetadata::read_footer_payload_length(file_read.as_ref(), input_file_length).await?; - let footer_bytes = FileMetadata::read_footer_bytes( - file_read.as_ref(), - input_file_length, - footer_payload_length, - ) - .await?; + FileMetadata::read_footer_payload_length(file_read, file_length).await?; + let footer_bytes = + FileMetadata::read_footer_bytes(file_read, file_length, footer_payload_length).await?; let magic_length = FileMetadata::MAGIC_LENGTH as usize; // check first four bytes of footer @@ -312,16 +305,14 @@ impl FileMetadata { /// read option. #[allow(dead_code)] pub(crate) async fn read_with_prefetch( - input_file: &InputFile, + file_read: &dyn FileRead, + file_length: u64, prefetch_hint: u8, ) -> Result { if prefetch_hint > 16 { - let input_file_length = input_file.metadata().await?.size; - let file_read = input_file.reader().await?; - // Hint cannot be larger than input file - if prefetch_hint as u64 > input_file_length { - return FileMetadata::read(input_file).await; + if prefetch_hint as u64 > file_length { + return FileMetadata::read(file_read, file_length).await; } // Validate file header magic @@ -329,8 +320,8 @@ impl FileMetadata { FileMetadata::check_magic(&first_four_bytes)?; // Read footer based on prefetch hint - let start = input_file_length - prefetch_hint as u64; - let end = input_file_length; + let start = file_length - prefetch_hint as u64; + let end = file_length; let footer_bytes = file_read.read(start..end).await?; let payload_length_start = @@ -350,7 +341,7 @@ impl FileMetadata { + FileMetadata::FOOTER_STRUCT_LENGTH as usize + FileMetadata::MAGIC_LENGTH as usize; if footer_length > prefetch_hint as usize { - return FileMetadata::read(input_file).await; + return FileMetadata::read(file_read, file_length).await; } // Read footer bytes @@ -369,7 +360,7 @@ impl FileMetadata { return FileMetadata::from_json_str(&footer_payload_str); } - FileMetadata::read(input_file).await + FileMetadata::read(file_read, file_length).await } #[inline] @@ -392,7 +383,6 @@ mod tests { use bytes::Bytes; use tempfile::TempDir; - use crate::ErrorKind; use crate::io::{FileIO, InputFile}; use crate::puffin::metadata::{BlobMetadata, CompressionCodec, FileMetadata}; use crate::puffin::test_utils::{ @@ -401,9 +391,27 @@ mod tests { java_zstd_compressed_metric_input_file, uncompressed_metric_file_metadata, zstd_compressed_metric_file_metadata, }; + use crate::{ErrorKind, Result}; const INVALID_MAGIC_VALUE: [u8; 4] = [80, 70, 65, 0]; + /// Reads file metadata from an [`InputFile`], resolving its reader and length. + async fn read_file_metadata(input_file: &InputFile) -> Result { + let file_read = input_file.reader().await?; + let file_length = input_file.metadata().await?.size; + FileMetadata::read(file_read.as_ref(), file_length).await + } + + /// Reads file metadata with a prefetch hint from an [`InputFile`]. + async fn read_file_metadata_with_prefetch( + input_file: &InputFile, + prefetch_hint: u8, + ) -> Result { + let file_read = input_file.reader().await?; + let file_length = input_file.metadata().await?.size; + FileMetadata::read_with_prefetch(file_read.as_ref(), file_length, prefetch_hint).await + } + async fn input_file_with_bytes(temp_dir: &TempDir, slice: &[u8]) -> InputFile { let file_io = FileIO::new_with_fs(); @@ -448,7 +456,7 @@ mod tests { let input_file = input_file_with_bytes(&temp_dir, &bytes).await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -471,7 +479,7 @@ mod tests { let input_file = input_file_with_bytes(&temp_dir, &bytes).await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -494,7 +502,7 @@ mod tests { let input_file = input_file_with_bytes(&temp_dir, &bytes).await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -519,7 +527,7 @@ mod tests { let input_file = input_file_with_bytes(&temp_dir, &bytes).await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -544,7 +552,7 @@ mod tests { let input_file = input_file_with_bytes(&temp_dir, &bytes).await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -567,7 +575,7 @@ mod tests { let input_file = input_file_with_bytes(&temp_dir, &bytes).await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -590,7 +598,7 @@ mod tests { let input_file = input_file_with_bytes(&temp_dir, &bytes).await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -616,7 +624,7 @@ mod tests { let input_file = input_file_with_bytes(&temp_dir, &bytes).await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -639,7 +647,7 @@ mod tests { let input_file = input_file_with_bytes(&temp_dir, &bytes).await; assert_eq!( - FileMetadata::read(&input_file).await.unwrap(), + read_file_metadata(&input_file).await.unwrap(), FileMetadata { blobs: vec![], properties: HashMap::new(), @@ -663,7 +671,7 @@ mod tests { .await; assert_eq!( - FileMetadata::read(&input_file).await.unwrap(), + read_file_metadata(&input_file).await.unwrap(), FileMetadata { blobs: vec![], properties: { @@ -692,7 +700,7 @@ mod tests { .await; assert_eq!( - FileMetadata::read(&input_file).await.unwrap(), + read_file_metadata(&input_file).await.unwrap(), FileMetadata { blobs: vec![], properties: { @@ -718,7 +726,7 @@ mod tests { .await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -741,7 +749,7 @@ mod tests { .await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -781,7 +789,7 @@ mod tests { .await; assert_eq!( - FileMetadata::read(&input_file).await.unwrap(), + read_file_metadata(&input_file).await.unwrap(), FileMetadata { blobs: vec![ BlobMetadata { @@ -835,7 +843,7 @@ mod tests { .await; assert_eq!( - FileMetadata::read(&input_file).await.unwrap(), + read_file_metadata(&input_file).await.unwrap(), FileMetadata { blobs: vec![BlobMetadata { r#type: "type-a".to_string(), @@ -882,7 +890,7 @@ mod tests { .await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -899,7 +907,7 @@ mod tests { let input_file = input_file_with_payload(&temp_dir, r#""blobs" = []"#).await; assert_eq!( - FileMetadata::read(&input_file) + read_file_metadata(&input_file) .await .unwrap_err() .to_string(), @@ -911,7 +919,7 @@ mod tests { async fn test_read_file_metadata_of_uncompressed_empty_file() { let input_file = java_empty_uncompressed_input_file(); - let file_metadata = FileMetadata::read(&input_file).await.unwrap(); + let file_metadata = read_file_metadata(&input_file).await.unwrap(); assert_eq!(file_metadata, empty_footer_payload()) } @@ -919,7 +927,7 @@ mod tests { async fn test_read_file_metadata_of_uncompressed_metric_data() { let input_file = java_uncompressed_metric_input_file(); - let file_metadata = FileMetadata::read(&input_file).await.unwrap(); + let file_metadata = read_file_metadata(&input_file).await.unwrap(); assert_eq!(file_metadata, uncompressed_metric_file_metadata()) } @@ -927,7 +935,7 @@ mod tests { async fn test_read_file_metadata_of_zstd_compressed_metric_data() { let input_file = java_zstd_compressed_metric_input_file(); - let file_metadata = FileMetadata::read_with_prefetch(&input_file, 64) + let file_metadata = read_file_metadata_with_prefetch(&input_file, 64) .await .unwrap(); assert_eq!(file_metadata, zstd_compressed_metric_file_metadata()) @@ -936,7 +944,7 @@ mod tests { #[tokio::test] async fn test_read_file_metadata_of_empty_file_with_prefetching() { let input_file = java_empty_uncompressed_input_file(); - let file_metadata = FileMetadata::read_with_prefetch(&input_file, 64) + let file_metadata = read_file_metadata_with_prefetch(&input_file, 64) .await .unwrap(); @@ -946,7 +954,7 @@ mod tests { #[tokio::test] async fn test_read_file_metadata_of_uncompressed_metric_data_with_prefetching() { let input_file = java_uncompressed_metric_input_file(); - let file_metadata = FileMetadata::read_with_prefetch(&input_file, 64) + let file_metadata = read_file_metadata_with_prefetch(&input_file, 64) .await .unwrap(); @@ -956,7 +964,7 @@ mod tests { #[tokio::test] async fn test_read_file_metadata_of_zstd_compressed_metric_data_with_prefetching() { let input_file = java_zstd_compressed_metric_input_file(); - let file_metadata = FileMetadata::read_with_prefetch(&input_file, 64) + let file_metadata = read_file_metadata_with_prefetch(&input_file, 64) .await .unwrap(); @@ -983,11 +991,11 @@ mod tests { let input_file = input_file_with_bytes(&temp_dir, &bytes).await; assert_eq!( - FileMetadata::read(&input_file).await.unwrap_err().kind(), + read_file_metadata(&input_file).await.unwrap_err().kind(), ErrorKind::DataInvalid, ); assert_eq!( - FileMetadata::read_with_prefetch(&input_file, prefetch_hint) + read_file_metadata_with_prefetch(&input_file, prefetch_hint) .await .unwrap_err() .kind(), @@ -1018,7 +1026,7 @@ mod tests { let input_file = input_file_with_payload(&temp_dir, payload).await; // Reading metadata should succeed (lazy validation) - let result = FileMetadata::read(&input_file).await; + let result = read_file_metadata(&input_file).await; assert!(result.is_ok()); let metadata = result.unwrap(); assert_eq!(metadata.blobs.len(), 1); diff --git a/crates/iceberg/src/puffin/reader.rs b/crates/iceberg/src/puffin/reader.rs index 0aced4186f..01601f48e0 100644 --- a/crates/iceberg/src/puffin/reader.rs +++ b/crates/iceberg/src/puffin/reader.rs @@ -19,21 +19,41 @@ use tokio::sync::OnceCell; use super::validate_puffin_compression; use crate::Result; -use crate::io::InputFile; +use crate::encryption::EncryptedInputFile; +use crate::io::{FileRead, InputFile}; use crate::puffin::blob::Blob; use crate::puffin::metadata::{BlobMetadata, FileMetadata}; /// Puffin reader pub struct PuffinReader { - input_file: InputFile, + file_read: Box, + file_length: u64, file_metadata: OnceCell, } impl PuffinReader { - /// Returns a new Puffin reader - pub fn new(input_file: InputFile) -> Self { + /// Returns a new Puffin reader for an unencrypted file. + pub async fn new(input_file: InputFile) -> Result { + let file_length = input_file.metadata().await?.size; + let file_read = input_file.reader().await?; + Ok(Self::from_parts(file_read, file_length)) + } + + /// Returns a new Puffin reader from an [`EncryptedInputFile`]. + /// + /// Use this when reading Puffin files with transparent decryption. The + /// reader operates over plaintext offsets and length, so all blob and + /// footer positions match those written to the unencrypted file. + pub async fn new_from_encrypted(encrypted_input: EncryptedInputFile) -> Result { + let file_length = encrypted_input.metadata().await?.size; + let file_read = encrypted_input.reader().await?; + Ok(Self::from_parts(file_read, file_length)) + } + + fn from_parts(file_read: Box, file_length: u64) -> Self { Self { - input_file, + file_read, + file_length, file_metadata: OnceCell::new(), } } @@ -41,7 +61,7 @@ impl PuffinReader { /// Returns file metadata pub async fn file_metadata(&self) -> Result<&FileMetadata> { self.file_metadata - .get_or_try_init(|| FileMetadata::read(&self.input_file)) + .get_or_try_init(|| FileMetadata::read(self.file_read.as_ref(), self.file_length)) .await } @@ -49,10 +69,9 @@ impl PuffinReader { pub async fn blob(&self, blob_metadata: &BlobMetadata) -> Result { validate_puffin_compression(blob_metadata.compression_codec)?; - let file_read = self.input_file.reader().await?; let start = blob_metadata.offset; let end = start + blob_metadata.length; - let bytes = file_read.read(start..end).await?; + let bytes = self.file_read.read(start..end).await?; let data = blob_metadata.compression_codec.decompress(bytes.to_vec())?; Ok(Blob { @@ -83,7 +102,7 @@ mod tests { #[tokio::test] async fn test_puffin_reader_uncompressed_metric_data() { let input_file = java_uncompressed_metric_input_file(); - let puffin_reader = PuffinReader::new(input_file); + let puffin_reader = PuffinReader::new(input_file).await.unwrap(); let file_metadata = puffin_reader.file_metadata().await.unwrap().clone(); assert_eq!(file_metadata, uncompressed_metric_file_metadata()); @@ -108,7 +127,7 @@ mod tests { #[tokio::test] async fn test_puffin_reader_zstd_compressed_metric_data() { let input_file = java_zstd_compressed_metric_input_file(); - let puffin_reader = PuffinReader::new(input_file); + let puffin_reader = PuffinReader::new(input_file).await.unwrap(); let file_metadata = puffin_reader.file_metadata().await.unwrap().clone(); assert_eq!(file_metadata, zstd_compressed_metric_file_metadata()); @@ -134,7 +153,7 @@ mod tests { async fn test_gzip_compression_rejected_on_blob_access() { // Use a real puffin file let input_file = java_uncompressed_metric_input_file(); - let reader = PuffinReader::new(input_file); + let reader = PuffinReader::new(input_file).await.unwrap(); // Create a BlobMetadata with Gzip compression let gzip_blob_metadata = BlobMetadata { diff --git a/crates/iceberg/src/puffin/writer.rs b/crates/iceberg/src/puffin/writer.rs index 4af4970b04..17c5a9298d 100644 --- a/crates/iceberg/src/puffin/writer.rs +++ b/crates/iceberg/src/puffin/writer.rs @@ -22,6 +22,7 @@ use bytes::Bytes; use super::validate_puffin_compression; use crate::Result; use crate::compression::CompressionCodec; +use crate::encryption::EncryptedOutputFile; use crate::io::{FileWrite, OutputFile}; use crate::puffin::blob::Blob; use crate::puffin::metadata::{BlobMetadata, FileMetadata, Flag}; @@ -38,12 +39,41 @@ pub struct PuffinWriter { } impl PuffinWriter { - /// Returns a new Puffin writer + /// Returns a new Puffin writer for an unencrypted file. pub async fn new( output_file: &OutputFile, properties: HashMap, compress_footer: bool, ) -> Result { + Ok(Self::from_writer( + output_file.writer().await?, + properties, + compress_footer, + )) + } + + /// Returns a new Puffin writer from an [`EncryptedOutputFile`]. + /// + /// Use this when writing Puffin files with transparent encryption. Blob + /// and footer offsets are recorded as plaintext positions, matching what + /// an unencrypted writer would produce. + pub async fn new_from_encrypted( + encrypted_output: &EncryptedOutputFile, + properties: HashMap, + compress_footer: bool, + ) -> Result { + Ok(Self::from_writer( + encrypted_output.writer().await?, + properties, + compress_footer, + )) + } + + fn from_writer( + writer: Box, + properties: HashMap, + compress_footer: bool, + ) -> Self { let mut flags = HashSet::::new(); let footer_compression_codec = if compress_footer { flags.insert(Flag::FooterPayloadCompressed); @@ -52,15 +82,15 @@ impl PuffinWriter { CompressionCodec::None }; - Ok(Self { - writer: output_file.writer().await?, + Self { + writer, is_header_written: false, num_bytes_written: 0, written_blobs_metadata: Vec::new(), properties, footer_compression_codec, flags, - }) + } } /// Adds blob to Puffin file @@ -184,8 +214,16 @@ mod tests { Ok(output_file) } + async fn read_file_metadata(input_file: &InputFile) -> FileMetadata { + let file_read = input_file.reader().await.unwrap(); + let file_length = input_file.metadata().await.unwrap().size; + FileMetadata::read(file_read.as_ref(), file_length) + .await + .unwrap() + } + async fn read_all_blobs_from_puffin_file(input_file: InputFile) -> Vec { - let puffin_reader = PuffinReader::new(input_file); + let puffin_reader = PuffinReader::new(input_file).await.unwrap(); let mut blobs = Vec::new(); let blobs_metadata = puffin_reader.file_metadata().await.unwrap().clone().blobs; for blob_metadata in blobs_metadata { @@ -204,7 +242,7 @@ mod tests { .to_input_file(); assert_eq!( - FileMetadata::read(&input_file).await.unwrap(), + read_file_metadata(&input_file).await, empty_footer_payload() ); @@ -240,7 +278,7 @@ mod tests { .to_input_file(); assert_eq!( - FileMetadata::read(&input_file).await.unwrap(), + read_file_metadata(&input_file).await, uncompressed_metric_file_metadata() ); @@ -260,7 +298,7 @@ mod tests { .to_input_file(); assert_eq!( - FileMetadata::read(&input_file).await.unwrap(), + read_file_metadata(&input_file).await, zstd_compressed_metric_file_metadata() ); @@ -354,4 +392,51 @@ mod tests { .contains("is not supported for Puffin files") ); } + + #[tokio::test] + async fn test_encrypted_write_read_roundtrip() { + use crate::encryption::{EncryptedInputFile, EncryptedOutputFile, StandardKeyMetadata}; + + let key_metadata = + || StandardKeyMetadata::new(b"0123456789abcdef").with_aad_prefix(b"test-aad-prefix!"); + + let file_io = FileIO::new_with_memory(); + let path = "memory:///test/encrypted.puffin"; + let blobs = vec![blob_0(), blob_1()]; + + // Write through the encrypting writer. + let encrypted_output = + EncryptedOutputFile::new(file_io.new_output(path).unwrap(), key_metadata()); + let mut writer = + PuffinWriter::new_from_encrypted(&encrypted_output, file_properties(), false) + .await + .unwrap(); + for blob in blobs.clone() { + writer.add(blob, CompressionCodec::None).await.unwrap(); + } + writer.close().await.unwrap(); + + // The ciphertext on disk must not equal a plaintext puffin file. + let raw = file_io.new_input(path).unwrap().read().await.unwrap(); + assert_ne!( + &raw[..FileMetadata::MAGIC_LENGTH as usize], + FileMetadata::MAGIC + ); + + // Read back through the decrypting reader over plaintext offsets. + let encrypted_input = + EncryptedInputFile::new(file_io.new_input(path).unwrap(), key_metadata()); + let reader = PuffinReader::new_from_encrypted(encrypted_input) + .await + .unwrap(); + + let file_metadata = reader.file_metadata().await.unwrap().clone(); + assert_eq!(file_metadata, uncompressed_metric_file_metadata()); + + let mut read_blobs = Vec::new(); + for blob_metadata in &file_metadata.blobs { + read_blobs.push(reader.blob(blob_metadata).await.unwrap()); + } + assert_eq!(read_blobs, blobs); + } }