Add ability to prepare the next operation and it's associated nonce for OpeningKey and SealingKey

aws-lc-rs/src/aead.rs

@@ -55,7 +55,7 @@
 //! # Nonce Sequence APIs
 //!
 //! The [`UnboundKey`], [`OpeningKey`], [`SealingKey`], and [`LessSafeKey`] types are the
-//! AEAD API's provided for compatability with the original *ring* API.
+//! AEAD API's provided for compatibility with the original *ring* API.
 //!
 //! Users should prefer [`RandomizedNonceKey`] which provides a simplified experience around
 //! Nonce construction.
@@ -116,7 +116,8 @@
 
 use crate::{derive_debug_via_id, error::Unspecified, hkdf};
 use aead_ctx::AeadCtx;
-use core::{fmt::Debug, ops::RangeFrom};
+use core::{fmt::Debug, ops::RangeFrom, stringify};
+use paste::paste;
 
 mod aead_ctx;
 mod aes_gcm;
@@ -307,6 +308,13 @@ impl<N: NonceSequence> OpeningKey<N> {
             ciphertext_and_tag,
         )
     }
+
+    /// Returns a `OpeningKeyOpMut` with the next computed `Nonce` from the `NonceSequence` for the next operation.
+    /// # Errors
+    /// `Unspecified` if there is a failure computing the nonce for the next operation.
+    pub fn prepare_operation(&mut self) -> Result<OpeningKeyOpMut<'_, N>, Unspecified> {
+        OpeningKeyOpMut::new(self)
+    }
 }
 
 /// An AEAD key for encrypting and signing ("sealing"), bound to a nonce
@@ -433,6 +441,149 @@ impl<N: NonceSequence> SealingKey<N> {
             .seal_in_place_separate_tag(Some(self.nonce_sequence.advance()?), aad.as_ref(), in_out)
             .map(|(_, tag)| tag)
     }
+
+    /// Returns a `SealingKeyOpMut` with the next computed `Nonce` from the `NonceSequence` for the next operation.
+    /// # Errors
+    /// `Unspecified` if there is a failure computing the nonce for the next operation.
+    pub fn prepare_operation(&mut self) -> Result<SealingKeyOpMut<'_, N>, Unspecified> {
+        SealingKeyOpMut::new(self)
+    }
+}
+
+macro_rules! nonce_seq_key_op_mut {
+    ($name:ident) => {
+        paste! {
+        /// A key operation with a precomputed nonce from a key's associated `NonceSequence`.
+        pub struct [<$name OpMut>]<'a, N: NonceSequence> {
+            key: &'a mut $name<N>,
+            nonce: Nonce,
+        }
+
+        impl<'a, N: NonceSequence> [<$name OpMut>]<'a, N> {
+            fn new(key: &'a mut $name<N>) -> Result<Self, Unspecified> {
+                let nonce = key.nonce_sequence.advance()?;
+                Ok(Self {
+                    key,
+                    nonce,
+                })
+            }
+        }
+
+        impl<'a, N: NonceSequence> Debug for [<$name OpMut>]<'a, N> {
+            fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
+                f.debug_struct(stringify!([<$name OpMut>])).finish_non_exhaustive()
+            }
+        }
+        }
+    };
+}
+
+nonce_seq_key_op_mut!(OpeningKey);
+nonce_seq_key_op_mut!(SealingKey);
+
+impl<N: NonceSequence> OpeningKeyOpMut<'_, N> {
+    /// Returns the Nonce that will be used for this operation.
+    #[must_use]
+    pub fn nonce(&self) -> Nonce {
+        let nonce_bytes = self.nonce.0.as_ref();
+        let nonce: Nonce = Nonce(nonce_bytes.into());
+        nonce
+    }
+
+    /// Authenticates and decrypts (“opens”) data in place.
+    ///
+    /// See [OpeningKey::open_in_place] for additional API information.
+    ///
+    /// # Errors
+    /// `error::Unspecified` when ciphertext is invalid. In this case, `in_out` may have been
+    /// overwritten in an unspecified way.
+    #[inline]
+    #[allow(clippy::needless_pass_by_value)]
+    pub fn open_in_place<A>(self, aad: Aad<A>, in_out: &mut [u8]) -> Result<&mut [u8], Unspecified>
+    where
+        A: AsRef<[u8]>,
+    {
+        self.open_within(aad, in_out, 0..)
+    }
+
+    /// Authenticates and decrypts (“opens”) data in place, with a shift.
+    ///
+    /// See [OpeningKey::open_within] for additional API information.
+    ///
+    /// # Errors
+    /// `error::Unspecified` when ciphertext is invalid. In this case, `in_out` may have been
+    /// overwritten in an unspecified way.
+    #[inline]
+    #[allow(clippy::needless_pass_by_value)]
+    pub fn open_within<A>(
+        self,
+        aad: Aad<A>,
+        in_out: &mut [u8],
+        ciphertext_and_tag: RangeFrom<usize>,
+    ) -> Result<&mut [u8], Unspecified>
+    where
+        A: AsRef<[u8]>,
+    {
+        self.key
+            .key
+            .open_within(self.nonce, aad.as_ref(), in_out, ciphertext_and_tag)
+    }
+}
+
+impl<N: NonceSequence> SealingKeyOpMut<'_, N> {
+    /// Returns the Nonce that will be used for this operation.
+    #[must_use]
+    pub fn nonce(&self) -> Nonce {
+        let nonce_bytes = self.nonce.0.as_ref();
+        let nonce: Nonce = Nonce(nonce_bytes.into());
+        nonce
+    }
+
+    /// Encrypts and signs (“seals”) data in place, appending the tag to the
+    /// resulting ciphertext.
+    ///
+    /// See [SealingKey::seal_in_place_append_tag] for additional API information.
+    ///
+    /// # Errors
+    /// `error::Unspecified` when `nonce_sequence` cannot be advanced.
+    #[inline]
+    #[allow(clippy::needless_pass_by_value)]
+    pub fn seal_in_place_append_tag<A, InOut>(
+        self,
+        aad: Aad<A>,
+        in_out: &mut InOut,
+    ) -> Result<(), Unspecified>
+    where
+        A: AsRef<[u8]>,
+        InOut: AsMut<[u8]> + for<'in_out> Extend<&'in_out u8>,
+    {
+        self.key
+            .key
+            .seal_in_place_append_tag(Some(self.nonce), aad.as_ref(), in_out)
+            .map(|_| ())
+    }
+
+    /// Encrypts and signs (“seals”) data in place.
+    ///
+    /// See [`SealingKey::seal_in_place_separate_tag`] for additional API information.
+    ///
+    /// # Errors
+    /// `error::Unspecified` when `nonce_sequence` cannot be advanced.
+    #[inline]
+    #[allow(clippy::needless_pass_by_value)]
+    pub fn seal_in_place_separate_tag<A>(
+        self,
+        aad: Aad<A>,
+        in_out: &mut [u8],
+    ) -> Result<Tag, Unspecified>
+    where
+        A: AsRef<[u8]>,
+    {
+        self.key
+            .key
+            .seal_in_place_separate_tag(Some(self.nonce), aad.as_ref(), in_out)
+            .map(|(_, tag)| tag)
+    }
 }
 
 /// The additionally authenticated data (AAD) for an opening or sealing
@@ -786,6 +937,12 @@ impl AsRef<[u8]> for Tag {
     }
 }
 
+impl core::fmt::Debug for Tag {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_tuple("Tag").finish()
+    }
+}
+
 #[allow(dead_code)]
 const MAX_KEY_LEN: usize = 32;
 
@@ -797,6 +954,8 @@ pub const MAX_TAG_LEN: usize = TAG_LEN;
 
 #[cfg(test)]
 mod tests {
+    use nonce_sequence::Counter32Builder;
+
     use super::*;
     use crate::{iv::FixedLength, test::from_hex};
 
@@ -853,4 +1012,26 @@ mod tests {
 
         assert_eq!(plaintext, in_out[..plaintext.len()]);
     }
+
+    #[test]
+    fn debug_key_op_mut() {
+        let mut sk = SealingKey::new(
+            UnboundKey::new(&AES_128_GCM, &[0u8; 16]).unwrap(),
+            Counter32Builder::new().build(),
+        );
+        let mut ok = OpeningKey::new(
+            UnboundKey::new(&AES_128_GCM, &[0u8; 16]).unwrap(),
+            Counter32Builder::new().build(),
+        );
+        let so = sk.prepare_operation().unwrap();
+        let oo = ok.prepare_operation().unwrap();
+        assert_eq!("SealingKeyOpMut { .. }", format!("{so:?}"));
+        assert_eq!("OpeningKeyOpMut { .. }", format!("{oo:?}"));
+    }
+
+    #[test]
+    fn debug_tag() {
+        let tag = Tag([0u8; MAX_TAG_LEN], MAX_TAG_LEN);
+        assert_eq!("Tag", format!("{tag:?}"));
+    }
 }

aws-lc-rs/tests/aead_test.rs

@@ -3,9 +3,12 @@
 // Modifications copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0 OR ISC
 
+use aws_lc_rs::aead::nonce_sequence::Counter32Builder;
 use aws_lc_rs::{aead, error, test, test_file};
 
-use aws_lc_rs::aead::{Nonce, NONCE_LEN};
+use aws_lc_rs::aead::{
+    Aad, BoundKey, Nonce, OpeningKey, SealingKey, UnboundKey, AES_128_GCM, NONCE_LEN,
+};
 use core::ops::RangeFrom;
 
 #[test]
@@ -660,3 +663,90 @@ impl aead::NonceSequence for OneNonceSequence {
         self.0.take().ok_or(error::Unspecified)
     }
 }
+
+#[test]
+fn prepare_operation() {
+    const KEY: &[u8] = &[
+        0x52, 0x05, 0x19, 0x7a, 0xcc, 0x88, 0xdb, 0x78, 0x39, 0x59, 0xbc, 0x03, 0xb8, 0x1d, 0x4a,
+        0x6c,
+    ];
+    const MESSAGE: &[u8] = &[
+        0x52, 0x61, 0x63, 0x63, 0x6f, 0x6f, 0x6e, 0x20, 0x4d, 0x69, 0x73, 0x63, 0x68, 0x69, 0x65,
+        0x66,
+    ];
+    const LIMIT: u32 = 10;
+
+    let mut sk = SealingKey::new(
+        UnboundKey::new(&AES_128_GCM, KEY).unwrap(),
+        Counter32Builder::new().limit(LIMIT).build(),
+    );
+    let mut ok = OpeningKey::new(
+        UnboundKey::new(&AES_128_GCM, KEY).unwrap(),
+        Counter32Builder::new().limit(LIMIT).build(),
+    );
+
+    let mut nonces: Vec<Vec<u8>> = vec![];
+
+    for _ in 0..(LIMIT / 2) {
+        let so = sk.prepare_operation().unwrap();
+        let oo = ok.prepare_operation().unwrap();
+        let so_nonce = Vec::from(so.nonce().as_ref());
+        let oo_nonce = Vec::from(oo.nonce().as_ref());
+
+        assert_eq!(so_nonce.as_slice(), oo_nonce.as_slice());
+        assert!(!nonces.contains(&so_nonce));
+        nonces.push(so_nonce);
+        nonces.push(oo_nonce);
+
+        let mut message: Vec<u8> = vec![];
+        message.extend_from_slice(MESSAGE);
+
+        so.seal_in_place_append_tag(Aad::empty(), &mut message)
+            .unwrap();
+        assert_ne!(MESSAGE, message.as_slice());
+
+        let message = oo.open_in_place(Aad::empty(), &mut message).unwrap();
+        assert_eq!(MESSAGE, message);
+
+        let so = sk.prepare_operation().unwrap();
+        let oo = ok.prepare_operation().unwrap();
+        let so_nonce = Vec::from(so.nonce().as_ref());
+        let oo_nonce = Vec::from(oo.nonce().as_ref());
+
+        assert_eq!(so_nonce.as_slice(), oo_nonce.as_slice());
+        assert!(!nonces.contains(&so_nonce));
+        nonces.push(so_nonce);
+        nonces.push(oo_nonce);
+
+        let mut message: Vec<u8> = vec![];
+        message.extend_from_slice(MESSAGE);
+
+        let tag = so
+            .seal_in_place_separate_tag(Aad::empty(), &mut message)
+            .unwrap();
+        assert_ne!(MESSAGE, message.as_slice());
+        message.extend_from_slice(tag.as_ref());
+
+        let message = oo.open_within(Aad::empty(), &mut message, 0..).unwrap();
+        assert_eq!(MESSAGE, message);
+    }
+
+    let nonce_chunks = nonces.chunks_exact(2);
+    assert_eq!(0, nonce_chunks.remainder().len());
+    for chunk in nonce_chunks {
+        assert_eq!(chunk[0].as_slice(), chunk[1].as_slice());
+    }
+
+    let mut message: Vec<u8> = vec![];
+    message.extend_from_slice(MESSAGE);
+
+    // Subsequent usage should fail now since the sequence is exhausted in each key.
+    sk.prepare_operation().expect_err("sequence limit reached");
+    ok.prepare_operation().expect_err("sequence limit reached");
+    sk.seal_in_place_append_tag(Aad::empty(), &mut message)
+        .expect_err("sequence limit reached");
+    sk.seal_in_place_separate_tag(Aad::empty(), &mut message)
+        .expect_err("sequence limit reached");
+    ok.open_in_place(Aad::empty(), &mut message)
+        .expect_err("sequence limit reached");
+}