Implement untrusted IO buffer for each thread

src/libos/Cargo.lock

@@ -36,6 +36,7 @@ dependencies = [
  "sgx_tse",
  "sgx_tstd",
  "sgx_types",
+ "spin 0.7.1",
 ]
 
 [[package]]
@@ -257,7 +258,7 @@ version = "1.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646"
 dependencies = [
- "spin",
+ "spin 0.5.2",
 ]
 
 [[package]]
@@ -444,7 +445,7 @@ name = "rcore-fs"
 version = "0.1.0"
 dependencies = [
  "bitflags",
- "spin",
+ "spin 0.5.2",
 ]
 
 [[package]]
@@ -453,7 +454,7 @@ version = "0.1.0"
 dependencies = [
  "log",
  "rcore-fs",
- "spin",
+ "spin 0.5.2",
 ]
 
 [[package]]
@@ -462,7 +463,7 @@ version = "0.1.0"
 dependencies = [
  "log",
  "rcore-fs",
- "spin",
+ "spin 0.5.2",
 ]
 
 [[package]]
@@ -471,7 +472,7 @@ version = "0.1.0"
 dependencies = [
  "log",
  "rcore-fs",
- "spin",
+ "spin 0.5.2",
 ]
 
 [[package]]
@@ -481,7 +482,7 @@ dependencies = [
  "bitvec",
  "log",
  "rcore-fs",
- "spin",
+ "spin 0.5.2",
  "static_assertions",
  "uuid",
 ]
@@ -492,7 +493,7 @@ version = "0.1.0"
 dependencies = [
  "log",
  "rcore-fs",
- "spin",
+ "spin 0.5.2",
 ]
 
 [[package]]
@@ -700,6 +701,12 @@ version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d"
 
+[[package]]
+name = "spin"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "13287b4da9d1207a4f4929ac390916d64eacfe236a487e9a9f5b3be392be5162"
+
 [[package]]
 name = "stable_deref_trait"
 version = "1.2.0"

src/libos/Cargo.toml

@@ -32,6 +32,7 @@ ctor = "0.1"
 regex = { git = "https://github.com/mesalock-linux/regex-sgx", default-features = false, features = ["std", "unicode", "mesalock_sgx"] }
 goblin = { version = "0.5.4", default-features = false, features = ["elf64", "elf32", "endian_fd"] }
 intrusive-collections = "0.9"
+spin = "0.7"
 
 [patch.'https://github.com/apache/teaclave-sgx-sdk.git']
 sgx_tstd = { path = "../../deps/rust-sgx-sdk/sgx_tstd" }

src/libos/src/net/socket/host/mod.rs

@@ -4,30 +4,25 @@ use std::mem;
 
 use atomic::Atomic;
 
-use self::recv::Receiver;
-use self::send::Sender;
-
 use super::*;
 use crate::fs::{
     occlum_ocall_ioctl, AccessMode, CreationFlags, File, FileRef, HostFd, IoEvents, IoNotifier,
     IoctlCmd, StatusFlags,
 };
 
+use crate::process::IO_BUF_SIZE;
+
 mod ioctl_impl;
 mod recv;
 mod send;
 mod socket_file;
 
-pub const SEND_BUF_SIZE: usize = 128 * 1024;
-pub const RECV_BUF_SIZE: usize = 128 * 1024;
 /// Native linux socket
 #[derive(Debug)]
 pub struct HostSocket {
     host_fd: HostFd,
     host_events: Atomic<IoEvents>,
     notifier: IoNotifier,
-    sender: SgxMutex<Sender>,
-    receiver: SgxMutex<Receiver>,
 }
 
 impl HostSocket {
@@ -49,14 +44,10 @@ impl HostSocket {
     fn from_host_fd(host_fd: HostFd) -> Result<HostSocket> {
         let host_events = Atomic::new(IoEvents::empty());
         let notifier = IoNotifier::new();
-        let sender = SgxMutex::new(Sender::new()?);
-        let receiver = SgxMutex::new(Receiver::new()?);
         Ok(Self {
             host_fd,
             host_events,
             notifier,
-            sender,
-            receiver,
         })
     }
 

src/libos/src/net/socket/host/recv.rs

@@ -1,18 +1,6 @@
 use super::*;
 use crate::untrusted::{SliceAsMutPtrAndLen, SliceAsPtrAndLen, UntrustedSliceAlloc};
 
-#[derive(Debug)]
-pub struct Receiver {
-    alloc: UntrustedSliceAlloc,
-}
-
-impl Receiver {
-    pub fn new() -> Result<Self> {
-        let alloc = UntrustedSliceAlloc::new(RECV_BUF_SIZE)?;
-        Ok(Self { alloc })
-    }
-}
-
 impl HostSocket {
     pub fn recv(&self, buf: &mut [u8], flags: RecvFlags) -> Result<usize> {
         let (bytes_recvd, _) = self.recvfrom(buf, flags)?;
@@ -43,18 +31,17 @@ impl HostSocket {
         mut name: Option<&mut [u8]>,
         mut control: Option<&mut [u8]>,
     ) -> Result<(usize, usize, usize, MsgHdrFlags)> {
+        let current = current!();
         let data_length = data.iter().map(|s| s.len()).sum();
-        let mut receiver: SgxMutexGuard<'_, Receiver>;
         let mut ocall_alloc;
         // Allocated slice in untrusted memory region
-        let u_allocator = if data_length > RECV_BUF_SIZE {
+        let u_allocator = if data_length > IO_BUF_SIZE {
             // Ocall allocator
             ocall_alloc = UntrustedSliceAlloc::new(data_length)?;
-            &mut ocall_alloc
+            ocall_alloc.guard()
         } else {
-            // Inner allocator, lock buffer until recv ocall completion
-            receiver = self.receiver.lock().unwrap();
-            &mut receiver.alloc
+            // IO buffer per thread
+            current.io_buffer()
         };
 
         let mut u_data = {
@@ -77,7 +64,6 @@ impl HostSocket {
                 break;
             }
         }
-        u_allocator.reset();
         Ok(retval)
     }
 

src/libos/src/net/socket/host/send.rs

@@ -1,17 +1,5 @@
 use super::*;
 
-#[derive(Debug)]
-pub struct Sender {
-    alloc: UntrustedSliceAlloc,
-}
-
-impl Sender {
-    pub fn new() -> Result<Self> {
-        let alloc = UntrustedSliceAlloc::new(SEND_BUF_SIZE)?;
-        Ok(Self { alloc })
-    }
-}
-
 impl HostSocket {
     pub fn send(&self, buf: &[u8], flags: SendFlags) -> Result<usize> {
         self.sendto(buf, flags, &None)
@@ -35,18 +23,17 @@ impl HostSocket {
         name: Option<&[u8]>,
         control: Option<&[u8]>,
     ) -> Result<usize> {
+        let current = current!();
         let data_length = data.iter().map(|s| s.len()).sum();
-        let mut sender: SgxMutexGuard<'_, Sender>;
         let mut ocall_alloc;
         // Allocated slice in untrusted memory region
-        let u_allocator = if data_length > SEND_BUF_SIZE {
+        let u_allocator = if data_length > IO_BUF_SIZE {
             // Ocall allocator
             ocall_alloc = UntrustedSliceAlloc::new(data_length)?;
-            &mut ocall_alloc
+            ocall_alloc.guard()
         } else {
-            // Inner allocator, lock buffer until send ocall completion
-            sender = self.sender.lock().unwrap();
-            &mut sender.alloc
+            // IO buffer per thread
+            current.io_buffer()
         };
 
         let u_data = {
@@ -59,7 +46,6 @@ impl HostSocket {
         };
 
         let retval = self.do_sendmsg_untrusted_data(&u_data, flags, name, control);
-        u_allocator.reset();
         retval
     }
 

src/libos/src/process/mod.rs

@@ -33,7 +33,7 @@ pub use self::spawn_attribute::SpawnAttr;
 pub use self::syscalls::*;
 pub use self::task::Task;
 pub use self::term_status::{ForcedExitStatus, TermStatus};
-pub use self::thread::{Thread, ThreadStatus};
+pub use self::thread::{Thread, ThreadStatus, IO_BUF_SIZE};
 
 mod do_arch_prctl;
 mod do_clone;

src/libos/src/process/thread/builder.rs

@@ -1,3 +1,4 @@
+use spin::Once;
 use std::ptr::NonNull;
 
 use super::{
@@ -143,6 +144,7 @@ impl ThreadBuilder {
         };
         let host_eventfd = Arc::new(HostEventFd::new()?);
         let raw_ptr = RwLock::new(0);
+        let io_buffer = Once::new();
 
         let new_thread = Arc::new(Thread {
             task,
@@ -165,6 +167,7 @@ impl ThreadBuilder {
             profiler,
             host_eventfd,
             raw_ptr,
+            io_buffer,
         });
 
         let mut inner = new_thread.process().inner();

src/libos/src/process/thread/mod.rs

@@ -1,3 +1,4 @@
+use spin::Once;
 use std::fmt;
 use std::ptr::NonNull;
 
@@ -12,6 +13,7 @@ use crate::net::THREAD_NOTIFIERS;
 use crate::prelude::*;
 use crate::signal::{SigQueues, SigSet, SigStack};
 use crate::time::ThreadProfiler;
+use crate::untrusted::{UntrustedSliceAlloc, UntrustedSliceAllocGuard};
 
 pub use self::builder::ThreadBuilder;
 pub use self::id::ThreadId;
@@ -21,6 +23,8 @@ mod builder;
 mod id;
 mod name;
 
+pub const IO_BUF_SIZE: usize = 128 * 1024;
+
 pub struct Thread {
     // Low-level info
     task: Task,
@@ -50,6 +54,8 @@ pub struct Thread {
     // Misc
     host_eventfd: Arc<HostEventFd>,
     raw_ptr: RwLock<usize>,
+    // Thread ocall buffer
+    io_buffer: Once<UntrustedSliceAlloc>,
 }
 
 #[derive(Debug, PartialEq, Clone, Copy)]
@@ -216,6 +222,12 @@ impl Thread {
         &self.host_eventfd
     }
 
+    pub fn io_buffer(&self) -> UntrustedSliceAllocGuard<'_> {
+        self.io_buffer
+            .call_once(|| UntrustedSliceAlloc::new(IO_BUF_SIZE).unwrap())
+            .guard()
+    }
+
     pub(super) fn start(&self, host_tid: pid_t) {
         self.sched().lock().unwrap().attach(host_tid);
         let mut raw_ptr = self.raw_ptr.write().unwrap();

src/libos/src/untrusted/mod.rs

@@ -6,5 +6,5 @@ mod slice_ext;
 use super::*;
 
 pub use self::alloc::UNTRUSTED_ALLOC;
-pub use self::slice_alloc::{UntrustedSlice, UntrustedSliceAlloc};
+pub use self::slice_alloc::{UntrustedSlice, UntrustedSliceAlloc, UntrustedSliceAllocGuard};
 pub use self::slice_ext::{SliceAsMutPtrAndLen, SliceAsPtrAndLen};

src/libos/src/untrusted/slice_alloc.rs

@@ -46,13 +46,17 @@ impl UntrustedSliceAlloc {
         })
     }
 
-    pub fn new_slice(&self, src_slice: &[u8]) -> Result<UntrustedSlice> {
+    pub fn guard(&self) -> UntrustedSliceAllocGuard<'_> {
+        UntrustedSliceAllocGuard { alloc: self }
+    }
+
+    fn new_slice(&self, src_slice: &[u8]) -> Result<UntrustedSlice> {
         let mut new_slice = self.new_slice_mut(src_slice.len())?;
         new_slice.read_from_slice(src_slice)?;
         Ok(new_slice)
     }
 
-    pub fn new_slice_mut(&self, new_slice_len: usize) -> Result<UntrustedSlice> {
+    fn new_slice_mut(&self, new_slice_len: usize) -> Result<UntrustedSlice> {
         let new_slice_ptr = {
             // Move self.buf_pos forward if enough space _atomically_.
             let old_pos = self
@@ -72,7 +76,7 @@ impl UntrustedSliceAlloc {
         Ok(UntrustedSlice { slice: new_slice })
     }
 
-    pub fn reset(&mut self) {
+    fn reset(&self) {
         self.buf_pos.store(0, Ordering::Relaxed);
     }
 }
@@ -100,6 +104,25 @@ impl Debug for UntrustedSliceAlloc {
     }
 }
 
+pub struct UntrustedSliceAllocGuard<'a> {
+    alloc: &'a UntrustedSliceAlloc,
+}
+
+impl<'a> UntrustedSliceAllocGuard<'a> {
+    pub fn new_slice(&self, src_slice: &[u8]) -> Result<UntrustedSlice> {
+        self.alloc.new_slice(src_slice)
+    }
+    pub fn new_slice_mut(&self, new_slice_len: usize) -> Result<UntrustedSlice> {
+        self.alloc.new_slice_mut(new_slice_len)
+    }
+}
+
+impl<'a> Drop for UntrustedSliceAllocGuard<'a> {
+    fn drop(&mut self) {
+        self.alloc.reset();
+    }
+}
+
 pub struct UntrustedSlice<'a> {
     slice: &'a mut [u8],
 }