Implement enclave APIs that facilitate Fog View Router dynamic store …

…discovery (#2200)

* Add client_init and client_connect enclave methods

* Move multiview query construction to ake enclave

* Implement Remoun's suggestions

* Implement second round of remoun's suggestions

crypto/ake/enclave/src/lib.rs

@@ -8,7 +8,7 @@ use alloc::{string::ToString, vec::Vec};
 use digest::Digest;
 use mc_attest_ake::{
     AuthPending, AuthRequestOutput, AuthResponseInput, AuthResponseOutput, ClientAuthRequestInput,
-    NodeAuthRequestInput, NodeInitiate, Ready, Start, Transition,
+    ClientInitiate, NodeAuthRequestInput, NodeInitiate, Ready, Start, Transition,
 };
 use mc_attest_core::{
     IasNonce, Nonce, NonceError, Quote, QuoteNonce, Report, ReportData, TargetInfo,
@@ -36,9 +36,15 @@ const MAX_AUTH_PENDING_REQUESTS: usize = 64;
 /// Max number of peer sessions.
 const MAX_PEER_SESSIONS: usize = 64;
 
+/// Max number of backends that this enclave can connect to as a client.
+const MAX_BACKEND_CONNECTIONS: usize = 10000;
+
 /// Max number of client sessions.
 const MAX_CLIENT_SESSIONS: usize = 10000;
 
+/// Max number of auth requests for enclave backends.
+const MAX_BACKEND_AUTH_PENDING_REQUESTS: usize = 10000;
+
 /// Any additional "identities" (e.g. key material) for a given enclave that
 /// needs to become a part of the report. We provide some simple identities, and
 /// a trait to allow extensions
@@ -75,6 +81,10 @@ pub struct AkeEnclaveState<EI: EnclaveIdentity> {
     /// A map of responder-ID to incomplete, outbound, AKE state.
     initiator_auth_pending: Mutex<LruCache<ResponderId, AuthPending<X25519, Aes256Gcm, Sha512>>>,
 
+    /// A map of responder-ID to incomplete, outbound AKE state for connections
+    /// to enclaves that serve as backends to the current enclave.
+    backend_auth_pending: Mutex<LruCache<ResponderId, AuthPending<X25519, Aes256Gcm, Sha512>>>,
+
     /// A map of channel ID outbound connection state.
     peer_outbound: Mutex<LruCache<PeerSession, Ready<Aes256Gcm>>>,
 
@@ -83,6 +93,10 @@ pub struct AkeEnclaveState<EI: EnclaveIdentity> {
 
     /// A map of channel ID to connection state
     clients: Mutex<LruCache<ClientSession, Ready<Aes256Gcm>>>,
+
+    /// A map of ResponderIds for each enclave that serves as a backend to the
+    /// current enclave.
+    backends: Mutex<LruCache<ResponderId, Ready<Aes256Gcm>>>,
 }
 
 impl<EI: EnclaveIdentity + Default> Default for AkeEnclaveState<EI> {
@@ -103,9 +117,11 @@ impl<EI: EnclaveIdentity> AkeEnclaveState<EI> {
             ias_pending: Mutex::new(LruCache::new(MAX_PENDING_QUOTES)),
             current_ias_report: Mutex::new(None),
             initiator_auth_pending: Mutex::new(LruCache::new(MAX_AUTH_PENDING_REQUESTS)),
+            backend_auth_pending: Mutex::new(LruCache::new(MAX_BACKEND_AUTH_PENDING_REQUESTS)),
             peer_outbound: Mutex::new(LruCache::new(MAX_PEER_SESSIONS)),
             peer_inbound: Mutex::new(LruCache::new(MAX_PEER_SESSIONS)),
             clients: Mutex::new(LruCache::new(MAX_CLIENT_SESSIONS)),
+            backends: Mutex::new(LruCache::new(MAX_BACKEND_CONNECTIONS)),
         }
     }
 
@@ -165,6 +181,53 @@ impl<EI: EnclaveIdentity> AkeEnclaveState<EI> {
         }
     }
 
+    /// Constructs a ClientAuthRequest to be sent to an enclave backend.
+    ///
+    /// Differs from peer_init in that this enclave does not establish a peer
+    /// connection to the enclave described by `backend_id`. Rather, this
+    /// enclave serves as a client to this other backend enclave.
+    pub fn backend_init(&self, backend_id: ResponderId) -> Result<ClientAuthRequest> {
+        let mut csprng = McRng::default();
+
+        let initiator = Start::new(backend_id.to_string());
+
+        let init_input = ClientInitiate::<X25519, Aes256Gcm, Sha512>::default();
+        let (initiator, auth_request_output) = initiator.try_next(&mut csprng, init_input)?;
+        self.backend_auth_pending.lock()?.put(backend_id, initiator);
+        let client_auth_request_data: Vec<u8> = auth_request_output.into();
+        Ok(client_auth_request_data.into())
+    }
+
+    /// Connect to an enclave backend as a client.
+    ///
+    /// This establishes the client to backend enclave connection, see
+    /// `backend_init` for more details on how this differs from a peer
+    /// connection.
+    pub fn backend_connect(
+        &self,
+        backend_id: ResponderId,
+        backend_auth_response: ClientAuthResponse,
+    ) -> Result<()> {
+        let initiator = self
+            .backend_auth_pending
+            .lock()?
+            .pop(&backend_id)
+            .ok_or(Error::NotFound)?;
+
+        let mut csprng = McRng::default();
+
+        let auth_response_output_bytes: Vec<u8> = backend_auth_response.into();
+        let auth_response_event =
+            AuthResponseInput::new(auth_response_output_bytes.into(), self.get_verifier()?);
+        let (initiator, _verification_report) =
+            initiator.try_next(&mut csprng, auth_response_event)?;
+
+        let mut backends = self.backends.lock()?;
+        backends.put(backend_id, initiator);
+
+        Ok(())
+    }
+
     /// Accept a client connection
     pub fn client_accept(
         &self,
@@ -374,6 +437,36 @@ impl<EI: EnclaveIdentity> AkeEnclaveState<EI> {
         })
     }
 
+    /// Transforms an incoming client message, i.e. a message sent from a client
+    /// to the current enclave, into a list of outbound  messages for
+    /// other enclaves that serve as backends to the current enclave.
+    ///
+    ///                               / --> Backend Enclave 1
+    ///   Client -> Current Enclave ---> Backend Enclave 2
+    ///                              \ --> Backend Enclave N
+    pub fn reencrypt_client_message_for_backends(
+        &self,
+        incoming_client_message: EnclaveMessage<ClientSession>,
+    ) -> Result<Vec<EnclaveMessage<ClientSession>>> {
+        let client_query_bytes: Vec<u8> = self.client_decrypt(incoming_client_message.clone())?;
+        let mut backends = self.backends.lock()?;
+        let backend_messages = backends
+            .iter_mut()
+            .map(|(_, encryptor)| {
+                let aad = incoming_client_message.aad.clone();
+                let data = encryptor.encrypt(&aad, &client_query_bytes)?;
+                let channel_id = incoming_client_message.channel_id.clone();
+                Ok(EnclaveMessage {
+                    aad,
+                    channel_id,
+                    data,
+                })
+            })
+            .collect::<Result<_>>()?;
+
+        Ok(backend_messages)
+    }
+
     //
     // IAS related
     //

fog/view/enclave/api/src/lib.rs

@@ -87,6 +87,12 @@ pub enum ViewEnclaveRequest {
     /// Takes an encrypted fog_types::view::QueryRequest and returns a list of
     /// fog_types::view::QueryRequest.
     CreateMultiViewStoreQuery(EnclaveMessage<ClientSession>),
+    /// Begin a client connection to a Fog View Store discovered after
+    /// initialization.
+    ViewStoreInit(ResponderId),
+    /// Complete the client connection to a Fog View store that accepted our
+    /// client auth request. This is meant to be called after [ViewStoreInit].
+    ViewStoreConnect(ResponderId, ClientAuthResponse),
 }
 
 /// The parameters needed to initialize the view enclave
@@ -126,6 +132,19 @@ pub trait ViewEnclaveApi: ReportableEnclave {
     /// Destroy a peer association
     fn client_close(&self, channel_id: ClientSession) -> Result<()>;
 
+    /// Begin a connection to a Fog View Store. The enclave calling this method
+    /// will act as a client to the Fog View Store.
+    fn view_store_init(&self, view_store_id: ResponderId) -> Result<ClientAuthRequest>;
+
+    /// Complete the connection to a Fog View Store that has accepted our
+    /// ClientAuthRequest. This is meant to be called after the enclave has
+    /// initialized and discovers a new Fog View Store.
+    fn view_store_connect(
+        &self,
+        view_store_id: ResponderId,
+        view_store_auth_response: ClientAuthResponse,
+    ) -> Result<()>;
+
     /// Service a user's encrypted QueryRequest
     fn query(
         &self,

fog/view/enclave/impl/src/lib.rs

@@ -9,13 +9,13 @@ extern crate alloc;
 mod e_tx_out_store;
 use e_tx_out_store::{ETxOutStore, StorageDataSize, StorageMetaSize};
 
-use aes_gcm::Aes256Gcm;
 use alloc::vec::Vec;
-use core::ops::DerefMut;
-use mc_attest_ake::Ready;
 use mc_attest_core::{IasNonce, Quote, QuoteNonce, Report, TargetInfo, VerificationReport};
 use mc_attest_enclave_api::{ClientAuthRequest, ClientAuthResponse, ClientSession, EnclaveMessage};
-use mc_common::logger::{log, Logger};
+use mc_common::{
+    logger::{log, Logger},
+    ResponderId,
+};
 use mc_crypto_ake_enclave::{AkeEnclaveState, NullIdentity};
 use mc_crypto_keys::X25519Public;
 use mc_fog_recovery_db_iface::FogUserEvent;
@@ -42,13 +42,6 @@ where
 
     /// Logger object
     logger: Logger,
-
-    /// Encrypts a QueryRequest for each individual Fog View Store.
-    /// TODO: Use a BTreeMap<FogViewShardLoadBalancerID,
-    /// BTreeMap<FogViewStoreId, Ready<...>>>  when implement the cursoring
-    /// optimization. For right now, it's fine to leave as a Vec because a
-    /// follow up PR will implement cursoring.
-    store_encryptors: Mutex<Vec<Ready<Aes256Gcm>>>,
 }
 
 impl<OSC> ViewEnclave<OSC>
@@ -58,7 +51,6 @@ where
     pub fn new(logger: Logger) -> Self {
         Self {
             e_tx_out_store: Mutex::new(None),
-            store_encryptors: Mutex::new(Vec::new()),
             ake: Default::default(),
             logger,
         }
@@ -203,21 +195,22 @@ where
         &self,
         client_query: EnclaveMessage<ClientSession>,
     ) -> Result<Vec<EnclaveMessage<ClientSession>>> {
-        let client_query_bytes = self.ake.client_decrypt(client_query.clone())?;
-
-        let mut encryptors = self.store_encryptors.lock()?;
-        let mut results = Vec::with_capacity(encryptors.len());
-        for store_encryptor in encryptors.deref_mut() {
-            let aad = client_query.aad.clone();
-            let data = store_encryptor.encrypt(&aad, &client_query_bytes)?;
-            let channel_id = client_query.channel_id.clone();
-            results.push(EnclaveMessage {
-                aad,
-                channel_id,
-                data,
-            });
-        }
+        Ok(self
+            .ake
+            .reencrypt_client_message_for_backends(client_query)?)
+    }
 
-        Ok(results)
+    fn view_store_init(&self, view_store_id: ResponderId) -> Result<ClientAuthRequest> {
+        Ok(self.ake.backend_init(view_store_id)?)
+    }
+
+    fn view_store_connect(
+        &self,
+        view_store_id: ResponderId,
+        view_store_auth_response: ClientAuthResponse,
+    ) -> Result<()> {
+        Ok(self
+            .ake
+            .backend_connect(view_store_id, view_store_auth_response)?)
     }
 }

fog/view/enclave/src/lib.rs

@@ -161,6 +161,25 @@ impl ViewEnclaveApi for SgxViewEnclave {
         mc_util_serial::deserialize(&outbuf[..])?
     }
 
+    fn view_store_init(&self, view_store_id: ResponderId) -> Result<ClientAuthRequest> {
+        let inbuf = mc_util_serial::serialize(&ViewEnclaveRequest::ViewStoreInit(view_store_id))?;
+        let outbuf = self.enclave_call(&inbuf)?;
+        mc_util_serial::deserialize(&outbuf[..])?
+    }
+
+    fn view_store_connect(
+        &self,
+        view_store_id: ResponderId,
+        view_store_auth_response: ClientAuthResponse,
+    ) -> Result<()> {
+        let inbuf = mc_util_serial::serialize(&ViewEnclaveRequest::ViewStoreConnect(
+            view_store_id,
+            view_store_auth_response,
+        ))?;
+        let outbuf = self.enclave_call(&inbuf)?;
+        mc_util_serial::deserialize(&outbuf[..])?
+    }
+
     fn query(
         &self,
         payload: EnclaveMessage<ClientSession>,

fog/view/enclave/trusted/src/lib.rs

@@ -115,6 +115,8 @@ pub fn ecall_dispatcher(inbuf: &[u8]) -> Result<Vec<u8>, sgx_status_t> {
         }
         ViewEnclaveRequest::GetIasReport => serialize(&ENCLAVE.get_ias_report()),
         ViewEnclaveRequest::ClientAccept(msg) => serialize(&ENCLAVE.client_accept(msg)),
+        ViewEnclaveRequest::ViewStoreInit(view_store_id) => serialize(&ENCLAVE.view_store_init(view_store_id)),
+        ViewEnclaveRequest::ViewStoreConnect(view_store_id, msg) => serialize(&ENCLAVE.view_store_connect(view_store_id, msg)),
         ViewEnclaveRequest::ClientClose(session) => serialize(&ENCLAVE.client_close(session)),
         ViewEnclaveRequest::Query(req, untrusted_query_response) => {
             serialize(&ENCLAVE.query(req, untrusted_query_response))