track_aex_count feature for enclave-runner

Author: Jethro Beekman

intel-sgx/enclave-runner/Cargo.toml

@@ -42,3 +42,4 @@ futures = { version = "0.3", features = ["compat", "io-compat"] } # MIT/Apache-2
 [features]
 default = ["crypto-openssl"]
 crypto-openssl = ["openssl", "sgxs/crypto-openssl"]
+instrumentation = []

intel-sgx/enclave-runner/src/command.rs

@@ -25,6 +25,7 @@ pub struct Command {
     force_time_usercalls: bool,
     cmd_args: Vec<Vec<u8>>,
     num_worker_threads: usize,
+    track_aex_count: bool,
 }
 
 impl MappingInfo for Command {
@@ -61,14 +62,20 @@ impl Command {
             force_time_usercalls,
             cmd_args,
             num_worker_threads,
+            track_aex_count: false,
         }
     }
 
     pub fn new<P: AsRef<Path>, L: Load>(enclave_path: P, loader: &mut L) -> Result<Command, Error> {
         EnclaveBuilder::new(enclave_path.as_ref()).build(loader)
     }
 
+    #[cfg(feature = "instrumentation")]
+    pub fn track_aex_count(&mut self) {
+        self.track_aex_count = true;
+    }
+
     pub fn run(self) -> Result<(), Error> {
-        EnclaveState::main_entry(self.main, self.threads, self.usercall_ext, self.forward_panics, self.force_time_usercalls, self.cmd_args, self.num_worker_threads)
+        EnclaveState::main_entry(self.main, self.threads, self.usercall_ext, self.forward_panics, self.force_time_usercalls, self.track_aex_count, self.cmd_args, self.num_worker_threads)
     }
 }

intel-sgx/enclave-runner/src/tcs.rs

@@ -10,6 +10,9 @@ use std::convert::{TryFrom, TryInto};
 use std::fmt;
 use std::io::Write;
 use std::os::raw::c_void;
+use std::sync::atomic::AtomicUsize;
+#[cfg(feature = "instrumentation")]
+use std::sync::Arc;
 
 use sgx_isa::Enclu;
 use sgxs::loader::Tcs;
@@ -22,10 +25,41 @@ pub enum CoResult<Y, R> {
     Return(R),
 }
 
+#[cfg(feature = "instrumentation")]
+#[derive(Clone, Debug)]
+pub(crate) struct AexCount(Option<Arc<AtomicUsize>>);
+#[cfg(not(feature = "instrumentation"))]
+#[derive(Clone, Debug)]
+pub(crate) struct AexCount(());
+
+impl AexCount {
+    pub fn none() -> AexCount {
+        #[cfg(feature = "instrumentation")]
+        { AexCount(None) }
+        #[cfg(not(feature = "instrumentation"))]
+        { AexCount(()) }
+    }
+
+    fn get(&self) -> Option<&AtomicUsize> {
+        #[cfg(feature = "instrumentation")]
+        { self.0.as_ref().map(|p| &**p) }
+        #[cfg(not(feature = "instrumentation"))]
+        { None }
+    }
+}
+
+#[cfg(feature = "instrumentation")]
+impl From<Arc<AtomicUsize>> for AexCount {
+    fn from(v: Arc<AtomicUsize>) -> AexCount {
+        AexCount(Some(v))
+    }
+}
+
 #[derive(Debug)]
 pub struct Usercall<T: Tcs> {
     tcs: T,
     parameters: (u64, u64, u64, u64, u64),
+    aex_count: AexCount,
 }
 
 pub type ThreadResult<T> = CoResult<Usercall<T>, (T, u64, u64)>;
@@ -40,7 +74,7 @@ impl<T: Tcs> Usercall<T> {
         retval: (u64, u64),
         debug_buf: Option<&RefCell<DebugBuffer>>,
     ) -> ThreadResult<T> {
-        coenter(self.tcs, 0, retval.0, retval.1, 0, 0, debug_buf)
+        coenter(self.tcs, 0, retval.0, retval.1, 0, 0, debug_buf, self.aex_count)
     }
 
     pub fn tcs_address(&self) -> *mut c_void {
@@ -56,6 +90,7 @@ pub(crate) fn coenter<T: Tcs>(
     mut p4: u64,
     mut p5: u64,
     debug_buf: Option<&RefCell<DebugBuffer>>,
+    aex_count: AexCount,
 ) -> ThreadResult<T> {
     /// Check if __vdso_sgx_enter_enclave exists. We're using weak linkage, so
     /// it might not.
@@ -101,7 +136,33 @@ pub(crate) fn coenter<T: Tcs>(
                 uninit_debug_buf.as_mut_ptr() as *mut _
             }
         };
-        if has_vdso_sgx_enter_enclave() {
+        if let Some(aex_count) = aex_count.get() {
+            asm!("
+                    decq ({1})
+                    push {1}
+                    lea 1f(%rip), %rcx // set SGX AEP
+                    xchg {0}, %rbx
+1:                  mov (%rsp), %r12
+                    incq (%r12)
+                    enclu
+                    addq $8, %rsp
+                    xchg %rbx, {0}
+                ",
+                inout(reg) tcs.address() => _, // rbx is used internally by LLVM and cannot be used as an operand for inline asm (#84658)
+                in(reg) aex_count,
+                inout("eax") Enclu::EEnter as u32 => sgx_result,
+                out("rcx") _,
+                inout("rdx") p3,
+                inout("rdi") p1,
+                inout("rsi") p2,
+                inout("r8") p4,
+                inout("r9") p5,
+                inout("r10") debug_buf_ptr => _,
+                lateout("r11") _,
+                lateout("r12") _,
+                options(nostack, att_syntax)
+            );
+        } else if has_vdso_sgx_enter_enclave() {
             #[repr(C)]
             #[derive(Default)]
             struct SgxEnclaveRun {
@@ -172,8 +233,9 @@ pub(crate) fn coenter<T: Tcs>(
                             eprintln!("Enclave triggered exception, treating as panic: {:?}", run);
                         }
                         return CoResult::Yield(Usercall {
-                            tcs: tcs,
+                            tcs,
                             parameters: (crate::usercalls::abi::UsercallList::exit as _, true as _, 0, 0, 0),
+                            aex_count,
                         });
                     },
                     _ => panic!("Error entering enclave (VDSO): ret = success, run = {:?}", run),
@@ -211,8 +273,9 @@ pub(crate) fn coenter<T: Tcs>(
         CoResult::Return((tcs, p2, p3))
     } else {
         CoResult::Yield(Usercall {
-            tcs: tcs,
+            tcs,
             parameters: (p1, p2, p3, p4, p5),
+            aex_count,
         })
     }
 }

intel-sgx/enclave-runner/src/usercalls/interface.rs

@@ -178,7 +178,7 @@ impl<'future, 'ioinput: 'future, 'tcs: 'ioinput> Usercalls<'future> for Handler<
         self,
     ) -> std::pin::Pin<Box<dyn Future<Output = (Self, UsercallResult<Result>)> + 'future>> {
         async move {
-            let ret = Ok(self.0.launch_thread().to_sgx_result());
+            let ret = Ok(self.0.launch_thread().await.to_sgx_result());
             return (self, ret);
         }.boxed_local()
     }

intel-sgx/enclave-runner/src/usercalls/mod.rs

@@ -611,6 +611,7 @@ impl PendingEvents {
 struct RunningTcs {
     tcs_address: TcsAddress,
     mode: EnclaveEntry,
+    aex_count: tcs::AexCount,
 }
 
 enum EnclaveKind {
@@ -665,6 +666,7 @@ pub(crate) struct EnclaveState {
     // 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>>>,
+    aex_counts: Option<Mutex<FnvHashMap<TcsAddress, Arc<AtomicUsize>>>>,
 }
 
 struct Work {
@@ -682,7 +684,7 @@ impl Work {
         let buf = RefCell::new([0u8; 1024]);
         let usercall_send_data = match self.entry {
             CoEntry::Initial(erased_tcs, p1, p2, p3, p4, p5) => {
-                let coresult = tcs::coenter(erased_tcs, p1, p2, p3, p4, p5, Some(&buf));
+                let coresult = tcs::coenter(erased_tcs, p1, p2, p3, p4, p5, Some(&buf), self.tcs.aex_count.clone());
                 UsercallSendData::Sync(coresult, self.tcs, buf)
             }
             CoEntry::Resume(usercall, coresult) => {
@@ -736,6 +738,7 @@ impl EnclaveState {
         threads_vector: Vec<ErasedTcs>,
         forward_panics: bool,
         force_time_usercalls: bool,
+        aex_counts: Option<Mutex<FnvHashMap<TcsAddress, Arc<AtomicUsize>>>>,
     ) -> Arc<Self> {
         let mut fds = FnvHashMap::default();
 
@@ -779,6 +782,7 @@ impl EnclaveState {
             force_time_usercalls,
             fifo_guards: Mutex::new(None),
             return_queue_tx: Mutex::new(None),
+            aex_counts,
         })
     }
 
@@ -1053,7 +1057,10 @@ impl EnclaveState {
                 }
             });
 
-        local_set.block_on(&mut rt, select_fut.unit_error()).unwrap()
+        let ret = local_set.block_on(&mut rt, select_fut.unit_error()).unwrap();
+        #[cfg(feature = "instrumentation")]
+        tokio::runtime::Runtime::new().unwrap().block_on(enclave.print_aex_counts());
+        ret
     }
 
     fn run(
@@ -1105,12 +1112,18 @@ impl EnclaveState {
         usercall_ext: Option<Box<dyn UsercallExtension>>,
         forward_panics: bool,
         force_time_usercalls: bool,
+        track_aex_count: bool,
         cmd_args: Vec<Vec<u8>>,
         num_of_worker_threads: usize,
     ) -> StdResult<(), anyhow::Error> {
         assert!(num_of_worker_threads > 0, "worker_threads cannot be zero");
         let mut event_queues =
             FnvHashMap::with_capacity_and_hasher(threads.len() + 1, Default::default());
+        let mut aex_counts: Option<FnvHashMap<TcsAddress, Arc<AtomicUsize>>> = if track_aex_count {
+            Some(FnvHashMap::with_capacity_and_hasher(threads.len() + 1, Default::default()))
+        } else {
+            None
+        };
         let main = Self::event_queue_add_tcs(&mut event_queues, main);
 
         let mut args = Vec::with_capacity(cmd_args.len());
@@ -1123,10 +1136,16 @@ impl EnclaveState {
         let argc = args.len();
         let argv = Box::into_raw(args.into_boxed_slice()) as *const u8;
 
+        let aex_count = match &mut aex_counts {
+            #[cfg(feature = "instrumentation")]
+            Some(aex_counts) => tcs::AexCount::from(aex_counts.entry(main.tcs.address()).or_default().clone()),
+            _ => tcs::AexCount::none(),
+        };
         let main_work = Work {
             tcs: RunningTcs {
                 tcs_address: main.tcs.address(),
                 mode: EnclaveEntry::ExecutableMain,
+                aex_count,
             },
             entry: CoEntry::Initial(main.tcs, argv as _, argc as _, 0, 0, 0),
         };
@@ -1137,7 +1156,7 @@ impl EnclaveState {
                 other_reasons: vec![],
             }),
         });
-        let enclave = EnclaveState::new(kind, event_queues, usercall_ext, threads, forward_panics, force_time_usercalls);
+        let enclave = EnclaveState::new(kind, event_queues, usercall_ext, threads, forward_panics, force_time_usercalls, aex_counts.map(Mutex::new));
 
         let main_result = EnclaveState::run(enclave.clone(), num_of_worker_threads, main_work);
 
@@ -1194,7 +1213,7 @@ impl EnclaveState {
 
         let kind = EnclaveKind::Library(Library {});
 
-        let enclave = EnclaveState::new(kind, event_queues, usercall_ext, threads, forward_panics, force_time_usercalls);
+        let enclave = EnclaveState::new(kind, event_queues, usercall_ext, threads, forward_panics, force_time_usercalls, None);
         return enclave;
     }
 
@@ -1211,6 +1230,7 @@ impl EnclaveState {
             tcs: RunningTcs {
                 tcs_address: thread.tcs.address(),
                 mode: EnclaveEntry::Library,
+                aex_count: tcs::AexCount::none(),
             },
             entry: CoEntry::Initial(thread.tcs, p1, p2, p3, p4, p5),
         };
@@ -1244,6 +1264,17 @@ impl EnclaveState {
             pending_events.abort();
         }
     }
+
+    #[cfg(feature = "instrumentation")]
+    async fn print_aex_counts(&self) {
+        if let Some(aex_counts) = &self.aex_counts {
+            eprintln!("===AEX counts per TCS===");
+            for (addr, count) in aex_counts.lock().await.iter() {
+                eprintln!("{:p} {}", *addr, count.load(Ordering::Relaxed));
+            }
+            eprintln!("========================");
+        }
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -1540,7 +1571,7 @@ impl<'tcs> IOHandlerInput<'tcs> {
     }
 
     #[inline(always)]
-    fn launch_thread(&self) -> IoResult<()> {
+    async fn launch_thread(&self) -> IoResult<()> {
         // check if enclave is of type command
         self.enclave
             .kind
@@ -1553,10 +1584,16 @@ impl<'tcs> IOHandlerInput<'tcs> {
             }
         };
 
+        let aex_count = match &self.enclave.aex_counts {
+            #[cfg(feature = "instrumentation")]
+            Some(aex_counts) => tcs::AexCount::from(aex_counts.lock().await.entry(new_tcs.tcs.address()).or_default().clone()),
+            _ => tcs::AexCount::none(),
+        };
         let ret = self.work_sender.send(Work {
             tcs: RunningTcs {
                 tcs_address: new_tcs.tcs.address(),
                 mode: EnclaveEntry::ExecutableNonMain,
+                aex_count,
             },
             entry: CoEntry::Initial(new_tcs.tcs, 0, 0, 0, 0, 0),
         });