Fix the emulation of RDTSC instruction with ocall

src/Enclave.edl

@@ -54,6 +54,7 @@ enclave {
 
         void occlum_ocall_gettimeofday([out] struct timeval* tv);
         void occlum_ocall_clock_gettime(clockid_t clockid, [out] struct timespec* ts);
+        void occlum_ocall_rdtsc([out] uint32_t* low, [out] uint32_t* high);
 
         void occlum_ocall_nanosleep([in] const struct timespec* req);
 

src/libos/src/exception/rdtsc.rs

@@ -1,9 +1,15 @@
 use super::*;
+use process::Task;
 use sgx_types::*;
 
 const RDTSC_OPCODE: u16 = 0x310F;
-static mut FAKE_RDTSC_VALUE: u64 = 0;
-static FAKE_RDTSC_INC_VALUE: u64 = 1000;
+
+extern "C" {
+    fn occlum_ocall_rdtsc(low: *mut u32, high: *mut u32) -> sgx_status_t;
+    fn __get_current_task() -> *const Task;
+    fn switch_td_to_kernel(task: *const Task);
+    fn switch_td_to_user(task: *const Task);
+}
 
 #[no_mangle]
 pub extern "C" fn handle_rdtsc_exception(info: *mut sgx_exception_info_t) -> u32 {
@@ -15,11 +21,19 @@ pub extern "C" fn handle_rdtsc_exception(info: *mut sgx_exception_info_t) -> u32
     {
         return EXCEPTION_CONTINUE_SEARCH;
     }
-    // rdtsc support here is temporary, only for SKL, later CPU's will support this inside enclave
     unsafe {
-        FAKE_RDTSC_VALUE += FAKE_RDTSC_INC_VALUE;
-        info.cpu_context.rax = (FAKE_RDTSC_VALUE & 0xFFFFFFFF);
-        info.cpu_context.rdx = (FAKE_RDTSC_VALUE >> 32);
+        let task = __get_current_task();
+        switch_td_to_kernel(task);
+        let (low, high) = {
+            let mut low = 0;
+            let mut high = 0;
+            let sgx_status = occlum_ocall_rdtsc(&mut low, &mut high);
+            assert!(sgx_status == sgx_status_t::SGX_SUCCESS);
+            (low, high)
+        };
+        info.cpu_context.rax = low as u64;
+        info.cpu_context.rdx = high as u64;
+        switch_td_to_user(task);
     }
     info.cpu_context.rip += 2;
 

src/libos/src/process/mod.rs

@@ -7,7 +7,7 @@ pub use self::process::{Status, IDLE_PROCESS};
 pub use self::process_table::get;
 pub use self::sched::{do_sched_getaffinity, do_sched_setaffinity, do_sched_yield, CpuSet};
 pub use self::spawn::{do_spawn, do_spawn_without_exec, ElfFile, FileAction, ProgramHeaderExt};
-pub use self::task::{current_pid, get_current, run_task};
+pub use self::task::{current_pid, get_current, run_task, Task};
 pub use self::thread::{do_clone, do_set_tid_address, CloneFlags, ThreadGroup};
 pub use self::wait::{WaitQueue, Waiter};
 
@@ -82,7 +82,6 @@ mod task;
 mod thread;
 mod wait;
 
-use self::task::Task;
 use super::*;
 use fs::{File, FileRef, FileTable};
 use misc::ResourceLimitsRef;

src/libos/src/process/task.c

@@ -37,6 +37,28 @@ static uint64_t get_syscall_stack(struct Task* this_task) {
 #define RESET_CURRENT_TASK()                    \
     __set_stack_guard(stack_guard);
 
+void switch_td_to_kernel(const struct Task* task) {
+    thread_data_t* td = get_thread_data();
+
+    // TODO: do do not support stack expansion, need a new design on SGX2 platform.
+    // Set the stack_commit_addr to 0, as the result no stack expansion happens at any situations
+    __atomic_store_n(&td->stack_commit_addr, 0, __ATOMIC_RELAXED);
+    td->stack_base_addr = task->kernel_stack_base;
+    td->stack_limit_addr = task->kernel_stack_limit;
+    td->stack_commit_addr = task->kernel_stack_limit;
+}
+
+void switch_td_to_user(const struct Task* task) {
+    thread_data_t* td = get_thread_data();
+
+    // TODO: do do not support stack expansion, need a new design on SGX2 platform.
+    // Set the stack_commit_addr to 0, as the result no stack expansion happens at any situations
+    __atomic_store_n(&td->stack_commit_addr, 0, __ATOMIC_RELAXED);
+    td->stack_base_addr = task->user_stack_base;
+    td->stack_limit_addr = task->user_stack_limit;
+    td->stack_commit_addr = task->user_stack_limit;
+}
+
 int do_run_task(struct Task* task) {
     jmp_buf libos_state = {0};
     thread_data_t* td = get_thread_data();
@@ -45,12 +67,7 @@ int do_run_task(struct Task* task) {
     task->kernel_stack_base = td->stack_base_addr;
     task->kernel_stack_limit = td->stack_limit_addr;
 
-    //Do do not support stack expansion, need a new design on SGX2 platform.
-    //Set the stack_commit_addr to 0, as the result no stack expansion happens at any situations
-    __atomic_store_n(&td->stack_commit_addr, 0, __ATOMIC_RELAXED);
-    td->stack_base_addr = task->user_stack_base;
-    td->stack_limit_addr = task->user_stack_limit;
-    td->stack_commit_addr = task->user_stack_limit;
+    switch_td_to_user(task);
 
     SET_CURRENT_TASK(task);
 
@@ -67,13 +84,7 @@ int do_run_task(struct Task* task) {
 void do_exit_task(void) {
     struct Task* task = __get_current_task();
     jmp_buf* jb = task->saved_state;
-    thread_data_t* td = get_thread_data();
 
-    //Do do not support stack expansion, need a new design on SGX2 platform.
-    //Set the stack_commit_addr to 0, as the result no stack expansion happens at any situations
-    __atomic_store_n(&td->stack_commit_addr, 0, __ATOMIC_RELAXED);
-    td->stack_base_addr = task->kernel_stack_base;
-    td->stack_limit_addr = task->kernel_stack_limit;
-    td->stack_commit_addr = task->kernel_stack_limit;
+    switch_td_to_kernel(task);
     longjmp(*jb, 1);
 }

src/pal/src/ocalls/time.c

@@ -24,3 +24,10 @@ int occlum_ocall_thread_getcpuclock(struct timespec *tp) {
 
     return clock_gettime(thread_clock_id, tp);
 }
+
+void occlum_ocall_rdtsc(uint32_t* low, uint32_t* high) {
+    uint64_t rax, rdx;
+    asm volatile("rdtsc" : "=a"(rax), "=d"(rdx));
+    *low = (uint32_t)rax;
+    *high = (uint32_t)rdx;
+}

test/rdtsc/main.c

@@ -1,19 +1,41 @@
 #include <stdio.h>
 #include <stdint.h>
+#include "test.h"
+
+// ============================================================================
+// Helper functions for rdtsc
+// ============================================================================
 
 static inline uint64_t native_rdtsc() {
-    uint32_t hi, lo;
-    asm volatile("rdtsc" : "=a"(lo), "=d"(hi));
-    return (( (uint64_t)lo)|( ((uint64_t)hi)<<32 ));
+    uint64_t low, high;
+    asm volatile("rdtsc" : "=a"(low), "=d"(high));
+    return (high << 32) | low;
 }
 
-int main(int argc, char **argv)
-{
-    /* Gets rdtsc information and tests the SGX support of the rdtsc */
-    uint64_t r;
-
-    r = native_rdtsc();
-    printf("rdtsc: %lu\n", r);
+// ============================================================================
+// Test cases for rdtsc
+// ============================================================================
 
+int test_rdtsc() {
+    uint64_t start_count = native_rdtsc();
+    if (start_count == 0) {
+        THROW_ERROR("call rdtsc failed");
+    }
+    uint64_t end_count = native_rdtsc();
+    if (end_count <= start_count) {
+        THROW_ERROR("check rdtsc return value failed");
+    }
     return 0;
 }
+
+// ============================================================================
+// Test suite main
+// ============================================================================
+
+static test_case_t test_cases[] = {
+    TEST_CASE(test_rdtsc),
+};
+
+int main() {
+    return test_suite_run(test_cases, ARRAY_SIZE(test_cases));
+}