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) -> Result<Self> {
+        self.algorithm = algorithm;
+        Ok(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 {

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
         }))
     }
 }

parquet/src/file/metadata/mod.rs

@@ -2045,6 +2045,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

@@ -114,7 +114,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,
@@ -144,7 +144,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(),