Return encryptors as a Result<Box<dyn BlockEncryptor>>

Author: adamreeve

parquet/src/encryption/ciphers.rs

@@ -81,7 +81,7 @@ impl BlockDecryptor for RingGcmBlockDecryptor {
     }
 }
 
-pub trait BlockEncryptor: Debug + Send + Sync + Clone {
+pub trait BlockEncryptor: Debug + Send + Sync {
     fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Vec<u8>;
 }
 
@@ -92,19 +92,19 @@ struct CounterNonce {
 }
 
 impl CounterNonce {
-    pub fn new(rng: &SystemRandom) -> Self {
+    pub fn new(rng: &SystemRandom) -> Result<Self> {
         let mut buf = [0; 16];
-        rng.fill(&mut buf).unwrap();
+        rng.fill(&mut buf)?;
 
-        // Since this is a random seed value, endianess doesn't matter at all,
+        // Since this is a random seed value, endianness doesn't matter at all,
         // and we can use whatever is platform-native.
         let start = u128::from_ne_bytes(buf) & RIGHT_TWELVE;
         let counter = start.wrapping_add(1);
 
-        Self { start, counter }
+        Ok(Self { start, counter })
     }
 
-    /// One accessor for the nonce bytes to avoid potentially flipping endianess
+    /// One accessor for the nonce bytes to avoid potentially flipping endianness
     #[inline]
     pub fn get_bytes(&self) -> [u8; NONCE_LEN] {
         self.counter.to_le_bytes()[0..NONCE_LEN].try_into().unwrap()
@@ -132,22 +132,21 @@ pub(crate) struct RingGcmBlockEncryptor {
 }
 
 impl RingGcmBlockEncryptor {
-    // todo TBD: some KMS systems produce data keys, need to be able to pass them to Encryptor.
-    // todo TBD: for other KMSs, we will create data keys inside arrow-rs, making sure to use SystemRandom
     /// Create a new `RingGcmBlockEncryptor` with a given key and random nonce.
     /// The nonce will advance appropriately with each block encryption and
     /// return an error if it wraps around.
-    pub(crate) fn new(key_bytes: &[u8]) -> Self {
+    pub(crate) fn new(key_bytes: &[u8]) -> Result<Self> {
         let rng = SystemRandom::new();
 
         // todo support other key sizes
-        let key = UnboundKey::new(&AES_128_GCM, key_bytes).unwrap();
-        let nonce = CounterNonce::new(&rng);
+        let key = UnboundKey::new(&AES_128_GCM, key_bytes)
+            .map_err(|e| general_err!("Error creating AES key: {}", e))?;
+        let nonce = CounterNonce::new(&rng)?;
 
-        Self {
+        Ok(Self {
             key: LessSafeKey::new(key),
             nonce_sequence: nonce,
-        }
+        })
     }
 }
 
@@ -176,7 +175,7 @@ mod tests {
     #[test]
     fn test_round_trip() {
         let key = [0u8; 16];
-        let mut encryptor = RingGcmBlockEncryptor::new(&key);
+        let mut encryptor = RingGcmBlockEncryptor::new(&key).unwrap();
         let decryptor = RingGcmBlockDecryptor::new(&key);
 
         let plaintext = b"hello, world!";

parquet/src/encryption/encrypt.rs

@@ -166,19 +166,24 @@ impl FileEncryptor {
         &self.aad_file_unique
     }
 
-    pub(crate) fn get_footer_encryptor(&self) -> RingGcmBlockEncryptor {
-        RingGcmBlockEncryptor::new(&self.properties.footer_key.key)
+    pub(crate) fn get_footer_encryptor(&self) -> Result<Box<dyn BlockEncryptor>> {
+        Ok(Box::new(RingGcmBlockEncryptor::new(
+            &self.properties.footer_key.key,
+        )?))
     }
 
-    pub(crate) fn get_column_encryptor(&self, column_path: &str) -> RingGcmBlockEncryptor {
+    pub(crate) fn get_column_encryptor(
+        &self,
+        column_path: &str,
+    ) -> Result<Box<dyn BlockEncryptor>> {
         if self.properties.column_keys.is_empty() {
-            return RingGcmBlockEncryptor::new(self.properties.footer_key.key());
+            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) => RingGcmBlockEncryptor::new(column_key.key()),
+            Some(column_key) => Ok(Box::new(RingGcmBlockEncryptor::new(column_key.key())?)),
         }
     }
 }
@@ -197,7 +202,7 @@ pub(crate) fn encrypt_object<T: TSerializable, W: Write>(
 
     // TODO: Get correct encryptor (footer vs column, data vs metadata)
     let encrypted_buffer = encryptor
-        .get_footer_encryptor()
+        .get_footer_encryptor()?
         .encrypt(buffer.as_ref(), module_aad);
 
     sink.write_all(&encrypted_buffer)?;

parquet/src/encryption/mod.rs

@@ -18,8 +18,8 @@
 //! Encryption implementation specific to Parquet, as described
 //! in the [spec](https://github.com/apache/parquet-format/blob/master/Encryption.md).
 
-pub mod ciphers;
+pub(crate) mod ciphers;
 pub mod decryption;
 pub mod encrypt;
-pub mod modules;
+pub(crate) mod modules;
 pub(crate) mod page_encryptor;

parquet/src/encryption/page_encryptor.rs

@@ -16,7 +16,6 @@
 // under the License.
 
 use crate::column::page::CompressedPage;
-use crate::encryption::ciphers::BlockEncryptor;
 use crate::encryption::encrypt::{encrypt_object, FileEncryptor};
 use crate::encryption::modules::{create_module_aad, ModuleType};
 use crate::errors::ParquetError;
@@ -66,7 +65,9 @@ impl PageEncryptor {
             self.column_index,
             Some(self.page_index),
         )?;
-        let mut encryptor = self.file_encryptor.get_column_encryptor(&self.column_path);
+        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);