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