Address review comments part II

- DescriptorGuard
- Fix description of recv operation in ABI
- Closed variant in TrySendError/TryRecvError for future
- Check pointers in FifoDescriptor in SGX using User<[T]>

enclave-runner/src/usercalls/mod.rs

@@ -20,7 +20,7 @@ use fortanix_sgx_abi::*;
 use futures::future::{poll_fn, Either, Future, FutureExt};
 use futures::lock::Mutex;
 use futures::StreamExt;
-use ipc_queue::{self, QueueEvent};
+use ipc_queue::{self, DescriptorGuard, QueueEvent};
 use sgxs::loader::Tcs as SgxsTcs;
 use std::alloc::{GlobalAlloc, Layout, System};
 use std::cell::RefCell;
@@ -584,9 +584,10 @@ impl EnclaveKind {
     }
 }
 
-struct FifoDescriptors {
-    usercall_queue: FifoDescriptor<Usercall>,
-    return_queue: FifoDescriptor<Return>,
+struct FifoGuards {
+    usercall_queue: DescriptorGuard<Usercall>,
+    return_queue: DescriptorGuard<Return>,
+    async_queues_called: bool,
 }
 
 pub(crate) struct EnclaveState {
@@ -598,7 +599,8 @@ pub(crate) struct EnclaveState {
     usercall_ext: Box<dyn UsercallExtension>,
     threads_queue: crossbeam::queue::SegQueue<StoppedTcs>,
     forward_panics: bool,
-    fifo_descriptors: Mutex<Option<FifoDescriptors>>,
+    // Once set to Some, the guards should not be dropped for the lifetime of the enclave.
+    fifo_guards: Mutex<Option<FifoGuards>>,
     return_queue_tx: Mutex<Option<ipc_queue::AsyncSender<Return, QueueSynchronizer>>>,
 }
 
@@ -691,7 +693,7 @@ impl EnclaveState {
             usercall_ext,
             threads_queue,
             forward_panics,
-            fifo_descriptors: Mutex::new(None),
+            fifo_guards: Mutex::new(None),
             return_queue_tx: Mutex::new(None),
         })
     }
@@ -820,12 +822,13 @@ impl EnclaveState {
             let (usercall_queue_tx, usercall_queue_rx) = ipc_queue::bounded_async(USERCALL_QUEUE_SIZE, usercall_queue_synchronizer);
             let (return_queue_tx, return_queue_rx) = ipc_queue::bounded_async(RETURN_QUEUE_SIZE, return_queue_synchronizer);
 
-            let fifo_descriptors = FifoDescriptors {
-                usercall_queue: usercall_queue_tx.into_descriptor(),
-                return_queue: return_queue_rx.into_descriptor(),
+            let fifo_guards = FifoGuards {
+                usercall_queue: usercall_queue_tx.into_descriptor_guard(),
+                return_queue: return_queue_rx.into_descriptor_guard(),
+                async_queues_called: false,
             };
 
-            *enclave_clone.fifo_descriptors.lock().await = Some(fifo_descriptors);
+            *enclave_clone.fifo_guards.lock().await = Some(fifo_guards);
             *enclave_clone.return_queue_tx.lock().await = Some(return_queue_tx);
 
             tokio::task::spawn_local(async move {
@@ -1562,14 +1565,21 @@ impl<'tcs> IOHandlerInput<'tcs> {
         usercall_queue: &mut FifoDescriptor<Usercall>,
         return_queue: &mut FifoDescriptor<Return>,
     ) -> StdResult<(), EnclaveAbort<bool>> {
-        let fifo_descriptors = self.enclave.fifo_descriptors.lock().await.take();
-        match fifo_descriptors {
-            Some(fifo_descriptors) => {
-                *usercall_queue = fifo_descriptors.usercall_queue;
-                *return_queue = fifo_descriptors.return_queue;
+        let mut fifo_guards = self.enclave.fifo_guards.lock().await;
+        match &mut *fifo_guards {
+            Some(ref mut fifo_guards) if !fifo_guards.async_queues_called => {
+                *usercall_queue = fifo_guards.usercall_queue.fifo_descriptor();
+                *return_queue = fifo_guards.return_queue.fifo_descriptor();
+                fifo_guards.async_queues_called = true;
                 Ok(())
             }
-            None => Err(self.exit(true)),
+            Some(_) => {
+                drop(fifo_guards);
+                Err(self.exit(true))
+            },
+            // Enclave would not be able to call `async_queues()` before the
+            // fifo_guards is set to Some in `fn syscall_loop`.
+            None => unreachable!(),
         }
     }
 }

fortanix-sgx-abi/src/lib.rs

@@ -702,7 +702,7 @@ pub mod async {
     ///    expected to be written imminently.
     /// 6. Read the data, then store `0` in the `id`.
     /// 7. Store the new read offset.
-    /// 8. If the queue was full in step 1, signal the writer to wake up.
+    /// 8. If the queue was full before step 7, signal the writer to wake up.
     #[repr(C)]
     #[cfg_attr(feature = "rustc-dep-of-std", unstable(feature = "sgx_platform", issue = "56975"))]
     pub struct FifoDescriptor<T> {

ipc-queue/Cargo.toml

@@ -15,6 +15,8 @@ categories = ["asynchronous"]
 
 [dependencies]
 fortanix-sgx-abi = { version = "0.3.0", path = "../fortanix-sgx-abi" }
+
+[dev-dependencies]
 static_assertions = "1.1.0"
 
 [target.'cfg(not(target_env = "sgx"))'.dev-dependencies]

ipc-queue/src/fifo.rs

@@ -36,7 +36,7 @@ where
     (tx, rx)
 }
 
-struct Fifo<T> {
+pub(crate) struct Fifo<T> {
     data: Box<[T]>,
     offsets: Box<AtomicUsize>,
 }
@@ -92,6 +92,12 @@ impl<T: Identified> FifoInner<T> {
             descriptor.len.is_power_of_two(),
             "Fifo len should be a power of two"
         );
+        #[cfg(target_env = "sgx")] {
+            // check pointers are outside enclave range, etc.
+            use std::os::fortanix_sgx::usercalls::alloc::User;
+            let data = User::<[T]>::from_raw_parts(descriptor.data, descriptor.len);
+            mem::forget(data);
+        }
         let data_slice = std::slice::from_raw_parts_mut(descriptor.data, descriptor.len);
         Self {
             data: NonNull::new_unchecked(data_slice as *mut [T] as *mut UnsafeCell<[T]>),
@@ -114,20 +120,20 @@ impl<T: Identified> FifoInner<T> {
         }
     }
 
-    /// Consumes `self` and returns a FifoDescriptor. If `self` was created
-    /// using `from_arc`, it leaks the internal `Arc` copy to ensure the
-    /// resulting descriptor is valid for `'static` lifetime.
-    pub(crate) fn into_descriptor(self) -> FifoDescriptor<T> {
-        match self.storage {
-            Storage::Shared(arc) => mem::forget(arc),
-            Storage::Static => {}
-        }
+    /// Consumes `self` and returns a DescriptorGuard.
+    /// Panics if `self` was created using `from_descriptor`.
+    pub(crate) fn into_descriptor_guard(self) -> DescriptorGuard<T> {
+        let arc = match self.storage {
+            Storage::Shared(arc) => arc,
+            Storage::Static => panic!("Sender/Receiver created using `from_descriptor()` cannot be turned into DescriptorGuard."),
+        };
         let data_mut = unsafe { &mut *self.data.as_ref().get() };
-        FifoDescriptor {
+        let descriptor = FifoDescriptor {
             data: data_mut.as_mut_ptr(),
             len: data_mut.len(),
             offsets: self.offsets.as_ptr(),
-        }
+        };
+        DescriptorGuard { descriptor, _fifo: arc }
     }
 
     fn slot(&self, index: usize) -> &mut T {
@@ -177,7 +183,6 @@ impl<T: Identified> FifoInner<T> {
     pub(crate) fn try_recv_impl(&self) -> Result<(T, /*wake up writer:*/ bool), TryRecvError> {
         // 1. Load the current offsets.
         let current = Offsets::new(self.offsets().load(Ordering::SeqCst), self.data_len() as u32);
-        let was_full = current.is_full();
 
         // 2. If the queue is empty, wait, then go to step 1.
         if current.is_empty() {
@@ -204,16 +209,15 @@ impl<T: Identified> FifoInner<T> {
         slot.set_id(0);
 
         // 7. Store the new read offset.
-        let after = fetch_adjust(
+        let before = fetch_adjust(
             self.offsets(),
             new.read as isize - current.read as isize,
             Ordering::SeqCst,
         );
 
-        // 8. If the queue was full in step 1, signal the writer to wake up.
-        //    ... or became full during read
-        let became_full = Offsets::new(after, self.data_len() as u32).is_full();
-        Ok((val, was_full || became_full))
+        // 8. If the queue was full before step 7, signal the writer to wake up.
+        let was_full = Offsets::new(before, self.data_len() as u32).is_full();
+        Ok((val, was_full))
     }
 }
 

ipc-queue/src/interface_async.rs

@@ -5,7 +5,6 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
 use super::*;
-use fortanix_sgx_abi::FifoDescriptor;
 
 unsafe impl<T: Send, S: Send> Send for AsyncSender<T, S> {}
 unsafe impl<T: Send, S: Sync> Sync for AsyncSender<T, S> {}
@@ -35,15 +34,16 @@ impl<T: Identified, S: AsyncSynchronizer> AsyncSender<T, S> {
                         .map_err(|SynchronizationError::ChannelClosed| SendError::Closed)?;
                     val
                 }
+                Err((TrySendError::Closed, _)) => return Err(SendError::Closed),
             };
         }
     }
 
-    /// Consumes `self` and returns a FifoDescriptor.
-    /// **NOTE:** this function leaks the internal storage to ensure that the
-    /// pointers in the resulting descriptor remain valid.
-    pub fn into_descriptor(self) -> FifoDescriptor<T> {
-        self.inner.into_descriptor()
+    /// Consumes `self` and returns a DescriptorGuard.
+    /// The returned guard can be used to make `FifoDescriptor`s that remain
+    /// valid as long as the guard is not dropped.
+    pub fn into_descriptor_guard(self) -> DescriptorGuard<T> {
+        self.inner.into_descriptor_guard()
     }
 }
 
@@ -64,15 +64,16 @@ impl<T: Identified, S: AsyncSynchronizer> AsyncReceiver<T, S> {
                         .wait(QueueEvent::NotEmpty).await
                         .map_err(|SynchronizationError::ChannelClosed| RecvError::Closed)?;
                 }
+                Err(TryRecvError::Closed) => return Err(RecvError::Closed),
             }
         }
     }
 
-    /// Consumes `self` and returns a FifoDescriptor.
-    /// **NOTE:** this function leaks the internal storage to ensure that the
-    /// pointers in the resulting descriptor remain valid.
-    pub fn into_descriptor(self) -> FifoDescriptor<T> {
-        self.inner.into_descriptor()
+    /// Consumes `self` and returns a DescriptorGuard.
+    /// The returned guard can be used to make `FifoDescriptor`s that remain
+    /// valid as long as the guard is not dropped.
+    pub fn into_descriptor_guard(self) -> DescriptorGuard<T> {
+        self.inner.into_descriptor_guard()
     }
 }
 

ipc-queue/src/interface_sync.rs

@@ -65,6 +65,7 @@ impl<T: Identified, S: Synchronizer> Sender<T, S> {
                         .map_err(|SynchronizationError::ChannelClosed| SendError::Closed)?;
                     val
                 }
+                Err((TrySendError::Closed, _)) => return Err(SendError::Closed),
             };
         }
     }
@@ -115,6 +116,7 @@ impl<T: Identified, S: Synchronizer> Receiver<T, S> {
                         .wait(QueueEvent::NotEmpty)
                         .map_err(|SynchronizationError::ChannelClosed| RecvError::Closed)?;
                 }
+                Err(TryRecvError::Closed) => return Err(RecvError::Closed),
             }
         }
     }

ipc-queue/src/lib.rs

@@ -6,8 +6,10 @@
 
 #![cfg_attr(target_env = "sgx", feature(sgx_platform))]
 
+use fortanix_sgx_abi::FifoDescriptor;
 use std::future::Future;
 use std::pin::Pin;
+use std::sync::Arc;
 
 mod fifo;
 mod interface_sync;
@@ -16,7 +18,7 @@ mod sealed;
 #[cfg(test)]
 mod test_support;
 
-use self::fifo::FifoInner;
+use self::fifo::{Fifo, FifoInner};
 
 pub fn bounded<T, S>(len: usize, s: S) -> (Sender<T, S>, Receiver<T, S>)
 where
@@ -47,11 +49,13 @@ pub enum QueueEvent {
 #[derive(Debug, PartialEq, Eq)]
 pub enum TrySendError {
     QueueFull,
+    Closed,
 }
 
 #[derive(Debug, PartialEq, Eq)]
 pub enum TryRecvError {
     QueueEmpty,
+    Closed,
 }
 
 #[derive(Debug, PartialEq, Eq)]
@@ -104,3 +108,16 @@ pub struct AsyncReceiver<T, S> {
     inner: FifoInner<T>,
     synchronizer: S,
 }
+
+/// `DescriptorGuard<T>` can produce a `FifoDescriptor<T>` that is guaranteed
+/// to remain valid as long as the DescriptorGuard is not dropped.
+pub struct DescriptorGuard<T> {
+    descriptor: FifoDescriptor<T>,
+    _fifo: Arc<Fifo<T>>,
+}
+
+impl<T> DescriptorGuard<T> {
+    pub fn fifo_descriptor(&self) -> FifoDescriptor<T> {
+        self.descriptor
+    }
+}