Skip to content

feat(encryption) [14/N] Read / Write Encrypted puffin files #2822

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xanderbailey wants to merge 1 commit into
base: main
Choose a base branch
from
+183 −69
Open

feat(encryption) [14/N] Read / Write Encrypted puffin files #2822

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion crates/iceberg/public-api.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<iceberg::puffin::Blob>
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<Self>
pub async fn iceberg::puffin::PuffinReader::new_from_encrypted(encrypted_input: iceberg::encryption::EncryptedInputFile) -> iceberg::Result<Self>
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<alloc::string::String, alloc::string::String>, compress_footer: bool) -> iceberg::Result<Self>
pub async fn iceberg::puffin::PuffinWriter::new_from_encrypted(encrypted_output: &iceberg::encryption::EncryptedOutputFile, properties: std::collections::hash::map::HashMap<alloc::string::String, alloc::string::String>, compress_footer: bool) -> iceberg::Result<Self>
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
Expand Down
106 changes: 57 additions & 49 deletions crates/iceberg/src/puffin/metadata.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -277,21 +277,14 @@ impl FileMetadata {
}

/// Returns the file metadata about a Puffin file
pub(crate) async fn read(input_file: &InputFile) -> Result<FileMetadata> {
let file_read = input_file.reader().await?;

pub(crate) async fn read(file_read: &dyn FileRead, file_length: u64) -> Result<FileMetadata> {
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
Expand All @@ -312,25 +305,23 @@ 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<FileMetadata> {
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
let first_four_bytes = file_read.read(0..FileMetadata::MAGIC_LENGTH.into()).await?;
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 =
Expand All @@ -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
Expand All @@ -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]
Expand All @@ -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::{
Expand All @@ -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<FileMetadata> {
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<FileMetadata> {
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();

Expand Down Expand Up @@ -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(),
Expand All @@ -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(),
Expand All @@ -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(),
Expand All @@ -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(),
Expand All @@ -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(),
Expand All @@ -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(),
Expand All @@ -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(),
Expand All @@ -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(),
Expand All @@ -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(),
Expand All @@ -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: {
Expand Down Expand Up @@ -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: {
Expand All @@ -718,7 +726,7 @@ mod tests {
.await;

assert_eq!(
FileMetadata::read(&input_file)
read_file_metadata(&input_file)
.await
.unwrap_err()
.to_string(),
Expand All @@ -741,7 +749,7 @@ mod tests {
.await;

assert_eq!(
FileMetadata::read(&input_file)
read_file_metadata(&input_file)
.await
.unwrap_err()
.to_string(),
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -882,7 +890,7 @@ mod tests {
.await;

assert_eq!(
FileMetadata::read(&input_file)
read_file_metadata(&input_file)
.await
.unwrap_err()
.to_string(),
Expand All @@ -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(),
Expand All @@ -911,23 +919,23 @@ 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())
}

#[tokio::test]
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())
}

#[tokio::test]
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())
Expand All @@ -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();

Expand All @@ -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();

Expand All @@ -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();

Expand All @@ -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(),
Expand Down Expand Up @@ -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);
Expand Down
Loading
Loading