feat(DhSessions): add with_acquire_* helpers and AcquireSessionError

rtc_tenclave/src/dh/sessions.rs

@@ -5,6 +5,7 @@ use std::sync::Arc;
 
 use once_cell::sync::OnceCell;
 use rtc_types::dh::{ExchangeReportResult, SessionRequestResult};
+use rtc_types::enclave_messages::errors::AcquireSessionError;
 use secrecy::Secret;
 use sgx_types::*;
 
@@ -65,6 +66,32 @@ where
     TResp: RtcDhResponder,
     TInit: RtcDhInitiator,
 {
+    /// Acquire exclusive access to a session's channel, establishing a new session if necessary.
+    pub fn with_acquire_new_or_established<T>(
+        &self,
+        enclave_id: sgx_enclave_id_t,
+        f: impl FnOnce(&mut ProtectedChannel) -> T,
+    ) -> Result<T, AcquireSessionError> {
+        let channel_mutex = self.get_or_create_session(enclave_id)?;
+        let channel_guard = &mut channel_mutex.lock()?;
+        let result = f(channel_guard);
+        Ok(result)
+    }
+
+    /// Acquire exclusive access to an already-established session's channel.
+    pub fn with_acquire_established<T>(
+        &self,
+        enclave_id: sgx_enclave_id_t,
+        f: impl FnOnce(&mut ProtectedChannel) -> T,
+    ) -> Result<T, AcquireSessionError> {
+        let channel_mutex = self
+            .get_active(enclave_id)
+            .ok_or_else(|| AcquireSessionError::NoActiveSession(enclave_id))?;
+        let channel_guard = &mut channel_mutex.lock()?;
+        let result = f(channel_guard);
+        Ok(result)
+    }
+
     fn get(&self, enclave_id: sgx_enclave_id_t) -> Option<Arc<AtomicSession<TResp>>> {
         self.sessions
             .read()
@@ -332,4 +359,41 @@ mod test {
             _phantom_init: PhantomData::default(),
         })
     }
+
+    // TODO: Tests for [`DhSessions::with_acquire_new_or_established`]
+    //       (Trying these currently fail with linker errors.)
+
+    #[test]
+    fn with_acquire_established_no_active() {
+        let sessions = dh_sessions();
+        sessions.close_session(5);
+
+        let err = sessions
+            .with_acquire_established(5, |channel| {
+                channel.encrypt_message([1, 2, 3], [4, 5, 6]).unwrap()
+            })
+            .unwrap_err();
+
+        assert_eq!(err, AcquireSessionError::NoActiveSession(5));
+    }
+
+    #[test]
+    fn with_acquire_established_active() {
+        let sessions = dh_sessions();
+        sessions
+            .set_active(5, AlignedKey::new(sgx_align_key_128bit_t::default()))
+            .unwrap();
+
+        let sealed = sessions
+            .with_acquire_established(5, |channel| {
+                channel.encrypt_message([1, 2, 3], [4, 5, 6]).unwrap()
+            })
+            .unwrap();
+
+        let unsealed = sessions
+            .with_acquire_established(5, |channel| channel.decrypt_message(sealed).unwrap())
+            .unwrap();
+
+        assert_eq!(unsealed, (([1, 2, 3], [4, 5, 6])));
+    }
 }

rtc_types/src/enclave_messages/errors.rs

@@ -0,0 +1,38 @@
+//! Errors related to enclave messages and sealing.
+
+use std::sync::PoisonError;
+
+use sgx_types::{sgx_enclave_id_t, sgx_status_t};
+use thiserror::Error;
+
+/// Failed to acquire session / protected channel.
+///
+/// See: `rtc_tenclave::dh::sessions::DhSessions`
+#[derive(Debug, PartialEq)] // core
+#[derive(Error)] // thiserror
+pub enum AcquireSessionError {
+    /// This should generally be treated as an unrecoverable error.
+    #[error("Channel mutex poisoned")]
+    ChannelMutexPoisoned, // see impl From<PoisonError>
+
+    #[error("No active session for enclave ID {0}")]
+    NoActiveSession(sgx_enclave_id_t),
+
+    #[error("SGX error: {0:?}")]
+    Sgx(sgx_status_t),
+}
+
+/// [`PoisonError`] contains a lock guard, which requires lifetime propagation,
+/// but we're not interested in using or recovering from a poisoned lock,
+/// so we discard the guard here.
+impl<_Guard> From<PoisonError<_Guard>> for AcquireSessionError {
+    fn from(_: PoisonError<_Guard>) -> Self {
+        AcquireSessionError::ChannelMutexPoisoned
+    }
+}
+
+impl From<sgx_status_t> for AcquireSessionError {
+    fn from(err: sgx_status_t) -> Self {
+        AcquireSessionError::Sgx(err)
+    }
+}

rtc_types/src/enclave_messages/mod.rs

@@ -19,6 +19,8 @@ pub struct EncryptedEnclaveMessage<const MESSAGE_SIZE: usize, const AAD_SIZE: us
     pub nonce: RecommendedAesGcmIv,
 }
 
+pub mod errors;
+
 /// XXX: Ignore this module to work around cbindgen generic type handling
 ///
 /// Issues: