diff --git a/parquet/src/encryption/decryption.rs b/parquet/src/encryption/decryption.rs index b93fb19385a1..08a46b2ddd0b 100644 --- a/parquet/src/encryption/decryption.rs +++ b/parquet/src/encryption/decryption.rs @@ -18,6 +18,7 @@ use crate::errors::{ParquetError, Result}; use ring::aead::{Aad, LessSafeKey, UnboundKey, AES_128_GCM}; use std::collections::HashMap; +use std::io::Read; const NONCE_LEN: usize = 12; const TAG_LEN: usize = 16; @@ -25,6 +26,8 @@ const SIZE_LEN: usize = 4; pub trait BlockDecryptor { fn decrypt(&self, length_and_ciphertext: &[u8], aad: &[u8]) -> Result>; + + fn read_and_decrypt(&self, input: &mut T, aad: &[u8]) -> Result>; } #[derive(Debug, Clone)] @@ -59,6 +62,16 @@ impl BlockDecryptor for RingGcmBlockDecryptor { result.resize(result.len() - TAG_LEN, 0u8); Ok(result) } + + fn read_and_decrypt(&self, input: &mut T, aad: &[u8]) -> Result> { + let mut len_bytes = [0; 4]; + input.read_exact(&mut len_bytes)?; + let ciphertext_len = u32::from_le_bytes(len_bytes) as usize; + let mut ciphertext = vec![0; 4 + ciphertext_len]; + input.read_exact(&mut ciphertext[4..])?; + + self.decrypt(&ciphertext, aad.as_ref()) + } } #[derive(Debug, Clone, PartialEq)] diff --git a/parquet/src/file/serialized_reader.rs b/parquet/src/file/serialized_reader.rs index 702201688fab..dc9cc6a315cc 100644 --- a/parquet/src/file/serialized_reader.rs +++ b/parquet/src/file/serialized_reader.rs @@ -377,16 +377,10 @@ pub(crate) fn read_page_header( crypto_context.page_ordinal, )?; - let mut len_bytes = [0; 4]; - input.read_exact(&mut len_bytes)?; - let ciphertext_len = u32::from_le_bytes(len_bytes) as usize; - let mut ciphertext = vec![0; 4 + ciphertext_len]; - input.read_exact(&mut ciphertext[4..])?; - let buf = data_decryptor .footer_decryptor() .unwrap() - .decrypt(&ciphertext, aad.as_ref())?; + .read_and_decrypt(input, aad.as_ref())?; let mut prot = TCompactSliceInputProtocol::new(buf.as_slice()); let page_header = PageHeader::read_from_in_protocol(&mut prot)?;