WIP Add per-column encryption keys

Author: adamreeve

parquet/src/column/writer/mod.rs

@@ -1551,7 +1551,8 @@ mod tests {
     use core::str;
     use rand::distributions::uniform::SampleUniform;
     use std::{fs::File, sync::Arc};
-
+    #[cfg(feature = "encryption")]
+    use crate::encryption::encrypt::EncryptionKey;
     use super::*;
 
     #[test]
@@ -3481,6 +3482,101 @@ mod tests {
         assert_eq!(row_count, file_metadata.num_rows() as usize);
     }
 
+    #[cfg(feature = "encryption")]
+    #[test]
+    fn test_write_encrypted_column_non_uniform() {
+        let message_type = "
+            message test_schema {
+                OPTIONAL BYTE_ARRAY a (UTF8);
+            }
+        ";
+        let schema = Arc::new(parse_message_type(message_type).unwrap());
+        let data = vec![ByteArray::from(b"parquet".to_vec()); 7];
+        let def_levels = [1, 1, 1, 1, 0, 1, 0, 1, 0, 1];
+
+        let num_row_groups = 3;
+        let num_batches = 3;
+        let rows_per_batch = def_levels.len();
+        let valid_rows_per_batch = def_levels.iter().filter(|&level| *level > 0).count();
+
+        let file: File = tempfile::tempfile().unwrap();
+
+        let builder = WriterProperties::builder();
+        let footer_key: &[u8] = "0123456789012345".as_bytes();
+        let column_key = EncryptionKey::new(b"1234567890123450".to_vec());
+        let file_encryption_properties =
+            FileEncryptionProperties::builder(footer_key.to_vec())
+                .with_column_key(b"a".to_vec(), column_key.clone()).build();
+
+        let props = Arc::new(
+            builder
+                .with_file_encryption_properties(file_encryption_properties)
+                .set_data_page_row_count_limit(rows_per_batch)
+                .build(),
+        );
+        let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap();
+        for _ in 0..num_row_groups {
+            let mut row_group_writer = writer.next_row_group().unwrap();
+            let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
+
+            for _ in 0..num_batches {
+                col_writer
+                    .typed::<ByteArrayType>()
+                    .write_batch(&data, Some(&def_levels), None)
+                    .unwrap();
+            }
+
+            col_writer.close().unwrap();
+            row_group_writer.close().unwrap();
+        }
+
+        let _file_metadata = writer.close().unwrap();
+
+        let decryption_properties = FileDecryptionProperties::builder(footer_key.to_vec())
+            .with_column_key(b"a".to_vec(), column_key.key().clone())
+            .build()
+            .unwrap();
+        let options = ArrowReaderOptions::default()
+            .with_file_decryption_properties(decryption_properties.clone());
+        let metadata = ArrowReaderMetadata::load(&file, options.clone()).unwrap();
+        let file_metadata = metadata.metadata.file_metadata();
+
+        let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap();
+        let record_reader = builder.build().unwrap();
+
+        assert_eq!(
+            file_metadata.num_rows(),
+            (num_row_groups * num_batches * rows_per_batch) as i64
+        );
+        assert_eq!(file_metadata.schema_descr().num_columns(), 1);
+
+        assert_eq!(metadata.metadata.num_row_groups(), num_row_groups);
+        metadata.metadata.row_groups().iter().for_each(|rg| {
+            assert_eq!(rg.num_columns(), 1);
+            assert_eq!(rg.num_rows(), (num_batches * rows_per_batch) as i64);
+        });
+
+        let mut row_count = 0;
+        for batch in record_reader {
+            let batch = batch.unwrap();
+            row_count += batch.num_rows();
+
+            let string_col = batch.column(0).as_string_opt::<i32>().unwrap();
+
+            let mut valid_count = 0;
+            for x in string_col.iter().flatten() {
+                valid_count += 1;
+                assert_eq!(x, "parquet");
+            }
+            assert_eq!(
+                valid_count,
+                valid_rows_per_batch * num_batches * num_row_groups
+            );
+        }
+
+        assert_eq!(row_count, file_metadata.num_rows() as usize);
+    }
+
     #[test]
     fn test_increment_max_binary_chars() {
         let r = increment(vec![0xFF, 0xFE, 0xFD, 0xFF, 0xFF]);

parquet/src/encryption/encrypt.rs

@@ -37,8 +37,13 @@ impl EncryptionKey {
         }
     }
 
-    pub fn set_metadata(&mut self, metadata: Vec<u8>) {
+    pub fn with_metadata(mut self, metadata: Vec<u8>) -> Self {
         self.key_metadata = Some(metadata);
+        self
+    }
+
+    pub fn key(&self) -> &Vec<u8> {
+        &self.key
     }
 }
 
@@ -98,7 +103,12 @@ impl EncryptionPropertiesBuilder {
     }
 
     pub fn with_footer_key_metadata(mut self, metadata: Vec<u8>) -> Self {
-        self.footer_key.set_metadata(metadata);
+        self.footer_key = self.footer_key.with_metadata(metadata);
+        self
+    }
+
+    pub fn with_column_key(mut self, column_name: Vec<u8>, encryption_key: EncryptionKey) -> Self {
+        self.column_keys.insert(column_name, encryption_key);
         self
     }
 
@@ -155,17 +165,19 @@ impl FileEncryptor {
     pub fn aad_file_unique(&self) -> &Vec<u8> {
         &self.aad_file_unique
     }
-
-    // let footer_encryptor = RingGcmBlockEncryptor::new(&encryption_properties.footer_key.clone());
-    // let mut column_encryptors: HashMap<Vec<u8>, Arc<dyn BlockEncryptor>> = HashMap::new();
-    // if let Some(column_keys) = encryption_properties.column_keys.clone() {
-    // for (column_name, key) in column_keys.iter() {
-    // let column_encryptor = Arc::new(RingGcmBlockEncryptor::new(key));
-    // column_encryptors.insert(column_name.clone(), column_encryptor);
-    // }
-    // }
+    
     pub(crate) fn get_footer_encryptor(&self) -> RingGcmBlockEncryptor {
-        RingGcmBlockEncryptor::new(&self.properties.footer_key.key.clone())
+        RingGcmBlockEncryptor::new(&self.properties.footer_key.key)
+    }
+
+    pub(crate) fn get_column_encryptor(&self, column_path: &Vec<u8>) -> RingGcmBlockEncryptor {
+        if self.properties.column_keys.is_empty() {
+            return RingGcmBlockEncryptor::new(&self.properties.footer_key.key());
+        }
+        match self.properties.column_keys.get(column_path) {
+            None => todo!("Handle unencrypted columns"),
+            Some(column_key) => RingGcmBlockEncryptor::new(&column_key.key())
+        }
     }
 }
 

parquet/src/encryption/page_encryptor.rs

@@ -30,19 +30,22 @@ pub struct PageEncryptor {
     row_group_index: usize,
     column_index: usize,
     page_index: usize,
+    column_path: Vec<u8>,
 }
 
 impl PageEncryptor {
     pub fn new(
         file_encryptor: Arc<FileEncryptor>,
         row_group_index: usize,
         column_index: usize,
+        column_path: Vec<u8>,
     ) -> Self {
         Self {
             file_encryptor,
             row_group_index,
             column_index,
             page_index: 0,
+            column_path,
         }
     }
 

parquet/src/file/writer.rs

@@ -579,20 +579,28 @@ impl<'a, W: Write + Send> SerializedRowGroupWriter<'a, W> {
         self.assert_previous_writer_closed()?;
 
         #[cfg(feature = "encryption")]
-        let page_encryptor = match self.file_encryptor.as_ref() {
-            None => None,
-            Some(file_encryptor) => Some(PageEncryptor::new(
-                file_encryptor.clone(),
-                self.row_group_index as usize,
-                self.column_index,
-            )),
-        };
+        let file_encryptor = self.file_encryptor.clone();
+        #[cfg(feature = "encryption")]
+        let row_group_index = self.row_group_index as usize;
+        #[cfg(feature = "encryption")]
+        let column_index = self.column_index;
 
         Ok(match self.next_column_desc() {
             Some(column) => {
                 let props = self.props.clone();
                 let (buf, on_close) = self.get_on_close();
 
+                #[cfg(feature = "encryption")]
+                let page_encryptor = match file_encryptor {
+                    None => None,
+                    Some(file_encryptor) => Some(PageEncryptor::new(
+                        file_encryptor,
+                        row_group_index,
+                        column_index,
+                        column.path().string().into_bytes(),
+                    )),
+                };
+
                 #[cfg(feature = "encryption")]
                 let page_writer =
                     SerializedPageWriter::new(buf).with_page_encryptor(page_encryptor);