Get per-column encryption working and various tidy ups

Author: adamreeve

parquet/src/arrow/arrow_writer/mod.rs

@@ -510,7 +510,7 @@ impl PageWriter for ArrowPageWriter {
         match self.page_encryptor.as_ref() {
             #[cfg(feature = "encryption")]
             Some(encryptor) => {
-                encryptor.encrypt_page_header(page_header, &mut header)?;
+                encryptor.encrypt_page_header(&page_header, &mut header)?;
             }
             _ => {
                 let mut protocol = TCompactOutputProtocol::new(&mut header);
@@ -3824,8 +3824,8 @@ mod tests {
         let column_1_key = EncryptionKey::new(column_1_key.as_bytes().to_vec());
         let column_2_key = EncryptionKey::new(column_2_key.as_bytes().to_vec());
         let file_encryption_properties = FileEncryptionProperties::builder(footer_key.to_vec())
-            .with_column_key("double_field".as_bytes().to_vec(), column_1_key)
-            .with_column_key("float_field".as_bytes().to_vec(), column_2_key)
+            .with_column_key("double_field".into(), column_1_key)
+            .with_column_key("float_field".into(), column_2_key)
             .build();
 
         let temp_file =

parquet/src/column/writer/mod.rs

@@ -32,6 +32,8 @@ use crate::compression::{create_codec, Codec, CodecOptionsBuilder};
 use crate::data_type::private::ParquetValueType;
 use crate::data_type::*;
 use crate::encodings::levels::LevelEncoder;
+#[cfg(feature = "encryption")]
+use crate::encryption::encrypt::get_column_crypto_metadata;
 use crate::errors::{ParquetError, Result};
 use crate::file::metadata::{ColumnIndexBuilder, LevelHistogram, OffsetIndexBuilder};
 use crate::file::properties::EnabledStatistics;
@@ -1199,6 +1201,14 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
                 );
         }
 
+        #[cfg(feature = "encryption")]
+        if let Some(encryption_properties) = self.props.file_encryption_properties.as_ref() {
+            builder = builder.set_column_crypto_metadata(get_column_crypto_metadata(
+                encryption_properties,
+                &self.descr,
+            ));
+        }
+
         let metadata = builder.build()?;
         Ok(metadata)
     }
@@ -3505,7 +3515,7 @@ mod tests {
         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())
+            .with_column_key("a".into(), column_key.clone())
             .build();
 
         let props = Arc::new(

parquet/src/encryption/ciphers.rs

@@ -82,7 +82,7 @@ impl BlockDecryptor for RingGcmBlockDecryptor {
 }
 
 pub trait BlockEncryptor: Debug + Send + Sync {
-    fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Vec<u8>;
+    fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Result<Vec<u8>>;
 }
 
 #[derive(Debug, Clone)]
@@ -151,20 +151,28 @@ impl RingGcmBlockEncryptor {
 }
 
 impl BlockEncryptor for RingGcmBlockEncryptor {
-    fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Vec<u8> {
-        let mut ciphertext = Vec::with_capacity(plaintext.len() + TAG_LEN);
+    fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Result<Vec<u8>> {
+        // Create encrypted buffer.
+        // Format is: [ciphertext size, nonce, ciphertext, authentication tag]
+        let ciphertext_length = NONCE_LEN + plaintext.len() + TAG_LEN;
+        let mut ciphertext = Vec::with_capacity(SIZE_LEN + ciphertext_length);
+        ciphertext.extend((ciphertext_length as u32).to_le_bytes());
+
+        let nonce = self.nonce_sequence.advance()?;
+        ciphertext.extend(nonce.as_ref());
         ciphertext.extend(plaintext);
-        let nonce = self.nonce_sequence.advance().unwrap();
-        let nonce_bytes = *nonce.as_ref();
-        self.key
-            .seal_in_place_append_tag(nonce, Aad::from(aad), &mut ciphertext)
-            .unwrap();
-
-        let mut out = Vec::with_capacity(ciphertext.len() + SIZE_LEN + NONCE_LEN);
-        out.extend(((ciphertext.len() + NONCE_LEN) as u32).to_le_bytes());
-        out.extend(nonce_bytes);
-        out.extend_from_slice(ciphertext.as_ref());
-        out
+
+        let tag = self.key.seal_in_place_separate_tag(
+            nonce,
+            Aad::from(aad),
+            &mut ciphertext[SIZE_LEN + NONCE_LEN..],
+        )?;
+
+        ciphertext.extend(tag.as_ref());
+
+        debug_assert_eq!(SIZE_LEN + ciphertext_length, ciphertext.len());
+
+        Ok(ciphertext)
     }
 }
 
@@ -181,7 +189,7 @@ mod tests {
         let plaintext = b"hello, world!";
         let aad = b"some aad";
 
-        let ciphertext = encryptor.encrypt(plaintext, aad);
+        let ciphertext = encryptor.encrypt(plaintext, aad).unwrap();
         let decrypted = decryptor.decrypt(&ciphertext, aad).unwrap();
 
         assert_eq!(plaintext, decrypted.as_slice());

parquet/src/encryption/encrypt.rs

@@ -17,6 +17,8 @@
 
 use crate::encryption::ciphers::{BlockEncryptor, RingGcmBlockEncryptor};
 use crate::errors::Result;
+use crate::file::column_crypto_metadata::{ColumnCryptoMetaData, EncryptionWithColumnKey};
+use crate::schema::types::ColumnDescPtr;
 use crate::thrift::TSerializable;
 use ring::rand::{SecureRandom, SystemRandom};
 use std::collections::HashMap;
@@ -51,7 +53,7 @@ impl EncryptionKey {
 pub struct FileEncryptionProperties {
     encrypt_footer: bool,
     footer_key: EncryptionKey,
-    column_keys: HashMap<Vec<u8>, EncryptionKey>,
+    column_keys: HashMap<String, EncryptionKey>,
     aad_prefix: Option<Vec<u8>>,
     store_aad_prefix: bool,
 }
@@ -80,7 +82,7 @@ impl FileEncryptionProperties {
 
 pub struct EncryptionPropertiesBuilder {
     footer_key: EncryptionKey,
-    column_keys: HashMap<Vec<u8>, EncryptionKey>,
+    column_keys: HashMap<String, EncryptionKey>,
     aad_prefix: Option<Vec<u8>>,
     encrypt_footer: bool,
     store_aad_prefix: bool,
@@ -107,8 +109,8 @@ impl EncryptionPropertiesBuilder {
         self
     }
 
-    pub fn with_column_key(mut self, column_name: Vec<u8>, encryption_key: EncryptionKey) -> Self {
-        self.column_keys.insert(column_name, encryption_key);
+    pub fn with_column_key(mut self, column_path: String, encryption_key: EncryptionKey) -> Self {
+        self.column_keys.insert(column_path, encryption_key);
         self
     }
 
@@ -179,18 +181,17 @@ impl FileEncryptor {
         if self.properties.column_keys.is_empty() {
             return self.get_footer_encryptor();
         }
-        // TODO: Column paths should be stored as String
-        let column_path = column_path.as_bytes();
         match self.properties.column_keys.get(column_path) {
             None => todo!("Handle unencrypted columns"),
             Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new(column_key.key())?)),
         }
     }
 }
 
+/// Encrypt a Thrift serializable object
 pub(crate) fn encrypt_object<T: TSerializable, W: Write>(
-    object: T,
-    encryptor: &FileEncryptor,
+    object: &T,
+    encryptor: &mut Box<dyn BlockEncryptor>,
     sink: &mut W,
     module_aad: &[u8],
 ) -> Result<()> {
@@ -200,11 +201,45 @@ pub(crate) fn encrypt_object<T: TSerializable, W: Write>(
         object.write_to_out_protocol(&mut unencrypted_protocol)?;
     }
 
-    // TODO: Get correct encryptor (footer vs column, data vs metadata)
-    let encrypted_buffer = encryptor
-        .get_footer_encryptor()?
-        .encrypt(buffer.as_ref(), module_aad);
+    let encrypted_buffer = encryptor.encrypt(buffer.as_ref(), module_aad)?;
 
     sink.write_all(&encrypted_buffer)?;
     Ok(())
 }
+
+/// Encrypt a Thrift serializable object
+pub(crate) fn encrypt_object_to_vec<T: TSerializable>(
+    object: &T,
+    encryptor: &mut Box<dyn BlockEncryptor>,
+    module_aad: &[u8],
+) -> Result<Vec<u8>> {
+    let mut buffer: Vec<u8> = vec![];
+    {
+        let mut unencrypted_protocol = TCompactOutputProtocol::new(&mut buffer);
+        object.write_to_out_protocol(&mut unencrypted_protocol)?;
+    }
+
+    encryptor.encrypt(buffer.as_ref(), module_aad)
+}
+
+/// Get the crypto metadata for a column from the file encryption properties
+pub fn get_column_crypto_metadata(
+    properties: &FileEncryptionProperties,
+    column: &ColumnDescPtr,
+) -> Option<ColumnCryptoMetaData> {
+    if properties.column_keys.is_empty() {
+        // Uniform encryption
+        Some(ColumnCryptoMetaData::EncryptionWithFooterKey)
+    } else {
+        properties
+            .column_keys
+            .get(&column.path().string())
+            .map(|encryption_key| {
+                // Column is encrypted with a column specific key
+                ColumnCryptoMetaData::EncryptionWithColumnKey(EncryptionWithColumnKey {
+                    path_in_schema: column.path().parts().to_vec(),
+                    key_metadata: encryption_key.key_metadata.clone(),
+                })
+            })
+    }
+}

parquet/src/encryption/page_encryptor.rs

@@ -68,16 +68,14 @@ impl PageEncryptor {
         let mut encryptor = self
             .file_encryptor
             .get_column_encryptor(&self.column_path)?;
-        // todo: use column encryptor when needed
-        // self.file_encryptor.get_column_encryptor(self.column_path.as_ref())
-        let encrypted_buffer = encryptor.encrypt(page.data(), &aad);
+        let encrypted_buffer = encryptor.encrypt(page.data(), &aad)?;
 
         Ok(encrypted_buffer)
     }
 
     pub fn encrypt_page_header<W: Write>(
         &self,
-        page_header: PageHeader,
+        page_header: &PageHeader,
         sink: &mut W,
     ) -> crate::errors::Result<()> {
         let module_type = match page_header.type_ {
@@ -99,6 +97,10 @@ impl PageEncryptor {
             Some(self.page_index),
         )?;
 
-        encrypt_object(page_header, &self.file_encryptor, sink, &aad)
+        let mut encryptor = self
+            .file_encryptor
+            .get_column_encryptor(&self.column_path)?;
+
+        encrypt_object(page_header, &mut encryptor, sink, &aad)
     }
 }