Move cyphertext reading into decryptor

Author: rok

parquet/src/encryption/decryption.rs

@@ -18,13 +18,16 @@
 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;
 const SIZE_LEN: usize = 4;
 
 pub trait BlockDecryptor {
     fn decrypt(&self, length_and_ciphertext: &[u8], aad: &[u8]) -> Result<Vec<u8>>;
+
+    fn read_and_decrypt<T: Read>(&self, input: &mut T, aad: &[u8]) -> Result<Vec<u8>>;
 }
 
 #[derive(Debug, Clone)]
@@ -59,6 +62,16 @@ impl BlockDecryptor for RingGcmBlockDecryptor {
         result.resize(result.len() - TAG_LEN, 0u8);
         Ok(result)
     }
+
+    fn read_and_decrypt<T: Read>(&self, input: &mut T, aad: &[u8]) -> Result<Vec<u8>> {
+        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)]
@@ -102,7 +115,9 @@ impl DecryptionPropertiesBuilder {
 
     pub fn build(self) -> Result<FileDecryptionProperties> {
         if self.footer_key.is_none() && self.column_keys.is_none() {
-            return Err(ParquetError::General("Footer or at least one column key is required".to_string()))
+            return Err(ParquetError::General(
+                "Footer or at least one column key is required".to_string(),
+            ));
         }
 
         Ok(FileDecryptionProperties {

parquet/src/file/serialized_reader.rs

@@ -377,16 +377,10 @@ pub(crate) fn read_page_header<T: Read>(
             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)?;