From 51f2ea3eadf7512054fec575986812e0083f504f Mon Sep 17 00:00:00 2001
From: James Cape <james@mobilecoin.com>
Date: Fri, 16 Sep 2022 12:33:03 -0700
Subject: [PATCH] Re-apply 2463.

---
 attest/api/src/conversions.rs | 20 ++++----
 attest/enclave-api/src/lib.rs | 88 ++++++++++++++++++++++++++++++++++-
 2 files changed, 98 insertions(+), 10 deletions(-)

diff --git a/attest/api/src/conversions.rs b/attest/api/src/conversions.rs
index c3ca7f8f5d..e20a825bfe 100644
--- a/attest/api/src/conversions.rs
+++ b/attest/api/src/conversions.rs
@@ -5,7 +5,7 @@
 use crate::attest::{AuthMessage, Message, NonceMessage};
 use mc_attest_ake::{AuthRequestOutput, AuthResponseOutput};
 use mc_attest_enclave_api::{
-    ClientAuthRequest, ClientAuthResponse, EnclaveMessage, EnclaveNonceMessage, PeerAuthRequest,
+    ClientAuthRequest, ClientAuthResponse, EnclaveMessage, NonceSession, PeerAuthRequest,
     PeerAuthResponse, Session,
 };
 use mc_crypto_keys::Kex;
@@ -109,24 +109,26 @@ impl<S: Session> From<EnclaveMessage<S>> for Message {
     }
 }
 
-impl<S: Session> From<NonceMessage> for EnclaveNonceMessage<S> {
-    fn from(src: NonceMessage) -> EnclaveNonceMessage<S> {
-        EnclaveNonceMessage {
+impl From<NonceMessage> for EnclaveMessage<NonceSession> {
+    fn from(src: NonceMessage) -> Self {
+        let channel_id = NonceSession::new(src.channel_id, src.nonce);
+        Self {
             aad: src.aad,
-            channel_id: S::from(&src.channel_id),
+            channel_id,
             data: src.data,
-            nonce: src.nonce,
         }
     }
 }
 
-impl<S: Session> From<EnclaveNonceMessage<S>> for NonceMessage {
-    fn from(src: EnclaveNonceMessage<S>) -> NonceMessage {
+impl From<EnclaveMessage<NonceSession>> for NonceMessage {
+    fn from(src: EnclaveMessage<NonceSession>) -> NonceMessage {
         let mut retval = NonceMessage::default();
         retval.set_aad(src.aad);
+        // it doesn't matter if we don't bump the nonce when retrieving it,
+        // src.channel_id will be discarded anyways.
+        retval.set_nonce(src.channel_id.peek_nonce());
         retval.set_channel_id(src.channel_id.into());
         retval.set_data(src.data);
-        retval.set_nonce(src.nonce);
         retval
     }
 }
diff --git a/attest/enclave-api/src/lib.rs b/attest/enclave-api/src/lib.rs
index efc4372ee2..bb1950cae4 100644
--- a/attest/enclave-api/src/lib.rs
+++ b/attest/enclave-api/src/lib.rs
@@ -12,7 +12,7 @@ mod error;
 pub use error::{Error, Result};
 
 use alloc::vec::Vec;
-use core::hash::Hash;
+use core::hash::{Hash, Hasher};
 use mc_attest_core::{IntelSealed, QuoteNonce, Report};
 use serde::{Deserialize, Serialize};
 
@@ -35,6 +35,7 @@ macro_rules! impl_newtype_vec_inout {
 impl_newtype_vec_inout! {
     ClientAuthRequest; ClientAuthResponse; ClientSession;
     PeerAuthRequest; PeerAuthResponse; PeerSession;
+    NonceAuthRequest; NonceAuthResponse;
 }
 
 /// The raw authentication request message, sent from an initiator to a
@@ -57,6 +58,16 @@ pub struct PeerAuthRequest(Vec<u8>);
 #[derive(Clone, Debug, Default, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
 pub struct PeerAuthResponse(Vec<u8>);
 
+/// The raw authentication request message, sent from an initiator to a
+/// responder.
+#[derive(Clone, Debug, Default, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
+pub struct NonceAuthRequest(Vec<u8>);
+
+/// The raw authentication response message, sent from a responder to an
+/// initiator.
+#[derive(Clone, Debug, Default, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
+pub struct NonceAuthResponse(Vec<u8>);
+
 /// Inbound and outbound messages to/from an enclave.
 #[derive(Clone, Debug, Default, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
 pub struct EnclaveMessage<S: Session> {
@@ -143,3 +154,78 @@ impl Session for PeerSession {
     type Request = PeerAuthRequest;
     type Response = PeerAuthResponse;
 }
+
+/// An opaque bytestream used as a session ID for a session which uses explicit
+/// nonces.
+#[derive(Clone, Debug, Default, Deserialize, PartialOrd, Serialize)]
+pub struct NonceSession {
+    channel_id: Vec<u8>,
+    nonce: u64,
+}
+
+impl NonceSession {
+    /// Create a new nonce session from a vector and nonce.
+    ///
+    /// This takes a pre-created Vec in order to remove an extra allocation that
+    /// would be required when converting from a NonceMessage to an
+    /// [`EnclaveMessage`]`<`[`NonceSession`]`>`.
+    pub fn new(channel_id: Vec<u8>, nonce: u64) -> Self {
+        Self { channel_id, nonce }
+    }
+
+    /// Retrieves the nonce for this session
+    pub fn peek_nonce(&self) -> u64 {
+        self.nonce
+    }
+
+    /// Retrieves a copy of the nonce, and increments it for the next time.
+    pub fn get_nonce(&mut self) -> u64 {
+        let retval = self.nonce;
+        self.nonce += 1;
+        retval
+    }
+}
+
+impl AsRef<[u8]> for NonceSession {
+    fn as_ref(&self) -> &[u8] {
+        self.channel_id.as_ref()
+    }
+}
+
+impl<'bytes> From<&'bytes [u8]> for NonceSession {
+    fn from(src: &'bytes [u8]) -> Self {
+        Self::from(Vec::from(src))
+    }
+}
+
+impl From<Vec<u8>> for NonceSession {
+    fn from(channel_id: Vec<u8>) -> Self {
+        NonceSession {
+            channel_id,
+            nonce: 0,
+        }
+    }
+}
+
+impl From<NonceSession> for Vec<u8> {
+    fn from(src: NonceSession) -> Self {
+        src.channel_id
+    }
+}
+
+impl Hash for NonceSession {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.channel_id.hash(state)
+    }
+}
+
+impl PartialEq for NonceSession {
+    fn eq(&self, other: &Self) -> bool {
+        self.channel_id == other.channel_id
+    }
+}
+
+impl Session for NonceSession {
+    type Request = NonceAuthRequest;
+    type Response = NonceAuthResponse;
+}