rok commented on code in PR #7111:
URL: https://github.com/apache/arrow-rs/pull/7111#discussion_r2004521719


##########
parquet/src/encryption/ciphers.rs:
##########
@@ -61,3 +64,118 @@ impl BlockDecryptor for RingGcmBlockDecryptor {
         Ok(result)
     }
 }
+
+pub trait BlockEncryptor: Debug + Send + Sync {
+    fn encrypt(&mut self, plaintext: &[u8], aad: &[u8]) -> Result<Vec<u8>>;
+}
+
+#[derive(Debug, Clone)]
+struct CounterNonce {
+    start: u128,
+    counter: u128,
+}
+
+impl CounterNonce {
+    pub fn new(rng: &SystemRandom) -> Result<Self> {
+        let mut buf = [0; 16];
+        rng.fill(&mut buf)?;
+
+        // 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);
+
+        Ok(Self { start, counter })
+    }
+
+    /// 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()
+    }
+}
+
+impl NonceSequence for CounterNonce {
+    fn advance(&mut self) -> Result<ring::aead::Nonce, 
ring::error::Unspecified> {
+        // If we've wrapped around, we've exhausted this nonce sequence
+        if (self.counter & RIGHT_TWELVE) == (self.start & RIGHT_TWELVE) {
+            Err(ring::error::Unspecified)
+        } else {
+            // Otherwise, just advance and return the new value
+            let buf: [u8; NONCE_LEN] = self.get_bytes();
+            self.counter = self.counter.wrapping_add(1);
+            Ok(ring::aead::Nonce::assume_unique_for_key(buf))
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct RingGcmBlockEncryptor {
+    key: LessSafeKey,
+    nonce_sequence: CounterNonce,
+}
+
+impl RingGcmBlockEncryptor {
+    /// 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]) -> Result<Self> {
+        let rng = SystemRandom::new();
+
+        // todo support other key sizes
+        let key = UnboundKey::new(&AES_128_GCM, key_bytes)
+            .map_err(|e| general_err!("Error creating AES key: {}", e))?;
+        let nonce = CounterNonce::new(&rng)?;
+
+        Ok(Self {
+            key: LessSafeKey::new(key),
+            nonce_sequence: nonce,
+        })
+    }
+}
+
+impl BlockEncryptor for RingGcmBlockEncryptor {
+    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());

Review Comment:
   Fair point! I've changed this to:
   ```rust
           let ciphertext_length : u32 = (NONCE_LEN + plaintext.len() + 
TAG_LEN).try_into()
               .map_err(|err|
                   General(format!("Plaintext data too long. {:?}", err)))?;
           // Not checking for overflow here because we've already checked for 
it with ciphertext_length
           let mut ciphertext = Vec::with_capacity(SIZE_LEN + ciphertext_length 
as usize);
           ciphertext.extend((ciphertext_length).to_le_bytes());
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org

Reply via email to