add update_aad

apache · Dec 17, 2024 · cf4c54d · cf4c54d
1 parent bb6d3e8
commit cf4c54d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/parquet/src/encryption/ciphers.rs b/parquet/src/encryption/ciphers.rs
@@ -35,6 +35,7 @@ const RIGHT_TWELVE: u128 = 0x0000_0000_ffff_ffff_ffff_ffff_ffff_ffff;
 const NONCE_LEN: usize = 12;
 const TAG_LEN: usize = 16;
 const SIZE_LEN: usize = 4;
+const NON_PAGE_ORDINAL: i32 = -1;
 
 struct CounterNonce {
     start: u128,
@@ -152,10 +153,12 @@ impl BlockDecryptor for RingGcmBlockDecryptor {
             .open_in_place(nonce, Aad::from(aad), &mut result)
             .unwrap();
 
+        result.resize(result.len() - TAG_LEN, 0u8);
         result
     }
 }
 
+#[derive(PartialEq)]
 pub(crate) enum ModuleType {
     Footer = 0,
     ColumnMetaData = 1,
@@ -170,7 +173,7 @@ pub(crate) enum ModuleType {
 }
 
 pub fn create_footer_aad(file_aad: &[u8]) -> Result<Vec<u8>> {
-    create_module_aad(file_aad, ModuleType::Footer, -1, -1, -1)
+    create_module_aad(file_aad, ModuleType::Footer, -1, -1, NON_PAGE_ORDINAL)
 }
 
 pub fn create_page_aad(file_aad: &[u8], module_type: ModuleType, row_group_ordinal: i16, column_ordinal: i16, page_ordinal: i32) -> Result<Vec<u8>> {
@@ -314,6 +317,14 @@ impl FileDecryptor {
     pub(crate) fn aad_prefix(&self) -> &Vec<u8> {
         &self.aad_prefix
     }
+
+    pub fn update_aad(&mut self, aad: Vec<u8>, row_group_ordinal: i16, column_ordinal: i16, module_type: ModuleType) {
+        // todo decr: update aad
+        debug_assert!(!self.aad_file_unique().is_empty(), "AAD is empty");
+
+        let aad = create_module_aad(self.aad_file_unique(), module_type, row_group_ordinal, column_ordinal, NON_PAGE_ORDINAL).unwrap();
+        self.aad_file_unique = aad;
+    }
 }
 
 #[derive(Debug, Clone)]

diff --git a/parquet/src/file/serialized_reader.rs b/parquet/src/file/serialized_reader.rs
@@ -430,7 +430,7 @@ pub(crate) fn decode_page(
         can_decompress = header_v2.is_compressed.unwrap_or(true);
     }
 
-    let buffer = if crypto_context.is_some() {
+    let buffer : Bytes = if crypto_context.is_some() {
         let crypto_context = crypto_context.as_ref().unwrap();
         let decryptor = crypto_context.data_decryptor();
         let file_decryptor = decryptor.footer_decryptor();
@@ -463,7 +463,6 @@ pub(crate) fn decode_page(
                 &mut decompressed,
                 Some(uncompressed_size - offset),
             )?;
-            todo!("page decompressed!");
 
             if decompressed.len() != uncompressed_size {
                 return Err(general_err!(