mobilecoinfoundation · jcape · Sep 13, 2022 · Sep 2, 2022 · Sep 2, 2022 · Sep 7, 2022
diff --git a/attest/ake/src/event.rs b/attest/ake/src/event.rs
@@ -354,32 +354,45 @@ impl MealyOutput for Vec<u8> {}
 
 /// A type similar to [`aead::Payload`] used to distinguish writer inputs from
 /// outputs when there's an explicit nonce.
-pub struct NoncePlaintext<'aad, 'msg> {
-    pub aad: &'aad [u8],
-    pub msg: &'msg [u8],
-    pub nonce: u64,
-}
+pub struct NoncePlaintext<'aad, 'msg>(Plaintext<'aad, 'msg>);
 
 impl<'aad, 'msg> NoncePlaintext<'aad, 'msg> {
-    pub fn new(aad: &'aad [u8], msg: &'msg [u8], nonce: u64) -> Self {
-        Self { aad, msg, nonce }
+    /// Create a new NoncePlaintext object from the given slices.
+    pub fn new(aad: &'aad [u8], msg: &'msg [u8]) -> Self {
+        Self(Plaintext::new(aad, msg))
+    }
+
+    /// Grab a reference to the internal `aad` slice.
+    pub fn aad(&self) -> &[u8] {
+        self.0.aad
+    }
+
+    /// Grab a reference to the internal `msg` slice.
+    pub fn msg(&self) -> &[u8] {
+        self.0.msg
     }
 }
 
 /// Plaintext may be provided to an FST for encryption into a vector
 impl MealyInput for NoncePlaintext<'_, '_> {}
 
+/// A tuple of bytes and a u64 can be output from an FST for the
+/// encrypt-for-explicit nonce case.
+impl MealyOutput for (Vec<u8>, u64) {}
+
 /// A type similar to [`aead::Payload`] used to distinguish reader inputs from
 /// outputs when there's an explicit nonce.
 pub struct NonceCiphertext<'aad, 'msg> {
-    pub aad: &'aad [u8],
-    pub msg: &'msg [u8],
+    pub ciphertext: Ciphertext<'aad, 'msg>,
     pub nonce: u64,
 }
 
 impl<'aad, 'msg> NonceCiphertext<'aad, 'msg> {
     pub fn new(aad: &'aad [u8], msg: &'msg [u8], nonce: u64) -> Self {
-        Self { aad, msg, nonce }
+        Self {
+            ciphertex
8000
t: Ciphertext::new(aad, msg),
+            nonce,
+        }
     }
 }
 

diff --git a/attest/ake/src/shared.rs b/attest/ake/src/shared.rs
@@ -60,13 +60,14 @@ where
         input: NonceCiphertext<'_, '_>,
     ) -> Result<(Ready<Cipher>, Vec<u8>), Self::Error> {
         let mut retval = self;
-        let plaintext = retval.decrypt_with_nonce(input.aad, input.msg, input.nonce)?;
+        let plaintext =
+            retval.decrypt_with_nonce(input.ciphertext.aad, input.ciphertext.msg, input.nonce)?;
         Ok((retval, plaintext))
     }
 }
 
-/// Ready + NoncePlaintext => Ready + Vec
-impl<Cipher> Transition<Ready<Cipher>, NoncePlaintext<'_, '_>, Vec<u8>> for Ready<Cipher>
+/// Ready + NoncePlaintext => Ready + (Vec, u64)
+impl<Cipher> Transition<Ready<Cipher>, NoncePlaintext<'_, '_>, (Vec<u8>, u64)> for Ready<Cipher>
 where
     Cipher: NoiseCipher,
 {
@@ -76,9 +77,9 @@ where
         self,
         _csprng: &mut R,
         input: NoncePlaintext<'_, '_>,
-    ) -> Result<(Ready<Cipher>, Vec<u8>), Self::Error> {
+    ) -> Result<(Ready<Cipher>, (Vec<u8>, u64)), Self::Error> {
         let mut retval = self;
-        let ciphertext = retval.encrypt_with_nonce(input.aad, input.msg, input.nonce)?;
-        Ok((retval, ciphertext))
+        let output = retval.encrypt_with_nonce(input.aad(), input.msg())?;
+        Ok((retval, output))
     }
 }
diff --git a/attest/ake/src/state.rs b/attest/ake/src/state.rs
@@ -85,15 +85,16 @@ where
         self.reader.decrypt_with_ad(aad, ciphertext)
     }
 
-    /// Using the writer cipher, encrypt the given plaintext for the nonce.
+    /// Using the writer cipher, encrypt the given plaintext and return the
+    /// nonce.
     pub fn encrypt_with_nonce(
         &mut self,
         aad: &[u8],
         plaintext: &[u8],
-        nonce: u64,
-    ) -> Result<Vec<u8>, CipherError> {
-        self.writer.set_nonce(nonce);
-        self.encrypt(aad, plaintext)
+    ) -> Result<(Vec<u8>, u64), CipherError> {
+        let nonce = self.writer.next_nonce();
+        let ciphertext = self.encrypt(aad, plaintext)?;
+        Ok((ciphertext, nonce))
     }
 
     /// Using the reader cipher, decrypt the provided ciphertext for the nonce.

diff --git a/crypto/noise/src/cipher_state.rs b/crypto/noise/src/cipher_state.rs
@@ -161,12 +161,19 @@ impl<Cipher: NoiseCipher> CipherState<Cipher> {
         self.cipher.is_some()
     }
 
+    /// Retrieve the nonce value which will be used in the next operation.
+    ///
+    /// This is an extension of the noise protocol to allow for implicit-nonce
+    /// writers with explicit-nonce readers to co-exist in the same stream.
+    pub fn next_nonce(&self) -> u64 {
+        self.nonce
+    }
+
     /// The noise protocol `SetNonce()` operation.
     ///
     /// This will irrevocably override the current nonce value.
     pub fn set_nonce(&mut self, nonce: u64) {
         self.nonce = nonce;
-        // TODO: return current nonce? We don't provide any access otherwise...
     }
 
     /// The noise protocol `EncryptWithAd()` operation.
@@ -393,4 +400,14 @@ mod test {
         assert_eq!(encryptor.nonce, 2);
         assert_eq!(encryptor.bytes_sent, key.len() as u64);
     }
+
+    /// Try to set the nonce, and retrieve it.
+    #[test]
+    fn set_nonce() {
+        let mut encryptor = CipherState::<Aes256Gcm>::default();
+        let expected = 1234;
+        encryptor.set_nonce(expected);
+        let actual = encryptor.next_nonce();
+        assert_eq!(expected, actual);
+    }
 }