Add encryption algorithm to file_metadata before writing

Author: rok

parquet/src/encryption/encrypt.rs

@@ -385,16 +385,19 @@ pub(crate) fn sign_and_write_object<T: TSerializable, W: Write>(
         let mut protocol = TCompactOutputProtocol::new(&mut buffer);
         object.write_to_out_protocol(&mut protocol)?;
     }
-    let plaintext_len = buffer.len();
     sink.write_all(&buffer)?;
+    buffer = encryptor.encrypt(buffer.as_ref(), module_aad)?;
+
+    let ciphertext_length : u32 = buffer.len()
+        .try_into()
+        .map_err(|err| general_err!("Plaintext data too long. {:?}", err))?;
 
-    let encrypted_buffer = encryptor.encrypt(buffer.as_ref(), module_aad)?;
- 
     // Format is: [ciphertext size, nonce, ciphertext, authentication tag]
-    let nonce = encrypted_buffer[SIZE_LEN..SIZE_LEN + NONCE_LEN].to_vec();
-    let tag = encrypted_buffer[SIZE_LEN + NONCE_LEN + plaintext_len..SIZE_LEN + NONCE_LEN + plaintext_len + TAG_LEN].to_vec();
+    let nonce = buffer[SIZE_LEN..SIZE_LEN + NONCE_LEN].to_vec();
+    let tag = buffer[(ciphertext_length - TAG_LEN as u32) as usize..].to_vec();
     sink.write_all(&nonce)?;
     sink.write_all(&tag)?;
+
     Ok(())
 }
 

parquet/src/file/metadata/writer.rs

@@ -16,9 +16,10 @@
 // under the License.
 
 #[cfg(feature = "encryption")]
-use crate::encryption::encrypt::{encrypt_object, encrypt_object_to_vec, sign_and_write_object, FileEncryptor};
-#[cfg(feature = "encryption")]
-use crate::encryption::modules::{create_footer_aad, create_module_aad, ModuleType};
+use crate::encryption::{
+    encrypt::{encrypt_object, encrypt_object_to_vec, sign_and_write_object, FileEncryptor},
+    modules::{create_footer_aad, create_module_aad, ModuleType},
+};
 #[cfg(feature = "encryption")]
 use crate::errors::ParquetError;
 use crate::errors::Result;
@@ -141,6 +142,25 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
             .object_writer
             .apply_row_group_encryption(self.row_groups)?;
 
+        let mut encryption_algorithm = None;
+        if let Some(file_encryptor) = self.object_writer.file_encryptor.clone() {
+            let properties = file_encryptor.properties();
+            if !properties.encrypt_footer() {
+                let supply_aad_prefix = properties
+                    .aad_prefix()
+                    .map(|_| !properties.store_aad_prefix());
+                encryption_algorithm = Some(EncryptionAlgorithm::AESGCMV1(AesGcmV1 {
+                    aad_prefix: if properties.store_aad_prefix() {
+                        properties.aad_prefix().cloned()
+                    } else {
+                        None
+                    },
+                    aad_file_unique: Some(file_encryptor.aad_file_unique().clone()),
+                    supply_aad_prefix,
+                }));
+            }
+        };
+
         let mut file_metadata = FileMetaData {
             num_rows,
             row_groups,
@@ -149,7 +169,7 @@ impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
             schema: types::to_thrift(self.schema.as_ref())?,
             created_by: self.created_by.clone(),
             column_orders,
-            encryption_algorithm: None,
+            encryption_algorithm,
             footer_signing_key_metadata: None,
         };
 
@@ -503,8 +523,7 @@ impl MetadataObjectWriter {
                 let mut encryptor = file_encryptor.get_footer_encryptor()?;
                 encrypt_object(file_metadata, &mut encryptor, &mut sink, &aad)
             }
-            Some(file_encryptor) if !file_encryptor.properties().encrypt_footer() => {
-                // todo: should we also check for file_metadata.encryption_algorithm.is_some() ?
+            Some(file_encryptor) if file_metadata.encryption_algorithm.is_some() => {
                 let aad = create_footer_aad(file_encryptor.file_aad())?;
                 let mut encryptor = file_encryptor.get_footer_encryptor()?;
                 sign_and_write_object(file_metadata, &mut encryptor, &mut sink, &aad)

parquet/src/file/writer.rs

@@ -349,7 +349,6 @@ impl<W: Write + Send> SerializedFileWriter<W> {
         );
 
         #[cfg(feature = "encryption")]
-        // todo
         {
             encoder = encoder.with_file_encryptor(self.file_encryptor.clone());
         }