[parquet] Allow more encryption algorithms

parquet/src/encryption/ciphers.rs

@@ -19,7 +19,7 @@ use crate::errors::ParquetError;
 use crate::errors::ParquetError::General;
 use crate::errors::Result;
 use crate::file::metadata::HeapSize;
-use ring::aead::{AES_128_GCM, Aad, LessSafeKey, NonceSequence, UnboundKey};
+use ring::aead::{AES_128_GCM, Aad, LessSafeKey, NonceSequence, UnboundKey, Algorithm};
 use ring::rand::{SecureRandom, SystemRandom};
 use std::fmt::Debug;
 
@@ -41,9 +41,12 @@ pub(crate) struct RingGcmBlockDecryptor {
 
 impl RingGcmBlockDecryptor {
     pub(crate) fn new(key_bytes: &[u8]) -> Result<Self> {
-        // todo support other key sizes
-        let key = UnboundKey::new(&AES_128_GCM, key_bytes)
-            .map_err(|_| General("Failed to create AES key".to_string()))?;
+        Self::new_with_algorithm(&AES_128_GCM, key_bytes)
+    }
+
+    pub(crate) fn new_with_algorithm(algorithm: &'static Algorithm, key_bytes: &[u8]) -> Result<Self> {
+        let key = UnboundKey::new(algorithm, key_bytes)
+            .map_err(|_| general_err!("Failed to create {:?} key", algorithm))?;
 
         Ok(Self {
             key: LessSafeKey::new(key),
@@ -143,11 +146,13 @@ impl RingGcmBlockEncryptor {
     /// The nonce will advance appropriately with each block encryption and
     /// return an error if it wraps around.
     pub(crate) fn new(key_bytes: &[u8]) -> Result<Self> {
-        let rng = SystemRandom::new();
+        Self::new_with_algorithm(&AES_128_GCM, key_bytes)
+    }
 
-        // todo support other key sizes
-        let key = UnboundKey::new(&AES_128_GCM, key_bytes)
-            .map_err(|e| general_err!("Error creating AES key: {}", e))?;
+    pub(crate) fn new_with_algorithm(algorithm: &'static Algorithm, key_bytes: &[u8]) -> Result<Self> {
+        let rng = SystemRandom::new();
+        let key = UnboundKey::new(&algorithm, key_bytes)
+            .map_err(|e| general_err!("Error creating {:?} key: {}", algorithm, e))?;
         let nonce = CounterNonce::new(&rng)?;
 
         Ok(Self {
@@ -188,6 +193,7 @@ impl BlockEncryptor for RingGcmBlockEncryptor {
 
 #[cfg(test)]
 mod tests {
+    use ring::aead::{AES_256_GCM, CHACHA20_POLY1305};
     use super::*;
 
     #[test]
@@ -204,4 +210,26 @@ mod tests {
 
         assert_eq!(plaintext, decrypted.as_slice());
     }
+
+    #[test]
+    fn test_round_trip_with_algorithm() {
+        let key = [0u8; 32];
+        let mut encryptor = RingGcmBlockEncryptor::new_with_algorithm(&AES_256_GCM, &key).unwrap();
+        let decryptor = RingGcmBlockDecryptor::new_with_algorithm(&AES_256_GCM, &key).unwrap();
+
+        let plaintext = b"hello, world!";
+        let aad = b"some aad";
+
+        let ciphertext = encryptor.encrypt(plaintext, aad).unwrap();
+        let decrypted = decryptor.decrypt(&ciphertext, aad).unwrap();
+
+        assert_eq!(plaintext, decrypted.as_slice());
+    }
+
+    #[test]
+    fn test_round_trip_with_incorrect_key_length() {
+        let key = [0u8; 16];
+        assert!(RingGcmBlockEncryptor::new_with_algorithm(&CHACHA20_POLY1305, &key).is_err());
+        assert!(RingGcmBlockDecryptor::new_with_algorithm(&CHACHA20_POLY1305, &key).is_err());
+    }
 }

parquet/src/encryption/decrypt.rs

@@ -27,6 +27,7 @@ use std::collections::HashMap;
 use std::fmt::Formatter;
 use std::io::Read;
 use std::sync::Arc;
+use ring::aead::{Algorithm, AES_128_GCM};
 
 /// Trait for retrieving an encryption key using the key's metadata
 ///
@@ -352,6 +353,7 @@ pub struct FileDecryptionProperties {
     keys: DecryptionKeys,
     aad_prefix: Option<Vec<u8>>,
     footer_signature_verification: bool,
+    algorithm: &'static Algorithm
 }
 
 impl HeapSize for FileDecryptionProperties {
@@ -448,6 +450,7 @@ pub struct DecryptionPropertiesBuilder {
     column_keys: HashMap<String, Vec<u8>>,
     aad_prefix: Option<Vec<u8>>,
     footer_signature_verification: bool,
+    algorithm: &'static Algorithm
 }
 
 impl DecryptionPropertiesBuilder {
@@ -459,6 +462,7 @@ impl DecryptionPropertiesBuilder {
             column_keys: HashMap::default(),
             aad_prefix: None,
             footer_signature_verification: true,
+            algorithm: &AES_128_GCM
         }
     }
 
@@ -472,6 +476,7 @@ impl DecryptionPropertiesBuilder {
             keys,
             aad_prefix: self.aad_prefix,
             footer_signature_verification: self.footer_signature_verification,
+            algorithm: self.algorithm
         }))
     }
 
@@ -505,6 +510,11 @@ impl DecryptionPropertiesBuilder {
         Ok(self)
     }
 
+    pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self {
+        self.algorithm = algorithm;
+        self
+    }
+
     /// Disable verification of footer tags for files that use plaintext footers.
     /// Signature verification is enabled by default.
     pub fn disable_footer_signature_verification(mut self) -> Self {
@@ -520,6 +530,7 @@ pub struct DecryptionPropertiesBuilderWithRetriever {
     key_retriever: Arc<dyn KeyRetriever>,
     aad_prefix: Option<Vec<u8>>,
     footer_signature_verification: bool,
+    algorithm: &'static Algorithm
 }
 
 impl DecryptionPropertiesBuilderWithRetriever {
@@ -530,6 +541,7 @@ impl DecryptionPropertiesBuilderWithRetriever {
             key_retriever,
             aad_prefix: None,
             footer_signature_verification: true,
+            algorithm: &AES_128_GCM
         }
     }
 
@@ -540,6 +552,7 @@ impl DecryptionPropertiesBuilderWithRetriever {
             keys,
             aad_prefix: self.aad_prefix,
             footer_signature_verification: self.footer_signature_verification,
+            algorithm: self.algorithm
         }))
     }
 
@@ -551,6 +564,11 @@ impl DecryptionPropertiesBuilderWithRetriever {
         self
     }
 
+    pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self {
+        self.algorithm = algorithm;
+        self
+    }
+
     /// Disable verification of footer tags for files that use plaintext footers.
     /// Signature verification is enabled by default.
     pub fn disable_footer_signature_verification(mut self) -> Self {
@@ -596,7 +614,8 @@ impl FileDecryptor {
     ) -> Result<Self> {
         let file_aad = [aad_prefix.as_slice(), aad_file_unique.as_slice()].concat();
         let footer_key = decryption_properties.footer_key(footer_key_metadata)?;
-        let footer_decryptor = RingGcmBlockDecryptor::new(&footer_key).map_err(|e| {
+        let algorithm = decryption_properties.algorithm;
+        let footer_decryptor = RingGcmBlockDecryptor::new_with_algorithm(algorithm, &footer_key).map_err(|e| {
             general_err!(
                 "Invalid footer key. {}",
                 e.to_string().replace("Parquet error: ", "")
@@ -641,7 +660,8 @@ impl FileDecryptor {
         let column_key = self
             .decryption_properties
             .column_key(column_name, key_metadata)?;
-        Ok(Arc::new(RingGcmBlockDecryptor::new(&column_key)?))
+        let algorithm = self.decryption_properties.algorithm;
+        Ok(Arc::new(RingGcmBlockDecryptor::new_with_algorithm(algorithm, &column_key)?))
     }
 
     pub(crate) fn get_column_metadata_decryptor(

parquet/src/encryption/encrypt.rs

@@ -28,6 +28,7 @@ use ring::rand::{SecureRandom, SystemRandom};
 use std::collections::{HashMap, HashSet};
 use std::io::Write;
 use std::sync::Arc;
+use ring::aead::{Algorithm, AES_128_GCM};
 
 #[derive(Debug, Clone, PartialEq)]
 struct EncryptionKey {
@@ -96,6 +97,7 @@ pub struct FileEncryptionProperties {
     column_keys: HashMap<String, EncryptionKey>,
     aad_prefix: Option<Vec<u8>>,
     store_aad_prefix: bool,
+    algorithm: &'static Algorithm
 }
 
 impl FileEncryptionProperties {
@@ -186,6 +188,7 @@ pub struct EncryptionPropertiesBuilder {
     column_keys: HashMap<String, EncryptionKey>,
     aad_prefix: Option<Vec<u8>>,
     store_aad_prefix: bool,
+    algorithm: &'static Algorithm
 }
 
 impl EncryptionPropertiesBuilder {
@@ -197,6 +200,7 @@ impl EncryptionPropertiesBuilder {
             aad_prefix: None,
             encrypt_footer: true,
             store_aad_prefix: false,
+            algorithm: &AES_128_GCM
         }
     }
 
@@ -274,6 +278,11 @@ impl EncryptionPropertiesBuilder {
         self
     }
 
+    pub fn with_algorithm(mut self, algorithm: &'static Algorithm) -> Self {
+        self.algorithm = algorithm;
+        self
+    }
+
     /// Build the encryption properties
     pub fn build(self) -> Result<Arc<FileEncryptionProperties>> {
         Ok(Arc::new(FileEncryptionProperties {
@@ -282,6 +291,7 @@ impl EncryptionPropertiesBuilder {
             column_keys: self.column_keys,
             aad_prefix: self.aad_prefix,
             store_aad_prefix: self.store_aad_prefix,
+            algorithm: self.algorithm
         }))
     }
 }
@@ -343,8 +353,8 @@ impl FileEncryptor {
 
     /// Get the BlockEncryptor for the footer
     pub(crate) fn get_footer_encryptor(&self) -> Result<Box<dyn BlockEncryptor>> {
-        Ok(Box::new(RingGcmBlockEncryptor::new(
-            &self.properties.footer_key.key,
+        Ok(Box::new(RingGcmBlockEncryptor::new_with_algorithm(
+            &self.properties.algorithm, &self.properties.footer_key.key,
         )?))
     }
 
@@ -359,7 +369,7 @@ impl FileEncryptor {
         }
         match self.properties.column_keys.get(column_path) {
             None => Err(general_err!("Column '{}' is not encrypted", column_path)),
-            Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new(column_key.key())?)),
+            Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new_with_algorithm(&self.properties.algorithm, column_key.key())?)),
         }
     }
 }

parquet/src/file/metadata/mod.rs

@@ -2118,6 +2118,9 @@ mod tests {
             .set_file_decryptor(Some(decryptor))
             .build();
 
+        #[cfg(feature = "encryption")]
+        let expected_size_with_decryptor = 3080;
+        #[cfg(not(feature = "encryption"))]
         let expected_size_with_decryptor = 3072;
         assert!(expected_size_with_decryptor > base_expected_size);
 

parquet/tests/encryption/encryption_async.rs

@@ -115,7 +115,7 @@ async fn test_misspecified_encryption_keys() {
 
     // Too short footer key
     check_for_error(
-        "Parquet error: Invalid footer key. Failed to create AES key",
+        "Parquet error: Invalid footer key. Failed to create AES_128_GCM key",
         &path,
         "bad_pwd".as_bytes(),
         column_1_key,
@@ -145,7 +145,7 @@ async fn test_misspecified_encryption_keys() {
 
     // Too short column key
     check_for_error(
-        "Parquet error: Failed to create AES key",
+        "Parquet error: Failed to create AES_128_GCM key",
         &path,
         footer_key,
         "abc".as_bytes(),