From c0178049d8fdf83cfb4b2c4d657b549ae1482c57 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Tue, 6 Jan 2026 15:45:39 -0800 Subject: [PATCH 01/26] Allow more encryption algorithms --- parquet/src/encryption/ciphers.rs | 44 ++++++++++++++++---- parquet/src/encryption/decrypt.rs | 18 ++++++++ parquet/src/encryption/encrypt.rs | 10 +++++ parquet/src/file/metadata/mod.rs | 3 ++ parquet/tests/encryption/encryption_async.rs | 4 +- 5 files changed, 69 insertions(+), 10 deletions(-) diff --git a/parquet/src/encryption/ciphers.rs b/parquet/src/encryption/ciphers.rs index a94c72dcd5ec..6e632876279c 100644 --- a/parquet/src/encryption/ciphers.rs +++ b/parquet/src/encryption/ciphers.rs @@ -19,7 +19,7 @@ use crate::errors::ParquetError; use crate::errors::ParquetError::General; use crate::errors::Result; use crate::file::metadata::HeapSize; -use ring::aead::{AES_128_GCM, Aad, LessSafeKey, NonceSequence, UnboundKey}; +use ring::aead::{AES_128_GCM, Aad, LessSafeKey, NonceSequence, UnboundKey, Algorithm}; use ring::rand::{SecureRandom, SystemRandom}; use std::fmt::Debug; @@ -41,9 +41,12 @@ pub(crate) struct RingGcmBlockDecryptor { impl RingGcmBlockDecryptor { pub(crate) fn new(key_bytes: &[u8]) -> Result { - // todo support other key sizes - let key = UnboundKey::new(&AES_128_GCM, key_bytes) - .map_err(|_| General("Failed to create AES key".to_string()))?; + Self::new_with_algorithm(&AES_128_GCM, key_bytes) + } + + pub(crate) fn new_with_algorithm(algorithm: &'static Algorithm, key_bytes: &[u8]) -> Result { + let key = UnboundKey::new(algorithm, key_bytes) + .map_err(|_| general_err!("Failed to create {:?} key", algorithm))?; Ok(Self { key: LessSafeKey::new(key), @@ -143,11 +146,13 @@ impl RingGcmBlockEncryptor { /// The nonce will advance appropriately with each block encryption and /// return an error if it wraps around. pub(crate) fn new(key_bytes: &[u8]) -> Result { - let rng = SystemRandom::new(); + Self::new_with_algorithm(&AES_128_GCM, key_bytes) + } - // todo support other key sizes - let key = UnboundKey::new(&AES_128_GCM, key_bytes) - .map_err(|e| general_err!("Error creating AES key: {}", e))?; + pub(crate) fn new_with_algorithm(algorithm: &'static Algorithm, key_bytes: &[u8]) -> Result { + let rng = SystemRandom::new(); + let key = UnboundKey::new(&algorithm, key_bytes) + .map_err(|e| general_err!("Error creating {:?} key: {}", algorithm, e))?; let nonce = CounterNonce::new(&rng)?; Ok(Self { @@ -188,6 +193,7 @@ impl BlockEncryptor for RingGcmBlockEncryptor { #[cfg(test)] mod tests { + use ring::aead::{AES_256_GCM, CHACHA20_POLY1305}; use super::*; #[test] @@ -204,4 +210,26 @@ mod tests { assert_eq!(plaintext, decrypted.as_slice()); } + + #[test] + fn test_round_trip_with_algorithm() { + let key = [0u8; 32]; + let mut encryptor = RingGcmBlockEncryptor::new_with_algorithm(&AES_256_GCM, &key).unwrap(); + let decryptor = RingGcmBlockDecryptor::new_with_algorithm(&AES_256_GCM, &key).unwrap(); + + let plaintext = b"hello, world!"; + let aad = b"some aad"; + + let ciphertext = encryptor.encrypt(plaintext, aad).unwrap(); + let decrypted = decryptor.decrypt(&ciphertext, aad).unwrap(); + + assert_eq!(plaintext, decrypted.as_slice()); + } + + #[test] + fn test_round_trip_with_incorrect_key_length() { + let key = [0u8; 16]; + assert!(RingGcmBlockEncryptor::new_with_algorithm(&CHACHA20_POLY1305, &key).is_err()); + assert!(RingGcmBlockDecryptor::new_with_algorithm(&CHACHA20_POLY1305, &key).is_err()); + } } diff --git a/parquet/src/encryption/decrypt.rs b/parquet/src/encryption/decrypt.rs index 0066523419de..210e5ee80cab 100644 --- a/parquet/src/encryption/decrypt.rs +++ b/parquet/src/encryption/decrypt.rs @@ -27,6 +27,7 @@ use std::collections::HashMap; use std::fmt::Formatter; use std::io::Read; use std::sync::Arc; +use ring::aead::{Algorithm, AES_128_GCM}; /// Trait for retrieving an encryption key using the key's metadata /// @@ -352,6 +353,7 @@ pub struct FileDecryptionProperties { keys: DecryptionKeys, aad_prefix: Option>, footer_signature_verification: bool, + algorithm: &'static Algorithm } impl HeapSize for FileDecryptionProperties { @@ -448,6 +450,7 @@ pub struct DecryptionPropertiesBuilder { column_keys: HashMap>, aad_prefix: Option>, footer_signature_verification: bool, + algorithm: &'static Algorithm } impl DecryptionPropertiesBuilder { @@ -459,6 +462,7 @@ impl DecryptionPropertiesBuilder { column_keys: HashMap::default(), aad_prefix: None, footer_signature_verification: true, + algorithm: &AES_128_GCM } } @@ -472,6 +476,7 @@ impl DecryptionPropertiesBuilder { keys, aad_prefix: self.aad_prefix, footer_signature_verification: self.footer_signature_verification, + algorithm: self.algorithm })) } @@ -505,6 +510,11 @@ impl DecryptionPropertiesBuilder { Ok(self) } + pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Result { + self.algorithm = algorithm; + Ok(self) + } + /// Disable verification of footer tags for files that use plaintext footers. /// Signature verification is enabled by default. pub fn disable_footer_signature_verification(mut self) -> Self { @@ -520,6 +530,7 @@ pub struct DecryptionPropertiesBuilderWithRetriever { key_retriever: Arc, aad_prefix: Option>, footer_signature_verification: bool, + algorithm: &'static Algorithm } impl DecryptionPropertiesBuilderWithRetriever { @@ -530,6 +541,7 @@ impl DecryptionPropertiesBuilderWithRetriever { key_retriever, aad_prefix: None, footer_signature_verification: true, + algorithm: &AES_128_GCM } } @@ -540,6 +552,7 @@ impl DecryptionPropertiesBuilderWithRetriever { keys, aad_prefix: self.aad_prefix, footer_signature_verification: self.footer_signature_verification, + algorithm: self.algorithm })) } @@ -551,6 +564,11 @@ impl DecryptionPropertiesBuilderWithRetriever { self } + pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self { + self.algorithm = algorithm; + self + } + /// Disable verification of footer tags for files that use plaintext footers. /// Signature verification is enabled by default. pub fn disable_footer_signature_verification(mut self) -> Self { diff --git a/parquet/src/encryption/encrypt.rs b/parquet/src/encryption/encrypt.rs index d69e3c02500a..d12dd5aeef6c 100644 --- a/parquet/src/encryption/encrypt.rs +++ b/parquet/src/encryption/encrypt.rs @@ -28,6 +28,7 @@ use ring::rand::{SecureRandom, SystemRandom}; use std::collections::{HashMap, HashSet}; use std::io::Write; use std::sync::Arc; +use ring::aead::{Algorithm, AES_128_GCM}; #[derive(Debug, Clone, PartialEq)] struct EncryptionKey { @@ -96,6 +97,7 @@ pub struct FileEncryptionProperties { column_keys: HashMap, aad_prefix: Option>, store_aad_prefix: bool, + algorithm: &'static Algorithm } impl FileEncryptionProperties { @@ -186,6 +188,7 @@ pub struct EncryptionPropertiesBuilder { column_keys: HashMap, aad_prefix: Option>, store_aad_prefix: bool, + algorithm: &'static Algorithm } impl EncryptionPropertiesBuilder { @@ -197,6 +200,7 @@ impl EncryptionPropertiesBuilder { aad_prefix: None, encrypt_footer: true, store_aad_prefix: false, + algorithm: &AES_128_GCM } } @@ -274,6 +278,11 @@ impl EncryptionPropertiesBuilder { self } + pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self { + self.algorithm = algorithm; + self + } + /// Build the encryption properties pub fn build(self) -> Result> { Ok(Arc::new(FileEncryptionProperties { @@ -282,6 +291,7 @@ impl EncryptionPropertiesBuilder { column_keys: self.column_keys, aad_prefix: self.aad_prefix, store_aad_prefix: self.store_aad_prefix, + algorithm: self.algorithm })) } } diff --git a/parquet/src/file/metadata/mod.rs b/parquet/src/file/metadata/mod.rs index 6bd426ee677f..62bd11354b4f 100644 --- a/parquet/src/file/metadata/mod.rs +++ b/parquet/src/file/metadata/mod.rs @@ -2045,6 +2045,9 @@ mod tests { .set_file_decryptor(Some(decryptor)) .build(); + #[cfg(feature = "encryption")] + let expected_size_with_decryptor = 3080; + #[cfg(not(feature = "encryption"))] let expected_size_with_decryptor = 3072; assert!(expected_size_with_decryptor > base_expected_size); diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index 51acd7374879..1c2e7e0fe15b 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -114,7 +114,7 @@ async fn test_misspecified_encryption_keys() { // Too short footer key check_for_error( - "Parquet error: Invalid footer key. Failed to create AES key", + "Parquet error: Invalid footer key. Failed to create AES_128_GCM key", &path, "bad_pwd".as_bytes(), column_1_key, @@ -144,7 +144,7 @@ async fn test_misspecified_encryption_keys() { // Too short column key check_for_error( - "Parquet error: Failed to create AES key", + "Parquet error: Failed to create AES_128_GCM key", &path, footer_key, "abc".as_bytes(), From e8da4b952866a2e0d7688852ec89a7a2b1002230 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Tue, 27 Jan 2026 11:20:09 -0800 Subject: [PATCH 02/26] Expose new API --- parquet/src/encryption/decrypt.rs | 6 ++++-- parquet/src/encryption/encrypt.rs | 6 +++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/parquet/src/encryption/decrypt.rs b/parquet/src/encryption/decrypt.rs index 210e5ee80cab..5a0426995f03 100644 --- a/parquet/src/encryption/decrypt.rs +++ b/parquet/src/encryption/decrypt.rs @@ -614,7 +614,8 @@ impl FileDecryptor { ) -> Result { let file_aad = [aad_prefix.as_slice(), aad_file_unique.as_slice()].concat(); let footer_key = decryption_properties.footer_key(footer_key_metadata)?; - let footer_decryptor = RingGcmBlockDecryptor::new(&footer_key).map_err(|e| { + let algorithm = decryption_properties.algorithm; + let footer_decryptor = RingGcmBlockDecryptor::new_with_algorithm(algorithm, &footer_key).map_err(|e| { general_err!( "Invalid footer key. {}", e.to_string().replace("Parquet error: ", "") @@ -659,7 +660,8 @@ impl FileDecryptor { let column_key = self .decryption_properties .column_key(column_name, key_metadata)?; - Ok(Arc::new(RingGcmBlockDecryptor::new(&column_key)?)) + let algorithm = self.decryption_properties.algorithm; + Ok(Arc::new(RingGcmBlockDecryptor::new_with_algorithm(algorithm, &column_key)?)) } pub(crate) fn get_column_metadata_decryptor( diff --git a/parquet/src/encryption/encrypt.rs b/parquet/src/encryption/encrypt.rs index d12dd5aeef6c..2762f7254ed2 100644 --- a/parquet/src/encryption/encrypt.rs +++ b/parquet/src/encryption/encrypt.rs @@ -353,8 +353,8 @@ impl FileEncryptor { /// Get the BlockEncryptor for the footer pub(crate) fn get_footer_encryptor(&self) -> Result> { - Ok(Box::new(RingGcmBlockEncryptor::new( - &self.properties.footer_key.key, + Ok(Box::new(RingGcmBlockEncryptor::new_with_algorithm( + &self.properties.algorithm, &self.properties.footer_key.key, )?)) } @@ -369,7 +369,7 @@ impl FileEncryptor { } match self.properties.column_keys.get(column_path) { None => Err(general_err!("Column '{}' is not encrypted", column_path)), - Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new(column_key.key())?)), + Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new_with_algorithm(&self.properties.algorithm, column_key.key())?)), } } } From 3e1b8a72e46600b7044d1f9f6fae8f5acaa9600f Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Tue, 27 Jan 2026 11:56:45 -0800 Subject: [PATCH 03/26] Fix API --- parquet/src/encryption/decrypt.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/parquet/src/encryption/decrypt.rs b/parquet/src/encryption/decrypt.rs index 5a0426995f03..c088c00c4618 100644 --- a/parquet/src/encryption/decrypt.rs +++ b/parquet/src/encryption/decrypt.rs @@ -510,9 +510,9 @@ impl DecryptionPropertiesBuilder { Ok(self) } - pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Result { + pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self { self.algorithm = algorithm; - Ok(self) + self } /// Disable verification of footer tags for files that use plaintext footers. From 1f1f0da24e01417e50bc6fc69fdd07c521146946 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Wed, 4 Feb 2026 20:36:05 -0800 Subject: [PATCH 04/26] Fix format and clippy --- parquet/src/encryption/ciphers.rs | 18 +++++++++++----- parquet/src/encryption/decrypt.rs | 34 +++++++++++++++++-------------- parquet/src/encryption/encrypt.rs | 18 +++++++++------- 3 files changed, 43 insertions(+), 27 deletions(-) diff --git a/parquet/src/encryption/ciphers.rs b/parquet/src/encryption/ciphers.rs index 6e632876279c..7421ad591a33 100644 --- a/parquet/src/encryption/ciphers.rs +++ b/parquet/src/encryption/ciphers.rs @@ -19,7 +19,7 @@ use crate::errors::ParquetError; use crate::errors::ParquetError::General; use crate::errors::Result; use crate::file::metadata::HeapSize; -use ring::aead::{AES_128_GCM, Aad, LessSafeKey, NonceSequence, UnboundKey, Algorithm}; +use ring::aead::{AES_128_GCM, Aad, Algorithm, LessSafeKey, NonceSequence, UnboundKey}; use ring::rand::{SecureRandom, SystemRandom}; use std::fmt::Debug; @@ -40,11 +40,15 @@ pub(crate) struct RingGcmBlockDecryptor { } impl RingGcmBlockDecryptor { + #[allow(dead_code)] pub(crate) fn new(key_bytes: &[u8]) -> Result { Self::new_with_algorithm(&AES_128_GCM, key_bytes) } - pub(crate) fn new_with_algorithm(algorithm: &'static Algorithm, key_bytes: &[u8]) -> Result { + pub(crate) fn new_with_algorithm( + algorithm: &'static Algorithm, + key_bytes: &[u8], + ) -> Result { let key = UnboundKey::new(algorithm, key_bytes) .map_err(|_| general_err!("Failed to create {:?} key", algorithm))?; @@ -145,13 +149,17 @@ impl RingGcmBlockEncryptor { /// Create a new `RingGcmBlockEncryptor` with a given key and random nonce. /// The nonce will advance appropriately with each block encryption and /// return an error if it wraps around. + #[allow(dead_code)] pub(crate) fn new(key_bytes: &[u8]) -> Result { Self::new_with_algorithm(&AES_128_GCM, key_bytes) } - pub(crate) fn new_with_algorithm(algorithm: &'static Algorithm, key_bytes: &[u8]) -> Result { + pub(crate) fn new_with_algorithm( + algorithm: &'static Algorithm, + key_bytes: &[u8], + ) -> Result { let rng = SystemRandom::new(); - let key = UnboundKey::new(&algorithm, key_bytes) + let key = UnboundKey::new(algorithm, key_bytes) .map_err(|e| general_err!("Error creating {:?} key: {}", algorithm, e))?; let nonce = CounterNonce::new(&rng)?; @@ -193,8 +201,8 @@ impl BlockEncryptor for RingGcmBlockEncryptor { #[cfg(test)] mod tests { - use ring::aead::{AES_256_GCM, CHACHA20_POLY1305}; use super::*; + use ring::aead::{AES_256_GCM, CHACHA20_POLY1305}; #[test] fn test_round_trip() { diff --git a/parquet/src/encryption/decrypt.rs b/parquet/src/encryption/decrypt.rs index c088c00c4618..d524042d6b92 100644 --- a/parquet/src/encryption/decrypt.rs +++ b/parquet/src/encryption/decrypt.rs @@ -22,12 +22,12 @@ use crate::encryption::modules::{ModuleType, create_footer_aad, create_module_aa use crate::errors::{ParquetError, Result}; use crate::file::column_crypto_metadata::ColumnCryptoMetaData; use crate::file::metadata::HeapSize; +use ring::aead::{AES_128_GCM, Algorithm}; use std::borrow::Cow; use std::collections::HashMap; use std::fmt::Formatter; use std::io::Read; use std::sync::Arc; -use ring::aead::{Algorithm, AES_128_GCM}; /// Trait for retrieving an encryption key using the key's metadata /// @@ -353,7 +353,7 @@ pub struct FileDecryptionProperties { keys: DecryptionKeys, aad_prefix: Option>, footer_signature_verification: bool, - algorithm: &'static Algorithm + algorithm: &'static Algorithm, } impl HeapSize for FileDecryptionProperties { @@ -450,7 +450,7 @@ pub struct DecryptionPropertiesBuilder { column_keys: HashMap>, aad_prefix: Option>, footer_signature_verification: bool, - algorithm: &'static Algorithm + algorithm: &'static Algorithm, } impl DecryptionPropertiesBuilder { @@ -462,7 +462,7 @@ impl DecryptionPropertiesBuilder { column_keys: HashMap::default(), aad_prefix: None, footer_signature_verification: true, - algorithm: &AES_128_GCM + algorithm: &AES_128_GCM, } } @@ -476,7 +476,7 @@ impl DecryptionPropertiesBuilder { keys, aad_prefix: self.aad_prefix, footer_signature_verification: self.footer_signature_verification, - algorithm: self.algorithm + algorithm: self.algorithm, })) } @@ -530,7 +530,7 @@ pub struct DecryptionPropertiesBuilderWithRetriever { key_retriever: Arc, aad_prefix: Option>, footer_signature_verification: bool, - algorithm: &'static Algorithm + algorithm: &'static Algorithm, } impl DecryptionPropertiesBuilderWithRetriever { @@ -541,7 +541,7 @@ impl DecryptionPropertiesBuilderWithRetriever { key_retriever, aad_prefix: None, footer_signature_verification: true, - algorithm: &AES_128_GCM + algorithm: &AES_128_GCM, } } @@ -552,7 +552,7 @@ impl DecryptionPropertiesBuilderWithRetriever { keys, aad_prefix: self.aad_prefix, footer_signature_verification: self.footer_signature_verification, - algorithm: self.algorithm + algorithm: self.algorithm, })) } @@ -615,12 +615,13 @@ impl FileDecryptor { let file_aad = [aad_prefix.as_slice(), aad_file_unique.as_slice()].concat(); let footer_key = decryption_properties.footer_key(footer_key_metadata)?; let algorithm = decryption_properties.algorithm; - let footer_decryptor = RingGcmBlockDecryptor::new_with_algorithm(algorithm, &footer_key).map_err(|e| { - general_err!( - "Invalid footer key. {}", - e.to_string().replace("Parquet error: ", "") - ) - })?; + let footer_decryptor = RingGcmBlockDecryptor::new_with_algorithm(algorithm, &footer_key) + .map_err(|e| { + general_err!( + "Invalid footer key. {}", + e.to_string().replace("Parquet error: ", "") + ) + })?; Ok(Self { footer_decryptor: Arc::new(footer_decryptor), @@ -661,7 +662,10 @@ impl FileDecryptor { .decryption_properties .column_key(column_name, key_metadata)?; let algorithm = self.decryption_properties.algorithm; - Ok(Arc::new(RingGcmBlockDecryptor::new_with_algorithm(algorithm, &column_key)?)) + Ok(Arc::new(RingGcmBlockDecryptor::new_with_algorithm( + algorithm, + &column_key, + )?)) } pub(crate) fn get_column_metadata_decryptor( diff --git a/parquet/src/encryption/encrypt.rs b/parquet/src/encryption/encrypt.rs index 2762f7254ed2..2cc38d7fe260 100644 --- a/parquet/src/encryption/encrypt.rs +++ b/parquet/src/encryption/encrypt.rs @@ -24,11 +24,11 @@ use crate::errors::{ParquetError, Result}; use crate::file::column_crypto_metadata::{ColumnCryptoMetaData, EncryptionWithColumnKey}; use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift}; use crate::schema::types::{ColumnDescPtr, SchemaDescriptor}; +use ring::aead::{AES_128_GCM, Algorithm}; use ring::rand::{SecureRandom, SystemRandom}; use std::collections::{HashMap, HashSet}; use std::io::Write; use std::sync::Arc; -use ring::aead::{Algorithm, AES_128_GCM}; #[derive(Debug, Clone, PartialEq)] struct EncryptionKey { @@ -97,7 +97,7 @@ pub struct FileEncryptionProperties { column_keys: HashMap, aad_prefix: Option>, store_aad_prefix: bool, - algorithm: &'static Algorithm + algorithm: &'static Algorithm, } impl FileEncryptionProperties { @@ -188,7 +188,7 @@ pub struct EncryptionPropertiesBuilder { column_keys: HashMap, aad_prefix: Option>, store_aad_prefix: bool, - algorithm: &'static Algorithm + algorithm: &'static Algorithm, } impl EncryptionPropertiesBuilder { @@ -200,7 +200,7 @@ impl EncryptionPropertiesBuilder { aad_prefix: None, encrypt_footer: true, store_aad_prefix: false, - algorithm: &AES_128_GCM + algorithm: &AES_128_GCM, } } @@ -291,7 +291,7 @@ impl EncryptionPropertiesBuilder { column_keys: self.column_keys, aad_prefix: self.aad_prefix, store_aad_prefix: self.store_aad_prefix, - algorithm: self.algorithm + algorithm: self.algorithm, })) } } @@ -354,7 +354,8 @@ impl FileEncryptor { /// Get the BlockEncryptor for the footer pub(crate) fn get_footer_encryptor(&self) -> Result> { Ok(Box::new(RingGcmBlockEncryptor::new_with_algorithm( - &self.properties.algorithm, &self.properties.footer_key.key, + self.properties.algorithm, + &self.properties.footer_key.key, )?)) } @@ -369,7 +370,10 @@ impl FileEncryptor { } match self.properties.column_keys.get(column_path) { None => Err(general_err!("Column '{}' is not encrypted", column_path)), - Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new_with_algorithm(&self.properties.algorithm, column_key.key())?)), + Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new_with_algorithm( + self.properties.algorithm, + column_key.key(), + )?)), } } } From 92b40944397da314c3670a7bfbe84e1c208232d2 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Wed, 4 Feb 2026 21:00:39 -0800 Subject: [PATCH 05/26] Document new api --- parquet/src/encryption/ciphers.rs | 7 ++++--- parquet/src/encryption/decrypt.rs | 2 ++ parquet/src/encryption/encrypt.rs | 1 + 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/parquet/src/encryption/ciphers.rs b/parquet/src/encryption/ciphers.rs index 7421ad591a33..9a93ec409444 100644 --- a/parquet/src/encryption/ciphers.rs +++ b/parquet/src/encryption/ciphers.rs @@ -45,6 +45,7 @@ impl RingGcmBlockDecryptor { Self::new_with_algorithm(&AES_128_GCM, key_bytes) } + /// Create a new `RingGcmBlockDecryptor` with an AEAD algorithm and a given key. pub(crate) fn new_with_algorithm( algorithm: &'static Algorithm, key_bytes: &[u8], @@ -146,14 +147,14 @@ pub(crate) struct RingGcmBlockEncryptor { } impl RingGcmBlockEncryptor { - /// Create a new `RingGcmBlockEncryptor` with a given key and random nonce. - /// The nonce will advance appropriately with each block encryption and - /// return an error if it wraps around. #[allow(dead_code)] pub(crate) fn new(key_bytes: &[u8]) -> Result { Self::new_with_algorithm(&AES_128_GCM, key_bytes) } + /// Create a new `RingGcmBlockEncryptor` with an AEAD algorithm, a given key and random nonce. + /// The nonce will advance appropriately with each block encryption and + /// return an error if it wraps around. pub(crate) fn new_with_algorithm( algorithm: &'static Algorithm, key_bytes: &[u8], diff --git a/parquet/src/encryption/decrypt.rs b/parquet/src/encryption/decrypt.rs index d524042d6b92..37e0444a7587 100644 --- a/parquet/src/encryption/decrypt.rs +++ b/parquet/src/encryption/decrypt.rs @@ -510,6 +510,7 @@ impl DecryptionPropertiesBuilder { Ok(self) } + /// The AEAD decryption algorithm to be used. pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self { self.algorithm = algorithm; self @@ -564,6 +565,7 @@ impl DecryptionPropertiesBuilderWithRetriever { self } + /// Specify the AEAD decryption algorithm to be used for decryption. pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self { self.algorithm = algorithm; self diff --git a/parquet/src/encryption/encrypt.rs b/parquet/src/encryption/encrypt.rs index 2cc38d7fe260..a67266962699 100644 --- a/parquet/src/encryption/encrypt.rs +++ b/parquet/src/encryption/encrypt.rs @@ -278,6 +278,7 @@ impl EncryptionPropertiesBuilder { self } + /// The AEAD encryption algorithm to be used pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self { self.algorithm = algorithm; self From 12c18882df630838290bcaae0455c5aa2ec199b7 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Fri, 13 Feb 2026 16:32:31 -0800 Subject: [PATCH 06/26] Pick algorithm by key size --- parquet/src/encryption/ciphers.rs | 49 +++++++++----------- parquet/src/encryption/decrypt.rs | 27 +---------- parquet/src/encryption/encrypt.rs | 17 +------ parquet/tests/encryption/encryption_async.rs | 8 +++- 4 files changed, 32 insertions(+), 69 deletions(-) diff --git a/parquet/src/encryption/ciphers.rs b/parquet/src/encryption/ciphers.rs index 9a93ec409444..eab24ea35380 100644 --- a/parquet/src/encryption/ciphers.rs +++ b/parquet/src/encryption/ciphers.rs @@ -19,7 +19,7 @@ use crate::errors::ParquetError; use crate::errors::ParquetError::General; use crate::errors::Result; use crate::file::metadata::HeapSize; -use ring::aead::{AES_128_GCM, Aad, Algorithm, LessSafeKey, NonceSequence, UnboundKey}; +use ring::aead::{AES_128_GCM, AES_256_GCM, Aad, LessSafeKey, NonceSequence, UnboundKey}; use ring::rand::{SecureRandom, SystemRandom}; use std::fmt::Debug; @@ -40,16 +40,17 @@ pub(crate) struct RingGcmBlockDecryptor { } impl RingGcmBlockDecryptor { - #[allow(dead_code)] - pub(crate) fn new(key_bytes: &[u8]) -> Result { - Self::new_with_algorithm(&AES_128_GCM, key_bytes) - } - - /// Create a new `RingGcmBlockDecryptor` with an AEAD algorithm and a given key. - pub(crate) fn new_with_algorithm( - algorithm: &'static Algorithm, + /// Create a new `RingGcmBlockDecryptor` with a given key. + pub(crate) fn new( key_bytes: &[u8], ) -> Result { + let algorithm = if key_bytes.len() == AES_128_GCM.key_len() { + &AES_128_GCM + } else if key_bytes.len() == AES_256_GCM.key_len() { + &AES_256_GCM + } else { + return Err(general_err!("Error creating RingGcmBlockDecryptor with unsupported key length: {}", key_bytes.len())); + }; let key = UnboundKey::new(algorithm, key_bytes) .map_err(|_| general_err!("Failed to create {:?} key", algorithm))?; @@ -147,19 +148,21 @@ pub(crate) struct RingGcmBlockEncryptor { } impl RingGcmBlockEncryptor { - #[allow(dead_code)] - pub(crate) fn new(key_bytes: &[u8]) -> Result { - Self::new_with_algorithm(&AES_128_GCM, key_bytes) - } - - /// Create a new `RingGcmBlockEncryptor` with an AEAD algorithm, a given key and random nonce. + /// Create a new `RingGcmBlockEncryptor` with a given key and random nonce. /// The nonce will advance appropriately with each block encryption and /// return an error if it wraps around. - pub(crate) fn new_with_algorithm( - algorithm: &'static Algorithm, + pub(crate) fn new( key_bytes: &[u8], ) -> Result { let rng = SystemRandom::new(); + let algorithm = if key_bytes.len() == AES_128_GCM.key_len() { + &AES_128_GCM + } else if key_bytes.len() == AES_256_GCM.key_len() { + &AES_256_GCM + } else { + return Err(general_err!("Error creating RingGcmBlockEncryptor with unsupported key length: {}", key_bytes.len())); + }; + let key = UnboundKey::new(algorithm, key_bytes) .map_err(|e| general_err!("Error creating {:?} key: {}", algorithm, e))?; let nonce = CounterNonce::new(&rng)?; @@ -203,7 +206,6 @@ impl BlockEncryptor for RingGcmBlockEncryptor { #[cfg(test)] mod tests { use super::*; - use ring::aead::{AES_256_GCM, CHACHA20_POLY1305}; #[test] fn test_round_trip() { @@ -223,8 +225,8 @@ mod tests { #[test] fn test_round_trip_with_algorithm() { let key = [0u8; 32]; - let mut encryptor = RingGcmBlockEncryptor::new_with_algorithm(&AES_256_GCM, &key).unwrap(); - let decryptor = RingGcmBlockDecryptor::new_with_algorithm(&AES_256_GCM, &key).unwrap(); + let mut encryptor = RingGcmBlockEncryptor::new(&key).unwrap(); + let decryptor = RingGcmBlockDecryptor::new(&key).unwrap(); let plaintext = b"hello, world!"; let aad = b"some aad"; @@ -234,11 +236,4 @@ mod tests { assert_eq!(plaintext, decrypted.as_slice()); } - - #[test] - fn test_round_trip_with_incorrect_key_length() { - let key = [0u8; 16]; - assert!(RingGcmBlockEncryptor::new_with_algorithm(&CHACHA20_POLY1305, &key).is_err()); - assert!(RingGcmBlockDecryptor::new_with_algorithm(&CHACHA20_POLY1305, &key).is_err()); - } } diff --git a/parquet/src/encryption/decrypt.rs b/parquet/src/encryption/decrypt.rs index 37e0444a7587..5ae8dd620680 100644 --- a/parquet/src/encryption/decrypt.rs +++ b/parquet/src/encryption/decrypt.rs @@ -22,7 +22,6 @@ use crate::encryption::modules::{ModuleType, create_footer_aad, create_module_aa use crate::errors::{ParquetError, Result}; use crate::file::column_crypto_metadata::ColumnCryptoMetaData; use crate::file::metadata::HeapSize; -use ring::aead::{AES_128_GCM, Algorithm}; use std::borrow::Cow; use std::collections::HashMap; use std::fmt::Formatter; @@ -353,7 +352,6 @@ pub struct FileDecryptionProperties { keys: DecryptionKeys, aad_prefix: Option>, footer_signature_verification: bool, - algorithm: &'static Algorithm, } impl HeapSize for FileDecryptionProperties { @@ -450,7 +448,6 @@ pub struct DecryptionPropertiesBuilder { column_keys: HashMap>, aad_prefix: Option>, footer_signature_verification: bool, - algorithm: &'static Algorithm, } impl DecryptionPropertiesBuilder { @@ -462,7 +459,6 @@ impl DecryptionPropertiesBuilder { column_keys: HashMap::default(), aad_prefix: None, footer_signature_verification: true, - algorithm: &AES_128_GCM, } } @@ -476,7 +472,6 @@ impl DecryptionPropertiesBuilder { keys, aad_prefix: self.aad_prefix, footer_signature_verification: self.footer_signature_verification, - algorithm: self.algorithm, })) } @@ -510,12 +505,6 @@ impl DecryptionPropertiesBuilder { Ok(self) } - /// The AEAD decryption algorithm to be used. - pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self { - self.algorithm = algorithm; - self - } - /// Disable verification of footer tags for files that use plaintext footers. /// Signature verification is enabled by default. pub fn disable_footer_signature_verification(mut self) -> Self { @@ -531,7 +520,6 @@ pub struct DecryptionPropertiesBuilderWithRetriever { key_retriever: Arc, aad_prefix: Option>, footer_signature_verification: bool, - algorithm: &'static Algorithm, } impl DecryptionPropertiesBuilderWithRetriever { @@ -542,7 +530,6 @@ impl DecryptionPropertiesBuilderWithRetriever { key_retriever, aad_prefix: None, footer_signature_verification: true, - algorithm: &AES_128_GCM, } } @@ -553,7 +540,6 @@ impl DecryptionPropertiesBuilderWithRetriever { keys, aad_prefix: self.aad_prefix, footer_signature_verification: self.footer_signature_verification, - algorithm: self.algorithm, })) } @@ -565,12 +551,6 @@ impl DecryptionPropertiesBuilderWithRetriever { self } - /// Specify the AEAD decryption algorithm to be used for decryption. - pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self { - self.algorithm = algorithm; - self - } - /// Disable verification of footer tags for files that use plaintext footers. /// Signature verification is enabled by default. pub fn disable_footer_signature_verification(mut self) -> Self { @@ -616,8 +596,7 @@ impl FileDecryptor { ) -> Result { let file_aad = [aad_prefix.as_slice(), aad_file_unique.as_slice()].concat(); let footer_key = decryption_properties.footer_key(footer_key_metadata)?; - let algorithm = decryption_properties.algorithm; - let footer_decryptor = RingGcmBlockDecryptor::new_with_algorithm(algorithm, &footer_key) + let footer_decryptor = RingGcmBlockDecryptor::new(&footer_key) .map_err(|e| { general_err!( "Invalid footer key. {}", @@ -663,9 +642,7 @@ impl FileDecryptor { let column_key = self .decryption_properties .column_key(column_name, key_metadata)?; - let algorithm = self.decryption_properties.algorithm; - Ok(Arc::new(RingGcmBlockDecryptor::new_with_algorithm( - algorithm, + Ok(Arc::new(RingGcmBlockDecryptor::new( &column_key, )?)) } diff --git a/parquet/src/encryption/encrypt.rs b/parquet/src/encryption/encrypt.rs index a67266962699..12d59ce89712 100644 --- a/parquet/src/encryption/encrypt.rs +++ b/parquet/src/encryption/encrypt.rs @@ -24,7 +24,6 @@ use crate::errors::{ParquetError, Result}; use crate::file::column_crypto_metadata::{ColumnCryptoMetaData, EncryptionWithColumnKey}; use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift}; use crate::schema::types::{ColumnDescPtr, SchemaDescriptor}; -use ring::aead::{AES_128_GCM, Algorithm}; use ring::rand::{SecureRandom, SystemRandom}; use std::collections::{HashMap, HashSet}; use std::io::Write; @@ -97,7 +96,6 @@ pub struct FileEncryptionProperties { column_keys: HashMap, aad_prefix: Option>, store_aad_prefix: bool, - algorithm: &'static Algorithm, } impl FileEncryptionProperties { @@ -188,7 +186,6 @@ pub struct EncryptionPropertiesBuilder { column_keys: HashMap, aad_prefix: Option>, store_aad_prefix: bool, - algorithm: &'static Algorithm, } impl EncryptionPropertiesBuilder { @@ -200,7 +197,6 @@ impl EncryptionPropertiesBuilder { aad_prefix: None, encrypt_footer: true, store_aad_prefix: false, - algorithm: &AES_128_GCM, } } @@ -278,12 +274,6 @@ impl EncryptionPropertiesBuilder { self } - /// The AEAD encryption algorithm to be used - pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self { - self.algorithm = algorithm; - self - } - /// Build the encryption properties pub fn build(self) -> Result> { Ok(Arc::new(FileEncryptionProperties { @@ -292,7 +282,6 @@ impl EncryptionPropertiesBuilder { column_keys: self.column_keys, aad_prefix: self.aad_prefix, store_aad_prefix: self.store_aad_prefix, - algorithm: self.algorithm, })) } } @@ -354,8 +343,7 @@ impl FileEncryptor { /// Get the BlockEncryptor for the footer pub(crate) fn get_footer_encryptor(&self) -> Result> { - Ok(Box::new(RingGcmBlockEncryptor::new_with_algorithm( - self.properties.algorithm, + Ok(Box::new(RingGcmBlockEncryptor::new( &self.properties.footer_key.key, )?)) } @@ -371,8 +359,7 @@ impl FileEncryptor { } match self.properties.column_keys.get(column_path) { None => Err(general_err!("Column '{}' is not encrypted", column_path)), - Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new_with_algorithm( - self.properties.algorithm, + Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new( column_key.key(), )?)), } diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index e76980a74c78..34cd4e5b6465 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -115,7 +115,7 @@ async fn test_misspecified_encryption_keys() { // Too short footer key check_for_error( - "Parquet error: Invalid footer key. Failed to create AES_128_GCM key", + format!("Parquet error: Invalid footer key. Error creating RingGcmBlockDecryptor with unsupported key length: {}", "bad_pwd".len()).as_str(), &path, "bad_pwd".as_bytes(), column_1_key, @@ -145,7 +145,11 @@ async fn test_misspecified_encryption_keys() { // Too short column key check_for_error( - "Parquet error: Failed to create AES_128_GCM key", + format!( + "Parquet error: Error creating RingGcmBlockDecryptor with unsupported key length: {}", + "abc".len() + ) + .as_str(), &path, footer_key, "abc".as_bytes(), From 3ec6b708c32404641cc7e1f77f81f88a15f0e27e Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Fri, 13 Feb 2026 16:51:50 -0800 Subject: [PATCH 07/26] Remove unnecessary tests --- parquet/src/encryption/ciphers.rs | 15 --------------- parquet/src/file/metadata/mod.rs | 3 --- 2 files changed, 18 deletions(-) diff --git a/parquet/src/encryption/ciphers.rs b/parquet/src/encryption/ciphers.rs index eab24ea35380..726b292d1c70 100644 --- a/parquet/src/encryption/ciphers.rs +++ b/parquet/src/encryption/ciphers.rs @@ -221,19 +221,4 @@ mod tests { assert_eq!(plaintext, decrypted.as_slice()); } - - #[test] - fn test_round_trip_with_algorithm() { - let key = [0u8; 32]; - let mut encryptor = RingGcmBlockEncryptor::new(&key).unwrap(); - let decryptor = RingGcmBlockDecryptor::new(&key).unwrap(); - - let plaintext = b"hello, world!"; - let aad = b"some aad"; - - let ciphertext = encryptor.encrypt(plaintext, aad).unwrap(); - let decrypted = decryptor.decrypt(&ciphertext, aad).unwrap(); - - assert_eq!(plaintext, decrypted.as_slice()); - } } diff --git a/parquet/src/file/metadata/mod.rs b/parquet/src/file/metadata/mod.rs index 36226721f04c..ca3a9e10978b 100644 --- a/parquet/src/file/metadata/mod.rs +++ b/parquet/src/file/metadata/mod.rs @@ -2118,9 +2118,6 @@ mod tests { .set_file_decryptor(Some(decryptor)) .build(); - #[cfg(feature = "encryption")] - let expected_size_with_decryptor = 3080; - #[cfg(not(feature = "encryption"))] let expected_size_with_decryptor = 3072; assert!(expected_size_with_decryptor > base_expected_size); From deec81ea0705945e0337988240a3d6747e4edc41 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Wed, 18 Feb 2026 16:52:32 -0800 Subject: [PATCH 08/26] Round-trip tests w/ AES-256 encrypted parquet --- parquet/tests/encryption/encryption_async.rs | 61 ++++++++++++++++++++ parquet/tests/encryption/encryption_util.rs | 14 +++++ 2 files changed, 75 insertions(+) diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index 34cd4e5b6465..4c246397ef1e 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -20,6 +20,7 @@ use crate::encryption_util::{ TestKeyRetriever, read_encrypted_file, verify_column_indexes, verify_encryption_double_test_data, verify_encryption_test_data, + verify_encryption_test_data_from_java, }; use arrow_array::RecordBatch; use arrow_schema::Schema; @@ -436,6 +437,53 @@ async fn test_decrypt_page_index_non_uniform() { .unwrap(); } +#[tokio::test] +async fn test_aes_256_non_uniform_encryption_plaintext_footer() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); + let mut file = File::open(&path).await.unwrap(); + + // There is always a footer key even with a plaintext footer, + // but this is used for signing the footer. + let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 128bit/16 + let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); + let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); + let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); + let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); + let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key) + .with_column_key("double_field", column_1_key) + .with_column_key("float_field", column_2_key) + .with_column_key("boolean_field", column_3_key) + .with_column_key("int32_field", column_4_key) + .with_column_key("ba_field", column_5_key) + .with_column_key("flba_field", column_6_key) + .build() + .unwrap(); + + verify_aes256_encryption_java_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_aes_256_uniform_encryption() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"); + let mut file = File::open(&path).await.unwrap(); + + let key_code: &[u8] = "01234567890123456789012345678901".as_bytes(); + let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) + .build() + .unwrap(); + + verify_aes256_encryption_java_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); +} + async fn test_decrypt_page_index( path: &str, decryption_properties: Arc, @@ -453,6 +501,19 @@ async fn test_decrypt_page_index( Ok(()) } +async fn verify_aes256_encryption_java_test_file_read_async( + file: &mut tokio::fs::File, + decryption_properties: Arc, +) -> Result<(), ParquetError> { + let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); + + let arrow_metadata = ArrowReaderMetadata::load_async(file, options).await?; + let metadata = arrow_metadata.metadata(); + + verify_encryption_test_data_from_java(metadata); + Ok(()) +} + async fn verify_encryption_test_file_read_async( file: &mut tokio::fs::File, decryption_properties: Arc, diff --git a/parquet/tests/encryption/encryption_util.rs b/parquet/tests/encryption/encryption_util.rs index 7f4cc5a9da45..ce3dcc3c60d7 100644 --- a/parquet/tests/encryption/encryption_util.rs +++ b/parquet/tests/encryption/encryption_util.rs @@ -104,6 +104,20 @@ pub(crate) fn verify_encryption_double_test_data( assert_eq!(row_count, file_metadata.num_rows() as usize); } +/// Verifies parquet-java generated AES-256 data from the parquet-testing repository +pub(crate) fn verify_encryption_test_data_from_java(metadata: &ParquetMetaData) { + let file_metadata = metadata.file_metadata(); + assert_eq!(file_metadata.num_rows(), 10000); + assert_eq!(file_metadata.schema_descr().num_columns(), 7); + + let mut total_rows = 0; + metadata.row_groups().iter().for_each(|rg| { + assert_eq!(rg.num_columns(), 7); + total_rows += rg.num_rows(); + }); + assert_eq!(total_rows, 10000); +} + /// Verifies data read from an encrypted file from the parquet-testing repository pub(crate) fn verify_encryption_test_data( record_batches: Vec, From f55fe3d8905d6a2acfe28ca3ce6dda42dda7cb1d Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Wed, 18 Feb 2026 17:30:05 -0800 Subject: [PATCH 09/26] Fix comment --- parquet/tests/encryption/encryption_async.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index eee12df5ff86..8c7fa8b6428a 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -445,7 +445,7 @@ async fn test_aes_256_non_uniform_encryption_plaintext_footer() { // There is always a footer key even with a plaintext footer, // but this is used for signing the footer. - let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 128bit/16 + let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 256bit/32 let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); From d92014ba3e0980e451a5934b5670344cb151e3fb Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Sun, 1 Mar 2026 11:07:18 +0800 Subject: [PATCH 10/26] More test cases --- parquet/tests/encryption/encryption.rs | 822 +++++++++++++++++++ parquet/tests/encryption/encryption_async.rs | 470 ++++++++++- parquet/tests/encryption/encryption_util.rs | 14 - 3 files changed, 1274 insertions(+), 32 deletions(-) diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index c26a0edee682..2c304736b25f 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -1247,6 +1247,828 @@ fn test_decrypt_page_index_non_uniform() { test_decrypt_page_index(&path, decryption_properties).unwrap(); } +#[test] +fn test_aes256_non_uniform_encryption_plaintext_footer() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let footer_key = "01234567890123456789012345678901".as_bytes(); // 256bit/32 + let column_1_key = "12345678901234567890123456789012".as_bytes(); + let column_2_key = "12345678901234567890123456789013".as_bytes(); + let column_3_key = "12345678901234567890123456789014".as_bytes(); + let column_4_key = "12345678901234567890123456789015".as_bytes(); + let column_5_key = "12345678901234567890123456789016".as_bytes(); + let column_6_key = "12345678901234567890123456789017".as_bytes(); + let column_7_key = "12345678901234567890123456789018".as_bytes(); + let column_8_key = "12345678901234567890123456789019".as_bytes(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_column_key("double_field", column_1_key.to_vec()) + .with_column_key("float_field", column_2_key.to_vec()) + .with_column_key("boolean_field", column_3_key.to_vec()) + .with_column_key("int32_field", column_4_key.to_vec()) + .with_column_key("ba_field", column_5_key.to_vec()) + .with_column_key("flba_field", column_6_key.to_vec()) + .with_column_key("int64_list_field", column_7_key.to_vec()) + .with_column_key("int64_field", column_8_key.to_vec()) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); +} + +#[test] +fn test_aes256_plaintext_footer_signature_verification() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); + let file = File::open(path.clone()).unwrap(); + + let wrong_footer_key = "00000000000000000000000000000000".as_bytes(); // 256bit/32 + let column_1_key = "12345678901234567890123456789012".as_bytes(); + let column_2_key = "12345678901234567890123456789013".as_bytes(); + let column_3_key = "12345678901234567890123456789014".as_bytes(); + let column_4_key = "12345678901234567890123456789015".as_bytes(); + let column_5_key = "12345678901234567890123456789016".as_bytes(); + let column_6_key = "12345678901234567890123456789017".as_bytes(); + let column_7_key = "12345678901234567890123456789018".as_bytes(); + let column_8_key = "12345678901234567890123456789019".as_bytes(); + + let decryption_properties = FileDecryptionProperties::builder(wrong_footer_key.to_vec()) + .disable_footer_signature_verification() + .with_column_key("double_field", column_1_key.to_vec()) + .with_column_key("float_field", column_2_key.to_vec()) + .with_column_key("boolean_field", column_3_key.to_vec()) + .with_column_key("int32_field", column_4_key.to_vec()) + .with_column_key("ba_field", column_5_key.to_vec()) + .with_column_key("flba_field", column_6_key.to_vec()) + .with_column_key("int64_list_field", column_7_key.to_vec()) + .with_column_key("int64_field", column_8_key.to_vec()) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + + let file = File::open(path.clone()).unwrap(); + + let decryption_properties = FileDecryptionProperties::builder(wrong_footer_key.to_vec()) + .with_column_key("double_field", column_1_key.to_vec()) + .with_column_key("float_field", column_2_key.to_vec()) + .with_column_key("boolean_field", column_3_key.to_vec()) + .with_column_key("int32_field", column_4_key.to_vec()) + .with_column_key("ba_field", column_5_key.to_vec()) + .with_column_key("flba_field", column_6_key.to_vec()) + .with_column_key("int64_list_field", column_7_key.to_vec()) + .with_column_key("int64_field", column_8_key.to_vec()) + .build() + .unwrap(); + + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .starts_with("Parquet error: Footer signature verification failed. Computed: [") + ); +} + +#[test] +fn test_aes256_non_uniform_encryption_disabled_aad_storage() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!( + "{test_data}/aes256/encrypt_columns_and_footer_disable_aad_storage.parquet.encrypted" + ); + let file = File::open(path.clone()).unwrap(); + + let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 256bit/32 + let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); + let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); + let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); + let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); + let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); + let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); + let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); + + // Can read successfully when providing the correct AAD prefix + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_key("double_field", column_1_key.clone()) + .with_column_key("float_field", column_2_key.clone()) + .with_column_key("boolean_field", column_3_key.clone()) + .with_column_key("int32_field", column_4_key.clone()) + .with_column_key("ba_field", column_5_key.clone()) + .with_column_key("flba_field", column_6_key.clone()) + .with_column_key("int64_list_field", column_7_key.clone()) + .with_column_key("int64_field", column_8_key.clone()) + .with_aad_prefix(b"tester".to_vec()) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + + // Using wrong AAD prefix should fail + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_key("double_field", column_1_key.clone()) + .with_column_key("float_field", column_2_key.clone()) + .with_column_key("boolean_field", column_3_key.clone()) + .with_column_key("int32_field", column_4_key.clone()) + .with_column_key("ba_field", column_5_key.clone()) + .with_column_key("flba_field", column_6_key.clone()) + .with_column_key("int64_list_field", column_7_key.clone()) + .with_column_key("int64_field", column_8_key.clone()) + .with_aad_prefix(b"wrong_aad_prefix".to_vec()) + .build() + .unwrap(); + + let file = File::open(path.clone()).unwrap(); + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer" + ); + + // Not providing any AAD prefix should fail as it isn't stored in the file + let decryption_properties = FileDecryptionProperties::builder(footer_key) + .with_column_key("double_field", column_1_key.clone()) + .with_column_key("float_field", column_2_key.clone()) + .with_column_key("boolean_field", column_3_key.clone()) + .with_column_key("int32_field", column_4_key.clone()) + .with_column_key("ba_field", column_5_key.clone()) + .with_column_key("flba_field", column_6_key.clone()) + .with_column_key("int64_list_field", column_7_key.clone()) + .with_column_key("int64_field", column_8_key.clone()) + .build() + .unwrap(); + + let file = File::open(path).unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file was encrypted with an AAD prefix that is not stored in the file, \ + but no AAD prefix was provided in the file decryption properties" + ); +} +#[test] +fn test_aes256_non_uniform_encryption() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); + let file = File::open(path).unwrap(); + + // There is always a footer key even with a plaintext footer, + // but this is used for signing the footer. + let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 256bit/32 + let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); + let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); + let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); + let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); + let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); + let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); + let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key) + .with_column_key("double_field", column_1_key) + .with_column_key("float_field", column_2_key) + .with_column_key("boolean_field", column_3_key) + .with_column_key("int32_field", column_4_key) + .with_column_key("ba_field", column_5_key) + .with_column_key("flba_field", column_6_key) + .with_column_key("int64_list_field", column_7_key) + .with_column_key("int64_field", column_8_key) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); +} + +#[test] +fn test_aes256_uniform_encryption() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let key_code: &[u8] = "01234567890123456789012345678901".as_bytes(); + let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); +} + +#[test] +fn test_aes256_ctr_encryption() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/encrypt_columns_and_footer_ctr.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); + let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key) + .with_column_key("double_field", column_1_key.clone()) + .with_column_key("float_field", column_1_key) + .build() + .unwrap(); + + let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); + let metadata = ArrowReaderMetadata::load(&file, options); + + match metadata { + Err(ParquetError::NYI(s)) => { + assert!(s.contains("AES_GCM_CTR_V1")); + } + _ => { + panic!("Expected ParquetError::NYI"); + } + }; +} + +#[test] +fn test_aes256_non_uniform_encryption_plaintext_footer_with_key_retriever() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let key_retriever = TestKeyRetriever::new() + .with_key( + "kf".to_owned(), + "01234567890123456789012345678901".as_bytes().to_vec(), + ) + .with_key( + "kc1".to_owned(), + "12345678901234567890123456789012".as_bytes().to_vec(), + ) + .with_key( + "kc2".to_owned(), + "12345678901234567890123456789013".as_bytes().to_vec(), + ) + .with_key( + "kc3".to_owned(), + "12345678901234567890123456789014".as_bytes().to_vec(), + ) + .with_key( + "kc4".to_owned(), + "12345678901234567890123456789015".as_bytes().to_vec(), + ) + .with_key( + "kc5".to_owned(), + "12345678901234567890123456789016".as_bytes().to_vec(), + ) + .with_key( + "kc6".to_owned(), + "12345678901234567890123456789017".as_bytes().to_vec(), + ) + .with_key( + "kc7".to_owned(), + "12345678901234567890123456789018".as_bytes().to_vec(), + ) + .with_key( + "kc8".to_owned(), + "12345678901234567890123456789019".as_bytes().to_vec(), + ); + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); +} + +#[test] +fn test_aes256_uniform_encryption_plaintext_footer_with_key_retriever() { + let test_data = arrow::util::test_util::parquet_test_data(); + + // Read example data with key retriever + let path = format!("{test_data}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let key_retriever = Arc::new( + TestKeyRetriever::new() + .with_key( + "kf".to_owned(), + "01234567890123456789012345678901".as_bytes().to_vec(), + ) + .with_key( + "kc1".to_owned(), + "12345678901234567890123456789012".as_bytes().to_vec(), + ) + .with_key( + "kc2".to_owned(), + "12345678901234567890123456789013".as_bytes().to_vec(), + ) + .with_key( + "kc3".to_owned(), + "12345678901234567890123456789014".as_bytes().to_vec(), + ) + .with_key( + "kc4".to_owned(), + "12345678901234567890123456789015".as_bytes().to_vec(), + ) + .with_key( + "kc5".to_owned(), + "12345678901234567890123456789016".as_bytes().to_vec(), + ) + .with_key( + "kc6".to_owned(), + "12345678901234567890123456789017".as_bytes().to_vec(), + ) + .with_key( + "kc7".to_owned(), + "12345678901234567890123456789018".as_bytes().to_vec(), + ) + .with_key( + "kc8".to_owned(), + "12345678901234567890123456789019".as_bytes().to_vec(), + ), + ); + + let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever.clone()) + .build() + .unwrap(); + + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + + // Write data into temporary file with plaintext footer and footer key metadata + let temp_file = tempfile::tempfile().unwrap(); + let encryption_properties = + FileEncryptionProperties::builder(b"01234567890123456789012345678901".to_vec()) + .with_footer_key_metadata("kf".into()) + .with_column_key_and_metadata( + "double_field", + b"12345678901234567890123456789012".to_vec(), + b"kc1".into(), + ) + .with_column_key_and_metadata( + "float_field", + b"12345678901234567890123456789013".to_vec(), + b"kc2".into(), + ) + .with_column_key_and_metadata( + "boolean_field", + b"12345678901234567890123456789014".to_vec(), + b"kc3".into(), + ) + .with_column_key_and_metadata( + "int32_field", + b"12345678901234567890123456789015".to_vec(), + b"kc4".into(), + ) + .with_column_key_and_metadata( + "ba_field", + b"12345678901234567890123456789016".to_vec(), + b"kc5".into(), + ) + .with_column_key_and_metadata( + "flba_field", + b"12345678901234567890123456789017".to_vec(), + b"kc6".into(), + ) + // .with_column_key_and_metadata("int64_list_field", b"12345678901234567890123456789018".to_vec(), b"kc7".into()) + .with_column_key_and_metadata( + "int64_list_field.list.int64_list_field", + b"12345678901234567890123456789018".to_vec(), + b"kc7".into(), + ) + .with_column_key_and_metadata( + "int64_field", + b"12345678901234567890123456789019".to_vec(), + b"kc8".into(), + ) + .with_plaintext_footer(true) + .build() + .unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); + let batch_reader = builder.build().unwrap(); + let batches = batch_reader + .collect::, _>>() + .unwrap(); + + let props = WriterProperties::builder() + .with_file_encryption_properties(encryption_properties) + .build(); + + let mut writer = ArrowWriter::try_new( + temp_file.try_clone().unwrap(), + metadata.schema().clone(), + Some(props), + ) + .unwrap(); + for batch in batches { + writer.write(&batch).unwrap(); + } + + writer.close().unwrap(); + + // Read temporary file with plaintext metadata using key retriever + let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever) + .build() + .unwrap(); + + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let _ = ArrowReaderMetadata::load(&temp_file, options.clone()).unwrap(); + + // Read temporary file with plaintext metadata using key retriever with invalid key + let key_retriever = Arc::new( + TestKeyRetriever::new() + .with_key( + "kf".to_owned(), + b"01334567890123456789012345678901".to_vec(), + ) // wrong footer key + .with_key( + "kc1".to_owned(), + b"12345678901234567890123456789012".to_vec(), + ) + .with_key( + "kc2".to_owned(), + b"12345678901234567890123456789013".to_vec(), + ) + .with_key( + "kc3".to_owned(), + b"12345678901234567890123456789014".to_vec(), + ) + .with_key( + "kc4".to_owned(), + b"12345678901234567890123456789015".to_vec(), + ) + .with_key( + "kc5".to_owned(), + b"12345678901234567890123456789016".to_vec(), + ) + .with_key( + "kc6".to_owned(), + b"12345678901234567890123456789017".to_vec(), + ) + .with_key( + "kc7".to_owned(), + b"12345678901234567890123456789018".to_vec(), + ) + .with_key( + "kc8".to_owned(), + b"12345678901234567890123456789019".to_vec(), + ), + ); + let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever) + .build() + .unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&temp_file, options.clone()); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .starts_with("Parquet error: Footer signature verification failed. Computed: [") + ); +} + +#[test] +fn test_aes256_non_uniform_encryption_with_key_retriever() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/encrypt_columns_and_footer.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let key_retriever = TestKeyRetriever::new() + .with_key( + "kf".to_owned(), + "01234567890123456789012345678901".as_bytes().to_vec(), + ) + .with_key( + "kc1".to_owned(), + "12345678901234567890123456789012".as_bytes().to_vec(), + ) + .with_key( + "kc2".to_owned(), + "12345678901234567890123456789013".as_bytes().to_vec(), + ) + .with_key( + "kc3".to_owned(), + "12345678901234567890123456789014".as_bytes().to_vec(), + ) + .with_key( + "kc4".to_owned(), + "12345678901234567890123456789015".as_bytes().to_vec(), + ) + .with_key( + "kc5".to_owned(), + "12345678901234567890123456789016".as_bytes().to_vec(), + ) + .with_key( + "kc6".to_owned(), + "12345678901234567890123456789017".as_bytes().to_vec(), + ) + .with_key( + "kc7".to_owned(), + "12345678901234567890123456789018".as_bytes().to_vec(), + ) + .with_key( + "kc8".to_owned(), + "12345678901234567890123456789019".as_bytes().to_vec(), + ); + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); +} + +#[test] +fn test_aes256_uniform_encryption_with_key_retriever() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/uniform_encryption.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 + + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .build() + .unwrap(); + + let file_encryption_properties = FileEncryptionProperties::builder(footer_key) + .build() + .unwrap(); + + read_and_roundtrip_to_encrypted_file(&file, decryption_properties, file_encryption_properties); +} + +#[test] +fn test_aes256_write_non_uniform_encryption() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/encrypt_columns_and_footer.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 + let column_names = vec![ + "double_field", + "float_field", + "boolean_field", + "int32_field", + "ba_field", + "flba_field", + "int64_list_field", + "int64_list_field.list.int64_list_field", + "int64_field", + ]; + let column_keys = vec![ + b"12345678901234567890123456789012".to_vec(), + b"12345678901234567890123456789013".to_vec(), + b"12345678901234567890123456789014".to_vec(), + b"12345678901234567890123456789015".to_vec(), + b"12345678901234567890123456789016".to_vec(), + b"12345678901234567890123456789017".to_vec(), + b"12345678901234567890123456789018".to_vec(), + b"12345678901234567890123456789018".to_vec(), + b"12345678901234567890123456789019".to_vec(), + ]; + + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_keys(column_names.clone(), column_keys.clone()) + .unwrap() + .build() + .unwrap(); + + let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); + let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); + let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); + let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); + let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); + let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); + let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); + + let file_encryption_properties = FileEncryptionProperties::builder(footer_key) + .with_column_key("double_field", column_1_key) + .with_column_key("float_field", column_2_key) + .with_column_key("boolean_field", column_3_key) + .with_column_key("int32_field", column_4_key) + .with_column_key("ba_field", column_5_key) + .with_column_key("flba_field", column_6_key.clone()) + .with_column_key("int64_list_field.list.int64_list_field", column_7_key) + .with_column_key("int64_field", column_8_key) + .build() + .unwrap(); + + read_and_roundtrip_to_encrypted_file(&file, decryption_properties, file_encryption_properties); +} + +#[test] +fn test_aes256_write_uniform_encryption_plaintext_footer() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 + let wrong_footer_key = b"00000000000000000000000000000000".to_vec(); // 256bit/32 + let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); + let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); + let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); + let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); + let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); + let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); + let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_key("double_field", column_1_key.clone()) + .with_column_key("float_field", column_2_key.clone()) + .with_column_key("boolean_field", column_3_key.clone()) + .with_column_key("int32_field", column_4_key.clone()) + .with_column_key("ba_field", column_5_key.clone()) + .with_column_key("flba_field", column_6_key.clone()) + .with_column_key("int64_list_field", column_7_key.clone()) + .with_column_key("int64_field", column_8_key.clone()) + .build() + .unwrap(); + + let wrong_decryption_properties = FileDecryptionProperties::builder(wrong_footer_key) + .with_column_key("double_field", column_1_key) + .with_column_key("float_field", column_2_key) + .with_column_key("boolean_field", column_3_key) + .with_column_key("int32_field", column_4_key) + .with_column_key("ba_field", column_5_key) + .with_column_key("flba_field", column_6_key) + .with_column_key("int64_list_field", column_7_key) + .with_column_key("int64_field", column_8_key) + .build() + .unwrap(); + + let file_encryption_properties = FileEncryptionProperties::builder(footer_key) + .with_plaintext_footer(true) + .build() + .unwrap(); + + // Try writing plaintext footer and then reading it with the correct footer key + read_and_roundtrip_to_encrypted_file( + &file, + Arc::clone(&decryption_properties), + file_encryption_properties.clone(), + ); + + // Try writing plaintext footer and then reading it with the wrong footer key + let temp_file = tempfile::tempfile().unwrap(); + + // read example data + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); + let batch_reader = builder.build().unwrap(); + let batches = batch_reader + .collect::, _>>() + .unwrap(); + + // write example data + let props = WriterProperties::builder() + .with_file_encryption_properties(file_encryption_properties) + .build(); + + let mut writer = ArrowWriter::try_new( + temp_file.try_clone().unwrap(), + metadata.schema().clone(), + Some(props), + ) + .unwrap(); + for batch in batches { + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + + // Try reading plaintext footer and with the wrong footer key + let options = + ArrowReaderOptions::default().with_file_decryption_properties(wrong_decryption_properties); + let result = ArrowReaderMetadata::load(&temp_file, options.clone()); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .starts_with("Parquet error: Footer signature verification failed. Computed: [") + ); +} + +#[test] +fn test_aes256_write_non_uniform_encryption_column_missmatch() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/encrypt_columns_and_footer.parquet.encrypted"); + + let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 + let column_1_key = b"12345678901234567890123456789012".to_vec(); + let column_2_key = b"12345678901234567890123456789013".to_vec(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_key("double_field", column_1_key.clone()) + .with_column_key("float_field", column_2_key.clone()) + .build() + .unwrap(); + + let file_encryption_properties = FileEncryptionProperties::builder(footer_key) + .with_column_key("double_field", column_1_key.clone()) + .with_column_key("other_field", column_1_key) + .with_column_key("yet_another_field", column_2_key) + .build() + .unwrap(); + + let temp_file = tempfile::tempfile().unwrap(); + + // read example data + let file = File::open(path).unwrap(); + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + let props = WriterProperties::builder() + .with_file_encryption_properties(file_encryption_properties) + .build(); + + let result = ArrowWriter::try_new( + temp_file.try_clone().unwrap(), + metadata.schema().clone(), + Some(props), + ); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: The following columns with encryption keys specified were not found in the schema: other_field, yet_another_field" + ); +} + +#[test] +fn test_aes256_write_uniform_encryption() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/uniform_encryption.parquet.encrypted"); + let file = File::open(path).unwrap(); + + let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 + + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .build() + .unwrap(); + + let file_encryption_properties = FileEncryptionProperties::builder(footer_key) + .build() + .unwrap(); + + read_and_roundtrip_to_encrypted_file(&file, decryption_properties, file_encryption_properties); +} + +#[test] +fn test_aes256_decrypt_page_index_uniform() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"); + + let key_code: &[u8] = "01234567890123456789012345678901".as_bytes(); + let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) + .build() + .unwrap(); + + test_decrypt_page_index(&path, decryption_properties).unwrap(); +} + +#[test] +fn test_aes256_decrypt_page_index_non_uniform() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/encrypt_columns_and_footer.parquet.encrypted"); + + let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); + let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); + let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); + let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); + let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); + let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); + let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); + let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_column_key("double_field", column_1_key) + .with_column_key("float_field", column_2_key) + .with_column_key("boolean_field", column_3_key) + .with_column_key("int32_field", column_4_key) + .with_column_key("ba_field", column_5_key) + .with_column_key("flba_field", column_6_key) + .with_column_key("int64_list_field", column_7_key) + .with_column_key("int64_field", column_8_key) + .build() + .unwrap(); + + test_decrypt_page_index(&path, decryption_properties).unwrap(); +} + fn test_decrypt_page_index( path: &str, decryption_properties: Arc, diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index 8c7fa8b6428a..9a47b808af32 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -20,7 +20,6 @@ use crate::encryption_util::{ TestKeyRetriever, read_encrypted_file, verify_column_indexes, verify_encryption_double_test_data, verify_encryption_test_data, - verify_encryption_test_data_from_java, }; use arrow_array::RecordBatch; use arrow_schema::Schema; @@ -438,7 +437,7 @@ async fn test_decrypt_page_index_non_uniform() { } #[tokio::test] -async fn test_aes_256_non_uniform_encryption_plaintext_footer() { +async fn test_aes256_non_uniform_encryption_plaintext_footer() { let test_data = arrow::util::test_util::parquet_test_data(); let path = format!("{test_data}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); let mut file = File::open(&path).await.unwrap(); @@ -452,6 +451,8 @@ async fn test_aes_256_non_uniform_encryption_plaintext_footer() { let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); + let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); + let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); let decryption_properties = FileDecryptionProperties::builder(footer_key) .with_column_key("double_field", column_1_key) @@ -460,16 +461,178 @@ async fn test_aes_256_non_uniform_encryption_plaintext_footer() { .with_column_key("int32_field", column_4_key) .with_column_key("ba_field", column_5_key) .with_column_key("flba_field", column_6_key) + .with_column_key("int64_list_field", column_7_key) + .with_column_key("int64_field", column_8_key) .build() .unwrap(); - verify_aes256_encryption_java_test_file_read_async(&mut file, decryption_properties) + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_aes256_misspecified_encryption_keys() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/encrypt_columns_and_footer.parquet.encrypted"); + + // There is always a footer key even with a plaintext footer, + // but this is used for signing the footer. + let footer_key = "01234567890123456789012345678901".as_bytes(); // 256bit/32 + let column_1_key = "12345678901234567890123456789012".as_bytes(); + let column_2_key = "12345678901234567890123456789013".as_bytes(); + + // read file with keys and check for expected error message + async fn check_for_error( + expected_message: &str, + path: &String, + footer_key: &[u8], + column_1_key: &[u8], + column_2_key: &[u8], + ) { + let mut file = File::open(&path).await.unwrap(); + + let mut decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()); + + if !column_1_key.is_empty() { + decryption_properties = + decryption_properties.with_column_key("double_field", column_1_key.to_vec()); + } + + if !column_2_key.is_empty() { + decryption_properties = + decryption_properties.with_column_key("float_field", column_2_key.to_vec()); + } + + let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); + let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); + let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); + let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); + let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); + let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); + + let decryption_properties = decryption_properties + .with_column_key("boolean_field", column_3_key) + .with_column_key("int32_field", column_4_key) + .with_column_key("ba_field", column_5_key) + .with_column_key("flba_field", column_6_key) + .with_column_key("int64_list_field", column_7_key) + .with_column_key("int64_field", column_8_key) + .build() + .unwrap(); + + match verify_encryption_test_file_read_async(&mut file, decryption_properties).await { + Ok(_) => { + panic!("did not get expected error") + } + Err(e) => { + assert_eq!(e.to_string(), expected_message); + } + } + } + + // Too short footer key + check_for_error( + format!("Parquet error: Invalid footer key. Error creating RingGcmBlockDecryptor with unsupported key length: {}", "bad_pwd".len()).as_str(), + &path, + "bad_pwd".as_bytes(), + column_1_key, + column_2_key, + ) + .await; + + // Wrong footer key + check_for_error( + "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer", + &path, + "11234567890123456789012345678901".as_bytes(), + column_1_key, + column_2_key, + ) + .await; + + // Missing column key + check_for_error( + "Parquet error: No column decryption key set for encrypted column 'double_field'", + &path, + footer_key, + "".as_bytes(), + column_2_key, + ) + .await; + + // Too short column key + check_for_error( + format!( + "Parquet error: Error creating RingGcmBlockDecryptor with unsupported key length: {}", + "abc".len() + ) + .as_str(), + &path, + footer_key, + "abc".as_bytes(), + column_2_key, + ) + .await; + + // Wrong column key + check_for_error( + "Parquet error: Unable to decrypt column 'double_field', perhaps the column key is wrong?", + &path, + footer_key, + "11234567890123456789012345678901".as_bytes(), + column_2_key, + ) + .await; + + // Mixed up keys + check_for_error( + "Parquet error: Unable to decrypt column 'float_field', perhaps the column key is wrong?", + &path, + footer_key, + column_2_key, + column_1_key, + ) + .await; +} + +#[tokio::test] +async fn test_aes256_non_uniform_encryption() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/encrypt_columns_and_footer.parquet.encrypted"); + let mut file = File::open(&path).await.unwrap(); + + // There is always a footer key even with a plaintext footer, + // but this is used for signing the footer. + let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 256bit/32 + let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); + let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); + let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); + let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); + let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); + let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); + let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key) + .with_column_key("double_field", column_1_key) + .with_column_key("float_field", column_2_key) + .with_column_key("boolean_field", column_3_key) + .with_column_key("int32_field", column_4_key) + .with_column_key("ba_field", column_5_key) + .with_column_key("flba_field", column_6_key) + .with_column_key("int64_list_field", column_7_key) + .with_column_key("int64_field", column_8_key) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) .await .unwrap(); } #[tokio::test] -async fn test_aes_256_uniform_encryption() { +async fn test_aes256_uniform_encryption() { let test_data = arrow::util::test_util::parquet_test_data(); let path = format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"); let mut file = File::open(&path).await.unwrap(); @@ -479,7 +642,291 @@ async fn test_aes_256_uniform_encryption() { .build() .unwrap(); - verify_aes256_encryption_java_test_file_read_async(&mut file, decryption_properties) + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_aes256_ctr_encryption() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/encrypt_columns_and_footer_ctr.parquet.encrypted"); + let mut file = File::open(&path).await.unwrap(); + + let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); + let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key) + .with_column_key("double_field", column_1_key.clone()) + .with_column_key("float_field", column_1_key) + .build() + .unwrap(); + + let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); + let metadata = ArrowReaderMetadata::load_async(&mut file, options).await; + + match metadata { + Err(ParquetError::NYI(s)) => { + assert!(s.contains("AES_GCM_CTR_V1")); + } + _ => { + panic!("Expected ParquetError::NYI"); + } + }; +} + +#[tokio::test] +async fn test_aes256_decrypting_without_decryption_properties_fails() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"); + let mut file = File::open(&path).await.unwrap(); + + let options = ArrowReaderOptions::new(); + let result = ArrowReaderMetadata::load_async(&mut file, options).await; + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); +} + +#[tokio::test] +async fn test_aes256_write_non_uniform_encryption() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/encrypt_columns_and_footer.parquet.encrypted"); + + let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 + let column_names = vec![ + "double_field", + "float_field", + "boolean_field", + "int32_field", + "ba_field", + "flba_field", + "int64_list_field", + "int64_list_field.list.int64_list_field", + "int64_field", + ]; + let column_keys = vec![ + b"12345678901234567890123456789012".to_vec(), + b"12345678901234567890123456789013".to_vec(), + b"12345678901234567890123456789014".to_vec(), + b"12345678901234567890123456789015".to_vec(), + b"12345678901234567890123456789016".to_vec(), + b"12345678901234567890123456789017".to_vec(), + b"12345678901234567890123456789018".to_vec(), + b"12345678901234567890123456789018".to_vec(), + b"12345678901234567890123456789019".to_vec(), + ]; + + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_keys(column_names.clone(), column_keys.clone()) + .unwrap() + .build() + .unwrap(); + + let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); + let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); + let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); + let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); + let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); + let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); + let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); + + let file_encryption_properties = FileEncryptionProperties::builder(footer_key) + .with_column_key("double_field", column_1_key) + .with_column_key("float_field", column_2_key) + .with_column_key("boolean_field", column_3_key) + .with_column_key("int32_field", column_4_key) + .with_column_key("ba_field", column_5_key) + .with_column_key("flba_field", column_6_key.clone()) + .with_column_key("int64_list_field.list.int64_list_field", column_7_key) + .with_column_key("int64_field", column_8_key) + .build() + .unwrap(); + + read_and_roundtrip_to_encrypted_file_async( + &path, + decryption_properties, + file_encryption_properties, + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_aes256_non_uniform_encryption_plaintext_footer_with_key_retriever() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); + let mut file = File::open(&path).await.unwrap(); + + let key_retriever = TestKeyRetriever::new() + .with_key( + "kf".to_owned(), + "01234567890123456789012345678901".as_bytes().to_vec(), + ) + .with_key( + "kc1".to_owned(), + "12345678901234567890123456789012".as_bytes().to_vec(), + ) + .with_key( + "kc2".to_owned(), + "12345678901234567890123456789013".as_bytes().to_vec(), + ) + .with_key( + "kc3".to_owned(), + "12345678901234567890123456789014".as_bytes().to_vec(), + ) + .with_key( + "kc4".to_owned(), + "12345678901234567890123456789015".as_bytes().to_vec(), + ) + .with_key( + "kc5".to_owned(), + "12345678901234567890123456789016".as_bytes().to_vec(), + ) + .with_key( + "kc6".to_owned(), + "12345678901234567890123456789017".as_bytes().to_vec(), + ) + .with_key( + "kc7".to_owned(), + "12345678901234567890123456789018".as_bytes().to_vec(), + ) + .with_key( + "kc8".to_owned(), + "12345678901234567890123456789019".as_bytes().to_vec(), + ); + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_aes256_non_uniform_encryption_with_key_retriever() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/encrypt_columns_and_footer.parquet.encrypted"); + let mut file = File::open(&path).await.unwrap(); + + let key_retriever = TestKeyRetriever::new() + .with_key( + "kf".to_owned(), + "01234567890123456789012345678901".as_bytes().to_vec(), + ) + .with_key( + "kc1".to_owned(), + "12345678901234567890123456789012".as_bytes().to_vec(), + ) + .with_key( + "kc2".to_owned(), + "12345678901234567890123456789013".as_bytes().to_vec(), + ) + .with_key( + "kc3".to_owned(), + "12345678901234567890123456789014".as_bytes().to_vec(), + ) + .with_key( + "kc4".to_owned(), + "12345678901234567890123456789015".as_bytes().to_vec(), + ) + .with_key( + "kc5".to_owned(), + "12345678901234567890123456789016".as_bytes().to_vec(), + ) + .with_key( + "kc6".to_owned(), + "12345678901234567890123456789017".as_bytes().to_vec(), + ) + .with_key( + "kc7".to_owned(), + "12345678901234567890123456789018".as_bytes().to_vec(), + ) + .with_key( + "kc8".to_owned(), + "12345678901234567890123456789019".as_bytes().to_vec(), + ); + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_aes256_uniform_encryption_with_key_retriever() { + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/aes256/uniform_encryption.parquet.encrypted"); + let mut file = File::open(&path).await.unwrap(); + + let key_retriever = TestKeyRetriever::new().with_key( + "kf".to_owned(), + "01234567890123456789012345678901".as_bytes().to_vec(), + ); + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_aes256_decrypt_page_index_uniform() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"); + + let key_code: &[u8] = "01234567890123456789012345678901".as_bytes(); + let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) + .build() + .unwrap(); + + test_decrypt_page_index(&path, decryption_properties) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_aes256_decrypt_page_index_non_uniform() { + let test_data = arrow::util::test_util::parquet_test_data(); + let path = format!("{test_data}/aes256/encrypt_columns_and_footer.parquet.encrypted"); + + let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 256bit/32 + let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); + let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); + let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); + let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); + let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); + let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); + let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); + + let decryption_properties = FileDecryptionProperties::builder(footer_key) + .with_column_key("double_field", column_1_key) + .with_column_key("float_field", column_2_key) + .with_column_key("boolean_field", column_3_key) + .with_column_key("int32_field", column_4_key) + .with_column_key("ba_field", column_5_key) + .with_column_key("flba_field", column_6_key) + .with_column_key("int64_list_field", column_7_key) + .with_column_key("int64_field", column_8_key) + .build() + .unwrap(); + + test_decrypt_page_index(&path, decryption_properties) .await .unwrap(); } @@ -501,19 +948,6 @@ async fn test_decrypt_page_index( Ok(()) } -async fn verify_aes256_encryption_java_test_file_read_async( - file: &mut tokio::fs::File, - decryption_properties: Arc, -) -> Result<(), ParquetError> { - let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); - - let arrow_metadata = ArrowReaderMetadata::load_async(file, options).await?; - let metadata = arrow_metadata.metadata(); - - verify_encryption_test_data_from_java(metadata); - Ok(()) -} - async fn verify_encryption_test_file_read_async( file: &mut tokio::fs::File, decryption_properties: Arc, diff --git a/parquet/tests/encryption/encryption_util.rs b/parquet/tests/encryption/encryption_util.rs index ce3dcc3c60d7..7f4cc5a9da45 100644 --- a/parquet/tests/encryption/encryption_util.rs +++ b/parquet/tests/encryption/encryption_util.rs @@ -104,20 +104,6 @@ pub(crate) fn verify_encryption_double_test_data( assert_eq!(row_count, file_metadata.num_rows() as usize); } -/// Verifies parquet-java generated AES-256 data from the parquet-testing repository -pub(crate) fn verify_encryption_test_data_from_java(metadata: &ParquetMetaData) { - let file_metadata = metadata.file_metadata(); - assert_eq!(file_metadata.num_rows(), 10000); - assert_eq!(file_metadata.schema_descr().num_columns(), 7); - - let mut total_rows = 0; - metadata.row_groups().iter().for_each(|rg| { - assert_eq!(rg.num_columns(), 7); - total_rows += rg.num_rows(); - }); - assert_eq!(total_rows, 10000); -} - /// Verifies data read from an encrypted file from the parquet-testing repository pub(crate) fn verify_encryption_test_data( record_batches: Vec, From bd79c92ce36c4a670386576899b0833dfceec0c5 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Mon, 2 Mar 2026 14:43:01 -0800 Subject: [PATCH 11/26] Align field names --- parquet/tests/encryption/encryption.rs | 54 ++++++++++---------- parquet/tests/encryption/encryption_async.rs | 24 ++++----- 2 files changed, 39 insertions(+), 39 deletions(-) diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index 2c304736b25f..cb2762750e00 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -1270,8 +1270,8 @@ fn test_aes256_non_uniform_encryption_plaintext_footer() { .with_column_key("int32_field", column_4_key.to_vec()) .with_column_key("ba_field", column_5_key.to_vec()) .with_column_key("flba_field", column_6_key.to_vec()) - .with_column_key("int64_list_field", column_7_key.to_vec()) - .with_column_key("int64_field", column_8_key.to_vec()) + .with_column_key("int64_field", column_7_key.to_vec()) + .with_column_key("int96_field", column_8_key.to_vec()) .build() .unwrap(); @@ -1302,8 +1302,8 @@ fn test_aes256_plaintext_footer_signature_verification() { .with_column_key("int32_field", column_4_key.to_vec()) .with_column_key("ba_field", column_5_key.to_vec()) .with_column_key("flba_field", column_6_key.to_vec()) - .with_column_key("int64_list_field", column_7_key.to_vec()) - .with_column_key("int64_field", column_8_key.to_vec()) + .with_column_key("int64_field", column_7_key.to_vec()) + .with_column_key("int96_field", column_8_key.to_vec()) .build() .unwrap(); @@ -1318,8 +1318,8 @@ fn test_aes256_plaintext_footer_signature_verification() { .with_column_key("int32_field", column_4_key.to_vec()) .with_column_key("ba_field", column_5_key.to_vec()) .with_column_key("flba_field", column_6_key.to_vec()) - .with_column_key("int64_list_field", column_7_key.to_vec()) - .with_column_key("int64_field", column_8_key.to_vec()) + .with_column_key("int64_field", column_7_key.to_vec()) + .with_column_key("int96_field", column_8_key.to_vec()) .build() .unwrap(); @@ -1361,8 +1361,8 @@ fn test_aes256_non_uniform_encryption_disabled_aad_storage() { .with_column_key("int32_field", column_4_key.clone()) .with_column_key("ba_field", column_5_key.clone()) .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_list_field", column_7_key.clone()) - .with_column_key("int64_field", column_8_key.clone()) + .with_column_key("int64_field", column_7_key.clone()) + .with_column_key("int96_field", column_8_key.clone()) .with_aad_prefix(b"tester".to_vec()) .build() .unwrap(); @@ -1377,8 +1377,8 @@ fn test_aes256_non_uniform_encryption_disabled_aad_storage() { .with_column_key("int32_field", column_4_key.clone()) .with_column_key("ba_field", column_5_key.clone()) .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_list_field", column_7_key.clone()) - .with_column_key("int64_field", column_8_key.clone()) + .with_column_key("int64_field", column_7_key.clone()) + .with_column_key("int96_field", column_8_key.clone()) .with_aad_prefix(b"wrong_aad_prefix".to_vec()) .build() .unwrap(); @@ -1401,8 +1401,8 @@ fn test_aes256_non_uniform_encryption_disabled_aad_storage() { .with_column_key("int32_field", column_4_key.clone()) .with_column_key("ba_field", column_5_key.clone()) .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_list_field", column_7_key.clone()) - .with_column_key("int64_field", column_8_key.clone()) + .with_column_key("int64_field", column_7_key.clone()) + .with_column_key("int96_field", column_8_key.clone()) .build() .unwrap(); @@ -1442,8 +1442,8 @@ fn test_aes256_non_uniform_encryption() { .with_column_key("int32_field", column_4_key) .with_column_key("ba_field", column_5_key) .with_column_key("flba_field", column_6_key) - .with_column_key("int64_list_field", column_7_key) - .with_column_key("int64_field", column_8_key) + .with_column_key("int64_field", column_7_key) + .with_column_key("int96_field", column_8_key) .build() .unwrap(); @@ -1635,14 +1635,14 @@ fn test_aes256_uniform_encryption_plaintext_footer_with_key_retriever() { b"12345678901234567890123456789017".to_vec(), b"kc6".into(), ) - // .with_column_key_and_metadata("int64_list_field", b"12345678901234567890123456789018".to_vec(), b"kc7".into()) + // .with_column_key_and_metadata("int64_field", b"12345678901234567890123456789018".to_vec(), b"kc7".into()) .with_column_key_and_metadata( - "int64_list_field.list.int64_list_field", + "int64_field.list.int64_field", b"12345678901234567890123456789018".to_vec(), b"kc7".into(), ) .with_column_key_and_metadata( - "int64_field", + "int96_field", b"12345678901234567890123456789019".to_vec(), b"kc8".into(), ) @@ -1821,9 +1821,9 @@ fn test_aes256_write_non_uniform_encryption() { "int32_field", "ba_field", "flba_field", - "int64_list_field", - "int64_list_field.list.int64_list_field", "int64_field", + "int64_field.list.int64_field", + "int96_field", ]; let column_keys = vec![ b"12345678901234567890123456789012".to_vec(), @@ -1859,8 +1859,8 @@ fn test_aes256_write_non_uniform_encryption() { .with_column_key("int32_field", column_4_key) .with_column_key("ba_field", column_5_key) .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_list_field.list.int64_list_field", column_7_key) - .with_column_key("int64_field", column_8_key) + .with_column_key("int64_field.list.int64_field", column_7_key) + .with_column_key("int96_field", column_8_key) .build() .unwrap(); @@ -1891,8 +1891,8 @@ fn test_aes256_write_uniform_encryption_plaintext_footer() { .with_column_key("int32_field", column_4_key.clone()) .with_column_key("ba_field", column_5_key.clone()) .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_list_field", column_7_key.clone()) - .with_column_key("int64_field", column_8_key.clone()) + .with_column_key("int64_field", column_7_key.clone()) + .with_column_key("int96_field", column_8_key.clone()) .build() .unwrap(); @@ -1903,8 +1903,8 @@ fn test_aes256_write_uniform_encryption_plaintext_footer() { .with_column_key("int32_field", column_4_key) .with_column_key("ba_field", column_5_key) .with_column_key("flba_field", column_6_key) - .with_column_key("int64_list_field", column_7_key) - .with_column_key("int64_field", column_8_key) + .with_column_key("int64_field", column_7_key) + .with_column_key("int96_field", column_8_key) .build() .unwrap(); @@ -2061,8 +2061,8 @@ fn test_aes256_decrypt_page_index_non_uniform() { .with_column_key("int32_field", column_4_key) .with_column_key("ba_field", column_5_key) .with_column_key("flba_field", column_6_key) - .with_column_key("int64_list_field", column_7_key) - .with_column_key("int64_field", column_8_key) + .with_column_key("int64_field", column_7_key) + .with_column_key("int96_field", column_8_key) .build() .unwrap(); diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index 9a47b808af32..9474f3a38c03 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -461,8 +461,8 @@ async fn test_aes256_non_uniform_encryption_plaintext_footer() { .with_column_key("int32_field", column_4_key) .with_column_key("ba_field", column_5_key) .with_column_key("flba_field", column_6_key) - .with_column_key("int64_list_field", column_7_key) - .with_column_key("int64_field", column_8_key) + .with_column_key("int64_field", column_7_key) + .with_column_key("int96_field", column_8_key) .build() .unwrap(); @@ -516,8 +516,8 @@ async fn test_aes256_misspecified_encryption_keys() { .with_column_key("int32_field", column_4_key) .with_column_key("ba_field", column_5_key) .with_column_key("flba_field", column_6_key) - .with_column_key("int64_list_field", column_7_key) - .with_column_key("int64_field", column_8_key) + .with_column_key("int64_field", column_7_key) + .with_column_key("int96_field", column_8_key) .build() .unwrap(); @@ -621,8 +621,8 @@ async fn test_aes256_non_uniform_encryption() { .with_column_key("int32_field", column_4_key) .with_column_key("ba_field", column_5_key) .with_column_key("flba_field", column_6_key) - .with_column_key("int64_list_field", column_7_key) - .with_column_key("int64_field", column_8_key) + .with_column_key("int64_field", column_7_key) + .with_column_key("int96_field", column_8_key) .build() .unwrap(); @@ -703,9 +703,9 @@ async fn test_aes256_write_non_uniform_encryption() { "int32_field", "ba_field", "flba_field", - "int64_list_field", - "int64_list_field.list.int64_list_field", "int64_field", + "int64_field.list.int64_field", + "int96_field", ]; let column_keys = vec![ b"12345678901234567890123456789012".to_vec(), @@ -741,8 +741,8 @@ async fn test_aes256_write_non_uniform_encryption() { .with_column_key("int32_field", column_4_key) .with_column_key("ba_field", column_5_key) .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_list_field.list.int64_list_field", column_7_key) - .with_column_key("int64_field", column_8_key) + .with_column_key("int64_field.list.int64_field", column_7_key) + .with_column_key("int96_field", column_8_key) .build() .unwrap(); @@ -921,8 +921,8 @@ async fn test_aes256_decrypt_page_index_non_uniform() { .with_column_key("int32_field", column_4_key) .with_column_key("ba_field", column_5_key) .with_column_key("flba_field", column_6_key) - .with_column_key("int64_list_field", column_7_key) - .with_column_key("int64_field", column_8_key) + .with_column_key("int64_field", column_7_key) + .with_column_key("int96_field", column_8_key) .build() .unwrap(); From e95eb4d2774c384ef6a9e11958783e0618d41e4a Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Tue, 3 Mar 2026 17:08:58 -0800 Subject: [PATCH 12/26] Refactor async encryption tests --- parquet/tests/encryption/encryption_async.rs | 1996 ++++++++++-------- parquet/tests/encryption/encryption_util.rs | 12 + 2 files changed, 1111 insertions(+), 897 deletions(-) diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index 9474f3a38c03..93bc44e9163d 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -17,6 +17,7 @@ //! This module contains tests for reading encrypted Parquet files with the async Arrow API +use crate::encryption_util; use crate::encryption_util::{ TestKeyRetriever, read_encrypted_file, verify_column_indexes, verify_encryption_double_test_data, verify_encryption_test_data, @@ -32,6 +33,7 @@ use parquet::arrow::arrow_writer::{ use parquet::arrow::{ ArrowSchemaConverter, ArrowWriter, AsyncArrowWriter, ParquetRecordBatchStreamBuilder, }; +use parquet::data_type::AsBytes; use parquet::encryption::decrypt::FileDecryptionProperties; use parquet::encryption::encrypt::FileEncryptionProperties; use parquet::errors::ParquetError; @@ -47,61 +49,102 @@ use tokio::task::JoinHandle; #[tokio::test] async fn test_non_uniform_encryption_plaintext_footer() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); + async fn non_uniform_encryption_plaintext_footer( + footer_key: &[u8], + column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } - // There is always a footer key even with a plaintext footer, + // AES-128: there is always a footer key even with a plaintext footer, // but this is used for signing the footer. - let footer_key = "0123456789012345".as_bytes().to_vec(); // 128bit/16 - let column_1_key = "1234567890123450".as_bytes().to_vec(); - let column_2_key = "1234567890123451".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); + non_uniform_encryption_plaintext_footer( + b"0123456789012345", // 128bit/16 + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ) + .await; - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + // AES-256 + non_uniform_encryption_plaintext_footer( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + ) + .await; } #[tokio::test] async fn test_misspecified_encryption_keys() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); - - // There is always a footer key even with a plaintext footer, - // but this is used for signing the footer. - let footer_key = "0123456789012345".as_bytes(); // 128bit/16 - let column_1_key = "1234567890123450".as_bytes(); - let column_2_key = "1234567890123451".as_bytes(); - // read file with keys and check for expected error message async fn check_for_error( expected_message: &str, - path: &String, footer_key: &[u8], column_1_key: &[u8], column_2_key: &[u8], + additional_column_keys: &[(&str, &[u8])], ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); let mut file = File::open(&path).await.unwrap(); - let mut decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()); + let mut builder = FileDecryptionProperties::builder(footer_key.as_bytes().to_vec()); if !column_1_key.is_empty() { - decryption_properties = - decryption_properties.with_column_key("double_field", column_1_key.to_vec()); + builder = builder.with_column_key("double_field", column_1_key.to_vec()); } if !column_2_key.is_empty() { - decryption_properties = - decryption_properties.with_column_key("float_field", column_2_key.to_vec()); + builder = builder.with_column_key("float_field", column_2_key.to_vec()); } - let decryption_properties = decryption_properties.build().unwrap(); + for (column_name, key) in additional_column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + + let decryption_properties = builder.build().unwrap(); match verify_encryption_test_file_read_async(&mut file, decryption_properties).await { Ok(_) => { @@ -113,33 +156,40 @@ async fn test_misspecified_encryption_keys() { } } + // There is always a footer key even with a plaintext footer, + // but this is used for signing the footer. + let footer_key = "0123456789012345".as_bytes(); // 128bit/16 + let column_1_key = "1234567890123450".as_bytes(); + let column_2_key = "1234567890123451".as_bytes(); + let empty_column_key = &[]; + // Too short footer key check_for_error( format!("Parquet error: Invalid footer key. Error creating RingGcmBlockDecryptor with unsupported key length: {}", "bad_pwd".len()).as_str(), - &path, - "bad_pwd".as_bytes(), + b"bad_pwd", column_1_key, column_2_key, + empty_column_key ) .await; // Wrong footer key check_for_error( "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer", - &path, - "1123456789012345".as_bytes(), + b"1123456789012345", column_1_key, column_2_key, + empty_column_key, ) .await; // Missing column key check_for_error( "Parquet error: No column decryption key set for encrypted column 'double_field'", - &path, footer_key, "".as_bytes(), column_2_key, + empty_column_key, ) .await; @@ -150,785 +200,666 @@ async fn test_misspecified_encryption_keys() { "abc".len() ) .as_str(), - &path, footer_key, "abc".as_bytes(), column_2_key, + empty_column_key, ) .await; // Wrong column key check_for_error( "Parquet error: Unable to decrypt column 'double_field', perhaps the column key is wrong?", - &path, footer_key, "1123456789012345".as_bytes(), column_2_key, + empty_column_key, ) .await; // Mixed up keys check_for_error( "Parquet error: Unable to decrypt column 'float_field', perhaps the column key is wrong?", - &path, footer_key, column_2_key, column_1_key, + empty_column_key, ) .await; -} -#[tokio::test] -#[cfg(feature = "snap")] -async fn test_plaintext_footer_read_without_decryption() { - crate::encryption_agnostic::read_plaintext_footer_file_without_decryption_properties_async() - .await; -} - -#[tokio::test] -async fn test_non_uniform_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); - - let footer_key = "0123456789012345".as_bytes().to_vec(); // 128bit/16 - let column_1_key = "1234567890123450".as_bytes().to_vec(); - let column_2_key = "1234567890123451".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); - - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); -} - -#[tokio::test] -async fn test_uniform_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); - - let key_code: &[u8] = "0123456789012345".as_bytes(); - let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) - .build() - .unwrap(); - - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); -} - -#[tokio::test] -async fn test_aes_ctr_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer_ctr.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); - - let footer_key = "0123456789012345".as_bytes().to_vec(); - let column_1_key = "1234567890123450".as_bytes().to_vec(); - //let column_2_key = "1234567890123451".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_1_key) - .build() - .unwrap(); - - let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); - let metadata = ArrowReaderMetadata::load_async(&mut file, options).await; - - match metadata { - Err(ParquetError::NYI(s)) => { - assert!(s.contains("AES_GCM_CTR_V1")); - } - _ => { - panic!("Expected ParquetError::NYI"); - } - }; -} - -#[tokio::test] -async fn test_decrypting_without_decryption_properties_fails() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); - - let options = ArrowReaderOptions::new(); - let result = ArrowReaderMetadata::load_async(&mut file, options).await; - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" - ); -} - -#[tokio::test] -async fn test_write_non_uniform_encryption() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); - - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let column_names = vec!["double_field", "float_field"]; - let column_keys = vec![b"1234567890123450".to_vec(), b"1234567890123451".to_vec()]; - - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_keys(column_names.clone(), column_keys.clone()) - .unwrap() - .build() - .unwrap(); + let aes256_footer_key = "01234567890123456789012345678901".as_bytes(); // 256bit/32 + let aes256_column_1_key = "12345678901234567890123456789012".as_bytes(); + let aes256_column_2_key = "12345678901234567890123456789013".as_bytes(); + let additional_column_keys = &[ + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ]; - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .with_column_keys(column_names, column_keys) - .unwrap() - .build() - .unwrap(); + // Too short footer key + check_for_error( + format!("Parquet error: Invalid footer key. Error creating RingGcmBlockDecryptor with unsupported key length: {}", "bad_pwd".len()).as_str(), + b"bad_pwd", + aes256_column_1_key, + aes256_column_2_key, + additional_column_keys + ).await; - read_and_roundtrip_to_encrypted_file_async( - &path, - decryption_properties, - file_encryption_properties, + // Wrong footer key + check_for_error( + "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer", + b"11234567890123456789012345678901", + aes256_column_1_key, + aes256_column_2_key, + additional_column_keys, ) - .await - .unwrap(); -} + .await; -#[cfg(feature = "object_store")] -async fn get_encrypted_meta_store() -> ( - object_store::ObjectMeta, - std::sync::Arc, -) { - use object_store::local::LocalFileSystem; - use object_store::path::Path; - use object_store::{ObjectStore, ObjectStoreExt}; + // Missing column key + check_for_error( + "Parquet error: No column decryption key set for encrypted column 'double_field'", + aes256_footer_key, + "".as_bytes(), + aes256_column_2_key, + additional_column_keys, + ) + .await; - use std::sync::Arc; - let test_data = arrow::util::test_util::parquet_test_data(); - let store = LocalFileSystem::new_with_prefix(test_data).unwrap(); + // Too short column key + check_for_error( + format!( + "Parquet error: Error creating RingGcmBlockDecryptor with unsupported key length: {}", + "abc".len() + ) + .as_str(), + aes256_footer_key, + "abc".as_bytes(), + aes256_column_2_key, + additional_column_keys, + ) + .await; - let meta = store - .head(&Path::from("uniform_encryption.parquet.encrypted")) - .await - .unwrap(); + // Wrong column key + check_for_error( + "Parquet error: Unable to decrypt column 'double_field', perhaps the column key is wrong?", + aes256_footer_key, + "22345678901234567890123456789012".as_bytes(), + aes256_column_2_key, + additional_column_keys, + ) + .await; - (meta, Arc::new(store) as Arc) + // Mixed up keys + check_for_error( + "Parquet error: Unable to decrypt column 'float_field', perhaps the column key is wrong?", + aes256_footer_key, + aes256_column_2_key, + aes256_column_1_key, + additional_column_keys, + ) + .await; } #[tokio::test] -#[cfg(feature = "object_store")] -async fn test_read_encrypted_file_from_object_store() { - use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader}; - let (meta, store) = get_encrypted_meta_store().await; - - let key_code: &[u8] = "0123456789012345".as_bytes(); - let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) - .build() - .unwrap(); - let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); - - let mut reader = ParquetObjectReader::new(store, meta.location).with_file_size(meta.size); - let metadata = reader.get_metadata(Some(&options)).await.unwrap(); - let builder = ParquetRecordBatchStreamBuilder::new_with_options(reader, options) - .await - .unwrap(); - let batch_stream = builder.build().unwrap(); - let record_batches: Vec<_> = batch_stream.try_collect().await.unwrap(); - - verify_encryption_test_data(record_batches, &metadata); +#[cfg(feature = "snap")] +async fn test_plaintext_footer_read_without_decryption() { + crate::encryption_agnostic::read_plaintext_footer_file_without_decryption_properties_async() + .await; } #[tokio::test] -async fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); +async fn test_non_uniform_encryption() { + async fn non_uniform_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); - let key_retriever = TestKeyRetriever::new() - .with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()) - .with_key("kc1".to_owned(), "1234567890123450".as_bytes().to_vec()) - .with_key("kc2".to_owned(), "1234567890123451".as_bytes().to_vec()); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await .unwrap(); + } - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); -} - -#[tokio::test] -async fn test_non_uniform_encryption_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); - - let key_retriever = TestKeyRetriever::new() - .with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()) - .with_key("kc1".to_owned(), "1234567890123450".as_bytes().to_vec()) - .with_key("kc2".to_owned(), "1234567890123451".as_bytes().to_vec()); - - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() - .unwrap(); + // AES-128 + non_uniform_encryption( + b"0123456789012345", + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ) + .await; - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + // AES-256 + non_uniform_encryption( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + ) + .await; } #[tokio::test] -async fn test_uniform_encryption_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/uniform_encryption.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); +async fn test_uniform_encryption() { + async fn uniform_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); - let key_retriever = - TestKeyRetriever::new().with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await .unwrap(); + } - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + // AES-128: there is always a footer key even with a plaintext footer, + // but this is used for signing the footer. + uniform_encryption( + b"0123456789012345", // 128bit/16 + &[], + ) + .await; + + // AES-256 + uniform_encryption( + b"01234567890123456789012345678901", // 256bit/32 + &[], + ) + .await; } #[tokio::test] -async fn test_decrypt_page_index_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); - - let key_code: &[u8] = "0123456789012345".as_bytes(); - let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) - .build() - .unwrap(); +async fn test_aes_ctr_encryption() { + async fn aes_ctr_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer_ctr.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); - test_decrypt_page_index(&path, decryption_properties) - .await - .unwrap(); -} + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); -#[tokio::test] -async fn test_decrypt_page_index_non_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); + let options = + ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); + let metadata = ArrowReaderMetadata::load_async(&mut file, options).await; - let footer_key = "0123456789012345".as_bytes().to_vec(); - let column_1_key = "1234567890123450".as_bytes().to_vec(); - let column_2_key = "1234567890123451".as_bytes().to_vec(); + match metadata { + Err(ParquetError::NYI(s)) => { + assert!(s.contains("AES_GCM_CTR_V1")); + } + _ => { + panic!("Expected ParquetError::NYI"); + } + }; + } - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); + // AES-128 + aes_ctr_encryption( + b"0123456789012345", // 128bit/16 + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ) + .await; - test_decrypt_page_index(&path, decryption_properties) - .await - .unwrap(); + // AES-256 + aes_ctr_encryption( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ], + ) + .await; } #[tokio::test] -async fn test_aes256_non_uniform_encryption_plaintext_footer() { +async fn test_decrypting_without_decryption_properties_fails() { let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); + let paths = [ + format!("{test_data}/uniform_encryption.parquet.encrypted"), + format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"), + ]; - // There is always a footer key even with a plaintext footer, - // but this is used for signing the footer. - let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 256bit/32 - let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); - let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); - let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); - let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); - let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); - let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); - let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); - let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .with_column_key("boolean_field", column_3_key) - .with_column_key("int32_field", column_4_key) - .with_column_key("ba_field", column_5_key) - .with_column_key("flba_field", column_6_key) - .with_column_key("int64_field", column_7_key) - .with_column_key("int96_field", column_8_key) - .build() - .unwrap(); + for path in &paths { + let mut file = File::open(&path).await.unwrap(); - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + let options = ArrowReaderOptions::new(); + let result = ArrowReaderMetadata::load_async(&mut file, options).await; + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); + } } #[tokio::test] -async fn test_aes256_misspecified_encryption_keys() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/encrypt_columns_and_footer.parquet.encrypted"); - - // There is always a footer key even with a plaintext footer, - // but this is used for signing the footer. - let footer_key = "01234567890123456789012345678901".as_bytes(); // 256bit/32 - let column_1_key = "12345678901234567890123456789012".as_bytes(); - let column_2_key = "12345678901234567890123456789013".as_bytes(); - - // read file with keys and check for expected error message - async fn check_for_error( - expected_message: &str, - path: &String, +async fn test_write_non_uniform_encryption() { + async fn write_non_uniform_encryption( footer_key: &[u8], - column_1_key: &[u8], - column_2_key: &[u8], + column_names: Vec<&str>, + column_keys: Vec>, + encryption_column_keys: &[(&str, &[u8])], ) { - let mut file = File::open(&path).await.unwrap(); - - let mut decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()); - - if !column_1_key.is_empty() { - decryption_properties = - decryption_properties.with_column_key("double_field", column_1_key.to_vec()); - } - - if !column_2_key.is_empty() { - decryption_properties = - decryption_properties.with_column_key("float_field", column_2_key.to_vec()); - } + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); - let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); - let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); - let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); - let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); - let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); - let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); - - let decryption_properties = decryption_properties - .with_column_key("boolean_field", column_3_key) - .with_column_key("int32_field", column_4_key) - .with_column_key("ba_field", column_5_key) - .with_column_key("flba_field", column_6_key) - .with_column_key("int64_field", column_7_key) - .with_column_key("int96_field", column_8_key) + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_column_keys(column_names.to_vec(), column_keys.clone()) + .unwrap() .build() .unwrap(); - match verify_encryption_test_file_read_async(&mut file, decryption_properties).await { - Ok(_) => { - panic!("did not get expected error") - } - Err(e) => { - assert_eq!(e.to_string(), expected_message); - } - } - } - - // Too short footer key - check_for_error( - format!("Parquet error: Invalid footer key. Error creating RingGcmBlockDecryptor with unsupported key length: {}", "bad_pwd".len()).as_str(), - &path, - "bad_pwd".as_bytes(), - column_1_key, - column_2_key, - ) - .await; - - // Wrong footer key - check_for_error( - "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer", - &path, - "11234567890123456789012345678901".as_bytes(), - column_1_key, - column_2_key, - ) - .await; - - // Missing column key - check_for_error( - "Parquet error: No column decryption key set for encrypted column 'double_field'", - &path, - footer_key, - "".as_bytes(), - column_2_key, - ) - .await; - - // Too short column key - check_for_error( - format!( - "Parquet error: Error creating RingGcmBlockDecryptor with unsupported key length: {}", - "abc".len() - ) - .as_str(), - &path, - footer_key, - "abc".as_bytes(), - column_2_key, - ) - .await; + let mut builder = FileEncryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in encryption_column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let file_encryption_properties = builder.build().unwrap(); - // Wrong column key - check_for_error( - "Parquet error: Unable to decrypt column 'double_field', perhaps the column key is wrong?", - &path, - footer_key, - "11234567890123456789012345678901".as_bytes(), - column_2_key, + read_and_roundtrip_to_encrypted_file_async( + &path, + decryption_properties, + file_encryption_properties, + ) + .await + .unwrap(); + } + + write_non_uniform_encryption( + b"0123456789012345", // 128bit/16, + vec!["double_field", "float_field"], + vec![b"1234567890123450".to_vec(), b"1234567890123451".to_vec()], + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], ) .await; - // Mixed up keys - check_for_error( - "Parquet error: Unable to decrypt column 'float_field', perhaps the column key is wrong?", - &path, - footer_key, - column_2_key, - column_1_key, + write_non_uniform_encryption( + b"01234567890123456789012345678901", + vec![ + "double_field", + "float_field", + "boolean_field", + "int32_field", + "ba_field", + "flba_field", + "int64_field", + "int64_field.list.int64_field", + "int96_field", + ], + vec![ + b"12345678901234567890123456789012".to_vec(), + b"12345678901234567890123456789013".to_vec(), + b"12345678901234567890123456789014".to_vec(), + b"12345678901234567890123456789015".to_vec(), + b"12345678901234567890123456789016".to_vec(), + b"12345678901234567890123456789017".to_vec(), + b"12345678901234567890123456789018".to_vec(), + b"12345678901234567890123456789018".to_vec(), + b"12345678901234567890123456789019".to_vec(), + ], + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field.list.int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], ) .await; } -#[tokio::test] -async fn test_aes256_non_uniform_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/encrypt_columns_and_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); +#[cfg(feature = "object_store")] +async fn get_encrypted_meta_store() -> ( + object_store::ObjectMeta, + std::sync::Arc, +) { + use object_store::local::LocalFileSystem; + use object_store::path::Path; + use object_store::{ObjectStore, ObjectStoreExt}; - // There is always a footer key even with a plaintext footer, - // but this is used for signing the footer. - let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 256bit/32 - let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); - let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); - let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); - let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); - let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); - let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); - let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); - let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .with_column_key("boolean_field", column_3_key) - .with_column_key("int32_field", column_4_key) - .with_column_key("ba_field", column_5_key) - .with_column_key("flba_field", column_6_key) - .with_column_key("int64_field", column_7_key) - .with_column_key("int96_field", column_8_key) - .build() - .unwrap(); + use std::sync::Arc; + let test_data = arrow::util::test_util::parquet_test_data(); + let store = LocalFileSystem::new_with_prefix(test_data).unwrap(); - verify_encryption_test_file_read_async(&mut file, decryption_properties) + let meta = store + .head(&Path::from("uniform_encryption.parquet.encrypted")) .await .unwrap(); + + (meta, Arc::new(store) as Arc) } #[tokio::test] -async fn test_aes256_uniform_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); +#[cfg(feature = "object_store")] +async fn test_read_encrypted_file_from_object_store() { + use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader}; + let (meta, store) = get_encrypted_meta_store().await; - let key_code: &[u8] = "01234567890123456789012345678901".as_bytes(); + let key_code: &[u8] = "0123456789012345".as_bytes(); let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) .build() .unwrap(); + let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); - verify_encryption_test_file_read_async(&mut file, decryption_properties) + let mut reader = ParquetObjectReader::new(store, meta.location).with_file_size(meta.size); + let metadata = reader.get_metadata(Some(&options)).await.unwrap(); + let builder = ParquetRecordBatchStreamBuilder::new_with_options(reader, options) .await .unwrap(); + let batch_stream = builder.build().unwrap(); + let record_batches: Vec<_> = batch_stream.try_collect().await.unwrap(); + + verify_encryption_test_data(record_batches, &metadata); } #[tokio::test] -async fn test_aes256_ctr_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/encrypt_columns_and_footer_ctr.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); +async fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { + async fn non_uniform_encryption_plaintext_footer_with_key_retriever( + footer_key: &[u8], + keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); - let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); - let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_1_key) - .build() - .unwrap(); + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); - let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); - let metadata = ArrowReaderMetadata::load_async(&mut file, options).await; + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } - match metadata { - Err(ParquetError::NYI(s)) => { - assert!(s.contains("AES_GCM_CTR_V1")); - } - _ => { - panic!("Expected ParquetError::NYI"); - } - }; + // AES-128 + non_uniform_encryption_plaintext_footer_with_key_retriever( + b"0123456789012345", + &[ + ("kf", b"0123456789012345".as_slice()), + ("kc1", b"1234567890123450".as_slice()), + ("kc2", b"1234567890123451".as_slice()), + ], + ) + .await; + + // AES-256 + non_uniform_encryption_plaintext_footer_with_key_retriever( + b"01234567890123456789012345678901", + &[ + ("kf", b"01234567890123456789012345678901".as_slice()), + ("kc1", b"12345678901234567890123456789012".as_slice()), + ("kc2", b"12345678901234567890123456789013".as_slice()), + ("kc3", b"12345678901234567890123456789014".as_slice()), + ("kc4", b"12345678901234567890123456789015".as_slice()), + ("kc5", b"12345678901234567890123456789016".as_slice()), + ("kc6", b"12345678901234567890123456789017".as_slice()), + ("kc7", b"12345678901234567890123456789018".as_slice()), + ("kc8", b"12345678901234567890123456789019".as_slice()), + ], + ) + .await; } #[tokio::test] -async fn test_aes256_decrypting_without_decryption_properties_fails() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); +async fn test_non_uniform_encryption_with_key_retriever() { + async fn non_uniform_encryption_with_key_retriever(footer_key: &[u8], keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut file = File::open(&path).await.unwrap(); - let options = ArrowReaderOptions::new(); - let result = ArrowReaderMetadata::load_async(&mut file, options).await; - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" - ); -} + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } -#[tokio::test] -async fn test_aes256_write_non_uniform_encryption() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/encrypt_columns_and_footer.parquet.encrypted"); - - let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 - let column_names = vec![ - "double_field", - "float_field", - "boolean_field", - "int32_field", - "ba_field", - "flba_field", - "int64_field", - "int64_field.list.int64_field", - "int96_field", - ]; - let column_keys = vec![ - b"12345678901234567890123456789012".to_vec(), - b"12345678901234567890123456789013".to_vec(), - b"12345678901234567890123456789014".to_vec(), - b"12345678901234567890123456789015".to_vec(), - b"12345678901234567890123456789016".to_vec(), - b"12345678901234567890123456789017".to_vec(), - b"12345678901234567890123456789018".to_vec(), - b"12345678901234567890123456789018".to_vec(), - b"12345678901234567890123456789019".to_vec(), - ]; + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_keys(column_names.clone(), column_keys.clone()) - .unwrap() - .build() - .unwrap(); + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await + .unwrap(); + } - let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); - let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); - let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); - let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); - let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); - let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); - let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); - let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); - - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .with_column_key("boolean_field", column_3_key) - .with_column_key("int32_field", column_4_key) - .with_column_key("ba_field", column_5_key) - .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_field.list.int64_field", column_7_key) - .with_column_key("int96_field", column_8_key) - .build() - .unwrap(); + // AES-128 + non_uniform_encryption_with_key_retriever( + b"0123456789012345", + &[ + ("kf", b"0123456789012345".as_slice()), + ("kc1", b"1234567890123450".as_slice()), + ("kc2", b"1234567890123451".as_slice()), + ], + ) + .await; - read_and_roundtrip_to_encrypted_file_async( - &path, - decryption_properties, - file_encryption_properties, + // AES-256 + non_uniform_encryption_with_key_retriever( + b"01234567890123456789012345678901", + &[ + ("kf", b"01234567890123456789012345678901".as_slice()), + ("kc1", b"12345678901234567890123456789012".as_slice()), + ("kc2", b"12345678901234567890123456789013".as_slice()), + ("kc3", b"12345678901234567890123456789014".as_slice()), + ("kc4", b"12345678901234567890123456789015".as_slice()), + ("kc5", b"12345678901234567890123456789016".as_slice()), + ("kc6", b"12345678901234567890123456789017".as_slice()), + ("kc7", b"12345678901234567890123456789018".as_slice()), + ("kc8", b"12345678901234567890123456789019".as_slice()), + ], ) - .await - .unwrap(); + .await; } #[tokio::test] -async fn test_aes256_non_uniform_encryption_plaintext_footer_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); - - let key_retriever = TestKeyRetriever::new() - .with_key( - "kf".to_owned(), - "01234567890123456789012345678901".as_bytes().to_vec(), - ) - .with_key( - "kc1".to_owned(), - "12345678901234567890123456789012".as_bytes().to_vec(), - ) - .with_key( - "kc2".to_owned(), - "12345678901234567890123456789013".as_bytes().to_vec(), - ) - .with_key( - "kc3".to_owned(), - "12345678901234567890123456789014".as_bytes().to_vec(), - ) - .with_key( - "kc4".to_owned(), - "12345678901234567890123456789015".as_bytes().to_vec(), - ) - .with_key( - "kc5".to_owned(), - "12345678901234567890123456789016".as_bytes().to_vec(), - ) - .with_key( - "kc6".to_owned(), - "12345678901234567890123456789017".as_bytes().to_vec(), - ) - .with_key( - "kc7".to_owned(), - "12345678901234567890123456789018".as_bytes().to_vec(), - ) - .with_key( - "kc8".to_owned(), - "12345678901234567890123456789019".as_bytes().to_vec(), +async fn test_uniform_encryption_with_key_retriever() { + async fn uniform_encryption_with_key_retriever(key_name: &str, footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", ); + let mut file = File::open(&path).await.unwrap(); - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() + let key_retriever = + TestKeyRetriever::new().with_key(key_name.to_owned(), footer_key.to_vec()); + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read_async(&mut file, decryption_properties) + .await .unwrap(); + } + // AES-128 + uniform_encryption_with_key_retriever("kf", b"0123456789012345").await; - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + // AES-256 + uniform_encryption_with_key_retriever("kf", b"01234567890123456789012345678901").await; } #[tokio::test] -async fn test_aes256_non_uniform_encryption_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/encrypt_columns_and_footer.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); - - let key_retriever = TestKeyRetriever::new() - .with_key( - "kf".to_owned(), - "01234567890123456789012345678901".as_bytes().to_vec(), - ) - .with_key( - "kc1".to_owned(), - "12345678901234567890123456789012".as_bytes().to_vec(), - ) - .with_key( - "kc2".to_owned(), - "12345678901234567890123456789013".as_bytes().to_vec(), - ) - .with_key( - "kc3".to_owned(), - "12345678901234567890123456789014".as_bytes().to_vec(), - ) - .with_key( - "kc4".to_owned(), - "12345678901234567890123456789015".as_bytes().to_vec(), - ) - .with_key( - "kc5".to_owned(), - "12345678901234567890123456789016".as_bytes().to_vec(), - ) - .with_key( - "kc6".to_owned(), - "12345678901234567890123456789017".as_bytes().to_vec(), - ) - .with_key( - "kc7".to_owned(), - "12345678901234567890123456789018".as_bytes().to_vec(), - ) - .with_key( - "kc8".to_owned(), - "12345678901234567890123456789019".as_bytes().to_vec(), +async fn test_decrypt_page_index_uniform() { + async fn decrypt_page_index_uniform(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", ); - - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) .build() .unwrap(); - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); -} - -#[tokio::test] -async fn test_aes256_uniform_encryption_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/uniform_encryption.parquet.encrypted"); - let mut file = File::open(&path).await.unwrap(); - - let key_retriever = TestKeyRetriever::new().with_key( - "kf".to_owned(), - "01234567890123456789012345678901".as_bytes().to_vec(), - ); - - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() + test_decrypt_page_index(&path, decryption_properties) + .await .unwrap(); + } - verify_encryption_test_file_read_async(&mut file, decryption_properties) - .await - .unwrap(); + decrypt_page_index_uniform("0123456789012345".as_bytes()).await; + decrypt_page_index_uniform("01234567890123456789012345678901".as_bytes()).await; } #[tokio::test] -async fn test_aes256_decrypt_page_index_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"); - - let key_code: &[u8] = "01234567890123456789012345678901".as_bytes(); - let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) - .build() - .unwrap(); - - test_decrypt_page_index(&path, decryption_properties) - .await - .unwrap(); -} +async fn test_decrypt_page_index_non_uniform() { + async fn decrypt_page_index_non_uniform(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + test_decrypt_page_index(&path, decryption_properties) + .await + .unwrap(); + } -#[tokio::test] -async fn test_aes256_decrypt_page_index_non_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/encrypt_columns_and_footer.parquet.encrypted"); - - let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 256bit/32 - let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); - let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); - let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); - let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); - let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); - let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); - let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); - let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .with_column_key("boolean_field", column_3_key) - .with_column_key("int32_field", column_4_key) - .with_column_key("ba_field", column_5_key) - .with_column_key("flba_field", column_6_key) - .with_column_key("int64_field", column_7_key) - .with_column_key("int96_field", column_8_key) - .build() - .unwrap(); + decrypt_page_index_non_uniform( + b"0123456789012345", // 128bit/16 + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ) + .await; - test_decrypt_page_index(&path, decryption_properties) - .await - .unwrap(); + decrypt_page_index_non_uniform( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + ) + .await; } async fn test_decrypt_page_index( @@ -1170,263 +1101,534 @@ async fn concatenate_parallel_row_groups( #[tokio::test] async fn test_concurrent_encrypted_writing_over_multiple_row_groups() { // Read example data and set up encryption/decryption properties - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); - let file = std::fs::File::open(path).unwrap(); - - let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) - .build() - .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) - .build() - .unwrap(); + async fn concurrent_encrypted_writing_over_multiple_row_groups( + footer_key: &[u8], + dec_column_keys: &[(&str, &[u8])], + enc_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = std::fs::File::open(path).unwrap(); - let (record_batches, metadata) = - read_encrypted_file(&file, decryption_properties.clone()).unwrap(); - let schema = metadata.schema(); + let mut enc_builder = + parquet::encryption::encrypt::FileEncryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in enc_column_keys { + enc_builder = enc_builder.with_column_key(column_name, key.to_vec()); + } + let file_encryption_properties = enc_builder.build().unwrap(); - // Create a channel to send RecordBatches to the writer and send row groups - let (record_batch_tx, data) = tokio::sync::mpsc::channel::(100); - let data_generator = tokio::spawn(async move { - for record_batch in record_batches { - record_batch_tx.send(record_batch).await.unwrap(); + let mut dec_builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in dec_column_keys { + dec_builder = dec_builder.with_column_key(column_name, key.to_vec()); } - }); - - let props = Arc::new( - WriterPropertiesBuilder::default() - .with_file_encryption_properties(file_encryption_properties) - .build(), - ); - let parquet_schema = ArrowSchemaConverter::new() - .with_coerce_types(props.coerce_types()) - .convert(schema) - .unwrap(); + let decryption_properties = dec_builder.build().unwrap(); - // Create a temporary file to write the encrypted data - let temp_file = tempfile::tempfile().unwrap(); + let (record_batches, metadata) = + read_encrypted_file(&file, decryption_properties.clone()).unwrap(); + let schema = metadata.schema(); + + // Create a channel to send RecordBatches to the writer and send row groups + let (record_batch_tx, data) = tokio::sync::mpsc::channel::(100); + let data_generator = tokio::spawn(async move { + for record_batch in record_batches { + record_batch_tx.send(record_batch).await.unwrap(); + } + }); + + let props = Arc::new( + WriterPropertiesBuilder::default() + .with_file_encryption_properties(file_encryption_properties) + .build(), + ); + let parquet_schema = ArrowSchemaConverter::new() + .with_coerce_types(props.coerce_types()) + .convert(schema) + .unwrap(); - let writer = - SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); - let row_group_writer_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(schema)); - let max_row_groups = 1; + // Create a temporary file to write the encrypted data + let temp_file = tempfile::tempfile().unwrap(); - let (serialize_tx, serialize_rx) = - tokio::sync::mpsc::channel::>(max_row_groups); + let writer = + SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); + let row_group_writer_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(schema)); + let max_row_groups = 1; - let launch_serialization_task = spawn_parquet_parallel_serialization_task( - row_group_writer_factory, - data, - serialize_tx, - schema.clone(), - ); + let (serialize_tx, serialize_rx) = + tokio::sync::mpsc::channel::>(max_row_groups); - let _file_metadata = concatenate_parallel_row_groups(writer, serialize_rx) - .await - .unwrap(); + let launch_serialization_task = spawn_parquet_parallel_serialization_task( + row_group_writer_factory, + data, + serialize_tx, + schema.clone(), + ); - data_generator.await.unwrap(); - launch_serialization_task.await.unwrap().unwrap(); + let _file_metadata = concatenate_parallel_row_groups(writer, serialize_rx) + .await + .unwrap(); + + data_generator.await.unwrap(); + launch_serialization_task.await.unwrap().unwrap(); + + // Check that the file was written correctly + let (read_record_batches, read_metadata) = + read_encrypted_file(&temp_file, decryption_properties.clone()).unwrap(); - // Check that the file was written correctly - let (read_record_batches, read_metadata) = - read_encrypted_file(&temp_file, decryption_properties.clone()).unwrap(); + assert_eq!(read_metadata.metadata().file_metadata().num_rows(), 50); + verify_encryption_test_data(read_record_batches, read_metadata.metadata()); + } + + // AES-128 + concurrent_encrypted_writing_over_multiple_row_groups( + b"0123456789012345", // 128bit/16 + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ) + .await; - assert_eq!(read_metadata.metadata().file_metadata().num_rows(), 50); - verify_encryption_test_data(read_record_batches, read_metadata.metadata()); + // AES-256 + concurrent_encrypted_writing_over_multiple_row_groups( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int64_field.list.int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field.list.int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + ) + .await; } #[tokio::test] async fn test_multi_threaded_encrypted_writing() { // Read example data and set up encryption/decryption properties - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); - let file = std::fs::File::open(path).unwrap(); - - let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) - .build() - .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) - .build() - .unwrap(); + async fn multi_threaded_encrypted_writing( + footer_key: &[u8], + dec_column_keys: &[(&str, &[u8])], + enc_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = std::fs::File::open(path).unwrap(); - let (record_batches, metadata) = - read_encrypted_file(&file, Arc::clone(&decryption_properties)).unwrap(); - let schema = metadata.schema().clone(); + let mut enc_builder = + parquet::encryption::encrypt::FileEncryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in enc_column_keys { + enc_builder = enc_builder.with_column_key(column_name, key.to_vec()); + } + let encryption_properties = enc_builder.build().unwrap(); - let props = Arc::new( - WriterPropertiesBuilder::default() - .with_file_encryption_properties(file_encryption_properties) - .build(), - ); + let mut dec_builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in dec_column_keys { + dec_builder = dec_builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = dec_builder.build().unwrap(); - let parquet_schema = ArrowSchemaConverter::new() - .with_coerce_types(props.coerce_types()) - .convert(&schema) - .unwrap(); + let (record_batches, metadata) = + read_encrypted_file(&file, Arc::clone(&decryption_properties)).unwrap(); + let schema = metadata.schema().clone(); - // Create a temporary file to write the encrypted data - let temp_file = tempfile::tempfile().unwrap(); - let mut writer = - SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); - let row_group_writer_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); + let props = Arc::new( + WriterPropertiesBuilder::default() + .with_file_encryption_properties(encryption_properties) + .build(), + ); - let (serialize_tx, mut serialize_rx) = - tokio::sync::mpsc::channel::>(1); + let parquet_schema = ArrowSchemaConverter::new() + .with_coerce_types(props.coerce_types()) + .convert(&schema) + .unwrap(); - // Create a channel to send RecordBatches to the writer and send row batches - let (record_batch_tx, mut data) = tokio::sync::mpsc::channel::(100); - let data_generator = tokio::spawn(async move { - for record_batch in record_batches { - record_batch_tx.send(record_batch).await.unwrap(); - } - }); + // Create a temporary file to write the encrypted data + let temp_file = tempfile::tempfile().unwrap(); + let mut writer = + SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); + let row_group_writer_factory = + ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); + + let (serialize_tx, mut serialize_rx) = + tokio::sync::mpsc::channel::>(1); + + // Create a channel to send RecordBatches to the writer and send row batches + let (record_batch_tx, mut data) = tokio::sync::mpsc::channel::(100); + let data_generator = tokio::spawn(async move { + for record_batch in record_batches { + record_batch_tx.send(record_batch).await.unwrap(); + } + }); - // Get column writers - let col_writers = row_group_writer_factory.create_column_writers(0).unwrap(); + // Get column writers + let col_writers = row_group_writer_factory.create_column_writers(0).unwrap(); - let (col_writer_tasks, col_array_channels) = - spawn_column_parallel_row_group_writer(col_writers, 10).unwrap(); + let (col_writer_tasks, col_array_channels) = + spawn_column_parallel_row_group_writer(col_writers, 10).unwrap(); - // Spawn serialization tasks for incoming RecordBatches - let launch_serialization_task = tokio::spawn(async move { - let Some(rb) = data.recv().await else { - panic!() - }; - send_arrays_to_column_writers(&col_array_channels, &rb, &schema) - .await - .unwrap(); - let finalize_rg_task = spawn_rg_join_and_finalize_task(col_writer_tasks, 10); + // Spawn serialization tasks for incoming RecordBatches + let launch_serialization_task = tokio::spawn(async move { + let Some(rb) = data.recv().await else { + panic!() + }; + send_arrays_to_column_writers(&col_array_channels, &rb, &schema) + .await + .unwrap(); + let finalize_rg_task = spawn_rg_join_and_finalize_task(col_writer_tasks, 10); - serialize_tx.send(finalize_rg_task).await.unwrap(); - drop(col_array_channels); - }); + serialize_tx.send(finalize_rg_task).await.unwrap(); + drop(col_array_channels); + }); - // Append the finalized row groups to the SerializedFileWriter - while let Some(task) = serialize_rx.recv().await { - let (arrow_column_chunks, _) = task.await.unwrap().unwrap(); - let mut row_group_writer = writer.next_row_group().unwrap(); - for chunk in arrow_column_chunks { - chunk.append_to_row_group(&mut row_group_writer).unwrap(); + // Append the finalized row groups to the SerializedFileWriter + while let Some(task) = serialize_rx.recv().await { + let (arrow_column_chunks, _) = task.await.unwrap().unwrap(); + let mut row_group_writer = writer.next_row_group().unwrap(); + for chunk in arrow_column_chunks { + chunk.append_to_row_group(&mut row_group_writer).unwrap(); + } + row_group_writer.close().unwrap(); } - row_group_writer.close().unwrap(); + + // Wait for data generator and serialization task to finish + data_generator.await.unwrap(); + launch_serialization_task.await.unwrap(); + let metadata = writer.close().unwrap(); + + // Close the file writer which writes the footer + assert_eq!(metadata.file_metadata().num_rows(), 50); + + // Check that the file was written correctly + let (read_record_batches, read_metadata) = + read_encrypted_file(&temp_file, decryption_properties).unwrap(); + verify_encryption_test_data(read_record_batches, read_metadata.metadata()); + + // Check that file was encrypted + let result = ArrowReaderMetadata::load(&temp_file, ArrowReaderOptions::default()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); } - // Wait for data generator and serialization task to finish - data_generator.await.unwrap(); - launch_serialization_task.await.unwrap(); - let metadata = writer.close().unwrap(); - - // Close the file writer which writes the footer - assert_eq!(metadata.file_metadata().num_rows(), 50); - - // Check that the file was written correctly - let (read_record_batches, read_metadata) = - read_encrypted_file(&temp_file, decryption_properties).unwrap(); - verify_encryption_test_data(read_record_batches, read_metadata.metadata()); - - // Check that file was encrypted - let result = ArrowReaderMetadata::load(&temp_file, ArrowReaderOptions::default()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" - ); + // AES-128 + multi_threaded_encrypted_writing( + b"0123456789012345", // 128bit/16 + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ) + .await; + + // AES-256 + multi_threaded_encrypted_writing( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int64_field.list.int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field.list.int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + ) + .await; } #[tokio::test] async fn test_multi_threaded_encrypted_writing_deprecated() { // Read example data and set up encryption/decryption properties - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); - let file = std::fs::File::open(path).unwrap(); + async fn multi_threaded_encrypted_writing_deprecated( + footer_key: &[u8], + dec_column_keys: &[(&str, &[u8])], + enc_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = std::fs::File::open(path).unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) - .build() - .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) - .build() - .unwrap(); + let mut enc_builder = + parquet::encryption::encrypt::FileEncryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in enc_column_keys { + enc_builder = enc_builder.with_column_key(column_name, key.to_vec()); + } + let encryption_properties = enc_builder.build().unwrap(); - let (record_batches, metadata) = - read_encrypted_file(&file, Arc::clone(&decryption_properties)).unwrap(); - let to_write: Vec<_> = record_batches - .iter() - .flat_map(|rb| rb.columns().to_vec()) - .collect(); - let schema = metadata.schema().clone(); - - let props = Some( - WriterPropertiesBuilder::default() - .with_file_encryption_properties(file_encryption_properties) - .build(), - ); - - // Create a temporary file to write the encrypted data - let temp_file = tempfile::tempfile().unwrap(); - let mut writer = ArrowWriter::try_new(&temp_file, schema.clone(), props).unwrap(); + let mut dec_builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in dec_column_keys { + dec_builder = dec_builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = dec_builder.build().unwrap(); + + let (record_batches, metadata) = + read_encrypted_file(&file, Arc::clone(&decryption_properties)).unwrap(); + let to_write: Vec<_> = record_batches + .iter() + .flat_map(|rb| rb.columns().to_vec()) + .collect(); + let schema = metadata.schema().clone(); + + let props = Some( + WriterPropertiesBuilder::default() + .with_file_encryption_properties(encryption_properties) + .build(), + ); - // LOW-LEVEL API: Use low level API to write into a file using multiple threads + // Create a temporary file to write the encrypted data + let temp_file = tempfile::tempfile().unwrap(); + let mut writer = ArrowWriter::try_new(&temp_file, schema.clone(), props).unwrap(); - // Get column writers - #[allow(deprecated)] - let col_writers = writer.get_column_writers().unwrap(); - let num_columns = col_writers.len(); + // LOW-LEVEL API: Use low level API to write into a file using multiple threads - let (col_writer_tasks, mut col_array_channels) = - spawn_column_parallel_row_group_writer(col_writers, 100).unwrap(); + // Get column writers + #[allow(deprecated)] + let col_writers = writer.get_column_writers().unwrap(); + let num_columns = col_writers.len(); - // Send the ArrowLeafColumn data to the respective column writer channels - let mut worker_iter = col_array_channels.iter_mut(); - for (array, field) in to_write.iter().zip(schema.fields()) { - for leaves in compute_leaves(field, array).unwrap() { - worker_iter.next().unwrap().send(leaves).await.unwrap(); + let (col_writer_tasks, mut col_array_channels) = + spawn_column_parallel_row_group_writer(col_writers, 100).unwrap(); + + // Send the ArrowLeafColumn data to the respective column writer channels + let mut worker_iter = col_array_channels.iter_mut(); + for (array, field) in to_write.iter().zip(schema.fields()) { + for leaves in compute_leaves(field, array).unwrap() { + worker_iter.next().unwrap().send(leaves).await.unwrap(); + } } - } - drop(col_array_channels); + drop(col_array_channels); - // Wait for all column writers to finish writing - let mut finalized_rg = Vec::with_capacity(num_columns); - for task in col_writer_tasks.into_iter() { - finalized_rg.push(task.await.unwrap().unwrap().close().unwrap()); - } + // Wait for all column writers to finish writing + let mut finalized_rg = Vec::with_capacity(num_columns); + for task in col_writer_tasks.into_iter() { + finalized_rg.push(task.await.unwrap().unwrap().close().unwrap()); + } - // Append the finalized row group to the SerializedFileWriter - #[allow(deprecated)] - writer.append_row_group(finalized_rg).unwrap(); + // Append the finalized row group to the SerializedFileWriter + #[allow(deprecated)] + writer.append_row_group(finalized_rg).unwrap(); - // HIGH-LEVEL API: Write RecordBatches into the file using ArrowWriter + // HIGH-LEVEL API: Write RecordBatches into the file using ArrowWriter - // Write individual RecordBatches into the file - for rb in record_batches { - writer.write(&rb).unwrap() + // Write individual RecordBatches into the file + for rb in record_batches { + writer.write(&rb).unwrap() + } + assert!(writer.flush().is_ok()); + + // Close the file writer which writes the footer + let metadata = writer.finish().unwrap(); + assert_eq!(metadata.file_metadata().num_rows(), 100); + + // Check that the file was written correctly + let (read_record_batches, read_metadata) = + read_encrypted_file(&temp_file, decryption_properties).unwrap(); + verify_encryption_double_test_data(read_record_batches, read_metadata.metadata()); + + // Check that file was encrypted + let result = ArrowReaderMetadata::load(&temp_file, ArrowReaderOptions::default()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); } - assert!(writer.flush().is_ok()); - - // Close the file writer which writes the footer - let metadata = writer.finish().unwrap(); - assert_eq!(metadata.file_metadata().num_rows(), 100); - - // Check that the file was written correctly - let (read_record_batches, read_metadata) = - read_encrypted_file(&temp_file, decryption_properties).unwrap(); - verify_encryption_double_test_data(read_record_batches, read_metadata.metadata()); - - // Check that file was encrypted - let result = ArrowReaderMetadata::load(&temp_file, ArrowReaderOptions::default()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" - ); + + // AES-128 + multi_threaded_encrypted_writing_deprecated( + b"0123456789012345", // 128bit/16 + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ) + .await; + + // AES-256 + multi_threaded_encrypted_writing_deprecated( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int64_field.list.int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field.list.int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + ) + .await; } diff --git a/parquet/tests/encryption/encryption_util.rs b/parquet/tests/encryption/encryption_util.rs index 7f4cc5a9da45..a940a47bbd69 100644 --- a/parquet/tests/encryption/encryption_util.rs +++ b/parquet/tests/encryption/encryption_util.rs @@ -26,6 +26,7 @@ use parquet::encryption::encrypt::FileEncryptionProperties; use parquet::errors::{ParquetError, Result}; use parquet::file::metadata::ParquetMetaData; use parquet::file::properties::WriterProperties; +use ring::aead::AES_256_GCM; use std::collections::HashMap; use std::fs::File; use std::sync::{Arc, Mutex}; @@ -317,3 +318,14 @@ impl KeyRetriever for TestKeyRetriever { } } } + +pub fn encrypted_data_path(footer_key: &[u8], file_name: &str) -> String { + let test_data = arrow::util::test_util::parquet_test_data(); + let subpath = if AES_256_GCM.key_len() == footer_key.len() { + "aes256/" + } else { + "" + }; + let path = format!("{test_data}/{subpath}/{file_name}"); + path +} From e8999d370a57e7e76de1c4c454ab130b4d5269a0 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Thu, 5 Mar 2026 13:00:21 -0800 Subject: [PATCH 13/26] Refactor sync encryption tests --- parquet/tests/encryption/encryption.rs | 2232 +++++++++++------------- 1 file changed, 1009 insertions(+), 1223 deletions(-) diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index cb2762750e00..4eed9c6a2a56 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -17,6 +17,7 @@ //! This module contains tests for reading encrypted Parquet files with the Arrow API +use crate::encryption_util; use crate::encryption_util::{ TestKeyRetriever, read_and_roundtrip_to_encrypted_file, verify_column_indexes, verify_encryption_test_file_read, @@ -43,119 +44,243 @@ use std::sync::Arc; #[test] fn test_non_uniform_encryption_plaintext_footer() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn non_uniform_encryption_plaintext_footer(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + verify_encryption_test_file_read(file, decryption_properties); + } - // There is always a footer key even with a plaintext footer, + // AES-128: there is always a footer key even with a plaintext footer, // but this is used for signing the footer. - let footer_key = "0123456789012345".as_bytes(); // 128bit/16 - let column_1_key = "1234567890123450".as_bytes(); - let column_2_key = "1234567890123451".as_bytes(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key.to_vec()) - .with_column_key("float_field", column_2_key.to_vec()) - .build() - .unwrap(); + non_uniform_encryption_plaintext_footer( + b"0123456789012345", // 128bit/16 + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ); - verify_encryption_test_file_read(file, decryption_properties); + // AES-256 + non_uniform_encryption_plaintext_footer( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + ); } #[test] fn test_plaintext_footer_signature_verification() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path.clone()).unwrap(); - - let footer_key = "0000000000000000".as_bytes(); // 128bit/16 - let column_1_key = "1234567890123450".as_bytes(); - let column_2_key = "1234567890123451".as_bytes(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .disable_footer_signature_verification() - .with_column_key("double_field", column_1_key.to_vec()) - .with_column_key("float_field", column_2_key.to_vec()) - .build() - .unwrap(); - - verify_encryption_test_file_read(file, decryption_properties); + fn plaintext_footer_signature_verification(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path.clone()).unwrap(); + let mut disable_footer_signature_verification_builder = + FileDecryptionProperties::builder(footer_key.to_vec()) + .disable_footer_signature_verification(); + for (column_name, key) in column_keys { + disable_footer_signature_verification_builder = + disable_footer_signature_verification_builder + .with_column_key(column_name, key.to_vec()); + } + let decryption_properties = disable_footer_signature_verification_builder + .build() + .unwrap(); + verify_encryption_test_file_read(file, decryption_properties); - let file = File::open(path.clone()).unwrap(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key.to_vec()) - .with_column_key("float_field", column_2_key.to_vec()) - .build() - .unwrap(); + let file = File::open(path.clone()).unwrap(); + let decryption_properties = builder.build().unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .starts_with("Parquet error: Footer signature verification failed. Computed: [") + ); + } - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let result = ArrowReaderMetadata::load(&file, options.clone()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .starts_with("Parquet error: Footer signature verification failed. Computed: [") + plaintext_footer_signature_verification( + // 128bit/16 + b"0000000000000000", // 128bit/16 + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], ); + + plaintext_footer_signature_verification( + b"00000000000000000000000000000000", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + ) } #[test] fn test_non_uniform_encryption_disabled_aad_storage() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = - format!("{test_data}/encrypt_columns_and_footer_disable_aad_storage.parquet.encrypted"); - let file = File::open(path.clone()).unwrap(); - - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let column_1_key = b"1234567890123450".to_vec(); - let column_2_key = b"1234567890123451".to_vec(); + fn non_uniform_encryption_disabled_aad_storage( + footer_key: &[u8], + column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer_disable_aad_storage.parquet.encrypted", + ); - // Can read successfully when providing the correct AAD prefix - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) - .with_aad_prefix(b"tester".to_vec()) - .build() - .unwrap(); + // Can read successfully when providing the correct AAD prefix + let file = File::open(path.clone()).unwrap(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_aad_prefix(b"tester".to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + verify_encryption_test_file_read(file, decryption_properties); + + // Using wrong AAD prefix should fail + let file = File::open(path.clone()).unwrap(); + let mut wrong_aad_builder = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_aad_prefix(b"wrong_aad_prefix".to_vec()); + for (column_name, key) in column_keys { + wrong_aad_builder = wrong_aad_builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = wrong_aad_builder.build().unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + std::assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer" + ); - verify_encryption_test_file_read(file, decryption_properties); + // Not providing any AAD prefix should fail as it isn't stored in the file + let mut no_aad_builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + no_aad_builder = no_aad_builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = no_aad_builder.build().unwrap(); - // Using wrong AAD prefix should fail - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) - .with_aad_prefix(b"wrong_aad_prefix".to_vec()) - .build() - .unwrap(); + let file = File::open(path).unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file was encrypted with an AAD prefix that is not stored in the file, \ + but no AAD prefix was provided in the file decryption properties" + ); + } - let file = File::open(path.clone()).unwrap(); - let options = ArrowReaderOptions::default() - .with_file_decryption_properties(decryption_properties.clone()); - let result = ArrowReaderMetadata::load(&file, options.clone()); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer" + // AES-128 + non_uniform_encryption_disabled_aad_storage( + b"0123456789012345", // 128bit/16 + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], ); - // Not providing any AAD prefix should fail as it isn't stored in the file - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); - - let file = File::open(path).unwrap(); - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let result = ArrowReaderMetadata::load(&file, options.clone()); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file was encrypted with an AAD prefix that is not stored in the file, \ - but no AAD prefix was provided in the file decryption properties" + // AES-256 + non_uniform_encryption_disabled_aad_storage( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], ); } @@ -167,33 +292,95 @@ fn test_plaintext_footer_read_without_decryption() { #[test] fn test_non_uniform_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn non_uniform_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let column_1_key = b"1234567890123450".to_vec(); - let column_2_key = b"1234567890123451".to_vec(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); + verify_encryption_test_file_read(file, decryption_properties); + } + + // AES-128 + non_uniform_encryption( + b"0123456789012345", + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ); - verify_encryption_test_file_read(file, decryption_properties); + // AES-256 + non_uniform_encryption( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + ); } #[test] fn test_uniform_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn uniform_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + } - let key_code = b"0123456789012345".to_vec(); - let decryption_properties = FileDecryptionProperties::builder(key_code).build().unwrap(); + // AES-128: there is always a footer key even with a plaintext footer, + // but this is used for signing the footer. + uniform_encryption( + b"0123456789012345", // 128bit/16 + &[], + ); - verify_encryption_test_file_read(file, decryption_properties); + // AES-256 + uniform_encryption( + b"01234567890123456789012345678901", // 256bit/32 + &[], + ); } #[test] @@ -213,173 +400,353 @@ fn test_decrypting_without_decryption_properties_fails() { #[test] fn test_aes_ctr_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer_ctr.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn aes_ctr_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer_ctr.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let footer_key = b"0123456789012345".to_vec(); - let column_1_key = b"1234567890123450".to_vec(); - let column_2_key = b"1234567890123451".to_vec(); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); + let options = + ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); + let metadata = ArrowReaderMetadata::load(&file, options); - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let metadata = ArrowReaderMetadata::load(&file, options); + match metadata { + Err(ParquetError::NYI(s)) => { + assert!(s.contains("AES_GCM_CTR_V1")); + } + _ => { + panic!("Expected ParquetError::NYI"); + } + }; + } + // AES-128 + aes_ctr_encryption( + b"0123456789012345", // 128bit/16 + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ); - match metadata { - Err(parquet::errors::ParquetError::NYI(s)) => { - assert!(s.contains("AES_GCM_CTR_V1")); - } - _ => { - panic!("Expected ParquetError::NYI"); - } - }; + // AES-256 + aes_ctr_encryption( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ], + ); } #[test] fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn non_uniform_encryption_plaintext_footer_with_key_retriever( + footer_key: &[u8], + keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); + + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } - let key_retriever = TestKeyRetriever::new() - .with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()) - .with_key("kc1".to_owned(), "1234567890123450".as_bytes().to_vec()) - .with_key("kc2".to_owned(), "1234567890123451".as_bytes().to_vec()); + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() - .unwrap(); + verify_encryption_test_file_read(file, decryption_properties); + } + + // AES-128 + non_uniform_encryption_plaintext_footer_with_key_retriever( + b"0123456789012345", + &[ + ("kf", b"0123456789012345".as_slice()), + ("kc1", b"1234567890123450".as_slice()), + ("kc2", b"1234567890123451".as_slice()), + ], + ); - verify_encryption_test_file_read(file, decryption_properties); + // AES-256 + non_uniform_encryption_plaintext_footer_with_key_retriever( + b"01234567890123456789012345678901", + &[ + ("kf", b"01234567890123456789012345678901".as_slice()), + ("kc1", b"12345678901234567890123456789012".as_slice()), + ("kc2", b"12345678901234567890123456789013".as_slice()), + ("kc3", b"12345678901234567890123456789014".as_slice()), + ("kc4", b"12345678901234567890123456789015".as_slice()), + ("kc5", b"12345678901234567890123456789016".as_slice()), + ("kc6", b"12345678901234567890123456789017".as_slice()), + ("kc7", b"12345678901234567890123456789018".as_slice()), + ("kc8", b"12345678901234567890123456789019".as_slice()), + ], + ); } #[test] fn test_uniform_encryption_plaintext_footer_with_key_retriever() { - let test_data = arrow::util::test_util::parquet_test_data(); + fn uniform_encryption_plaintext_footer_with_key_retriever( + footer_key: &[u8], + footer_key_metadata: &str, + wrong_footer_key: &[u8], + dec_column_keys: &[(&str, &[u8])], + enc_column_keys: &[(&str, &[u8], &str)], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_plaintext_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - // Read example data with key retriever - let path = format!("{test_data}/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + let mut key_retriever = + TestKeyRetriever::new().with_key(footer_key_metadata.to_owned(), footer_key.to_vec()); - let key_retriever = Arc::new( - TestKeyRetriever::new() - .with_key("kf".to_owned(), b"0123456789012345".to_vec()) - .with_key("kc1".to_owned(), b"1234567890123450".to_vec()) - .with_key("kc2".to_owned(), b"1234567890123451".to_vec()), - ); + for (key_name, key) in dec_column_keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } - let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever.clone()) - .build() - .unwrap(); + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + + // Write data into temporary file with plaintext footer and footer key metadata + let temp_file = tempfile::tempfile().unwrap(); + let mut encryption_properties_builder = + FileEncryptionProperties::builder(footer_key.to_vec()) + .with_footer_key_metadata(footer_key_metadata.into()); + for (column_name, key, metadata) in enc_column_keys { + encryption_properties_builder = encryption_properties_builder + .with_column_key_and_metadata(column_name, key.to_vec(), (*metadata).into()); + } + let encryption_properties = encryption_properties_builder + .with_plaintext_footer(true) + .build() + .unwrap(); - // Write data into temporary file with plaintext footer and footer key metadata - let temp_file = tempfile::tempfile().unwrap(); - let encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".to_vec()) - .with_footer_key_metadata("kf".into()) - .with_column_key_and_metadata("double_field", b"1234567890123450".to_vec(), b"kc1".into()) - .with_column_key_and_metadata("float_field", b"1234567890123451".to_vec(), b"kc2".into()) - .with_plaintext_footer(true) - .build() - .unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); + let batch_reader = builder.build().unwrap(); + let batches = batch_reader + .collect::, _>>() + .unwrap(); - let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); - let batch_reader = builder.build().unwrap(); - let batches = batch_reader - .collect::, _>>() + let props = WriterProperties::builder() + .with_file_encryption_properties(encryption_properties) + .build(); + + let mut writer = ArrowWriter::try_new( + temp_file.try_clone().unwrap(), + metadata.schema().clone(), + Some(props), + ) .unwrap(); + for batch in batches { + writer.write(&batch).unwrap(); + } - let props = WriterProperties::builder() - .with_file_encryption_properties(encryption_properties) - .build(); + writer.close().unwrap(); - let mut writer = ArrowWriter::try_new( - temp_file.try_clone().unwrap(), - metadata.schema().clone(), - Some(props), - ) - .unwrap(); - for batch in batches { - writer.write(&batch).unwrap(); - } + // Read temporary file with plaintext metadata using key retriever + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let _ = ArrowReaderMetadata::load(&temp_file, options.clone()).unwrap(); - writer.close().unwrap(); + // Read temporary file with plaintext metadata using key retriever with invalid key + let mut key_retriever = TestKeyRetriever::new() + .with_key(footer_key_metadata.to_owned(), wrong_footer_key.to_vec()); - // Read temporary file with plaintext metadata using key retriever - let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever) - .build() - .unwrap(); + for (key_name, key) in dec_column_keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + let options = + ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); + let result = ArrowReaderMetadata::load(&temp_file, options.clone()); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .starts_with("Parquet error: Footer signature verification failed. Computed: [") + ); + } - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let _ = ArrowReaderMetadata::load(&temp_file, options.clone()).unwrap(); - - // Read temporary file with plaintext metadata using key retriever with invalid key - let key_retriever = Arc::new( - TestKeyRetriever::new() - .with_key("kf".to_owned(), b"0133756789012345".to_vec()) - .with_key("kc1".to_owned(), b"1234567890123450".to_vec()) - .with_key("kc2".to_owned(), b"1234567890123451".to_vec()), + // AES-128 + uniform_encryption_plaintext_footer_with_key_retriever( + b"0123456789012345", + "kf", + b"0133456789012345", + &[ + ("kc1", b"1234567890123450".as_slice()), + ("kc2", b"1234567890123451".as_slice()), + ], + &[ + ("double_field", b"1234567890123450".as_slice(), "kc1"), + ("float_field", b"1234567890123451".as_slice(), "kc2"), + ], ); - let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever) - .build() - .unwrap(); - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let result = ArrowReaderMetadata::load(&temp_file, options.clone()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .starts_with("Parquet error: Footer signature verification failed. Computed: [") - ); -} -#[test] -fn test_non_uniform_encryption_with_key_retriever() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + // AES-256 + uniform_encryption_plaintext_footer_with_key_retriever( + b"01234567890123456789012345678901", + "kf", + b"01334567890123456789012345678901", + &[ + ("kc1", b"12345678901234567890123456789012".as_slice()), + ("kc2", b"12345678901234567890123456789013".as_slice()), + ("kc3", b"12345678901234567890123456789014".as_slice()), + ("kc4", b"12345678901234567890123456789015".as_slice()), + ("kc5", b"12345678901234567890123456789016".as_slice()), + ("kc6", b"12345678901234567890123456789017".as_slice()), + ("kc7", b"12345678901234567890123456789018".as_slice()), + ("kc8", b"12345678901234567890123456789019".as_slice()), + ], + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + "kc1", + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + "kc2", + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + "kc3", + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + "kc4", + ), + ( + "ba_field", + b"12345678901234567890123456789016".as_slice(), + "kc5", + ), + ( + "flba_field", + b"12345678901234567890123456789017".as_slice(), + "kc6", + ), + ( + "int64_field.list.int64_field", + b"12345678901234567890123456789018".as_slice(), + "kc7", + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + "kc8", + ), + ], + ); +} - let key_retriever = TestKeyRetriever::new() - .with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()) - .with_key("kc1".to_owned(), "1234567890123450".as_bytes().to_vec()) - .with_key("kc2".to_owned(), "1234567890123451".as_bytes().to_vec()); +#[test] +fn test_non_uniform_encryption_with_key_retriever() { + fn non_uniform_encryption_with_key_retriever(footer_key: &[u8], keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() - .unwrap(); + let mut key_retriever = TestKeyRetriever::new(); + for (key_name, key) in keys { + key_retriever = key_retriever.with_key((*key_name).to_owned(), (*key).to_vec()); + } + + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + } + + // AES-128 + non_uniform_encryption_with_key_retriever( + b"0123456789012345", + &[ + ("kf", b"0123456789012345".as_slice()), + ("kc1", b"1234567890123450".as_slice()), + ("kc2", b"1234567890123451".as_slice()), + ], + ); - verify_encryption_test_file_read(file, decryption_properties); + // AES-256 + non_uniform_encryption_with_key_retriever( + b"01234567890123456789012345678901", + &[ + ("kf", b"01234567890123456789012345678901".as_slice()), + ("kc1", b"12345678901234567890123456789012".as_slice()), + ("kc2", b"12345678901234567890123456789013".as_slice()), + ("kc3", b"12345678901234567890123456789014".as_slice()), + ("kc4", b"12345678901234567890123456789015".as_slice()), + ("kc5", b"12345678901234567890123456789016".as_slice()), + ("kc6", b"12345678901234567890123456789017".as_slice()), + ("kc7", b"12345678901234567890123456789018".as_slice()), + ("kc8", b"12345678901234567890123456789019".as_slice()), + ], + ); } #[test] fn test_uniform_encryption_with_key_retriever() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn uniform_encryption_with_key_retriever(key_name: &str, footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let key_retriever = - TestKeyRetriever::new().with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()); + let key_retriever = + TestKeyRetriever::new().with_key(key_name.to_owned(), footer_key.to_vec()); - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() - .unwrap(); + let decryption_properties = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + verify_encryption_test_file_read(file, decryption_properties); + } + + // AES-128 + uniform_encryption_with_key_retriever("kf", b"0123456789012345"); - verify_encryption_test_file_read(file, decryption_properties); + // AES-256 + uniform_encryption_with_key_retriever("kf", b"01234567890123456789012345678901"); } fn row_group_sizes(metadata: &ParquetMetaData) -> Vec { @@ -388,27 +755,32 @@ fn row_group_sizes(metadata: &ParquetMetaData) -> Vec { #[test] fn test_uniform_encryption_roundtrip() { - uniform_encryption_roundtrip(false, false).unwrap(); + uniform_encryption_roundtrip(false, false, b"0123456789012345").unwrap(); + uniform_encryption_roundtrip(false, false, b"01234567890123456789012345678901").unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_dictionary() { - uniform_encryption_roundtrip(false, true).unwrap(); + uniform_encryption_roundtrip(false, true, b"0123456789012345").unwrap(); + uniform_encryption_roundtrip(false, true, b"01234567890123456789012345678901").unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_page_index() { - uniform_encryption_roundtrip(true, false).unwrap(); + uniform_encryption_roundtrip(true, false, b"0123456789012345").unwrap(); + uniform_encryption_roundtrip(true, false, b"01234567890123456789012345678901").unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_page_index_and_dictionary() { - uniform_encryption_roundtrip(true, true).unwrap(); + uniform_encryption_roundtrip(true, true, b"0123456789012345").unwrap(); + uniform_encryption_roundtrip(true, true, b"01234567890123456789012345678901").unwrap(); } fn uniform_encryption_roundtrip( page_index: bool, dictionary_encoding: bool, + footer_key: &[u8], ) -> parquet::errors::Result<()> { let x0_arrays = [ Int32Array::from((0..100).collect::>()), @@ -426,7 +798,6 @@ fn uniform_encryption_roundtrip( let file = tempfile::tempfile()?; - let footer_key = b"0123456789012345"; let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()).build()?; @@ -506,15 +877,20 @@ fn uniform_encryption_roundtrip( #[test] fn test_uniform_encryption_page_skipping() { - uniform_encryption_page_skipping(false).unwrap(); + uniform_encryption_page_skipping(false, b"0123456789012345").unwrap(); + uniform_encryption_page_skipping(false, b"01234567890123456789012345678901").unwrap(); } #[test] fn test_uniform_encryption_page_skipping_with_page_index() { - uniform_encryption_page_skipping(true).unwrap(); + uniform_encryption_page_skipping(true, b"0123456789012345").unwrap(); + uniform_encryption_page_skipping(true, b"01234567890123456789012345678901").unwrap(); } -fn uniform_encryption_page_skipping(page_index: bool) -> parquet::errors::Result<()> { +fn uniform_encryption_page_skipping( + page_index: bool, + footer_key: &[u8], +) -> parquet::errors::Result<()> { let x0_arrays = [ Int32Array::from((0..100).collect::>()), Int32Array::from((100..150).collect::>()), @@ -531,7 +907,6 @@ fn uniform_encryption_page_skipping(page_index: bool) -> parquet::errors::Result let file = tempfile::tempfile()?; - let footer_key = b"0123456789012345"; let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()).build()?; @@ -618,29 +993,103 @@ fn uniform_encryption_page_skipping(page_index: bool) -> parquet::errors::Result #[test] fn test_write_non_uniform_encryption() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn write_non_uniform_encryption( + footer_key: &[u8], + column_names: Vec<&str>, + column_keys: Vec>, + encryption_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let column_names = vec!["double_field", "float_field"]; - let column_keys = vec![b"1234567890123450".to_vec(), b"1234567890123451".to_vec()]; + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_column_keys(column_names.to_vec(), column_keys.clone()) + .unwrap() + .build() + .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_keys(column_names.clone(), column_keys.clone()) - .unwrap() - .build() - .unwrap(); + let mut builder = FileEncryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in encryption_column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let file_encryption_properties = builder.build().unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .with_column_keys(column_names, column_keys) - .unwrap() - .build() - .unwrap(); + read_and_roundtrip_to_encrypted_file( + &file, + decryption_properties, + file_encryption_properties, + ); + } + + write_non_uniform_encryption( + b"0123456789012345", // 128bit/16, + vec!["double_field", "float_field"], + vec![b"1234567890123450".to_vec(), b"1234567890123451".to_vec()], + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ); - read_and_roundtrip_to_encrypted_file(&file, decryption_properties, file_encryption_properties); + write_non_uniform_encryption( + b"01234567890123456789012345678901", + vec![ + "double_field", + "float_field", + "boolean_field", + "int32_field", + "ba_field", + "flba_field", + "int64_field", + "int64_field.list.int64_field", + "int96_field", + ], + vec![ + b"12345678901234567890123456789012".to_vec(), + b"12345678901234567890123456789013".to_vec(), + b"12345678901234567890123456789014".to_vec(), + b"12345678901234567890123456789015".to_vec(), + b"12345678901234567890123456789016".to_vec(), + b"12345678901234567890123456789017".to_vec(), + b"12345678901234567890123456789018".to_vec(), + b"12345678901234567890123456789018".to_vec(), + b"12345678901234567890123456789019".to_vec(), + ], + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field.list.int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + ); } +// TODO #[test] fn test_write_uniform_encryption_plaintext_footer() { let testdata = arrow::util::test_util::parquet_test_data(); @@ -721,114 +1170,138 @@ fn test_write_uniform_encryption_plaintext_footer() { #[test] pub fn test_column_statistics_with_plaintext_footer() { - let footer_key = b"0123456789012345".to_vec(); - let column_key = b"1234567890123450".to_vec(); + let keys = [ + (b"0123456789012345".to_vec(), b"1234567890123450".to_vec()), + ( + b"01234567890123456789012345678901".to_vec(), + b"12345678901234567890123456789012".to_vec(), + ), + ]; - // Encrypt with a plaintext footer and column-specific keys - let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) - .with_plaintext_footer(true) - .with_column_key("x", column_key.clone()) - .with_column_key("y", column_key.clone()) - .with_column_key("s", column_key.clone()) - .build() - .unwrap(); + for (footer_key, column_key) in &keys { + // Encrypt with a plaintext footer and column-specific keys + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(true) + .with_column_key("x", column_key.clone()) + .with_column_key("y", column_key.clone()) + .with_column_key("s", column_key.clone()) + .build() + .unwrap(); - // Read with only the footer key and the key for one column - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("x", column_key.clone()) - .build() - .unwrap(); + // Read with only the footer key and the key for one column + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_key("x", column_key.clone()) + .build() + .unwrap(); - // Reader can read plaintext stats from the unencrypted column z - // and column x for which the key is provided, but not columns y and s - // for which no key is provided. - write_and_read_stats( - Arc::clone(&encryption_properties), - Some(decryption_properties), - &[true, false, true, false], - ); + // Reader can read plaintext stats from the unencrypted column z + // and column x for which the key is provided, but not columns y and s + // for which no key is provided. + write_and_read_stats( + Arc::clone(&encryption_properties), + Some(decryption_properties), + &[true, false, true, false], + ); - // Read without any decryption properties. - // Reader can only read plaintext stats from the unencrypted column z. - write_and_read_stats(encryption_properties, None, &[false, false, true, false]); + // Read without any decryption properties. + // Reader can only read plaintext stats from the unencrypted column z. + write_and_read_stats(encryption_properties, None, &[false, false, true, false]); + } } #[test] pub fn test_column_statistics_with_plaintext_footer_and_uniform_encryption() { - let footer_key = b"0123456789012345".to_vec(); + let footer_keys = [ + b"0123456789012345".to_vec(), + b"01234567890123456789012345678901".to_vec(), + ]; - // Write with uniform encryption and a plaintext footer. - let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) - .with_plaintext_footer(true) - .build() - .unwrap(); + for footer_key in &footer_keys { + // Write with uniform encryption and a plaintext footer. + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(true) + .build() + .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .build() - .unwrap(); + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .build() + .unwrap(); - // Reader can read stats from plaintext footer metadata if a footer key is provided - write_and_read_stats( - Arc::clone(&encryption_properties), - Some(decryption_properties), - &[true, true, true, true], - ); + // Reader can read stats from plaintext footer metadata if a footer key is provided + write_and_read_stats( + Arc::clone(&encryption_properties), + Some(decryption_properties), + &[true, true, true, true], + ); - // Reader can not read stats from plaintext footer metadata if no key is provided - write_and_read_stats(encryption_properties, None, &[false, false, false, false]); + // Reader can not read stats from plaintext footer metadata if no key is provided + write_and_read_stats(encryption_properties, None, &[false, false, false, false]); + } } #[test] pub fn test_column_statistics_with_encrypted_footer() { - let footer_key = b"0123456789012345".to_vec(); - let column_key = b"1234567890123450".to_vec(); + let keys = [ + (b"0123456789012345".to_vec(), b"1234567890123450".to_vec()), + ( + b"01234567890123456789012345678901".to_vec(), + b"12345678901234567890123456789012".to_vec(), + ), + ]; - // Encrypt with an encrypted footer and column-specific keys - let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) - .with_plaintext_footer(false) - .with_column_key("x", column_key.clone()) - .with_column_key("y", column_key.clone()) - .with_column_key("s", column_key.clone()) - .build() - .unwrap(); + for (footer_key, column_key) in &keys { + // Encrypt with an encrypted footer and column-specific keys + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(false) + .with_column_key("x", column_key.clone()) + .with_column_key("y", column_key.clone()) + .with_column_key("s", column_key.clone()) + .build() + .unwrap(); - // Read with only the footer key and the key for one column - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("x", column_key.clone()) - .build() - .unwrap(); + // Read with only the footer key and the key for one column + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_key("x", column_key.clone()) + .build() + .unwrap(); - // Reader can read plaintext stats from the unencrypted column z - // and column x for which the key is provided, but not columns y and s - // for which no key is provided. - write_and_read_stats( - encryption_properties, - Some(decryption_properties), - &[true, false, true, false], - ); + // Reader can read plaintext stats from the unencrypted column z + // and column x for which the key is provided, but not columns y and s + // for which no key is provided. + write_and_read_stats( + encryption_properties, + Some(decryption_properties), + &[true, false, true, false], + ); + } } #[test] pub fn test_column_statistics_with_encrypted_footer_and_uniform_encryption() { - let footer_key = b"0123456789012345".to_vec(); + let footer_keys = [ + b"0123456789012345".to_vec(), + b"01234567890123456789012345678901".to_vec(), + ]; - // Encrypt with an encrypted footer and uniform encryption - let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) - .with_plaintext_footer(false) - .build() - .unwrap(); + for footer_key in &footer_keys { + // Encrypt with an encrypted footer and uniform encryption + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(false) + .build() + .unwrap(); - // Read with the footer key - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .build() - .unwrap(); + // Read with the footer key + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .build() + .unwrap(); - // Reader can read stats for all columns. - write_and_read_stats( - encryption_properties, - Some(decryption_properties), - &[true, true, true, true], - ); + // Reader can read stats for all columns. + write_and_read_stats( + encryption_properties, + Some(decryption_properties), + &[true, true, true, true], + ); + } } /// Write a file with encryption and then verify whether statistics are readable with the provided decryption properties. @@ -934,64 +1407,156 @@ fn write_and_read_stats( #[test] fn test_write_uniform_encryption() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/uniform_encryption.parquet.encrypted"); - let file = File::open(path).unwrap(); + fn write_non_uniform_encryption(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let file = File::open(path).unwrap(); - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .build() - .unwrap(); + let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .build() - .unwrap(); + read_and_roundtrip_to_encrypted_file( + &file, + decryption_properties, + file_encryption_properties, + ); + } - read_and_roundtrip_to_encrypted_file(&file, decryption_properties, file_encryption_properties); + write_non_uniform_encryption(b"0123456789012345"); + write_non_uniform_encryption(b"01234567890123456789012345678901"); } #[test] fn test_write_non_uniform_encryption_column_missmatch() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); + fn write_non_uniform_encryption_column_missmatch( + footer_key: &[u8], + column_keys: &[(&str, &[u8])], + encryption_column_keys: &[(&str, &[u8])], + ) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let column_1_key = b"1234567890123450".to_vec(); - let column_2_key = b"1234567890123451".to_vec(); + let mut decryption_builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + decryption_builder = decryption_builder.with_column_key(column_name, key.to_vec()); + } - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) - .build() - .unwrap(); + let decryption_properties = decryption_builder.build().unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("other_field", column_1_key) - .with_column_key("yet_another_field", column_2_key) - .build() - .unwrap(); + let mut file_encryption_builder = FileEncryptionProperties::builder(footer_key.to_vec()) + .with_column_key("other_field", encryption_column_keys[0].1.to_vec()) + .with_column_key("yet_another_field", encryption_column_keys[1].1.to_vec()); - let temp_file = tempfile::tempfile().unwrap(); + for (column_name, key) in encryption_column_keys { + file_encryption_builder = + file_encryption_builder.with_column_key(column_name, key.to_vec()); + } - // read example data - let file = File::open(path).unwrap(); - let options = ArrowReaderOptions::default() - .with_file_decryption_properties(decryption_properties.clone()); - let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); - let props = WriterProperties::builder() - .with_file_encryption_properties(file_encryption_properties) - .build(); + let file_encryption_properties = file_encryption_builder.build().unwrap(); - let result = ArrowWriter::try_new( - temp_file.try_clone().unwrap(), - metadata.schema().clone(), - Some(props), + let temp_file = tempfile::tempfile().unwrap(); + + // read example data + let file = File::open(path).unwrap(); + let options = ArrowReaderOptions::default() + .with_file_decryption_properties(decryption_properties.clone()); + let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); + let props = WriterProperties::builder() + .with_file_encryption_properties(file_encryption_properties) + .build(); + + let result = ArrowWriter::try_new( + temp_file.try_clone().unwrap(), + metadata.schema().clone(), + Some(props), + ); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: The following columns with encryption keys specified were not found in the schema: other_field, yet_another_field" + ); + } + + // AES-128 + write_non_uniform_encryption_column_missmatch( + b"0123456789012345", + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], ); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: The following columns with encryption keys specified were not found in the schema: other_field, yet_another_field" + + // AES-256 + write_non_uniform_encryption_column_missmatch( + b"01234567890123456789012345678901", + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field.list.int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], + &[ + ( + "double_field", + b"12345678901234567890123456789012".as_slice(), + ), + ( + "float_field", + b"12345678901234567890123456789013".as_slice(), + ), + ( + "boolean_field", + b"12345678901234567890123456789014".as_slice(), + ), + ( + "int32_field", + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field.list.int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( + "int96_field", + b"12345678901234567890123456789019".as_slice(), + ), + ], ); } @@ -1218,855 +1783,76 @@ pub fn test_retrieve_row_group_statistics_after_encrypted_write() { #[test] fn test_decrypt_page_index_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); + fn decrypt_page_index_uniform(footer_key: &[u8]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "uniform_encryption.parquet.encrypted", + ); + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); - let key_code: &[u8] = "0123456789012345".as_bytes(); - let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) - .build() - .unwrap(); + test_decrypt_page_index(&path, decryption_properties).unwrap(); + } - test_decrypt_page_index(&path, decryption_properties).unwrap(); + decrypt_page_index_uniform("0123456789012345".as_bytes()); + decrypt_page_index_uniform("01234567890123456789012345678901".as_bytes()); } #[test] fn test_decrypt_page_index_non_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/encrypt_columns_and_footer.parquet.encrypted"); - - let footer_key = "0123456789012345".as_bytes().to_vec(); - let column_1_key = "1234567890123450".as_bytes().to_vec(); - let column_2_key = "1234567890123451".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .build() - .unwrap(); - - test_decrypt_page_index(&path, decryption_properties).unwrap(); -} - -#[test] -fn test_aes256_non_uniform_encryption_plaintext_footer() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); - - let footer_key = "01234567890123456789012345678901".as_bytes(); // 256bit/32 - let column_1_key = "12345678901234567890123456789012".as_bytes(); - let column_2_key = "12345678901234567890123456789013".as_bytes(); - let column_3_key = "12345678901234567890123456789014".as_bytes(); - let column_4_key = "12345678901234567890123456789015".as_bytes(); - let column_5_key = "12345678901234567890123456789016".as_bytes(); - let column_6_key = "12345678901234567890123456789017".as_bytes(); - let column_7_key = "12345678901234567890123456789018".as_bytes(); - let column_8_key = "12345678901234567890123456789019".as_bytes(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key.to_vec()) - .with_column_key("float_field", column_2_key.to_vec()) - .with_column_key("boolean_field", column_3_key.to_vec()) - .with_column_key("int32_field", column_4_key.to_vec()) - .with_column_key("ba_field", column_5_key.to_vec()) - .with_column_key("flba_field", column_6_key.to_vec()) - .with_column_key("int64_field", column_7_key.to_vec()) - .with_column_key("int96_field", column_8_key.to_vec()) - .build() - .unwrap(); + fn decrypt_page_index_non_uniform(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + let path = encryption_util::encrypted_data_path( + footer_key, + "encrypt_columns_and_footer.parquet.encrypted", + ); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); + for (column_name, key) in column_keys { + builder = builder.with_column_key(column_name, key.to_vec()); + } + let decryption_properties = builder.build().unwrap(); + test_decrypt_page_index(&path, decryption_properties).unwrap(); + } - verify_encryption_test_file_read(file, decryption_properties); -} + decrypt_page_index_non_uniform( + b"0123456789012345", // 128bit/16 + &[ + ("double_field", b"1234567890123450".as_slice()), + ("float_field", b"1234567890123451".as_slice()), + ], + ); -#[test] -fn test_aes256_plaintext_footer_signature_verification() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path.clone()).unwrap(); - - let wrong_footer_key = "00000000000000000000000000000000".as_bytes(); // 256bit/32 - let column_1_key = "12345678901234567890123456789012".as_bytes(); - let column_2_key = "12345678901234567890123456789013".as_bytes(); - let column_3_key = "12345678901234567890123456789014".as_bytes(); - let column_4_key = "12345678901234567890123456789015".as_bytes(); - let column_5_key = "12345678901234567890123456789016".as_bytes(); - let column_6_key = "12345678901234567890123456789017".as_bytes(); - let column_7_key = "12345678901234567890123456789018".as_bytes(); - let column_8_key = "12345678901234567890123456789019".as_bytes(); - - let decryption_properties = FileDecryptionProperties::builder(wrong_footer_key.to_vec()) - .disable_footer_signature_verification() - .with_column_key("double_field", column_1_key.to_vec()) - .with_column_key("float_field", column_2_key.to_vec()) - .with_column_key("boolean_field", column_3_key.to_vec()) - .with_column_key("int32_field", column_4_key.to_vec()) - .with_column_key("ba_field", column_5_key.to_vec()) - .with_column_key("flba_field", column_6_key.to_vec()) - .with_column_key("int64_field", column_7_key.to_vec()) - .with_column_key("int96_field", column_8_key.to_vec()) - .build() - .unwrap(); - - verify_encryption_test_file_read(file, decryption_properties); - - let file = File::open(path.clone()).unwrap(); - - let decryption_properties = FileDecryptionProperties::builder(wrong_footer_key.to_vec()) - .with_column_key("double_field", column_1_key.to_vec()) - .with_column_key("float_field", column_2_key.to_vec()) - .with_column_key("boolean_field", column_3_key.to_vec()) - .with_column_key("int32_field", column_4_key.to_vec()) - .with_column_key("ba_field", column_5_key.to_vec()) - .with_column_key("flba_field", column_6_key.to_vec()) - .with_column_key("int64_field", column_7_key.to_vec()) - .with_column_key("int96_field", column_8_key.to_vec()) - .build() - .unwrap(); - - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let result = ArrowReaderMetadata::load(&file, options.clone()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .starts_with("Parquet error: Footer signature verification failed. Computed: [") - ); -} - -#[test] -fn test_aes256_non_uniform_encryption_disabled_aad_storage() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!( - "{test_data}/aes256/encrypt_columns_and_footer_disable_aad_storage.parquet.encrypted" - ); - let file = File::open(path.clone()).unwrap(); - - let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 256bit/32 - let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); - let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); - let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); - let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); - let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); - let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); - let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); - let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); - - // Can read successfully when providing the correct AAD prefix - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) - .with_column_key("boolean_field", column_3_key.clone()) - .with_column_key("int32_field", column_4_key.clone()) - .with_column_key("ba_field", column_5_key.clone()) - .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_field", column_7_key.clone()) - .with_column_key("int96_field", column_8_key.clone()) - .with_aad_prefix(b"tester".to_vec()) - .build() - .unwrap(); - - verify_encryption_test_file_read(file, decryption_properties); - - // Using wrong AAD prefix should fail - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) - .with_column_key("boolean_field", column_3_key.clone()) - .with_column_key("int32_field", column_4_key.clone()) - .with_column_key("ba_field", column_5_key.clone()) - .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_field", column_7_key.clone()) - .with_column_key("int96_field", column_8_key.clone()) - .with_aad_prefix(b"wrong_aad_prefix".to_vec()) - .build() - .unwrap(); - - let file = File::open(path.clone()).unwrap(); - let options = ArrowReaderOptions::default() - .with_file_decryption_properties(decryption_properties.clone()); - let result = ArrowReaderMetadata::load(&file, options.clone()); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer" - ); - - // Not providing any AAD prefix should fail as it isn't stored in the file - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) - .with_column_key("boolean_field", column_3_key.clone()) - .with_column_key("int32_field", column_4_key.clone()) - .with_column_key("ba_field", column_5_key.clone()) - .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_field", column_7_key.clone()) - .with_column_key("int96_field", column_8_key.clone()) - .build() - .unwrap(); - - let file = File::open(path).unwrap(); - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let result = ArrowReaderMetadata::load(&file, options.clone()); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file was encrypted with an AAD prefix that is not stored in the file, \ - but no AAD prefix was provided in the file decryption properties" - ); -} -#[test] -fn test_aes256_non_uniform_encryption() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); - - // There is always a footer key even with a plaintext footer, - // but this is used for signing the footer. - let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); // 256bit/32 - let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); - let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); - let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); - let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); - let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); - let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); - let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); - let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .with_column_key("boolean_field", column_3_key) - .with_column_key("int32_field", column_4_key) - .with_column_key("ba_field", column_5_key) - .with_column_key("flba_field", column_6_key) - .with_column_key("int64_field", column_7_key) - .with_column_key("int96_field", column_8_key) - .build() - .unwrap(); - - verify_encryption_test_file_read(file, decryption_properties); -} - -#[test] -fn test_aes256_uniform_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"); - let file = File::open(path).unwrap(); - - let key_code: &[u8] = "01234567890123456789012345678901".as_bytes(); - let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) - .build() - .unwrap(); - - verify_encryption_test_file_read(file, decryption_properties); -} - -#[test] -fn test_aes256_ctr_encryption() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/encrypt_columns_and_footer_ctr.parquet.encrypted"); - let file = File::open(path).unwrap(); - - let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); - let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_1_key) - .build() - .unwrap(); - - let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption_properties); - let metadata = ArrowReaderMetadata::load(&file, options); - - match metadata { - Err(ParquetError::NYI(s)) => { - assert!(s.contains("AES_GCM_CTR_V1")); - } - _ => { - panic!("Expected ParquetError::NYI"); - } - }; -} - -#[test] -fn test_aes256_non_uniform_encryption_plaintext_footer_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); - - let key_retriever = TestKeyRetriever::new() - .with_key( - "kf".to_owned(), - "01234567890123456789012345678901".as_bytes().to_vec(), - ) - .with_key( - "kc1".to_owned(), - "12345678901234567890123456789012".as_bytes().to_vec(), - ) - .with_key( - "kc2".to_owned(), - "12345678901234567890123456789013".as_bytes().to_vec(), - ) - .with_key( - "kc3".to_owned(), - "12345678901234567890123456789014".as_bytes().to_vec(), - ) - .with_key( - "kc4".to_owned(), - "12345678901234567890123456789015".as_bytes().to_vec(), - ) - .with_key( - "kc5".to_owned(), - "12345678901234567890123456789016".as_bytes().to_vec(), - ) - .with_key( - "kc6".to_owned(), - "12345678901234567890123456789017".as_bytes().to_vec(), - ) - .with_key( - "kc7".to_owned(), - "12345678901234567890123456789018".as_bytes().to_vec(), - ) - .with_key( - "kc8".to_owned(), - "12345678901234567890123456789019".as_bytes().to_vec(), - ); - - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() - .unwrap(); - - verify_encryption_test_file_read(file, decryption_properties); -} - -#[test] -fn test_aes256_uniform_encryption_plaintext_footer_with_key_retriever() { - let test_data = arrow::util::test_util::parquet_test_data(); - - // Read example data with key retriever - let path = format!("{test_data}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); - - let key_retriever = Arc::new( - TestKeyRetriever::new() - .with_key( - "kf".to_owned(), - "01234567890123456789012345678901".as_bytes().to_vec(), - ) - .with_key( - "kc1".to_owned(), - "12345678901234567890123456789012".as_bytes().to_vec(), - ) - .with_key( - "kc2".to_owned(), - "12345678901234567890123456789013".as_bytes().to_vec(), - ) - .with_key( - "kc3".to_owned(), - "12345678901234567890123456789014".as_bytes().to_vec(), - ) - .with_key( - "kc4".to_owned(), - "12345678901234567890123456789015".as_bytes().to_vec(), - ) - .with_key( - "kc5".to_owned(), - "12345678901234567890123456789016".as_bytes().to_vec(), - ) - .with_key( - "kc6".to_owned(), - "12345678901234567890123456789017".as_bytes().to_vec(), - ) - .with_key( - "kc7".to_owned(), - "12345678901234567890123456789018".as_bytes().to_vec(), - ) - .with_key( - "kc8".to_owned(), - "12345678901234567890123456789019".as_bytes().to_vec(), - ), - ); - - let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever.clone()) - .build() - .unwrap(); - - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); - - // Write data into temporary file with plaintext footer and footer key metadata - let temp_file = tempfile::tempfile().unwrap(); - let encryption_properties = - FileEncryptionProperties::builder(b"01234567890123456789012345678901".to_vec()) - .with_footer_key_metadata("kf".into()) - .with_column_key_and_metadata( + decrypt_page_index_non_uniform( + b"01234567890123456789012345678901", // 256bit/32 + &[ + ( "double_field", - b"12345678901234567890123456789012".to_vec(), - b"kc1".into(), - ) - .with_column_key_and_metadata( + b"12345678901234567890123456789012".as_slice(), + ), + ( "float_field", - b"12345678901234567890123456789013".to_vec(), - b"kc2".into(), - ) - .with_column_key_and_metadata( + b"12345678901234567890123456789013".as_slice(), + ), + ( "boolean_field", - b"12345678901234567890123456789014".to_vec(), - b"kc3".into(), - ) - .with_column_key_and_metadata( + b"12345678901234567890123456789014".as_slice(), + ), + ( "int32_field", - b"12345678901234567890123456789015".to_vec(), - b"kc4".into(), - ) - .with_column_key_and_metadata( - "ba_field", - b"12345678901234567890123456789016".to_vec(), - b"kc5".into(), - ) - .with_column_key_and_metadata( - "flba_field", - b"12345678901234567890123456789017".to_vec(), - b"kc6".into(), - ) - // .with_column_key_and_metadata("int64_field", b"12345678901234567890123456789018".to_vec(), b"kc7".into()) - .with_column_key_and_metadata( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".to_vec(), - b"kc7".into(), - ) - .with_column_key_and_metadata( + b"12345678901234567890123456789015".as_slice(), + ), + ("ba_field", b"12345678901234567890123456789016".as_slice()), + ("flba_field", b"12345678901234567890123456789017".as_slice()), + ( + "int64_field", + b"12345678901234567890123456789018".as_slice(), + ), + ( "int96_field", - b"12345678901234567890123456789019".to_vec(), - b"kc8".into(), - ) - .with_plaintext_footer(true) - .build() - .unwrap(); - - let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); - let batch_reader = builder.build().unwrap(); - let batches = batch_reader - .collect::, _>>() - .unwrap(); - - let props = WriterProperties::builder() - .with_file_encryption_properties(encryption_properties) - .build(); - - let mut writer = ArrowWriter::try_new( - temp_file.try_clone().unwrap(), - metadata.schema().clone(), - Some(props), - ) - .unwrap(); - for batch in batches { - writer.write(&batch).unwrap(); - } - - writer.close().unwrap(); - - // Read temporary file with plaintext metadata using key retriever - let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever) - .build() - .unwrap(); - - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let _ = ArrowReaderMetadata::load(&temp_file, options.clone()).unwrap(); - - // Read temporary file with plaintext metadata using key retriever with invalid key - let key_retriever = Arc::new( - TestKeyRetriever::new() - .with_key( - "kf".to_owned(), - b"01334567890123456789012345678901".to_vec(), - ) // wrong footer key - .with_key( - "kc1".to_owned(), - b"12345678901234567890123456789012".to_vec(), - ) - .with_key( - "kc2".to_owned(), - b"12345678901234567890123456789013".to_vec(), - ) - .with_key( - "kc3".to_owned(), - b"12345678901234567890123456789014".to_vec(), - ) - .with_key( - "kc4".to_owned(), - b"12345678901234567890123456789015".to_vec(), - ) - .with_key( - "kc5".to_owned(), - b"12345678901234567890123456789016".to_vec(), - ) - .with_key( - "kc6".to_owned(), - b"12345678901234567890123456789017".to_vec(), - ) - .with_key( - "kc7".to_owned(), - b"12345678901234567890123456789018".to_vec(), - ) - .with_key( - "kc8".to_owned(), - b"12345678901234567890123456789019".to_vec(), + b"12345678901234567890123456789019".as_slice(), ), + ], ); - let decryption_properties = FileDecryptionProperties::with_key_retriever(key_retriever) - .build() - .unwrap(); - let options = - ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties); - let result = ArrowReaderMetadata::load(&temp_file, options.clone()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .starts_with("Parquet error: Footer signature verification failed. Computed: [") - ); -} - -#[test] -fn test_aes256_non_uniform_encryption_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/encrypt_columns_and_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); - - let key_retriever = TestKeyRetriever::new() - .with_key( - "kf".to_owned(), - "01234567890123456789012345678901".as_bytes().to_vec(), - ) - .with_key( - "kc1".to_owned(), - "12345678901234567890123456789012".as_bytes().to_vec(), - ) - .with_key( - "kc2".to_owned(), - "12345678901234567890123456789013".as_bytes().to_vec(), - ) - .with_key( - "kc3".to_owned(), - "12345678901234567890123456789014".as_bytes().to_vec(), - ) - .with_key( - "kc4".to_owned(), - "12345678901234567890123456789015".as_bytes().to_vec(), - ) - .with_key( - "kc5".to_owned(), - "12345678901234567890123456789016".as_bytes().to_vec(), - ) - .with_key( - "kc6".to_owned(), - "12345678901234567890123456789017".as_bytes().to_vec(), - ) - .with_key( - "kc7".to_owned(), - "12345678901234567890123456789018".as_bytes().to_vec(), - ) - .with_key( - "kc8".to_owned(), - "12345678901234567890123456789019".as_bytes().to_vec(), - ); - - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) - .build() - .unwrap(); - - verify_encryption_test_file_read(file, decryption_properties); -} - -#[test] -fn test_aes256_uniform_encryption_with_key_retriever() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/uniform_encryption.parquet.encrypted"); - let file = File::open(path).unwrap(); - - let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 - - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .build() - .unwrap(); - - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .build() - .unwrap(); - - read_and_roundtrip_to_encrypted_file(&file, decryption_properties, file_encryption_properties); -} - -#[test] -fn test_aes256_write_non_uniform_encryption() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/encrypt_columns_and_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); - - let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 - let column_names = vec![ - "double_field", - "float_field", - "boolean_field", - "int32_field", - "ba_field", - "flba_field", - "int64_field", - "int64_field.list.int64_field", - "int96_field", - ]; - let column_keys = vec![ - b"12345678901234567890123456789012".to_vec(), - b"12345678901234567890123456789013".to_vec(), - b"12345678901234567890123456789014".to_vec(), - b"12345678901234567890123456789015".to_vec(), - b"12345678901234567890123456789016".to_vec(), - b"12345678901234567890123456789017".to_vec(), - b"12345678901234567890123456789018".to_vec(), - b"12345678901234567890123456789018".to_vec(), - b"12345678901234567890123456789019".to_vec(), - ]; - - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_keys(column_names.clone(), column_keys.clone()) - .unwrap() - .build() - .unwrap(); - - let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); - let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); - let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); - let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); - let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); - let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); - let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); - let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); - - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .with_column_key("boolean_field", column_3_key) - .with_column_key("int32_field", column_4_key) - .with_column_key("ba_field", column_5_key) - .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_field.list.int64_field", column_7_key) - .with_column_key("int96_field", column_8_key) - .build() - .unwrap(); - - read_and_roundtrip_to_encrypted_file(&file, decryption_properties, file_encryption_properties); -} - -#[test] -fn test_aes256_write_uniform_encryption_plaintext_footer() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/encrypt_columns_plaintext_footer.parquet.encrypted"); - let file = File::open(path).unwrap(); - - let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 - let wrong_footer_key = b"00000000000000000000000000000000".to_vec(); // 256bit/32 - let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); - let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); - let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); - let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); - let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); - let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); - let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); - let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) - .with_column_key("boolean_field", column_3_key.clone()) - .with_column_key("int32_field", column_4_key.clone()) - .with_column_key("ba_field", column_5_key.clone()) - .with_column_key("flba_field", column_6_key.clone()) - .with_column_key("int64_field", column_7_key.clone()) - .with_column_key("int96_field", column_8_key.clone()) - .build() - .unwrap(); - - let wrong_decryption_properties = FileDecryptionProperties::builder(wrong_footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .with_column_key("boolean_field", column_3_key) - .with_column_key("int32_field", column_4_key) - .with_column_key("ba_field", column_5_key) - .with_column_key("flba_field", column_6_key) - .with_column_key("int64_field", column_7_key) - .with_column_key("int96_field", column_8_key) - .build() - .unwrap(); - - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .with_plaintext_footer(true) - .build() - .unwrap(); - - // Try writing plaintext footer and then reading it with the correct footer key - read_and_roundtrip_to_encrypted_file( - &file, - Arc::clone(&decryption_properties), - file_encryption_properties.clone(), - ); - - // Try writing plaintext footer and then reading it with the wrong footer key - let temp_file = tempfile::tempfile().unwrap(); - - // read example data - let options = ArrowReaderOptions::default() - .with_file_decryption_properties(decryption_properties.clone()); - let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); - - let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); - let batch_reader = builder.build().unwrap(); - let batches = batch_reader - .collect::, _>>() - .unwrap(); - - // write example data - let props = WriterProperties::builder() - .with_file_encryption_properties(file_encryption_properties) - .build(); - - let mut writer = ArrowWriter::try_new( - temp_file.try_clone().unwrap(), - metadata.schema().clone(), - Some(props), - ) - .unwrap(); - for batch in batches { - writer.write(&batch).unwrap(); - } - writer.close().unwrap(); - - // Try reading plaintext footer and with the wrong footer key - let options = - ArrowReaderOptions::default().with_file_decryption_properties(wrong_decryption_properties); - let result = ArrowReaderMetadata::load(&temp_file, options.clone()); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .starts_with("Parquet error: Footer signature verification failed. Computed: [") - ); -} - -#[test] -fn test_aes256_write_non_uniform_encryption_column_missmatch() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/encrypt_columns_and_footer.parquet.encrypted"); - - let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 - let column_1_key = b"12345678901234567890123456789012".to_vec(); - let column_2_key = b"12345678901234567890123456789013".to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) - .build() - .unwrap(); - - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("other_field", column_1_key) - .with_column_key("yet_another_field", column_2_key) - .build() - .unwrap(); - - let temp_file = tempfile::tempfile().unwrap(); - - // read example data - let file = File::open(path).unwrap(); - let options = ArrowReaderOptions::default() - .with_file_decryption_properties(decryption_properties.clone()); - let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap(); - let props = WriterProperties::builder() - .with_file_encryption_properties(file_encryption_properties) - .build(); - - let result = ArrowWriter::try_new( - temp_file.try_clone().unwrap(), - metadata.schema().clone(), - Some(props), - ); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: The following columns with encryption keys specified were not found in the schema: other_field, yet_another_field" - ); -} - -#[test] -fn test_aes256_write_uniform_encryption() { - let testdata = arrow::util::test_util::parquet_test_data(); - let path = format!("{testdata}/aes256/uniform_encryption.parquet.encrypted"); - let file = File::open(path).unwrap(); - - let footer_key = b"01234567890123456789012345678901".to_vec(); // 256bit/32 - - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .build() - .unwrap(); - - let file_encryption_properties = FileEncryptionProperties::builder(footer_key) - .build() - .unwrap(); - - read_and_roundtrip_to_encrypted_file(&file, decryption_properties, file_encryption_properties); -} - -#[test] -fn test_aes256_decrypt_page_index_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"); - - let key_code: &[u8] = "01234567890123456789012345678901".as_bytes(); - let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec()) - .build() - .unwrap(); - - test_decrypt_page_index(&path, decryption_properties).unwrap(); -} - -#[test] -fn test_aes256_decrypt_page_index_non_uniform() { - let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/aes256/encrypt_columns_and_footer.parquet.encrypted"); - - let footer_key = "01234567890123456789012345678901".as_bytes().to_vec(); - let column_1_key = "12345678901234567890123456789012".as_bytes().to_vec(); - let column_2_key = "12345678901234567890123456789013".as_bytes().to_vec(); - let column_3_key = "12345678901234567890123456789014".as_bytes().to_vec(); - let column_4_key = "12345678901234567890123456789015".as_bytes().to_vec(); - let column_5_key = "12345678901234567890123456789016".as_bytes().to_vec(); - let column_6_key = "12345678901234567890123456789017".as_bytes().to_vec(); - let column_7_key = "12345678901234567890123456789018".as_bytes().to_vec(); - let column_8_key = "12345678901234567890123456789019".as_bytes().to_vec(); - - let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) - .with_column_key("boolean_field", column_3_key) - .with_column_key("int32_field", column_4_key) - .with_column_key("ba_field", column_5_key) - .with_column_key("flba_field", column_6_key) - .with_column_key("int64_field", column_7_key) - .with_column_key("int96_field", column_8_key) - .build() - .unwrap(); - - test_decrypt_page_index(&path, decryption_properties).unwrap(); } fn test_decrypt_page_index( From 0abe16d8c17f9ef6245eac3da53376b87f0a13ec Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Thu, 5 Mar 2026 17:13:30 -0800 Subject: [PATCH 14/26] Comment & clean up --- parquet/tests/encryption/encryption.rs | 34 ++++++++++++-------- parquet/tests/encryption/encryption_async.rs | 7 ++-- 2 files changed, 26 insertions(+), 15 deletions(-) diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index 4eed9c6a2a56..420ea94b98c0 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -386,16 +386,22 @@ fn test_uniform_encryption() { #[test] fn test_decrypting_without_decryption_properties_fails() { let test_data = arrow::util::test_util::parquet_test_data(); - let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); - let file = File::open(path).unwrap(); + let paths = [ + format!("{test_data}/uniform_encryption.parquet.encrypted"), + format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"), + ]; - let options = ArrowReaderOptions::default(); - let result = ArrowReaderMetadata::load(&file, options.clone()); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" - ); + for path in &paths { + let file = File::open(path).unwrap(); + + let options = ArrowReaderOptions::default(); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); + } } #[test] @@ -614,6 +620,9 @@ fn test_uniform_encryption_plaintext_footer_with_key_retriever() { ); // AES-256 + // The asymmetric column names is because we check column paths in the validate_encrypted_column_names function of [encryption::encrypt] + // The column path of the repeated field (int64_field) is int64_field.list.int64_field + // Switching from `c.path().string()` to `c.name().to_string()` can fix the asymmetricity uniform_encryption_plaintext_footer_with_key_retriever( b"01234567890123456789012345678901", "kf", @@ -1089,7 +1098,6 @@ fn test_write_non_uniform_encryption() { ); } -// TODO #[test] fn test_write_uniform_encryption_plaintext_footer() { let testdata = arrow::util::test_util::parquet_test_data(); @@ -1407,7 +1415,7 @@ fn write_and_read_stats( #[test] fn test_write_uniform_encryption() { - fn write_non_uniform_encryption(footer_key: &[u8]) { + fn write_uniform_encryption(footer_key: &[u8]) { let path = encryption_util::encrypted_data_path( footer_key, "uniform_encryption.parquet.encrypted", @@ -1429,8 +1437,8 @@ fn test_write_uniform_encryption() { ); } - write_non_uniform_encryption(b"0123456789012345"); - write_non_uniform_encryption(b"01234567890123456789012345678901"); + write_uniform_encryption(b"0123456789012345"); + write_uniform_encryption(b"01234567890123456789012345678901"); } #[test] diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index 93bc44e9163d..f4076e3783af 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -33,7 +33,6 @@ use parquet::arrow::arrow_writer::{ use parquet::arrow::{ ArrowSchemaConverter, ArrowWriter, AsyncArrowWriter, ParquetRecordBatchStreamBuilder, }; -use parquet::data_type::AsBytes; use parquet::encryption::decrypt::FileDecryptionProperties; use parquet::encryption::encrypt::FileEncryptionProperties; use parquet::errors::ParquetError; @@ -130,7 +129,7 @@ async fn test_misspecified_encryption_keys() { ); let mut file = File::open(&path).await.unwrap(); - let mut builder = FileDecryptionProperties::builder(footer_key.as_bytes().to_vec()); + let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); if !column_1_key.is_empty() { builder = builder.with_column_key("double_field", column_1_key.to_vec()); @@ -545,6 +544,10 @@ async fn test_write_non_uniform_encryption() { ) .await; + // AES-256 + // The asymmetric column names is because we check column paths in the validate_encrypted_column_names function of [encryption::encrypt] + // The column path of the repeated field (int64_field) is int64_field.list.int64_field + // Switching from `c.path().string()` to `c.name().to_string()` can fix the asymmetricity write_non_uniform_encryption( b"01234567890123456789012345678901", vec![ From 206c6c63a2b51cbb0b73ae0f65d42802d2b2b3a9 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Wed, 22 Apr 2026 16:39:06 -0700 Subject: [PATCH 15/26] Update comment --- parquet/tests/encryption/encryption_async.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index f4076e3783af..e91fe3849276 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -545,9 +545,7 @@ async fn test_write_non_uniform_encryption() { .await; // AES-256 - // The asymmetric column names is because we check column paths in the validate_encrypted_column_names function of [encryption::encrypt] - // The column path of the repeated field (int64_field) is int64_field.list.int64_field - // Switching from `c.path().string()` to `c.name().to_string()` can fix the asymmetricity + // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name write_non_uniform_encryption( b"01234567890123456789012345678901", vec![ From efe04f615d6caa5231b43a8eafe731eaefb9e6ee Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Wed, 22 Apr 2026 16:51:53 -0700 Subject: [PATCH 16/26] Format --- parquet/src/encryption/ciphers.rs | 18 ++++++++++-------- parquet/src/encryption/decrypt.rs | 17 +++++++---------- parquet/src/encryption/encrypt.rs | 4 +--- 3 files changed, 18 insertions(+), 21 deletions(-) diff --git a/parquet/src/encryption/ciphers.rs b/parquet/src/encryption/ciphers.rs index 726b292d1c70..beb3acda268d 100644 --- a/parquet/src/encryption/ciphers.rs +++ b/parquet/src/encryption/ciphers.rs @@ -41,15 +41,16 @@ pub(crate) struct RingGcmBlockDecryptor { impl RingGcmBlockDecryptor { /// Create a new `RingGcmBlockDecryptor` with a given key. - pub(crate) fn new( - key_bytes: &[u8], - ) -> Result { + pub(crate) fn new(key_bytes: &[u8]) -> Result { let algorithm = if key_bytes.len() == AES_128_GCM.key_len() { &AES_128_GCM } else if key_bytes.len() == AES_256_GCM.key_len() { &AES_256_GCM } else { - return Err(general_err!("Error creating RingGcmBlockDecryptor with unsupported key length: {}", key_bytes.len())); + return Err(general_err!( + "Error creating RingGcmBlockDecryptor with unsupported key length: {}", + key_bytes.len() + )); }; let key = UnboundKey::new(algorithm, key_bytes) .map_err(|_| general_err!("Failed to create {:?} key", algorithm))?; @@ -151,16 +152,17 @@ impl RingGcmBlockEncryptor { /// Create a new `RingGcmBlockEncryptor` with a given key and random nonce. /// The nonce will advance appropriately with each block encryption and /// return an error if it wraps around. - pub(crate) fn new( - key_bytes: &[u8], - ) -> Result { + pub(crate) fn new(key_bytes: &[u8]) -> Result { let rng = SystemRandom::new(); let algorithm = if key_bytes.len() == AES_128_GCM.key_len() { &AES_128_GCM } else if key_bytes.len() == AES_256_GCM.key_len() { &AES_256_GCM } else { - return Err(general_err!("Error creating RingGcmBlockEncryptor with unsupported key length: {}", key_bytes.len())); + return Err(general_err!( + "Error creating RingGcmBlockEncryptor with unsupported key length: {}", + key_bytes.len() + )); }; let key = UnboundKey::new(algorithm, key_bytes) diff --git a/parquet/src/encryption/decrypt.rs b/parquet/src/encryption/decrypt.rs index 5ae8dd620680..0066523419de 100644 --- a/parquet/src/encryption/decrypt.rs +++ b/parquet/src/encryption/decrypt.rs @@ -596,13 +596,12 @@ impl FileDecryptor { ) -> Result { let file_aad = [aad_prefix.as_slice(), aad_file_unique.as_slice()].concat(); let footer_key = decryption_properties.footer_key(footer_key_metadata)?; - let footer_decryptor = RingGcmBlockDecryptor::new(&footer_key) - .map_err(|e| { - general_err!( - "Invalid footer key. {}", - e.to_string().replace("Parquet error: ", "") - ) - })?; + let footer_decryptor = RingGcmBlockDecryptor::new(&footer_key).map_err(|e| { + general_err!( + "Invalid footer key. {}", + e.to_string().replace("Parquet error: ", "") + ) + })?; Ok(Self { footer_decryptor: Arc::new(footer_decryptor), @@ -642,9 +641,7 @@ impl FileDecryptor { let column_key = self .decryption_properties .column_key(column_name, key_metadata)?; - Ok(Arc::new(RingGcmBlockDecryptor::new( - &column_key, - )?)) + Ok(Arc::new(RingGcmBlockDecryptor::new(&column_key)?)) } pub(crate) fn get_column_metadata_decryptor( diff --git a/parquet/src/encryption/encrypt.rs b/parquet/src/encryption/encrypt.rs index 12d59ce89712..d69e3c02500a 100644 --- a/parquet/src/encryption/encrypt.rs +++ b/parquet/src/encryption/encrypt.rs @@ -359,9 +359,7 @@ impl FileEncryptor { } match self.properties.column_keys.get(column_path) { None => Err(general_err!("Column '{}' is not encrypted", column_path)), - Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new( - column_key.key(), - )?)), + Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new(column_key.key())?)), } } } From d26f41b48bc582aeb04d38e43db2b4b737d6bb53 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Thu, 30 Apr 2026 13:19:33 -0700 Subject: [PATCH 17/26] Revert tests --- parquet/tests/encryption/encryption_async.rs | 717 ++++++------------- 1 file changed, 220 insertions(+), 497 deletions(-) diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index e91fe3849276..ba331d502042 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -481,22 +481,16 @@ async fn test_aes_ctr_encryption() { #[tokio::test] async fn test_decrypting_without_decryption_properties_fails() { let test_data = arrow::util::test_util::parquet_test_data(); - let paths = [ - format!("{test_data}/uniform_encryption.parquet.encrypted"), - format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"), - ]; - - for path in &paths { - let mut file = File::open(&path).await.unwrap(); + let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); + let mut file = File::open(&path).await.unwrap(); - let options = ArrowReaderOptions::new(); - let result = ArrowReaderMetadata::load_async(&mut file, options).await; - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" - ); - } + let options = ArrowReaderOptions::new(); + let result = ArrowReaderMetadata::load_async(&mut file, options).await; + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); } #[tokio::test] @@ -1102,534 +1096,263 @@ async fn concatenate_parallel_row_groups( #[tokio::test] async fn test_concurrent_encrypted_writing_over_multiple_row_groups() { // Read example data and set up encryption/decryption properties - async fn concurrent_encrypted_writing_over_multiple_row_groups( - footer_key: &[u8], - dec_column_keys: &[(&str, &[u8])], - enc_column_keys: &[(&str, &[u8])], - ) { - let path = encryption_util::encrypted_data_path( - footer_key, - "encrypt_columns_and_footer.parquet.encrypted", - ); - let file = std::fs::File::open(path).unwrap(); - - let mut enc_builder = - parquet::encryption::encrypt::FileEncryptionProperties::builder(footer_key.to_vec()); - for (column_name, key) in enc_column_keys { - enc_builder = enc_builder.with_column_key(column_name, key.to_vec()); - } - let file_encryption_properties = enc_builder.build().unwrap(); - - let mut dec_builder = FileDecryptionProperties::builder(footer_key.to_vec()); - for (column_name, key) in dec_column_keys { - dec_builder = dec_builder.with_column_key(column_name, key.to_vec()); - } - let decryption_properties = dec_builder.build().unwrap(); + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); + let file = std::fs::File::open(path).unwrap(); - let (record_batches, metadata) = - read_encrypted_file(&file, decryption_properties.clone()).unwrap(); - let schema = metadata.schema(); - - // Create a channel to send RecordBatches to the writer and send row groups - let (record_batch_tx, data) = tokio::sync::mpsc::channel::(100); - let data_generator = tokio::spawn(async move { - for record_batch in record_batches { - record_batch_tx.send(record_batch).await.unwrap(); - } - }); + let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) + .with_column_key("double_field", b"1234567890123450".into()) + .with_column_key("float_field", b"1234567890123451".into()) + .build() + .unwrap(); + let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) + .with_column_key("double_field", b"1234567890123450".into()) + .with_column_key("float_field", b"1234567890123451".into()) + .build() + .unwrap(); - let props = Arc::new( - WriterPropertiesBuilder::default() - .with_file_encryption_properties(file_encryption_properties) - .build(), - ); - let parquet_schema = ArrowSchemaConverter::new() - .with_coerce_types(props.coerce_types()) - .convert(schema) - .unwrap(); + let (record_batches, metadata) = + read_encrypted_file(&file, decryption_properties.clone()).unwrap(); + let schema = metadata.schema(); - // Create a temporary file to write the encrypted data - let temp_file = tempfile::tempfile().unwrap(); + // Create a channel to send RecordBatches to the writer and send row groups + let (record_batch_tx, data) = tokio::sync::mpsc::channel::(100); + let data_generator = tokio::spawn(async move { + for record_batch in record_batches { + record_batch_tx.send(record_batch).await.unwrap(); + } + }); + + let props = Arc::new( + WriterPropertiesBuilder::default() + .with_file_encryption_properties(file_encryption_properties) + .build(), + ); + let parquet_schema = ArrowSchemaConverter::new() + .with_coerce_types(props.coerce_types()) + .convert(schema) + .unwrap(); - let writer = - SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); - let row_group_writer_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(schema)); - let max_row_groups = 1; + // Create a temporary file to write the encrypted data + let temp_file = tempfile::tempfile().unwrap(); - let (serialize_tx, serialize_rx) = - tokio::sync::mpsc::channel::>(max_row_groups); + let writer = + SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); + let row_group_writer_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(schema)); + let max_row_groups = 1; - let launch_serialization_task = spawn_parquet_parallel_serialization_task( - row_group_writer_factory, - data, - serialize_tx, - schema.clone(), - ); + let (serialize_tx, serialize_rx) = + tokio::sync::mpsc::channel::>(max_row_groups); - let _file_metadata = concatenate_parallel_row_groups(writer, serialize_rx) - .await - .unwrap(); + let launch_serialization_task = spawn_parquet_parallel_serialization_task( + row_group_writer_factory, + data, + serialize_tx, + schema.clone(), + ); - data_generator.await.unwrap(); - launch_serialization_task.await.unwrap().unwrap(); + let _file_metadata = concatenate_parallel_row_groups(writer, serialize_rx) + .await + .unwrap(); - // Check that the file was written correctly - let (read_record_batches, read_metadata) = - read_encrypted_file(&temp_file, decryption_properties.clone()).unwrap(); + data_generator.await.unwrap(); + launch_serialization_task.await.unwrap().unwrap(); - assert_eq!(read_metadata.metadata().file_metadata().num_rows(), 50); - verify_encryption_test_data(read_record_batches, read_metadata.metadata()); - } + // Check that the file was written correctly + let (read_record_batches, read_metadata) = + read_encrypted_file(&temp_file, decryption_properties.clone()).unwrap(); - // AES-128 - concurrent_encrypted_writing_over_multiple_row_groups( - b"0123456789012345", // 128bit/16 - &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), - ], - &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), - ], - ) - .await; - - // AES-256 - concurrent_encrypted_writing_over_multiple_row_groups( - b"01234567890123456789012345678901", // 256bit/32 - &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), - ], - &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), - ], - ) - .await; + assert_eq!(read_metadata.metadata().file_metadata().num_rows(), 50); + verify_encryption_test_data(read_record_batches, read_metadata.metadata()); } #[tokio::test] async fn test_multi_threaded_encrypted_writing() { // Read example data and set up encryption/decryption properties - async fn multi_threaded_encrypted_writing( - footer_key: &[u8], - dec_column_keys: &[(&str, &[u8])], - enc_column_keys: &[(&str, &[u8])], - ) { - let path = encryption_util::encrypted_data_path( - footer_key, - "encrypt_columns_and_footer.parquet.encrypted", - ); - let file = std::fs::File::open(path).unwrap(); - - let mut enc_builder = - parquet::encryption::encrypt::FileEncryptionProperties::builder(footer_key.to_vec()); - for (column_name, key) in enc_column_keys { - enc_builder = enc_builder.with_column_key(column_name, key.to_vec()); - } - let encryption_properties = enc_builder.build().unwrap(); - - let mut dec_builder = FileDecryptionProperties::builder(footer_key.to_vec()); - for (column_name, key) in dec_column_keys { - dec_builder = dec_builder.with_column_key(column_name, key.to_vec()); - } - let decryption_properties = dec_builder.build().unwrap(); - - let (record_batches, metadata) = - read_encrypted_file(&file, Arc::clone(&decryption_properties)).unwrap(); - let schema = metadata.schema().clone(); + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); + let file = std::fs::File::open(path).unwrap(); - let props = Arc::new( - WriterPropertiesBuilder::default() - .with_file_encryption_properties(encryption_properties) - .build(), - ); - - let parquet_schema = ArrowSchemaConverter::new() - .with_coerce_types(props.coerce_types()) - .convert(&schema) - .unwrap(); + let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) + .with_column_key("double_field", b"1234567890123450".into()) + .with_column_key("float_field", b"1234567890123451".into()) + .build() + .unwrap(); + let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) + .with_column_key("double_field", b"1234567890123450".into()) + .with_column_key("float_field", b"1234567890123451".into()) + .build() + .unwrap(); - // Create a temporary file to write the encrypted data - let temp_file = tempfile::tempfile().unwrap(); - let mut writer = - SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); - let row_group_writer_factory = - ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); - - let (serialize_tx, mut serialize_rx) = - tokio::sync::mpsc::channel::>(1); - - // Create a channel to send RecordBatches to the writer and send row batches - let (record_batch_tx, mut data) = tokio::sync::mpsc::channel::(100); - let data_generator = tokio::spawn(async move { - for record_batch in record_batches { - record_batch_tx.send(record_batch).await.unwrap(); - } - }); + let (record_batches, metadata) = + read_encrypted_file(&file, Arc::clone(&decryption_properties)).unwrap(); + let schema = metadata.schema().clone(); - // Get column writers - let col_writers = row_group_writer_factory.create_column_writers(0).unwrap(); + let props = Arc::new( + WriterPropertiesBuilder::default() + .with_file_encryption_properties(file_encryption_properties) + .build(), + ); - let (col_writer_tasks, col_array_channels) = - spawn_column_parallel_row_group_writer(col_writers, 10).unwrap(); + let parquet_schema = ArrowSchemaConverter::new() + .with_coerce_types(props.coerce_types()) + .convert(&schema) + .unwrap(); - // Spawn serialization tasks for incoming RecordBatches - let launch_serialization_task = tokio::spawn(async move { - let Some(rb) = data.recv().await else { - panic!() - }; - send_arrays_to_column_writers(&col_array_channels, &rb, &schema) - .await - .unwrap(); - let finalize_rg_task = spawn_rg_join_and_finalize_task(col_writer_tasks, 10); + // Create a temporary file to write the encrypted data + let temp_file = tempfile::tempfile().unwrap(); + let mut writer = + SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); + let row_group_writer_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); - serialize_tx.send(finalize_rg_task).await.unwrap(); - drop(col_array_channels); - }); + let (serialize_tx, mut serialize_rx) = + tokio::sync::mpsc::channel::>(1); - // Append the finalized row groups to the SerializedFileWriter - while let Some(task) = serialize_rx.recv().await { - let (arrow_column_chunks, _) = task.await.unwrap().unwrap(); - let mut row_group_writer = writer.next_row_group().unwrap(); - for chunk in arrow_column_chunks { - chunk.append_to_row_group(&mut row_group_writer).unwrap(); - } - row_group_writer.close().unwrap(); + // Create a channel to send RecordBatches to the writer and send row batches + let (record_batch_tx, mut data) = tokio::sync::mpsc::channel::(100); + let data_generator = tokio::spawn(async move { + for record_batch in record_batches { + record_batch_tx.send(record_batch).await.unwrap(); } + }); - // Wait for data generator and serialization task to finish - data_generator.await.unwrap(); - launch_serialization_task.await.unwrap(); - let metadata = writer.close().unwrap(); + // Get column writers + let col_writers = row_group_writer_factory.create_column_writers(0).unwrap(); - // Close the file writer which writes the footer - assert_eq!(metadata.file_metadata().num_rows(), 50); + let (col_writer_tasks, col_array_channels) = + spawn_column_parallel_row_group_writer(col_writers, 10).unwrap(); - // Check that the file was written correctly - let (read_record_batches, read_metadata) = - read_encrypted_file(&temp_file, decryption_properties).unwrap(); - verify_encryption_test_data(read_record_batches, read_metadata.metadata()); + // Spawn serialization tasks for incoming RecordBatches + let launch_serialization_task = tokio::spawn(async move { + let Some(rb) = data.recv().await else { + panic!() + }; + send_arrays_to_column_writers(&col_array_channels, &rb, &schema) + .await + .unwrap(); + let finalize_rg_task = spawn_rg_join_and_finalize_task(col_writer_tasks, 10); - // Check that file was encrypted - let result = ArrowReaderMetadata::load(&temp_file, ArrowReaderOptions::default()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" - ); - } + serialize_tx.send(finalize_rg_task).await.unwrap(); + drop(col_array_channels); + }); - // AES-128 - multi_threaded_encrypted_writing( - b"0123456789012345", // 128bit/16 - &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), - ], - &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), - ], - ) - .await; + // Append the finalized row groups to the SerializedFileWriter + while let Some(task) = serialize_rx.recv().await { + let (arrow_column_chunks, _) = task.await.unwrap().unwrap(); + let mut row_group_writer = writer.next_row_group().unwrap(); + for chunk in arrow_column_chunks { + chunk.append_to_row_group(&mut row_group_writer).unwrap(); + } + row_group_writer.close().unwrap(); + } - // AES-256 - multi_threaded_encrypted_writing( - b"01234567890123456789012345678901", // 256bit/32 - &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), - ], - &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), - ], - ) - .await; + // Wait for data generator and serialization task to finish + data_generator.await.unwrap(); + launch_serialization_task.await.unwrap(); + let metadata = writer.close().unwrap(); + + // Close the file writer which writes the footer + assert_eq!(metadata.file_metadata().num_rows(), 50); + + // Check that the file was written correctly + let (read_record_batches, read_metadata) = + read_encrypted_file(&temp_file, decryption_properties).unwrap(); + verify_encryption_test_data(read_record_batches, read_metadata.metadata()); + + // Check that file was encrypted + let result = ArrowReaderMetadata::load(&temp_file, ArrowReaderOptions::default()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); } #[tokio::test] async fn test_multi_threaded_encrypted_writing_deprecated() { // Read example data and set up encryption/decryption properties - async fn multi_threaded_encrypted_writing_deprecated( - footer_key: &[u8], - dec_column_keys: &[(&str, &[u8])], - enc_column_keys: &[(&str, &[u8])], - ) { - let path = encryption_util::encrypted_data_path( - footer_key, - "encrypt_columns_and_footer.parquet.encrypted", - ); - let file = std::fs::File::open(path).unwrap(); - - let mut enc_builder = - parquet::encryption::encrypt::FileEncryptionProperties::builder(footer_key.to_vec()); - for (column_name, key) in enc_column_keys { - enc_builder = enc_builder.with_column_key(column_name, key.to_vec()); - } - let encryption_properties = enc_builder.build().unwrap(); + let testdata = arrow::util::test_util::parquet_test_data(); + let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); + let file = std::fs::File::open(path).unwrap(); - let mut dec_builder = FileDecryptionProperties::builder(footer_key.to_vec()); - for (column_name, key) in dec_column_keys { - dec_builder = dec_builder.with_column_key(column_name, key.to_vec()); - } - let decryption_properties = dec_builder.build().unwrap(); - - let (record_batches, metadata) = - read_encrypted_file(&file, Arc::clone(&decryption_properties)).unwrap(); - let to_write: Vec<_> = record_batches - .iter() - .flat_map(|rb| rb.columns().to_vec()) - .collect(); - let schema = metadata.schema().clone(); - - let props = Some( - WriterPropertiesBuilder::default() - .with_file_encryption_properties(encryption_properties) - .build(), - ); + let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) + .with_column_key("double_field", b"1234567890123450".into()) + .with_column_key("float_field", b"1234567890123451".into()) + .build() + .unwrap(); + let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) + .with_column_key("double_field", b"1234567890123450".into()) + .with_column_key("float_field", b"1234567890123451".into()) + .build() + .unwrap(); - // Create a temporary file to write the encrypted data - let temp_file = tempfile::tempfile().unwrap(); - let mut writer = ArrowWriter::try_new(&temp_file, schema.clone(), props).unwrap(); + let (record_batches, metadata) = + read_encrypted_file(&file, Arc::clone(&decryption_properties)).unwrap(); + let to_write: Vec<_> = record_batches + .iter() + .flat_map(|rb| rb.columns().to_vec()) + .collect(); + let schema = metadata.schema().clone(); + + let props = Some( + WriterPropertiesBuilder::default() + .with_file_encryption_properties(file_encryption_properties) + .build(), + ); + + // Create a temporary file to write the encrypted data + let temp_file = tempfile::tempfile().unwrap(); + let mut writer = ArrowWriter::try_new(&temp_file, schema.clone(), props).unwrap(); - // LOW-LEVEL API: Use low level API to write into a file using multiple threads + // LOW-LEVEL API: Use low level API to write into a file using multiple threads - // Get column writers - #[allow(deprecated)] - let col_writers = writer.get_column_writers().unwrap(); - let num_columns = col_writers.len(); + // Get column writers + #[allow(deprecated)] + let col_writers = writer.get_column_writers().unwrap(); + let num_columns = col_writers.len(); - let (col_writer_tasks, mut col_array_channels) = - spawn_column_parallel_row_group_writer(col_writers, 100).unwrap(); + let (col_writer_tasks, mut col_array_channels) = + spawn_column_parallel_row_group_writer(col_writers, 100).unwrap(); // Send the ArrowLeafColumn data to the respective column writer channels - let mut worker_iter = col_array_channels.iter_mut(); - for (array, field) in to_write.iter().zip(schema.fields()) { - for leaves in compute_leaves(field, array).unwrap() { - worker_iter.next().unwrap().send(leaves).await.unwrap(); - } + let mut worker_iter = col_array_channels.iter_mut(); + for (array, field) in to_write.iter().zip(schema.fields()) { + for leaves in compute_leaves(field, array).unwrap() { + worker_iter.next().unwrap().send(leaves).await.unwrap(); } - drop(col_array_channels); + } + drop(col_array_channels); - // Wait for all column writers to finish writing - let mut finalized_rg = Vec::with_capacity(num_columns); - for task in col_writer_tasks.into_iter() { - finalized_rg.push(task.await.unwrap().unwrap().close().unwrap()); - } + // Wait for all column writers to finish writing + let mut finalized_rg = Vec::with_capacity(num_columns); + for task in col_writer_tasks.into_iter() { + finalized_rg.push(task.await.unwrap().unwrap().close().unwrap()); + } - // Append the finalized row group to the SerializedFileWriter - #[allow(deprecated)] - writer.append_row_group(finalized_rg).unwrap(); + // Append the finalized row group to the SerializedFileWriter + #[allow(deprecated)] + writer.append_row_group(finalized_rg).unwrap(); - // HIGH-LEVEL API: Write RecordBatches into the file using ArrowWriter + // HIGH-LEVEL API: Write RecordBatches into the file using ArrowWriter - // Write individual RecordBatches into the file - for rb in record_batches { - writer.write(&rb).unwrap() - } - assert!(writer.flush().is_ok()); - - // Close the file writer which writes the footer - let metadata = writer.finish().unwrap(); - assert_eq!(metadata.file_metadata().num_rows(), 100); - - // Check that the file was written correctly - let (read_record_batches, read_metadata) = - read_encrypted_file(&temp_file, decryption_properties).unwrap(); - verify_encryption_double_test_data(read_record_batches, read_metadata.metadata()); - - // Check that file was encrypted - let result = ArrowReaderMetadata::load(&temp_file, ArrowReaderOptions::default()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" - ); + // Write individual RecordBatches into the file + for rb in record_batches { + writer.write(&rb).unwrap() } - - // AES-128 - multi_threaded_encrypted_writing_deprecated( - b"0123456789012345", // 128bit/16 - &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), - ], - &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), - ], - ) - .await; - - // AES-256 - multi_threaded_encrypted_writing_deprecated( - b"01234567890123456789012345678901", // 256bit/32 - &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), - ], - &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), - ], - ) - .await; + assert!(writer.flush().is_ok()); + + // Close the file writer which writes the footer + let metadata = writer.finish().unwrap(); + assert_eq!(metadata.file_metadata().num_rows(), 100); + + // Check that the file was written correctly + let (read_record_batches, read_metadata) = + read_encrypted_file(&temp_file, decryption_properties).unwrap(); + verify_encryption_double_test_data(read_record_batches, read_metadata.metadata()); + + // Check that file was encrypted + let result = ArrowReaderMetadata::load(&temp_file, ArrowReaderOptions::default()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); } From 1b3f2f7d290f0351ca78808d9a54793ae47d22cf Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Fri, 1 May 2026 16:06:36 -0700 Subject: [PATCH 18/26] Revert tests --- parquet/tests/encryption/encryption.rs | 24 +++++++++--------------- 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index 420ea94b98c0..6a1f538b6834 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -386,22 +386,16 @@ fn test_uniform_encryption() { #[test] fn test_decrypting_without_decryption_properties_fails() { let test_data = arrow::util::test_util::parquet_test_data(); - let paths = [ - format!("{test_data}/uniform_encryption.parquet.encrypted"), - format!("{test_data}/aes256/uniform_encryption.parquet.encrypted"), - ]; - - for path in &paths { - let file = File::open(path).unwrap(); + let path = format!("{test_data}/uniform_encryption.parquet.encrypted"); + let file = File::open(path).unwrap(); - let options = ArrowReaderOptions::default(); - let result = ArrowReaderMetadata::load(&file, options.clone()); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err().to_string(), - "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" - ); - } + let options = ArrowReaderOptions::default(); + let result = ArrowReaderMetadata::load(&file, options.clone()); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().to_string(), + "Parquet error: Parquet file has an encrypted footer but decryption properties were not provided" + ); } #[test] From 5b8b2441c43441f102fd75ad495ea8114123fd6f Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Fri, 1 May 2026 16:09:35 -0700 Subject: [PATCH 19/26] Fix comments --- parquet/tests/encryption/encryption.rs | 4 +--- parquet/tests/encryption/encryption_async.rs | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index 6a1f538b6834..d546837b0ca8 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -614,9 +614,7 @@ fn test_uniform_encryption_plaintext_footer_with_key_retriever() { ); // AES-256 - // The asymmetric column names is because we check column paths in the validate_encrypted_column_names function of [encryption::encrypt] - // The column path of the repeated field (int64_field) is int64_field.list.int64_field - // Switching from `c.path().string()` to `c.name().to_string()` can fix the asymmetricity + // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name uniform_encryption_plaintext_footer_with_key_retriever( b"01234567890123456789012345678901", "kf", diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index ba331d502042..ef31b167cf12 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -1313,7 +1313,7 @@ async fn test_multi_threaded_encrypted_writing_deprecated() { let (col_writer_tasks, mut col_array_channels) = spawn_column_parallel_row_group_writer(col_writers, 100).unwrap(); - // Send the ArrowLeafColumn data to the respective column writer channels + // Send the ArrowLeafColumn data to the respective column writer channels let mut worker_iter = col_array_channels.iter_mut(); for (array, field) in to_write.iter().zip(schema.fields()) { for leaves in compute_leaves(field, array).unwrap() { From f414b3e5dfd1170257268f8244bdaa76acad85a2 Mon Sep 17 00:00:00 2001 From: hsiang-c Date: Mon, 4 May 2026 10:39:31 -0700 Subject: [PATCH 20/26] Extract hardcoded strings to constants --- parquet-testing | 2 +- parquet/tests/encryption/encryption.rs | 646 +++++++------------ parquet/tests/encryption/encryption_async.rs | 389 ++++------- parquet/tests/encryption/encryption_util.rs | 37 ++ 4 files changed, 414 insertions(+), 660 deletions(-) diff --git a/parquet-testing b/parquet-testing index a3d96a65e11e..e74785d85a4e 160000 --- a/parquet-testing +++ b/parquet-testing @@ -1 +1 @@ -Subproject commit a3d96a65e11e2bbca7d22a894e8313ede90a33a3 +Subproject commit e74785d85a4ecee829e1e405444d6a1b24b8bc9c diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index d546837b0ca8..4a20d36421c7 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -19,6 +19,9 @@ use crate::encryption_util; use crate::encryption_util::{ + AES_128_COLUMN_KEYS, AES_128_COLUMN_NAMES, AES_128_FOOTER_KEY, AES_128_FOOTER_KEY_NAME, + AES_128_KEY_NAMES, AES_256_COLUMN_KEYS, AES_256_COLUMN_NAMES, AES_256_FOOTER_KEY, + AES_256_FOOTER_KEY_NAME, AES_256_KEY_NAMES, BAD_AES_128_FOOTER_KEY, BAD_AES_256_FOOTER_KEY, TestKeyRetriever, read_and_roundtrip_to_encrypted_file, verify_column_indexes, verify_encryption_test_file_read, }; @@ -61,43 +64,25 @@ fn test_non_uniform_encryption_plaintext_footer() { // AES-128: there is always a footer key even with a plaintext footer, // but this is used for signing the footer. non_uniform_encryption_plaintext_footer( - b"0123456789012345", // 128bit/16 + AES_128_FOOTER_KEY, &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ); // AES-256 non_uniform_encryption_plaintext_footer( - b"01234567890123456789012345678901", // 256bit/32 + AES_256_FOOTER_KEY, &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ); } @@ -143,43 +128,24 @@ fn test_plaintext_footer_signature_verification() { } plaintext_footer_signature_verification( - // 128bit/16 - b"0000000000000000", // 128bit/16 + BAD_AES_128_FOOTER_KEY, &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ); plaintext_footer_signature_verification( - b"00000000000000000000000000000000", // 256bit/32 + BAD_AES_256_FOOTER_KEY, &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ) } @@ -243,43 +209,25 @@ fn test_non_uniform_encryption_disabled_aad_storage() { // AES-128 non_uniform_encryption_disabled_aad_storage( - b"0123456789012345", // 128bit/16 + AES_128_FOOTER_KEY, &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ); // AES-256 non_uniform_encryption_disabled_aad_storage( - b"01234567890123456789012345678901", // 256bit/32 + AES_256_FOOTER_KEY, &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ); } @@ -310,77 +258,51 @@ fn test_non_uniform_encryption() { // AES-128 non_uniform_encryption( - b"0123456789012345", + AES_128_FOOTER_KEY, &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ); // AES-256 non_uniform_encryption( - b"01234567890123456789012345678901", // 256bit/32 + AES_256_FOOTER_KEY, &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ); } #[test] fn test_uniform_encryption() { - fn uniform_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + fn uniform_encryption(footer_key: &[u8]) { let path = encryption_util::encrypted_data_path( footer_key, "uniform_encryption.parquet.encrypted", ); let file = File::open(path).unwrap(); - let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); - for (column_name, key) in column_keys { - builder = builder.with_column_key(column_name, key.to_vec()); - } - let decryption_properties = builder.build().unwrap(); + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); verify_encryption_test_file_read(file, decryption_properties); } // AES-128: there is always a footer key even with a plaintext footer, // but this is used for signing the footer. - uniform_encryption( - b"0123456789012345", // 128bit/16 - &[], - ); + uniform_encryption(AES_128_FOOTER_KEY); // AES-256 - uniform_encryption( - b"01234567890123456789012345678901", // 256bit/32 - &[], - ); + uniform_encryption(AES_256_FOOTER_KEY); } #[test] @@ -428,25 +350,25 @@ fn test_aes_ctr_encryption() { } // AES-128 aes_ctr_encryption( - b"0123456789012345", // 128bit/16 + AES_128_FOOTER_KEY, &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ); // AES-256 aes_ctr_encryption( - b"01234567890123456789012345678901", // 256bit/32 + AES_256_FOOTER_KEY, &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ); } @@ -478,27 +400,27 @@ fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { // AES-128 non_uniform_encryption_plaintext_footer_with_key_retriever( - b"0123456789012345", + AES_128_FOOTER_KEY, &[ - ("kf", b"0123456789012345".as_slice()), - ("kc1", b"1234567890123450".as_slice()), - ("kc2", b"1234567890123451".as_slice()), + (AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY), + (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ); // AES-256 non_uniform_encryption_plaintext_footer_with_key_retriever( - b"01234567890123456789012345678901", + AES_256_FOOTER_KEY, &[ - ("kf", b"01234567890123456789012345678901".as_slice()), - ("kc1", b"12345678901234567890123456789012".as_slice()), - ("kc2", b"12345678901234567890123456789013".as_slice()), - ("kc3", b"12345678901234567890123456789014".as_slice()), - ("kc4", b"12345678901234567890123456789015".as_slice()), - ("kc5", b"12345678901234567890123456789016".as_slice()), - ("kc6", b"12345678901234567890123456789017".as_slice()), - ("kc7", b"12345678901234567890123456789018".as_slice()), - ("kc8", b"12345678901234567890123456789019".as_slice()), + (AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY), + (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ); } @@ -600,75 +522,83 @@ fn test_uniform_encryption_plaintext_footer_with_key_retriever() { // AES-128 uniform_encryption_plaintext_footer_with_key_retriever( - b"0123456789012345", - "kf", - b"0133456789012345", + AES_128_FOOTER_KEY, + AES_128_FOOTER_KEY_NAME, + BAD_AES_128_FOOTER_KEY, &[ - ("kc1", b"1234567890123450".as_slice()), - ("kc2", b"1234567890123451".as_slice()), + (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), ], &[ - ("double_field", b"1234567890123450".as_slice(), "kc1"), - ("float_field", b"1234567890123451".as_slice(), "kc2"), + ( + AES_128_COLUMN_NAMES[0], + AES_128_COLUMN_KEYS[0], + AES_128_KEY_NAMES[0], + ), + ( + AES_128_COLUMN_NAMES[1], + AES_128_COLUMN_KEYS[1], + AES_128_KEY_NAMES[1], + ), ], ); // AES-256 // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name uniform_encryption_plaintext_footer_with_key_retriever( - b"01234567890123456789012345678901", - "kf", - b"01334567890123456789012345678901", + AES_256_FOOTER_KEY, + AES_256_FOOTER_KEY_NAME, + BAD_AES_256_FOOTER_KEY, &[ - ("kc1", b"12345678901234567890123456789012".as_slice()), - ("kc2", b"12345678901234567890123456789013".as_slice()), - ("kc3", b"12345678901234567890123456789014".as_slice()), - ("kc4", b"12345678901234567890123456789015".as_slice()), - ("kc5", b"12345678901234567890123456789016".as_slice()), - ("kc6", b"12345678901234567890123456789017".as_slice()), - ("kc7", b"12345678901234567890123456789018".as_slice()), - ("kc8", b"12345678901234567890123456789019".as_slice()), + (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), ], &[ ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - "kc1", + AES_256_COLUMN_NAMES[0], + AES_256_COLUMN_KEYS[0], + AES_256_KEY_NAMES[0], ), ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - "kc2", + AES_256_COLUMN_NAMES[1], + AES_256_COLUMN_KEYS[1], + AES_256_KEY_NAMES[1], ), ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - "kc3", + AES_256_COLUMN_NAMES[2], + AES_256_COLUMN_KEYS[2], + AES_256_KEY_NAMES[2], ), ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - "kc4", + AES_256_COLUMN_NAMES[3], + AES_256_COLUMN_KEYS[3], + AES_256_KEY_NAMES[3], ), ( - "ba_field", - b"12345678901234567890123456789016".as_slice(), - "kc5", + AES_256_COLUMN_NAMES[4], + AES_256_COLUMN_KEYS[4], + AES_256_KEY_NAMES[4], ), ( - "flba_field", - b"12345678901234567890123456789017".as_slice(), - "kc6", + AES_256_COLUMN_NAMES[5], + AES_256_COLUMN_KEYS[5], + AES_256_KEY_NAMES[5], ), ( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".as_slice(), - "kc7", + AES_256_COLUMN_NAMES[7], + AES_256_COLUMN_KEYS[7], + AES_256_KEY_NAMES[7], ), ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - "kc8", + AES_256_COLUMN_NAMES[8], + AES_256_COLUMN_KEYS[8], + AES_256_KEY_NAMES[6], ), ], ); @@ -698,27 +628,27 @@ fn test_non_uniform_encryption_with_key_retriever() { // AES-128 non_uniform_encryption_with_key_retriever( - b"0123456789012345", + AES_128_FOOTER_KEY, &[ - ("kf", b"0123456789012345".as_slice()), - ("kc1", b"1234567890123450".as_slice()), - ("kc2", b"1234567890123451".as_slice()), + (AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY), + (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ); // AES-256 non_uniform_encryption_with_key_retriever( - b"01234567890123456789012345678901", + AES_256_FOOTER_KEY, &[ - ("kf", b"01234567890123456789012345678901".as_slice()), - ("kc1", b"12345678901234567890123456789012".as_slice()), - ("kc2", b"12345678901234567890123456789013".as_slice()), - ("kc3", b"12345678901234567890123456789014".as_slice()), - ("kc4", b"12345678901234567890123456789015".as_slice()), - ("kc5", b"12345678901234567890123456789016".as_slice()), - ("kc6", b"12345678901234567890123456789017".as_slice()), - ("kc7", b"12345678901234567890123456789018".as_slice()), - ("kc8", b"12345678901234567890123456789019".as_slice()), + (AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY), + (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ); } @@ -744,10 +674,10 @@ fn test_uniform_encryption_with_key_retriever() { } // AES-128 - uniform_encryption_with_key_retriever("kf", b"0123456789012345"); + uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY); // AES-256 - uniform_encryption_with_key_retriever("kf", b"01234567890123456789012345678901"); + uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY); } fn row_group_sizes(metadata: &ParquetMetaData) -> Vec { @@ -756,26 +686,26 @@ fn row_group_sizes(metadata: &ParquetMetaData) -> Vec { #[test] fn test_uniform_encryption_roundtrip() { - uniform_encryption_roundtrip(false, false, b"0123456789012345").unwrap(); - uniform_encryption_roundtrip(false, false, b"01234567890123456789012345678901").unwrap(); + uniform_encryption_roundtrip(false, false, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_roundtrip(false, false, AES_256_FOOTER_KEY).unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_dictionary() { - uniform_encryption_roundtrip(false, true, b"0123456789012345").unwrap(); - uniform_encryption_roundtrip(false, true, b"01234567890123456789012345678901").unwrap(); + uniform_encryption_roundtrip(false, true, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_roundtrip(false, true, AES_256_FOOTER_KEY).unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_page_index() { - uniform_encryption_roundtrip(true, false, b"0123456789012345").unwrap(); - uniform_encryption_roundtrip(true, false, b"01234567890123456789012345678901").unwrap(); + uniform_encryption_roundtrip(true, false, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_roundtrip(true, false, AES_256_FOOTER_KEY).unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_page_index_and_dictionary() { - uniform_encryption_roundtrip(true, true, b"0123456789012345").unwrap(); - uniform_encryption_roundtrip(true, true, b"01234567890123456789012345678901").unwrap(); + uniform_encryption_roundtrip(true, true, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_roundtrip(true, true, AES_256_FOOTER_KEY).unwrap(); } fn uniform_encryption_roundtrip( @@ -878,14 +808,14 @@ fn uniform_encryption_roundtrip( #[test] fn test_uniform_encryption_page_skipping() { - uniform_encryption_page_skipping(false, b"0123456789012345").unwrap(); - uniform_encryption_page_skipping(false, b"01234567890123456789012345678901").unwrap(); + uniform_encryption_page_skipping(false, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_page_skipping(false, AES_256_FOOTER_KEY).unwrap(); } #[test] fn test_uniform_encryption_page_skipping_with_page_index() { - uniform_encryption_page_skipping(true, b"0123456789012345").unwrap(); - uniform_encryption_page_skipping(true, b"01234567890123456789012345678901").unwrap(); + uniform_encryption_page_skipping(true, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_page_skipping(true, AES_256_FOOTER_KEY).unwrap(); } fn uniform_encryption_page_skipping( @@ -1026,66 +956,29 @@ fn test_write_non_uniform_encryption() { } write_non_uniform_encryption( - b"0123456789012345", // 128bit/16, - vec!["double_field", "float_field"], - vec![b"1234567890123450".to_vec(), b"1234567890123451".to_vec()], + AES_128_FOOTER_KEY, + AES_128_COLUMN_NAMES.to_vec(), + AES_128_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ); + // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name write_non_uniform_encryption( - b"01234567890123456789012345678901", - vec![ - "double_field", - "float_field", - "boolean_field", - "int32_field", - "ba_field", - "flba_field", - "int64_field", - "int64_field.list.int64_field", - "int96_field", - ], - vec![ - b"12345678901234567890123456789012".to_vec(), - b"12345678901234567890123456789013".to_vec(), - b"12345678901234567890123456789014".to_vec(), - b"12345678901234567890123456789015".to_vec(), - b"12345678901234567890123456789016".to_vec(), - b"12345678901234567890123456789017".to_vec(), - b"12345678901234567890123456789018".to_vec(), - b"12345678901234567890123456789018".to_vec(), - b"12345678901234567890123456789019".to_vec(), - ], + AES_256_FOOTER_KEY, + AES_256_COLUMN_NAMES.to_vec(), + AES_256_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), + (AES_256_COLUMN_NAMES[8], AES_256_COLUMN_KEYS[8]), ], ); } @@ -1096,20 +989,20 @@ fn test_write_uniform_encryption_plaintext_footer() { let path = format!("{testdata}/encrypt_columns_plaintext_footer.parquet.encrypted"); let file = File::open(path).unwrap(); - let footer_key = b"0123456789012345".to_vec(); // 128bit/16 - let wrong_footer_key = b"0000000000000000".to_vec(); // 128bit/16 - let column_1_key = b"1234567890123450".to_vec(); - let column_2_key = b"1234567890123451".to_vec(); + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let wrong_footer_key = BAD_AES_128_FOOTER_KEY.to_vec(); + let column_1_key = AES_128_COLUMN_KEYS[0].to_vec(); + let column_2_key = AES_128_COLUMN_KEYS[1].to_vec(); let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("double_field", column_1_key.clone()) - .with_column_key("float_field", column_2_key.clone()) + .with_column_key(AES_128_COLUMN_NAMES[0], column_1_key.clone()) + .with_column_key(AES_128_COLUMN_NAMES[1], column_2_key.clone()) .build() .unwrap(); let wrong_decryption_properties = FileDecryptionProperties::builder(wrong_footer_key) - .with_column_key("double_field", column_1_key) - .with_column_key("float_field", column_2_key) + .with_column_key(AES_128_COLUMN_NAMES[0], column_1_key) + .with_column_key(AES_128_COLUMN_NAMES[1], column_2_key) .build() .unwrap(); @@ -1171,11 +1064,8 @@ fn test_write_uniform_encryption_plaintext_footer() { #[test] pub fn test_column_statistics_with_plaintext_footer() { let keys = [ - (b"0123456789012345".to_vec(), b"1234567890123450".to_vec()), - ( - b"01234567890123456789012345678901".to_vec(), - b"12345678901234567890123456789012".to_vec(), - ), + (AES_128_FOOTER_KEY.to_vec(), AES_128_COLUMN_KEYS[0].to_vec()), + (AES_256_FOOTER_KEY.to_vec(), AES_256_COLUMN_KEYS[0].to_vec()), ]; for (footer_key, column_key) in &keys { @@ -1211,10 +1101,7 @@ pub fn test_column_statistics_with_plaintext_footer() { #[test] pub fn test_column_statistics_with_plaintext_footer_and_uniform_encryption() { - let footer_keys = [ - b"0123456789012345".to_vec(), - b"01234567890123456789012345678901".to_vec(), - ]; + let footer_keys = [AES_128_FOOTER_KEY.to_vec(), AES_256_FOOTER_KEY.to_vec()]; for footer_key in &footer_keys { // Write with uniform encryption and a plaintext footer. @@ -1242,11 +1129,8 @@ pub fn test_column_statistics_with_plaintext_footer_and_uniform_encryption() { #[test] pub fn test_column_statistics_with_encrypted_footer() { let keys = [ - (b"0123456789012345".to_vec(), b"1234567890123450".to_vec()), - ( - b"01234567890123456789012345678901".to_vec(), - b"12345678901234567890123456789012".to_vec(), - ), + (AES_128_FOOTER_KEY.to_vec(), AES_128_COLUMN_KEYS[0].to_vec()), + (AES_256_FOOTER_KEY.to_vec(), AES_256_COLUMN_KEYS[0].to_vec()), ]; for (footer_key, column_key) in &keys { @@ -1278,10 +1162,7 @@ pub fn test_column_statistics_with_encrypted_footer() { #[test] pub fn test_column_statistics_with_encrypted_footer_and_uniform_encryption() { - let footer_keys = [ - b"0123456789012345".to_vec(), - b"01234567890123456789012345678901".to_vec(), - ]; + let footer_keys = [AES_128_FOOTER_KEY.to_vec(), AES_256_FOOTER_KEY.to_vec()]; for footer_key in &footer_keys { // Encrypt with an encrypted footer and uniform encryption @@ -1429,8 +1310,8 @@ fn test_write_uniform_encryption() { ); } - write_uniform_encryption(b"0123456789012345"); - write_uniform_encryption(b"01234567890123456789012345678901"); + write_uniform_encryption(AES_128_FOOTER_KEY); + write_uniform_encryption(AES_256_FOOTER_KEY); } #[test] @@ -1487,75 +1368,40 @@ fn test_write_non_uniform_encryption_column_missmatch() { // AES-128 write_non_uniform_encryption_column_missmatch( - b"0123456789012345", + AES_128_FOOTER_KEY, &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ); // AES-256 + // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name write_non_uniform_encryption_column_missmatch( - b"01234567890123456789012345678901", + AES_256_FOOTER_KEY, &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), + (AES_256_COLUMN_NAMES[8], AES_256_COLUMN_KEYS[8]), ], &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), + (AES_256_COLUMN_NAMES[8], AES_256_COLUMN_KEYS[8]), ], ); } @@ -1579,7 +1425,7 @@ fn test_write_encrypted_column() { let file: File = tempfile::tempfile().unwrap(); let builder = WriterProperties::builder(); - let footer_key: &[u8] = "0123456789012345".as_bytes(); + let footer_key: &[u8] = AES_128_FOOTER_KEY; let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()) .build() .unwrap(); @@ -1682,9 +1528,9 @@ fn test_write_encrypted_struct_field() { // keys need to be specified for each leaf-level Parquet column using the full "." separated // column path. let builder = WriterProperties::builder(); - let footer_key = b"0123456789012345".to_vec(); - let column_key_1 = b"1234567890123450".to_vec(); - let column_key_2 = b"1234567890123451".to_vec(); + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key_1 = AES_128_COLUMN_KEYS[0].to_vec(); + let column_key_2 = AES_128_COLUMN_KEYS[1].to_vec(); let file_encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) .with_column_key("struct_col.int64_col", column_key_1.clone()) .with_column_key("struct_col.float64_col", column_key_2.clone()) @@ -1749,8 +1595,8 @@ pub fn test_retrieve_row_group_statistics_after_encrypted_write() { let temp_file = tempfile::tempfile().unwrap(); - let footer_key = b"0123456789012345".to_vec(); - let column_key = b"1234567890123450".to_vec(); + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key = AES_128_COLUMN_KEYS[0].to_vec(); let file_encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) .with_column_key("x", column_key.clone()) .build() @@ -1795,8 +1641,8 @@ fn test_decrypt_page_index_uniform() { test_decrypt_page_index(&path, decryption_properties).unwrap(); } - decrypt_page_index_uniform("0123456789012345".as_bytes()); - decrypt_page_index_uniform("01234567890123456789012345678901".as_bytes()); + decrypt_page_index_uniform(AES_128_FOOTER_KEY); + decrypt_page_index_uniform(AES_256_FOOTER_KEY); } #[test] @@ -1815,42 +1661,24 @@ fn test_decrypt_page_index_non_uniform() { } decrypt_page_index_non_uniform( - b"0123456789012345", // 128bit/16 + AES_128_FOOTER_KEY, &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ); decrypt_page_index_non_uniform( - b"01234567890123456789012345678901", // 256bit/32 + AES_256_FOOTER_KEY, &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ); } diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index ef31b167cf12..be8a75590f44 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -19,6 +19,9 @@ use crate::encryption_util; use crate::encryption_util::{ + AES_128_COLUMN_KEYS, AES_128_COLUMN_NAMES, AES_128_FOOTER_KEY, AES_128_FOOTER_KEY_NAME, + AES_128_KEY_NAMES, AES_256_COLUMN_KEYS, AES_256_COLUMN_NAMES, AES_256_FOOTER_KEY, + AES_256_FOOTER_KEY_NAME, AES_256_KEY_NAMES, BAD_AES_128_FOOTER_KEY, BAD_AES_256_FOOTER_KEY, TestKeyRetriever, read_encrypted_file, verify_column_indexes, verify_encryption_double_test_data, verify_encryption_test_data, }; @@ -70,44 +73,26 @@ async fn test_non_uniform_encryption_plaintext_footer() { // AES-128: there is always a footer key even with a plaintext footer, // but this is used for signing the footer. non_uniform_encryption_plaintext_footer( - b"0123456789012345", // 128bit/16 + AES_128_FOOTER_KEY, &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ) .await; // AES-256 non_uniform_encryption_plaintext_footer( - b"01234567890123456789012345678901", // 256bit/32 + AES_256_FOOTER_KEY, &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ) .await; @@ -132,11 +117,11 @@ async fn test_misspecified_encryption_keys() { let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); if !column_1_key.is_empty() { - builder = builder.with_column_key("double_field", column_1_key.to_vec()); + builder = builder.with_column_key(AES_128_COLUMN_NAMES[0], column_1_key.to_vec()); } if !column_2_key.is_empty() { - builder = builder.with_column_key("float_field", column_2_key.to_vec()); + builder = builder.with_column_key(AES_128_COLUMN_NAMES[1], column_2_key.to_vec()); } for (column_name, key) in additional_column_keys { @@ -157,9 +142,9 @@ async fn test_misspecified_encryption_keys() { // There is always a footer key even with a plaintext footer, // but this is used for signing the footer. - let footer_key = "0123456789012345".as_bytes(); // 128bit/16 - let column_1_key = "1234567890123450".as_bytes(); - let column_2_key = "1234567890123451".as_bytes(); + let footer_key = AES_128_FOOTER_KEY; + let column_1_key = AES_128_COLUMN_KEYS[0]; + let column_2_key = AES_128_COLUMN_KEYS[1]; let empty_column_key = &[]; // Too short footer key @@ -175,7 +160,7 @@ async fn test_misspecified_encryption_keys() { // Wrong footer key check_for_error( "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer", - b"1123456789012345", + BAD_AES_128_FOOTER_KEY, column_1_key, column_2_key, empty_column_key, @@ -226,28 +211,16 @@ async fn test_misspecified_encryption_keys() { ) .await; - let aes256_footer_key = "01234567890123456789012345678901".as_bytes(); // 256bit/32 - let aes256_column_1_key = "12345678901234567890123456789012".as_bytes(); - let aes256_column_2_key = "12345678901234567890123456789013".as_bytes(); + let aes256_footer_key = AES_256_FOOTER_KEY; + let aes256_column_1_key = AES_256_COLUMN_KEYS[0]; + let aes256_column_2_key = AES_256_COLUMN_KEYS[1]; let additional_column_keys = &[ - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), ]; // Too short footer key @@ -262,7 +235,7 @@ async fn test_misspecified_encryption_keys() { // Wrong footer key check_for_error( "Parquet error: Provided footer key and AAD were unable to decrypt parquet footer", - b"11234567890123456789012345678901", + BAD_AES_256_FOOTER_KEY, aes256_column_1_key, aes256_column_2_key, additional_column_keys, @@ -343,44 +316,26 @@ async fn test_non_uniform_encryption() { // AES-128 non_uniform_encryption( - b"0123456789012345", + AES_128_FOOTER_KEY, &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ) .await; // AES-256 non_uniform_encryption( - b"01234567890123456789012345678901", // 256bit/32 + AES_256_FOOTER_KEY, &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ) .await; @@ -388,18 +343,16 @@ async fn test_non_uniform_encryption() { #[tokio::test] async fn test_uniform_encryption() { - async fn uniform_encryption(footer_key: &[u8], column_keys: &[(&str, &[u8])]) { + async fn uniform_encryption(footer_key: &[u8]) { let path = encryption_util::encrypted_data_path( footer_key, "uniform_encryption.parquet.encrypted", ); let mut file = File::open(&path).await.unwrap(); - let mut builder = FileDecryptionProperties::builder(footer_key.to_vec()); - for (column_name, key) in column_keys { - builder = builder.with_column_key(column_name, key.to_vec()); - } - let decryption_properties = builder.build().unwrap(); + let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec()) + .build() + .unwrap(); verify_encryption_test_file_read_async(&mut file, decryption_properties) .await @@ -408,18 +361,10 @@ async fn test_uniform_encryption() { // AES-128: there is always a footer key even with a plaintext footer, // but this is used for signing the footer. - uniform_encryption( - b"0123456789012345", // 128bit/16 - &[], - ) - .await; + uniform_encryption(AES_128_FOOTER_KEY).await; // AES-256 - uniform_encryption( - b"01234567890123456789012345678901", // 256bit/32 - &[], - ) - .await; + uniform_encryption(AES_256_FOOTER_KEY).await; } #[tokio::test] @@ -453,26 +398,26 @@ async fn test_aes_ctr_encryption() { // AES-128 aes_ctr_encryption( - b"0123456789012345", // 128bit/16 + AES_128_FOOTER_KEY, &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ) .await; // AES-256 aes_ctr_encryption( - b"01234567890123456789012345678901", // 256bit/32 + AES_256_FOOTER_KEY, &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ) .await; @@ -528,12 +473,12 @@ async fn test_write_non_uniform_encryption() { } write_non_uniform_encryption( - b"0123456789012345", // 128bit/16, - vec!["double_field", "float_field"], - vec![b"1234567890123450".to_vec(), b"1234567890123451".to_vec()], + AES_128_FOOTER_KEY, + AES_128_COLUMN_NAMES.to_vec(), + AES_128_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ) .await; @@ -541,56 +486,18 @@ async fn test_write_non_uniform_encryption() { // AES-256 // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name write_non_uniform_encryption( - b"01234567890123456789012345678901", - vec![ - "double_field", - "float_field", - "boolean_field", - "int32_field", - "ba_field", - "flba_field", - "int64_field", - "int64_field.list.int64_field", - "int96_field", - ], - vec![ - b"12345678901234567890123456789012".to_vec(), - b"12345678901234567890123456789013".to_vec(), - b"12345678901234567890123456789014".to_vec(), - b"12345678901234567890123456789015".to_vec(), - b"12345678901234567890123456789016".to_vec(), - b"12345678901234567890123456789017".to_vec(), - b"12345678901234567890123456789018".to_vec(), - b"12345678901234567890123456789018".to_vec(), - b"12345678901234567890123456789019".to_vec(), - ], + AES_256_FOOTER_KEY, + AES_256_COLUMN_NAMES.to_vec(), + AES_256_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field.list.int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), + (AES_256_COLUMN_NAMES[8], AES_256_COLUMN_KEYS[8]), ], ) .await; @@ -669,28 +576,28 @@ async fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { // AES-128 non_uniform_encryption_plaintext_footer_with_key_retriever( - b"0123456789012345", + AES_128_FOOTER_KEY, &[ - ("kf", b"0123456789012345".as_slice()), - ("kc1", b"1234567890123450".as_slice()), - ("kc2", b"1234567890123451".as_slice()), + (AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY), + (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ) .await; // AES-256 non_uniform_encryption_plaintext_footer_with_key_retriever( - b"01234567890123456789012345678901", + AES_256_FOOTER_KEY, &[ - ("kf", b"01234567890123456789012345678901".as_slice()), - ("kc1", b"12345678901234567890123456789012".as_slice()), - ("kc2", b"12345678901234567890123456789013".as_slice()), - ("kc3", b"12345678901234567890123456789014".as_slice()), - ("kc4", b"12345678901234567890123456789015".as_slice()), - ("kc5", b"12345678901234567890123456789016".as_slice()), - ("kc6", b"12345678901234567890123456789017".as_slice()), - ("kc7", b"12345678901234567890123456789018".as_slice()), - ("kc8", b"12345678901234567890123456789019".as_slice()), + (AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY), + (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ) .await; @@ -722,28 +629,28 @@ async fn test_non_uniform_encryption_with_key_retriever() { // AES-128 non_uniform_encryption_with_key_retriever( - b"0123456789012345", + AES_128_FOOTER_KEY, &[ - ("kf", b"0123456789012345".as_slice()), - ("kc1", b"1234567890123450".as_slice()), - ("kc2", b"1234567890123451".as_slice()), + (AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY), + (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ) .await; // AES-256 non_uniform_encryption_with_key_retriever( - b"01234567890123456789012345678901", + AES_256_FOOTER_KEY, &[ - ("kf", b"01234567890123456789012345678901".as_slice()), - ("kc1", b"12345678901234567890123456789012".as_slice()), - ("kc2", b"12345678901234567890123456789013".as_slice()), - ("kc3", b"12345678901234567890123456789014".as_slice()), - ("kc4", b"12345678901234567890123456789015".as_slice()), - ("kc5", b"12345678901234567890123456789016".as_slice()), - ("kc6", b"12345678901234567890123456789017".as_slice()), - ("kc7", b"12345678901234567890123456789018".as_slice()), - ("kc8", b"12345678901234567890123456789019".as_slice()), + (AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY), + (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ) .await; @@ -771,10 +678,10 @@ async fn test_uniform_encryption_with_key_retriever() { .unwrap(); } // AES-128 - uniform_encryption_with_key_retriever("kf", b"0123456789012345").await; + uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY).await; // AES-256 - uniform_encryption_with_key_retriever("kf", b"01234567890123456789012345678901").await; + uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY).await; } #[tokio::test] @@ -793,8 +700,8 @@ async fn test_decrypt_page_index_uniform() { .unwrap(); } - decrypt_page_index_uniform("0123456789012345".as_bytes()).await; - decrypt_page_index_uniform("01234567890123456789012345678901".as_bytes()).await; + decrypt_page_index_uniform(AES_128_FOOTER_KEY).await; + decrypt_page_index_uniform(AES_256_FOOTER_KEY).await; } #[tokio::test] @@ -815,43 +722,25 @@ async fn test_decrypt_page_index_non_uniform() { } decrypt_page_index_non_uniform( - b"0123456789012345", // 128bit/16 + AES_128_FOOTER_KEY, &[ - ("double_field", b"1234567890123450".as_slice()), - ("float_field", b"1234567890123451".as_slice()), + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), ], ) .await; decrypt_page_index_non_uniform( - b"01234567890123456789012345678901", // 256bit/32 + AES_256_FOOTER_KEY, &[ - ( - "double_field", - b"12345678901234567890123456789012".as_slice(), - ), - ( - "float_field", - b"12345678901234567890123456789013".as_slice(), - ), - ( - "boolean_field", - b"12345678901234567890123456789014".as_slice(), - ), - ( - "int32_field", - b"12345678901234567890123456789015".as_slice(), - ), - ("ba_field", b"12345678901234567890123456789016".as_slice()), - ("flba_field", b"12345678901234567890123456789017".as_slice()), - ( - "int64_field", - b"12345678901234567890123456789018".as_slice(), - ), - ( - "int96_field", - b"12345678901234567890123456789019".as_slice(), - ), + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), ], ) .await; @@ -1100,14 +989,14 @@ async fn test_concurrent_encrypted_writing_over_multiple_row_groups() { let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); let file = std::fs::File::open(path).unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let file_encryption_properties = FileEncryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let decryption_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); @@ -1173,14 +1062,14 @@ async fn test_multi_threaded_encrypted_writing() { let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); let file = std::fs::File::open(path).unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let file_encryption_properties = FileEncryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let decryption_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); @@ -1274,14 +1163,14 @@ async fn test_multi_threaded_encrypted_writing_deprecated() { let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted"); let file = std::fs::File::open(path).unwrap(); - let file_encryption_properties = FileEncryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let file_encryption_properties = FileEncryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(b"0123456789012345".into()) - .with_column_key("double_field", b"1234567890123450".into()) - .with_column_key("float_field", b"1234567890123451".into()) + let decryption_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.into()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].into()) + .with_column_key(AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1].into()) .build() .unwrap(); diff --git a/parquet/tests/encryption/encryption_util.rs b/parquet/tests/encryption/encryption_util.rs index a940a47bbd69..ec495e6c9f04 100644 --- a/parquet/tests/encryption/encryption_util.rs +++ b/parquet/tests/encryption/encryption_util.rs @@ -31,6 +31,43 @@ use std::collections::HashMap; use std::fs::File; use std::sync::{Arc, Mutex}; +pub(crate) const AES_128_FOOTER_KEY: &[u8; 16] = b"0123456789012345"; // 128bit/16 +pub(crate) const BAD_AES_128_FOOTER_KEY: &[u8; 16] = b"0000000000000000"; +pub(crate) const AES_128_FOOTER_KEY_NAME: &str = "kf"; +pub(crate) const AES_128_KEY_NAMES: [&str; 2] = ["kc1", "kc2"]; +pub(crate) const AES_128_COLUMN_NAMES: [&str; 2] = ["double_field", "float_field"]; +pub(crate) const AES_128_COLUMN_KEYS: [&[u8; 16]; 2] = [b"1234567890123450", b"1234567890123451"]; + +pub(crate) const AES_256_FOOTER_KEY: &[u8; 32] = b"01234567890123456789012345678901"; // 256bit/32 +pub(crate) const BAD_AES_256_FOOTER_KEY: &[u8; 32] = b"00000000000000000000000000000000"; +pub(crate) const AES_256_FOOTER_KEY_NAME: &str = "kf"; +pub(crate) const AES_256_KEY_NAMES: [&str; 8] = + ["kc1", "kc2", "kc3", "kc4", "kc5", "kc6", "kc7", "kc8"]; + +// TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name +pub(crate) const AES_256_COLUMN_NAMES: [&str; 9] = [ + "double_field", + "float_field", + "boolean_field", + "int32_field", + "ba_field", + "flba_field", + "int64_field", + "int96_field", + "int64_field.list.int64_field", +]; +pub(crate) const AES_256_COLUMN_KEYS: [&[u8]; 9] = [ + b"12345678901234567890123456789012", + b"12345678901234567890123456789013", + b"12345678901234567890123456789014", + b"12345678901234567890123456789015", + b"12345678901234567890123456789016", + b"12345678901234567890123456789017", + b"12345678901234567890123456789018", + b"12345678901234567890123456789019", + b"12345678901234567890123456789018", +]; + pub(crate) fn verify_encryption_double_test_data( record_batches: Vec, metadata: &ParquetMetaData, From 4336eed735382f3b98a01bc5376d88741b31b82e Mon Sep 17 00:00:00 2001 From: hsiang-c Date: Mon, 4 May 2026 10:01:19 -0700 Subject: [PATCH 21/26] Clean up --- parquet/tests/encryption/encryption.rs | 157 +++++++++---------- parquet/tests/encryption/encryption_async.rs | 3 +- 2 files changed, 71 insertions(+), 89 deletions(-) diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index 4a20d36421c7..8e259c14e030 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -297,8 +297,7 @@ fn test_uniform_encryption() { verify_encryption_test_file_read(file, decryption_properties); } - // AES-128: there is always a footer key even with a plaintext footer, - // but this is used for signing the footer. + // AES-128 uniform_encryption(AES_128_FOOTER_KEY); // AES-256 @@ -426,8 +425,8 @@ fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { } #[test] -fn test_uniform_encryption_plaintext_footer_with_key_retriever() { - fn uniform_encryption_plaintext_footer_with_key_retriever( +fn test_roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever() { + fn roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( footer_key: &[u8], footer_key_metadata: &str, wrong_footer_key: &[u8], @@ -521,7 +520,7 @@ fn test_uniform_encryption_plaintext_footer_with_key_retriever() { } // AES-128 - uniform_encryption_plaintext_footer_with_key_retriever( + roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( AES_128_FOOTER_KEY, AES_128_FOOTER_KEY_NAME, BAD_AES_128_FOOTER_KEY, @@ -545,7 +544,7 @@ fn test_uniform_encryption_plaintext_footer_with_key_retriever() { // AES-256 // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name - uniform_encryption_plaintext_footer_with_key_retriever( + roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( AES_256_FOOTER_KEY, AES_256_FOOTER_KEY_NAME, BAD_AES_256_FOOTER_KEY, @@ -687,25 +686,21 @@ fn row_group_sizes(metadata: &ParquetMetaData) -> Vec { #[test] fn test_uniform_encryption_roundtrip() { uniform_encryption_roundtrip(false, false, AES_128_FOOTER_KEY).unwrap(); - uniform_encryption_roundtrip(false, false, AES_256_FOOTER_KEY).unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_dictionary() { uniform_encryption_roundtrip(false, true, AES_128_FOOTER_KEY).unwrap(); - uniform_encryption_roundtrip(false, true, AES_256_FOOTER_KEY).unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_page_index() { uniform_encryption_roundtrip(true, false, AES_128_FOOTER_KEY).unwrap(); - uniform_encryption_roundtrip(true, false, AES_256_FOOTER_KEY).unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_page_index_and_dictionary() { uniform_encryption_roundtrip(true, true, AES_128_FOOTER_KEY).unwrap(); - uniform_encryption_roundtrip(true, true, AES_256_FOOTER_KEY).unwrap(); } fn uniform_encryption_roundtrip( @@ -809,13 +804,11 @@ fn uniform_encryption_roundtrip( #[test] fn test_uniform_encryption_page_skipping() { uniform_encryption_page_skipping(false, AES_128_FOOTER_KEY).unwrap(); - uniform_encryption_page_skipping(false, AES_256_FOOTER_KEY).unwrap(); } #[test] fn test_uniform_encryption_page_skipping_with_page_index() { uniform_encryption_page_skipping(true, AES_128_FOOTER_KEY).unwrap(); - uniform_encryption_page_skipping(true, AES_256_FOOTER_KEY).unwrap(); } fn uniform_encryption_page_skipping( @@ -1063,40 +1056,36 @@ fn test_write_uniform_encryption_plaintext_footer() { #[test] pub fn test_column_statistics_with_plaintext_footer() { - let keys = [ - (AES_128_FOOTER_KEY.to_vec(), AES_128_COLUMN_KEYS[0].to_vec()), - (AES_256_FOOTER_KEY.to_vec(), AES_256_COLUMN_KEYS[0].to_vec()), - ]; + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key = AES_128_COLUMN_KEYS[0].to_vec(); - for (footer_key, column_key) in &keys { - // Encrypt with a plaintext footer and column-specific keys - let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) - .with_plaintext_footer(true) - .with_column_key("x", column_key.clone()) - .with_column_key("y", column_key.clone()) - .with_column_key("s", column_key.clone()) - .build() - .unwrap(); + // Encrypt with a plaintext footer and column-specific keys + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(true) + .with_column_key("x", column_key.clone()) + .with_column_key("y", column_key.clone()) + .with_column_key("s", column_key.clone()) + .build() + .unwrap(); - // Read with only the footer key and the key for one column - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("x", column_key.clone()) - .build() - .unwrap(); + // Read with only the footer key and the key for one column + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_key("x", column_key.clone()) + .build() + .unwrap(); - // Reader can read plaintext stats from the unencrypted column z - // and column x for which the key is provided, but not columns y and s - // for which no key is provided. - write_and_read_stats( - Arc::clone(&encryption_properties), - Some(decryption_properties), - &[true, false, true, false], - ); + // Reader can read plaintext stats from the unencrypted column z + // and column x for which the key is provided, but not columns y and s + // for which no key is provided. + write_and_read_stats( + Arc::clone(&encryption_properties), + Some(decryption_properties), + &[true, false, true, false], + ); - // Read without any decryption properties. - // Reader can only read plaintext stats from the unencrypted column z. - write_and_read_stats(encryption_properties, None, &[false, false, true, false]); - } + // Read without any decryption properties. + // Reader can only read plaintext stats from the unencrypted column z. + write_and_read_stats(encryption_properties, None, &[false, false, true, false]); } #[test] @@ -1128,61 +1117,55 @@ pub fn test_column_statistics_with_plaintext_footer_and_uniform_encryption() { #[test] pub fn test_column_statistics_with_encrypted_footer() { - let keys = [ - (AES_128_FOOTER_KEY.to_vec(), AES_128_COLUMN_KEYS[0].to_vec()), - (AES_256_FOOTER_KEY.to_vec(), AES_256_COLUMN_KEYS[0].to_vec()), - ]; + let footer_key = AES_128_FOOTER_KEY.to_vec(); + let column_key = AES_128_COLUMN_KEYS[0].to_vec(); - for (footer_key, column_key) in &keys { - // Encrypt with an encrypted footer and column-specific keys - let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) - .with_plaintext_footer(false) - .with_column_key("x", column_key.clone()) - .with_column_key("y", column_key.clone()) - .with_column_key("s", column_key.clone()) - .build() - .unwrap(); + // Encrypt with an encrypted footer and column-specific keys + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(false) + .with_column_key("x", column_key.clone()) + .with_column_key("y", column_key.clone()) + .with_column_key("s", column_key.clone()) + .build() + .unwrap(); - // Read with only the footer key and the key for one column - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .with_column_key("x", column_key.clone()) - .build() - .unwrap(); + // Read with only the footer key and the key for one column + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .with_column_key("x", column_key.clone()) + .build() + .unwrap(); - // Reader can read plaintext stats from the unencrypted column z - // and column x for which the key is provided, but not columns y and s - // for which no key is provided. - write_and_read_stats( - encryption_properties, - Some(decryption_properties), - &[true, false, true, false], - ); - } + // Reader can read plaintext stats from the unencrypted column z + // and column x for which the key is provided, but not columns y and s + // for which no key is provided. + write_and_read_stats( + encryption_properties, + Some(decryption_properties), + &[true, false, true, false], + ); } #[test] pub fn test_column_statistics_with_encrypted_footer_and_uniform_encryption() { - let footer_keys = [AES_128_FOOTER_KEY.to_vec(), AES_256_FOOTER_KEY.to_vec()]; + let footer_key = AES_128_FOOTER_KEY.to_vec(); - for footer_key in &footer_keys { - // Encrypt with an encrypted footer and uniform encryption - let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) - .with_plaintext_footer(false) - .build() - .unwrap(); + // Encrypt with an encrypted footer and uniform encryption + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(false) + .build() + .unwrap(); - // Read with the footer key - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .build() - .unwrap(); + // Read with the footer key + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .build() + .unwrap(); - // Reader can read stats for all columns. - write_and_read_stats( - encryption_properties, - Some(decryption_properties), - &[true, true, true, true], - ); - } + // Reader can read stats for all columns. + write_and_read_stats( + encryption_properties, + Some(decryption_properties), + &[true, true, true, true], + ); } /// Write a file with encryption and then verify whether statistics are readable with the provided decryption properties. diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index be8a75590f44..29230ff71b80 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -359,8 +359,7 @@ async fn test_uniform_encryption() { .unwrap(); } - // AES-128: there is always a footer key even with a plaintext footer, - // but this is used for signing the footer. + // AES-128 uniform_encryption(AES_128_FOOTER_KEY).await; // AES-256 From 44caeccad8dd75c031b7545e70bfa3c14f8da56c Mon Sep 17 00:00:00 2001 From: hsiang-c Date: Mon, 4 May 2026 10:35:30 -0700 Subject: [PATCH 22/26] Revert uncessary changes --- parquet/tests/encryption/encryption.rs | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index 8e259c14e030..734ce8a2151f 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -685,28 +685,27 @@ fn row_group_sizes(metadata: &ParquetMetaData) -> Vec { #[test] fn test_uniform_encryption_roundtrip() { - uniform_encryption_roundtrip(false, false, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_roundtrip(false, false).unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_dictionary() { - uniform_encryption_roundtrip(false, true, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_roundtrip(false, true).unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_page_index() { - uniform_encryption_roundtrip(true, false, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_roundtrip(true, false).unwrap(); } #[test] fn test_uniform_encryption_roundtrip_with_page_index_and_dictionary() { - uniform_encryption_roundtrip(true, true, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_roundtrip(true, true).unwrap(); } fn uniform_encryption_roundtrip( page_index: bool, dictionary_encoding: bool, - footer_key: &[u8], ) -> parquet::errors::Result<()> { let x0_arrays = [ Int32Array::from((0..100).collect::>()), @@ -724,6 +723,7 @@ fn uniform_encryption_roundtrip( let file = tempfile::tempfile()?; + let footer_key = AES_128_FOOTER_KEY; let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()).build()?; @@ -803,18 +803,15 @@ fn uniform_encryption_roundtrip( #[test] fn test_uniform_encryption_page_skipping() { - uniform_encryption_page_skipping(false, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_page_skipping(false).unwrap(); } #[test] fn test_uniform_encryption_page_skipping_with_page_index() { - uniform_encryption_page_skipping(true, AES_128_FOOTER_KEY).unwrap(); + uniform_encryption_page_skipping(true).unwrap(); } -fn uniform_encryption_page_skipping( - page_index: bool, - footer_key: &[u8], -) -> parquet::errors::Result<()> { +fn uniform_encryption_page_skipping(page_index: bool) -> parquet::errors::Result<()> { let x0_arrays = [ Int32Array::from((0..100).collect::>()), Int32Array::from((100..150).collect::>()), @@ -831,6 +828,7 @@ fn uniform_encryption_page_skipping( let file = tempfile::tempfile()?; + let footer_key = AES_128_FOOTER_KEY; let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec()).build()?; From 12d768d9625917904dd23f9dceceefa401f74d86 Mon Sep 17 00:00:00 2001 From: hsiang-c Date: Mon, 4 May 2026 11:17:12 -0700 Subject: [PATCH 23/26] Extract common data --- parquet/tests/encryption/encryption.rs | 219 ++----------------- parquet/tests/encryption/encryption_async.rs | 173 ++------------- parquet/tests/encryption/encryption_util.rs | 34 +++ 3 files changed, 76 insertions(+), 350 deletions(-) diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index 734ce8a2151f..d1e1f4812114 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -19,9 +19,10 @@ use crate::encryption_util; use crate::encryption_util::{ - AES_128_COLUMN_KEYS, AES_128_COLUMN_NAMES, AES_128_FOOTER_KEY, AES_128_FOOTER_KEY_NAME, - AES_128_KEY_NAMES, AES_256_COLUMN_KEYS, AES_256_COLUMN_NAMES, AES_256_FOOTER_KEY, - AES_256_FOOTER_KEY_NAME, AES_256_KEY_NAMES, BAD_AES_128_FOOTER_KEY, BAD_AES_256_FOOTER_KEY, + AES_128_COLUMN_KEYS, AES_128_COLUMN_NAME_KEYS, AES_128_COLUMN_NAMES, AES_128_FOOTER_KEY, + AES_128_FOOTER_KEY_NAME, AES_128_KEY_NAME_KEY, AES_128_KEY_NAMES, AES_256_COLUMN_KEYS, + AES_256_COLUMN_NAME_KEYS, AES_256_COLUMN_NAMES, AES_256_FOOTER_KEY, AES_256_FOOTER_KEY_NAME, + AES_256_KEY_NAME_KEY, AES_256_KEY_NAMES, BAD_AES_128_FOOTER_KEY, BAD_AES_256_FOOTER_KEY, TestKeyRetriever, read_and_roundtrip_to_encrypted_file, verify_column_indexes, verify_encryption_test_file_read, }; @@ -63,28 +64,10 @@ fn test_non_uniform_encryption_plaintext_footer() { // AES-128: there is always a footer key even with a plaintext footer, // but this is used for signing the footer. - non_uniform_encryption_plaintext_footer( - AES_128_FOOTER_KEY, - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ); + non_uniform_encryption_plaintext_footer(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); // AES-256 - non_uniform_encryption_plaintext_footer( - AES_256_FOOTER_KEY, - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ); + non_uniform_encryption_plaintext_footer(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); } #[test] @@ -127,27 +110,8 @@ fn test_plaintext_footer_signature_verification() { ); } - plaintext_footer_signature_verification( - BAD_AES_128_FOOTER_KEY, - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ); - - plaintext_footer_signature_verification( - BAD_AES_256_FOOTER_KEY, - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ) + plaintext_footer_signature_verification(BAD_AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + plaintext_footer_signature_verification(BAD_AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS) } #[test] @@ -207,29 +171,8 @@ fn test_non_uniform_encryption_disabled_aad_storage() { ); } - // AES-128 - non_uniform_encryption_disabled_aad_storage( - AES_128_FOOTER_KEY, - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ); - - // AES-256 - non_uniform_encryption_disabled_aad_storage( - AES_256_FOOTER_KEY, - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ); + non_uniform_encryption_disabled_aad_storage(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + non_uniform_encryption_disabled_aad_storage(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); } #[test] @@ -256,29 +199,8 @@ fn test_non_uniform_encryption() { verify_encryption_test_file_read(file, decryption_properties); } - // AES-128 - non_uniform_encryption( - AES_128_FOOTER_KEY, - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ); - - // AES-256 - non_uniform_encryption( - AES_256_FOOTER_KEY, - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ); + non_uniform_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + non_uniform_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); } #[test] @@ -297,10 +219,7 @@ fn test_uniform_encryption() { verify_encryption_test_file_read(file, decryption_properties); } - // AES-128 uniform_encryption(AES_128_FOOTER_KEY); - - // AES-256 uniform_encryption(AES_256_FOOTER_KEY); } @@ -347,29 +266,9 @@ fn test_aes_ctr_encryption() { } }; } - // AES-128 - aes_ctr_encryption( - AES_128_FOOTER_KEY, - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ); - // AES-256 - aes_ctr_encryption( - AES_256_FOOTER_KEY, - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ); + aes_ctr_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); + aes_ctr_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); } #[test] @@ -397,30 +296,14 @@ fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { verify_encryption_test_file_read(file, decryption_properties); } - // AES-128 non_uniform_encryption_plaintext_footer_with_key_retriever( AES_128_FOOTER_KEY, - &[ - (AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY), - (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], + AES_128_KEY_NAME_KEY, ); - // AES-256 non_uniform_encryption_plaintext_footer_with_key_retriever( AES_256_FOOTER_KEY, - &[ - (AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY), - (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], + AES_256_KEY_NAME_KEY, ); } @@ -519,7 +402,6 @@ fn test_roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever() { ); } - // AES-128 roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( AES_128_FOOTER_KEY, AES_128_FOOTER_KEY_NAME, @@ -542,7 +424,6 @@ fn test_roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever() { ], ); - // AES-256 // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( AES_256_FOOTER_KEY, @@ -625,31 +506,9 @@ fn test_non_uniform_encryption_with_key_retriever() { verify_encryption_test_file_read(file, decryption_properties); } - // AES-128 - non_uniform_encryption_with_key_retriever( - AES_128_FOOTER_KEY, - &[ - (AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY), - (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ); + non_uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY, AES_128_KEY_NAME_KEY); - // AES-256 - non_uniform_encryption_with_key_retriever( - AES_256_FOOTER_KEY, - &[ - (AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY), - (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ); + non_uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY, AES_256_KEY_NAME_KEY); } #[test] @@ -672,10 +531,7 @@ fn test_uniform_encryption_with_key_retriever() { verify_encryption_test_file_read(file, decryption_properties); } - // AES-128 uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY); - - // AES-256 uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY); } @@ -950,10 +806,7 @@ fn test_write_non_uniform_encryption() { AES_128_FOOTER_KEY, AES_128_COLUMN_NAMES.to_vec(), AES_128_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], + AES_128_COLUMN_NAME_KEYS, ); // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name @@ -1347,20 +1200,12 @@ fn test_write_non_uniform_encryption_column_missmatch() { ); } - // AES-128 write_non_uniform_encryption_column_missmatch( AES_128_FOOTER_KEY, - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], + AES_128_COLUMN_NAME_KEYS, + AES_128_COLUMN_NAME_KEYS, ); - // AES-256 // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name write_non_uniform_encryption_column_missmatch( AES_256_FOOTER_KEY, @@ -1641,27 +1486,9 @@ fn test_decrypt_page_index_non_uniform() { test_decrypt_page_index(&path, decryption_properties).unwrap(); } - decrypt_page_index_non_uniform( - AES_128_FOOTER_KEY, - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ); + decrypt_page_index_non_uniform(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS); - decrypt_page_index_non_uniform( - AES_256_FOOTER_KEY, - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ); + decrypt_page_index_non_uniform(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS); } fn test_decrypt_page_index( diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index 29230ff71b80..8832922bab64 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -19,11 +19,11 @@ use crate::encryption_util; use crate::encryption_util::{ - AES_128_COLUMN_KEYS, AES_128_COLUMN_NAMES, AES_128_FOOTER_KEY, AES_128_FOOTER_KEY_NAME, - AES_128_KEY_NAMES, AES_256_COLUMN_KEYS, AES_256_COLUMN_NAMES, AES_256_FOOTER_KEY, - AES_256_FOOTER_KEY_NAME, AES_256_KEY_NAMES, BAD_AES_128_FOOTER_KEY, BAD_AES_256_FOOTER_KEY, - TestKeyRetriever, read_encrypted_file, verify_column_indexes, - verify_encryption_double_test_data, verify_encryption_test_data, + AES_128_COLUMN_KEYS, AES_128_COLUMN_NAME_KEYS, AES_128_COLUMN_NAMES, AES_128_FOOTER_KEY, + AES_128_FOOTER_KEY_NAME, AES_128_KEY_NAME_KEY, AES_256_COLUMN_KEYS, AES_256_COLUMN_NAME_KEYS, + AES_256_COLUMN_NAMES, AES_256_FOOTER_KEY, AES_256_FOOTER_KEY_NAME, AES_256_KEY_NAME_KEY, + BAD_AES_128_FOOTER_KEY, BAD_AES_256_FOOTER_KEY, TestKeyRetriever, read_encrypted_file, + verify_column_indexes, verify_encryption_double_test_data, verify_encryption_test_data, }; use arrow_array::RecordBatch; use arrow_schema::Schema; @@ -72,30 +72,10 @@ async fn test_non_uniform_encryption_plaintext_footer() { // AES-128: there is always a footer key even with a plaintext footer, // but this is used for signing the footer. - non_uniform_encryption_plaintext_footer( - AES_128_FOOTER_KEY, - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ) - .await; + non_uniform_encryption_plaintext_footer(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; // AES-256 - non_uniform_encryption_plaintext_footer( - AES_256_FOOTER_KEY, - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ) - .await; + non_uniform_encryption_plaintext_footer(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; } #[tokio::test] @@ -314,31 +294,8 @@ async fn test_non_uniform_encryption() { .unwrap(); } - // AES-128 - non_uniform_encryption( - AES_128_FOOTER_KEY, - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ) - .await; - - // AES-256 - non_uniform_encryption( - AES_256_FOOTER_KEY, - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ) - .await; + non_uniform_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; + non_uniform_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; } #[tokio::test] @@ -359,10 +316,7 @@ async fn test_uniform_encryption() { .unwrap(); } - // AES-128 uniform_encryption(AES_128_FOOTER_KEY).await; - - // AES-256 uniform_encryption(AES_256_FOOTER_KEY).await; } @@ -395,31 +349,8 @@ async fn test_aes_ctr_encryption() { }; } - // AES-128 - aes_ctr_encryption( - AES_128_FOOTER_KEY, - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ) - .await; - - // AES-256 - aes_ctr_encryption( - AES_256_FOOTER_KEY, - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ) - .await; + aes_ctr_encryption(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; + aes_ctr_encryption(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; } #[tokio::test] @@ -475,10 +406,7 @@ async fn test_write_non_uniform_encryption() { AES_128_FOOTER_KEY, AES_128_COLUMN_NAMES.to_vec(), AES_128_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], + AES_128_COLUMN_NAME_KEYS, ) .await; @@ -573,31 +501,15 @@ async fn test_non_uniform_encryption_plaintext_footer_with_key_retriever() { .unwrap(); } - // AES-128 non_uniform_encryption_plaintext_footer_with_key_retriever( AES_128_FOOTER_KEY, - &[ - (AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY), - (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], + AES_128_KEY_NAME_KEY, ) .await; - // AES-256 non_uniform_encryption_plaintext_footer_with_key_retriever( AES_256_FOOTER_KEY, - &[ - (AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY), - (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], + AES_256_KEY_NAME_KEY, ) .await; } @@ -626,33 +538,8 @@ async fn test_non_uniform_encryption_with_key_retriever() { .unwrap(); } - // AES-128 - non_uniform_encryption_with_key_retriever( - AES_128_FOOTER_KEY, - &[ - (AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY), - (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ) - .await; - - // AES-256 - non_uniform_encryption_with_key_retriever( - AES_256_FOOTER_KEY, - &[ - (AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY), - (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ) - .await; + non_uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY, AES_128_KEY_NAME_KEY).await; + non_uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY, AES_256_KEY_NAME_KEY).await; } #[tokio::test] @@ -676,10 +563,8 @@ async fn test_uniform_encryption_with_key_retriever() { .await .unwrap(); } - // AES-128 - uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY).await; - // AES-256 + uniform_encryption_with_key_retriever(AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY).await; uniform_encryption_with_key_retriever(AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY).await; } @@ -720,29 +605,9 @@ async fn test_decrypt_page_index_non_uniform() { .unwrap(); } - decrypt_page_index_non_uniform( - AES_128_FOOTER_KEY, - &[ - (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), - (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), - ], - ) - .await; + decrypt_page_index_non_uniform(AES_128_FOOTER_KEY, AES_128_COLUMN_NAME_KEYS).await; - decrypt_page_index_non_uniform( - AES_256_FOOTER_KEY, - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - ], - ) - .await; + decrypt_page_index_non_uniform(AES_256_FOOTER_KEY, AES_256_COLUMN_NAME_KEYS).await; } async fn test_decrypt_page_index( diff --git a/parquet/tests/encryption/encryption_util.rs b/parquet/tests/encryption/encryption_util.rs index ec495e6c9f04..1c2322cebe30 100644 --- a/parquet/tests/encryption/encryption_util.rs +++ b/parquet/tests/encryption/encryption_util.rs @@ -38,6 +38,17 @@ pub(crate) const AES_128_KEY_NAMES: [&str; 2] = ["kc1", "kc2"]; pub(crate) const AES_128_COLUMN_NAMES: [&str; 2] = ["double_field", "float_field"]; pub(crate) const AES_128_COLUMN_KEYS: [&[u8; 16]; 2] = [b"1234567890123450", b"1234567890123451"]; +pub(crate) const AES_128_COLUMN_NAME_KEYS: &[(&str, &[u8]); 2] = &[ + (AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_COLUMN_NAMES[1], AES_128_COLUMN_KEYS[1]), +]; + +pub(crate) const AES_128_KEY_NAME_KEY: &[(&str, &[u8]); 3] = &[ + (AES_128_FOOTER_KEY_NAME, AES_128_FOOTER_KEY), + (AES_128_KEY_NAMES[0], AES_128_COLUMN_KEYS[0]), + (AES_128_KEY_NAMES[1], AES_128_COLUMN_KEYS[1]), +]; + pub(crate) const AES_256_FOOTER_KEY: &[u8; 32] = b"01234567890123456789012345678901"; // 256bit/32 pub(crate) const BAD_AES_256_FOOTER_KEY: &[u8; 32] = b"00000000000000000000000000000000"; pub(crate) const AES_256_FOOTER_KEY_NAME: &str = "kf"; @@ -68,6 +79,29 @@ pub(crate) const AES_256_COLUMN_KEYS: [&[u8]; 9] = [ b"12345678901234567890123456789018", ]; +pub(crate) const AES_256_COLUMN_NAME_KEYS: &[(&str, &[u8]); 8] = &[ + (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_COLUMN_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), +]; + +pub(crate) const AES_256_KEY_NAME_KEY: &[(&str, &[u8]); 9] = &[ + (AES_256_FOOTER_KEY_NAME, AES_256_FOOTER_KEY), + (AES_256_KEY_NAMES[0], AES_256_COLUMN_KEYS[0]), + (AES_256_KEY_NAMES[1], AES_256_COLUMN_KEYS[1]), + (AES_256_KEY_NAMES[2], AES_256_COLUMN_KEYS[2]), + (AES_256_KEY_NAMES[3], AES_256_COLUMN_KEYS[3]), + (AES_256_KEY_NAMES[4], AES_256_COLUMN_KEYS[4]), + (AES_256_KEY_NAMES[5], AES_256_COLUMN_KEYS[5]), + (AES_256_KEY_NAMES[6], AES_256_COLUMN_KEYS[6]), + (AES_256_KEY_NAMES[7], AES_256_COLUMN_KEYS[7]), +]; + pub(crate) fn verify_encryption_double_test_data( record_batches: Vec, metadata: &ParquetMetaData, From 06c6b73b4c3b9b570dfd82eb2ba8f0f09a184e5f Mon Sep 17 00:00:00 2001 From: hsiang-c Date: Mon, 4 May 2026 11:27:15 -0700 Subject: [PATCH 24/26] Restore one test --- parquet/tests/encryption/encryption.rs | 36 ++++++++++++-------------- 1 file changed, 17 insertions(+), 19 deletions(-) diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index d1e1f4812114..1848f80722c5 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -941,29 +941,27 @@ pub fn test_column_statistics_with_plaintext_footer() { #[test] pub fn test_column_statistics_with_plaintext_footer_and_uniform_encryption() { - let footer_keys = [AES_128_FOOTER_KEY.to_vec(), AES_256_FOOTER_KEY.to_vec()]; + let footer_key = AES_128_FOOTER_KEY.to_vec(); - for footer_key in &footer_keys { - // Write with uniform encryption and a plaintext footer. - let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) - .with_plaintext_footer(true) - .build() - .unwrap(); + // Write with uniform encryption and a plaintext footer. + let encryption_properties = FileEncryptionProperties::builder(footer_key.clone()) + .with_plaintext_footer(true) + .build() + .unwrap(); - let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) - .build() - .unwrap(); + let decryption_properties = FileDecryptionProperties::builder(footer_key.clone()) + .build() + .unwrap(); - // Reader can read stats from plaintext footer metadata if a footer key is provided - write_and_read_stats( - Arc::clone(&encryption_properties), - Some(decryption_properties), - &[true, true, true, true], - ); + // Reader can read stats from plaintext footer metadata if a footer key is provided + write_and_read_stats( + Arc::clone(&encryption_properties), + Some(decryption_properties), + &[true, true, true, true], + ); - // Reader can not read stats from plaintext footer metadata if no key is provided - write_and_read_stats(encryption_properties, None, &[false, false, false, false]); - } + // Reader can not read stats from plaintext footer metadata if no key is provided + write_and_read_stats(encryption_properties, None, &[false, false, false, false]); } #[test] From a91f2f90f452beddf0d37aaa0b0b7cdc9a4454fa Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Mon, 11 May 2026 14:12:03 -0700 Subject: [PATCH 25/26] Remove the workaround --- parquet/tests/encryption/encryption.rs | 46 ++++---------------- parquet/tests/encryption/encryption_async.rs | 12 +---- parquet/tests/encryption/encryption_util.rs | 9 ++-- 3 files changed, 12 insertions(+), 55 deletions(-) diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index 1848f80722c5..edd26f29619c 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -424,7 +424,6 @@ fn test_roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever() { ], ); - // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever( AES_256_FOOTER_KEY, AES_256_FOOTER_KEY_NAME, @@ -470,16 +469,16 @@ fn test_roundtrip_non_uniform_encryption_plaintext_footer_with_key_retriever() { AES_256_COLUMN_KEYS[5], AES_256_KEY_NAMES[5], ), + ( + AES_256_COLUMN_NAMES[6], + AES_256_COLUMN_KEYS[6], + AES_256_KEY_NAMES[6], + ), ( AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7], AES_256_KEY_NAMES[7], ), - ( - AES_256_COLUMN_NAMES[8], - AES_256_COLUMN_KEYS[8], - AES_256_KEY_NAMES[6], - ), ], ); } @@ -809,21 +808,11 @@ fn test_write_non_uniform_encryption() { AES_128_COLUMN_NAME_KEYS, ); - // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name write_non_uniform_encryption( AES_256_FOOTER_KEY, AES_256_COLUMN_NAMES.to_vec(), AES_256_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - (AES_256_COLUMN_NAMES[8], AES_256_COLUMN_KEYS[8]), - ], + AES_256_COLUMN_NAME_KEYS, ); } @@ -1204,29 +1193,10 @@ fn test_write_non_uniform_encryption_column_missmatch() { AES_128_COLUMN_NAME_KEYS, ); - // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name write_non_uniform_encryption_column_missmatch( AES_256_FOOTER_KEY, - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - (AES_256_COLUMN_NAMES[8], AES_256_COLUMN_KEYS[8]), - ], - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - (AES_256_COLUMN_NAMES[8], AES_256_COLUMN_KEYS[8]), - ], + AES_256_COLUMN_NAME_KEYS, + AES_256_COLUMN_NAME_KEYS, ); } diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index 8832922bab64..f86ab59bf755 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -411,21 +411,11 @@ async fn test_write_non_uniform_encryption() { .await; // AES-256 - // TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name write_non_uniform_encryption( AES_256_FOOTER_KEY, AES_256_COLUMN_NAMES.to_vec(), AES_256_COLUMN_KEYS.iter().map(|&s| s.to_vec()).collect(), - &[ - (AES_256_COLUMN_NAMES[0], AES_256_COLUMN_KEYS[0]), - (AES_256_COLUMN_NAMES[1], AES_256_COLUMN_KEYS[1]), - (AES_256_COLUMN_NAMES[2], AES_256_COLUMN_KEYS[2]), - (AES_256_COLUMN_NAMES[3], AES_256_COLUMN_KEYS[3]), - (AES_256_COLUMN_NAMES[4], AES_256_COLUMN_KEYS[4]), - (AES_256_COLUMN_NAMES[5], AES_256_COLUMN_KEYS[5]), - (AES_256_COLUMN_NAMES[7], AES_256_COLUMN_KEYS[7]), - (AES_256_COLUMN_NAMES[8], AES_256_COLUMN_KEYS[8]), - ], + AES_256_COLUMN_NAME_KEYS, ) .await; } diff --git a/parquet/tests/encryption/encryption_util.rs b/parquet/tests/encryption/encryption_util.rs index 1c2322cebe30..f0fe66651d11 100644 --- a/parquet/tests/encryption/encryption_util.rs +++ b/parquet/tests/encryption/encryption_util.rs @@ -55,19 +55,17 @@ pub(crate) const AES_256_FOOTER_KEY_NAME: &str = "kf"; pub(crate) const AES_256_KEY_NAMES: [&str; 8] = ["kc1", "kc2", "kc3", "kc4", "kc5", "kc6", "kc7", "kc8"]; -// TODO: Update the test files with 3-level list schema structure to avoid 'int64_field.list.int64_field' column name -pub(crate) const AES_256_COLUMN_NAMES: [&str; 9] = [ +pub(crate) const AES_256_COLUMN_NAMES: [&str; 8] = [ "double_field", "float_field", "boolean_field", "int32_field", "ba_field", "flba_field", - "int64_field", + "int64_field.list.element", "int96_field", - "int64_field.list.int64_field", ]; -pub(crate) const AES_256_COLUMN_KEYS: [&[u8]; 9] = [ +pub(crate) const AES_256_COLUMN_KEYS: [&[u8]; 8] = [ b"12345678901234567890123456789012", b"12345678901234567890123456789013", b"12345678901234567890123456789014", @@ -76,7 +74,6 @@ pub(crate) const AES_256_COLUMN_KEYS: [&[u8]; 9] = [ b"12345678901234567890123456789017", b"12345678901234567890123456789018", b"12345678901234567890123456789019", - b"12345678901234567890123456789018", ]; pub(crate) const AES_256_COLUMN_NAME_KEYS: &[(&str, &[u8]); 8] = &[ From ec912ab37fd9be910a77bc85459c3eb0261a87d7 Mon Sep 17 00:00:00 2001 From: Huang-Hsiang Cheng Date: Thu, 14 May 2026 13:21:19 -0700 Subject: [PATCH 26/26] Update test data --- parquet-testing | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/parquet-testing b/parquet-testing index e74785d85a4e..5a6cf84678df 160000 --- a/parquet-testing +++ b/parquet-testing @@ -1 +1 @@ -Subproject commit e74785d85a4ecee829e1e405444d6a1b24b8bc9c +Subproject commit 5a6cf84678df3af65c231109b78b86ba9bf495df