From cf4c54dbb6388f3f2613c94342a6e189089fd9f8 Mon Sep 17 00:00:00 2001
From: Rok Mihevc <rok@mihevc.org>
Date: Wed, 18 Dec 2024 00:17:43 +0100
Subject: [PATCH] add update_aad

---
 parquet/src/encryption/ciphers.rs     | 13 ++++++++++++-
 parquet/src/file/serialized_reader.rs |  3 +--
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/parquet/src/encryption/ciphers.rs b/parquet/src/encryption/ciphers.rs
index b347c292692..c0e7ce1a803 100644
--- 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
index 530caef57ce..4ec21316f0a 100644
--- 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!(