Change encrypt and decrypt to return Results

apache · Dec 18, 2024 · b8fdff1 · b8fdff1
1 parent cf4c54d
commit b8fdff1
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 16 deletions.
diff --git a/parquet/src/encryption/ciphers.rs b/parquet/src/encryption/ciphers.rs
@@ -24,11 +24,11 @@ use ring::rand::{SecureRandom, SystemRandom};
 use crate::errors::{ParquetError, Result};
 
 pub trait BlockEncryptor {
-    fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Vec<u8>;
+    fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Result<Vec<u8>>;
 }
 
 pub trait BlockDecryptor {
-    fn decrypt(&self, length_and_ciphertext: &[u8], aad: &[u8]) -> Vec<u8>;
+    fn decrypt(&self, length_and_ciphertext: &[u8], aad: &[u8]) -> Result<Vec<u8>>;
 }
 
 const RIGHT_TWELVE: u128 = 0x0000_0000_ffff_ffff_ffff_ffff_ffff_ffff;
@@ -102,8 +102,8 @@ impl RingGcmBlockEncryptor {
 }
 
 impl BlockEncryptor for RingGcmBlockEncryptor {
-    fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Vec<u8> {
-        let nonce = self.nonce_sequence.advance().unwrap();
+    fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Result<Vec<u8>> {
+        let nonce = self.nonce_sequence.advance()?;
         let ciphertext_len = plaintext.len() + NONCE_LEN + TAG_LEN;
         // todo TBD: add first 4 bytes with the length, per https://github.com/apache/parquet-format/blob/master/Encryption.md#51-encrypted-module-serialization
         let mut result = Vec::with_capacity(SIZE_LEN + ciphertext_len);
@@ -113,11 +113,10 @@ impl BlockEncryptor for RingGcmBlockEncryptor {
 
         let tag = self
             .key
-            .seal_in_place_separate_tag(nonce, Aad::from(aad), &mut result[SIZE_LEN + NONCE_LEN..])
-            .unwrap();
+            .seal_in_place_separate_tag(nonce, Aad::from(aad), &mut result[SIZE_LEN + NONCE_LEN..])?;
         result.extend_from_slice(tag.as_ref());
 
-        result
+        Ok(result)
     }
 }
 
@@ -138,23 +137,22 @@ impl RingGcmBlockDecryptor {
 }
 
 impl BlockDecryptor for RingGcmBlockDecryptor {
-    fn decrypt(&self, length_and_ciphertext: &[u8], aad: &[u8]) -> Vec<u8> {
+    fn decrypt(&self, length_and_ciphertext: &[u8], aad: &[u8]) -> Result<Vec<u8>> {
         let mut result = Vec::with_capacity(
             length_and_ciphertext.len() - SIZE_LEN - NONCE_LEN - TAG_LEN,
         );
         result.extend_from_slice(&length_and_ciphertext[SIZE_LEN + NONCE_LEN..]);
 
         let nonce = ring::aead::Nonce::try_assume_unique_for_key(
             &length_and_ciphertext[SIZE_LEN..SIZE_LEN + NONCE_LEN],
-        )
-        .unwrap();
+        )?;
 
         self.key
-            .open_in_place(nonce, Aad::from(aad), &mut result)
-            .unwrap();
+            .open_in_place(nonce, Aad::from(aad), &mut result)?;
 
+        // Truncate result to remove the tag
         result.resize(result.len() - TAG_LEN, 0u8);
-        result
+        Ok(result)
     }
 }
 

diff --git a/parquet/src/errors.rs b/parquet/src/errors.rs
@@ -120,6 +120,13 @@ impl From<object_store::Error> for ParquetError {
     }
 }
 
+//#[cfg(feature = "encryption")]
+impl From<ring::error::Unspecified> for ParquetError {
+    fn from(e: ring::error::Unspecified) -> ParquetError {
+        ParquetError::External(Box::new(e))
+    }
+}
+
 /// A specialized `Result` for Parquet errors.
 pub type Result<T, E = ParquetError> = result::Result<T, E>;
 

diff --git a/parquet/src/file/metadata/reader.rs b/parquet/src/file/metadata/reader.rs
@@ -671,7 +671,7 @@ impl ParquetMetaDataReader {
             // file_decryptor = Some(FileDecryptor::new(file_decryption_properties, aad, aad_prefix));
 
             decrypted_fmd_buf =
-                decryptor.decrypt(prot.as_slice().as_ref(), aad_footer.as_ref());
+                decryptor.decrypt(prot.as_slice().as_ref(), aad_footer.as_ref())?;
             prot = TCompactSliceInputProtocol::new(decrypted_fmd_buf.as_ref());
         }
 

diff --git a/parquet/src/file/serialized_reader.rs b/parquet/src/file/serialized_reader.rs
@@ -361,7 +361,7 @@ pub(crate) fn read_page_header<T: Read>(input: &mut T, crypto_context: Option<Ar
         let ciphertext_len = u32::from_le_bytes(len_bytes) as usize;
         let mut ciphertext = vec![0; 4 + ciphertext_len];
         input.read_exact(&mut ciphertext[4..])?;
-        let buf = file_decryptor.decrypt(&ciphertext, aad.as_ref());
+        let buf = file_decryptor.decrypt(&ciphertext, aad.as_ref())?;
 
         let mut prot = TCompactSliceInputProtocol::new(buf.as_slice());
         let page_header = PageHeader::read_from_in_protocol(&mut prot)?;
@@ -443,7 +443,7 @@ pub(crate) fn decode_page(
             crypto_context.column_ordinal,
             0,
         )?;
-        let decrypted = file_decryptor.decrypt(&buffer.as_ref(), &aad);
+        let decrypted = file_decryptor.decrypt(&buffer.as_ref(), &aad)?;
         Bytes::from(decrypted)
     } else {
         buffer