Save floating point registers in exception/interrupt flow

src/libos/Cargo.lock

@@ -4,6 +4,7 @@
 name = "Occlum"
 version = "0.14.0"
 dependencies = [
+ "aligned",
  "bitflags",
  "bitvec",
  "derive_builder",
@@ -25,6 +26,26 @@ dependencies = [
  "xmas-elf",
 ]
 
+[[package]]
+name = "aligned"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c19796bd8d477f1a9d4ac2465b464a8b1359474f06a96bb3cda650b4fca309bf"
+dependencies = [
+ "as-slice",
+]
+
+[[package]]
+name = "as-slice"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37dfb65bc03b2bc85ee827004f14a6817e04160e3b1a28931986a666a9290e70"
+dependencies = [
+ "generic-array 0.12.3",
+ "generic-array 0.13.2",
+ "stable_deref_trait",
+]
+
 [[package]]
 name = "autocfg"
 version = "0.1.7"
@@ -152,6 +173,24 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a06f77d526c1a601b7c4cdd98f54b5eaabffc14d5f2f0296febdc7f357c6d3ba"
 
+[[package]]
+name = "generic-array"
+version = "0.12.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c68f0274ae0e023facc3c97b2e00f076be70e254bc851d972503b328db79b2ec"
+dependencies = [
+ "typenum",
+]
+
+[[package]]
+name = "generic-array"
+version = "0.13.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ed1e761351b56f54eb9dcd0cfaca9fd0daecf93918e1cfc01c8a3d26ee7adcd"
+dependencies = [
+ "typenum",
+]
+
 [[package]]
 name = "hashbrown_tstd"
 version = "0.7.1"
@@ -529,6 +568,12 @@ version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d"
 
+[[package]]
+name = "stable_deref_trait"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3"
+
 [[package]]
 name = "static_assertions"
 version = "0.3.4"
@@ -563,6 +608,12 @@ dependencies = [
  "unicode-xid 0.2.1",
 ]
 
+[[package]]
+name = "typenum"
+version = "1.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "373c8a200f9e67a0c95e62a4f52fbf80c23b4381c05a17845531982fa99e6b33"
+
 [[package]]
 name = "unicode-xid"
 version = "0.1.0"

src/libos/Cargo.toml

@@ -11,6 +11,7 @@ crate-type = ["staticlib"]
 bitflags = "1.0"
 bitvec = { version = "0.17", default-features = false, features = ["alloc"]  }
 log = "0.4"
+aligned = "0.3.4"
 lazy_static = { version = "1.1.0", features = ["spin_no_std"] } # Implies nightly
 derive_builder = "0.7.2"
 ringbuf = { path = "../../deps/ringbuf" }

src/libos/src/exception/mod.rs

@@ -6,6 +6,8 @@ use self::syscall::{handle_syscall_exception, SYSCALL_OPCODE};
 use super::*;
 use crate::signal::{FaultSignal, SigSet};
 use crate::syscall::{CpuContext, SyscallNum};
+use aligned::{Aligned, A16};
+use core::arch::x86_64::_fxsave;
 use sgx_types::*;
 
 // Modules for instruction simulation
@@ -25,15 +27,26 @@ pub fn register_exception_handlers() {
 #[no_mangle]
 extern "C" fn handle_exception(info: *mut sgx_exception_info_t) -> i32 {
     extern "C" {
-        fn __occlum_syscall_c_abi(num: u32, info: *mut sgx_exception_info_t) -> u32;
+        fn __occlum_syscall_c_abi(
+            num: u32,
+            info: *mut sgx_exception_info_t,
+            fpregs: *mut u8,
+        ) -> u32;
     }
-    unsafe { __occlum_syscall_c_abi(SyscallNum::HandleException as u32, info) };
+
+    let mut fpregs: Aligned<A16, _> = Aligned([0u8; 512]);
+    let mut fpregs = fpregs.as_mut_ptr();
+    unsafe {
+        _fxsave(fpregs);
+        __occlum_syscall_c_abi(SyscallNum::HandleException as u32, info, fpregs)
+    };
     unreachable!();
 }
 
 /// Exceptions are handled as a special kind of system calls.
 pub fn do_handle_exception(
     info: *mut sgx_exception_info_t,
+    fpregs: *mut u8,
     user_context: *mut CpuContext,
 ) -> Result<isize> {
     let info = unsafe { &mut *info };
@@ -65,6 +78,7 @@ pub fn do_handle_exception(
     //
     // As the thread cannot proceed without handling the exception, we choose to force
     // delivering the signal regardless of the current signal mask.
+    user_context.fpregs = fpregs;
     let signal = Box::new(FaultSignal::new(info));
     crate::signal::force_signal(signal, user_context);
 

src/libos/src/interrupt/mod.rs

@@ -1,8 +1,9 @@
+pub use self::sgx::sgx_interrupt_info_t;
 use crate::prelude::*;
 use crate::process::ThreadRef;
 use crate::syscall::{CpuContext, SyscallNum};
-
+use aligned::{Aligned, A16};
-pub use self::sgx::sgx_interrupt_info_t;
+use core::arch::x86_64::_fxsave;
 
 mod sgx;
 
@@ -15,14 +16,25 @@ pub fn init() {
 
 extern "C" fn handle_interrupt(info: *mut sgx_interrupt_info_t) -> i32 {
     extern "C" {
-        fn __occlum_syscall_c_abi(num: u32, info: *mut sgx_interrupt_info_t) -> u32;
+        fn __occlum_syscall_c_abi(
+            num: u32,
+            info: *mut sgx_interrupt_info_t,
+            fpregs: *mut u8,
+        ) -> u32;
     }
-    unsafe { __occlum_syscall_c_abi(SyscallNum::HandleInterrupt as u32, info) };
+
+    let mut fpregs: Aligned<A16, _> = Aligned([0u8; 512]);
+    let mut fpregs = fpregs.as_mut_ptr();
+    unsafe {
+        _fxsave(fpregs);
+        __occlum_syscall_c_abi(SyscallNum::HandleInterrupt as u32, info, fpregs)
+    };
     unreachable!();
 }
 
 pub fn do_handle_interrupt(
     info: *mut sgx_interrupt_info_t,
+    fpregs: *mut u8,
     cpu_context: *mut CpuContext,
 ) -> Result<isize> {
     let info = unsafe { &*info };
@@ -30,6 +42,7 @@ pub fn do_handle_interrupt(
     // The cpu context is overriden so that it is as if the syscall is called from where the
     // interrupt happened
     *context = CpuContext::from_sgx(&info.cpu_context);
+    context.fpregs = fpregs;
     Ok(0)
 }
 

src/libos/src/signal/c_types.rs

@@ -207,7 +207,7 @@ pub struct ucontext_t {
     pub uc_stack: stack_t,
     pub uc_mcontext: mcontext_t,
     pub uc_sigmask: sigset_t,
-    __fpregs_mem: [u64; 64],
+    pub fpregs: [u8; 64 * 8], //fxsave structure
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -220,7 +220,7 @@ pub struct sigaltstack_t {
 
 pub type stack_t = sigaltstack_t;
 
-#[derive(Debug, Clone, Copy, Default)]
+#[derive(Debug, Clone, Copy)]
 #[repr(C)]
 pub struct mcontext_t {
     pub inner: CpuContext,

src/libos/src/signal/do_sigreturn.rs

@@ -2,9 +2,13 @@ use super::c_types::{mcontext_t, siginfo_t, ucontext_t};
 use super::constants::SIGKILL;
 use super::sig_stack::SigStackFlags;
 use super::{SigAction, SigActionFlags, SigDefaultAction, SigSet, Signal};
+use crate::lazy_static::__Deref;
 use crate::prelude::*;
 use crate::process::{ProcessRef, TermStatus, ThreadRef};
 use crate::syscall::CpuContext;
+use aligned::{Aligned, A16};
+use core::arch::x86_64::{_fxrstor, _fxsave};
+use std::{ptr, slice};
 
 pub fn do_rt_sigreturn(curr_user_ctxt: &mut CpuContext) -> Result<()> {
     debug!("do_rt_sigreturn");
@@ -23,10 +27,19 @@ pub fn do_rt_sigreturn(curr_user_ctxt: &mut CpuContext) -> Result<()> {
         }
         unsafe { &*last_ucontext.unwrap() }
     };
+
     // Restore sigmask
     *current!().sig_mask().write().unwrap() = SigSet::from_c(last_ucontext.uc_sigmask);
     // Restore user context
     *curr_user_ctxt = last_ucontext.uc_mcontext.inner;
+
+    // Restore the floating point registers to a temp area
+    // The floating point registers would be recoved just
+    // before return to user's code
+    let mut fpregs: Box<Aligned<A16, [u8]>> = Box::new(Aligned([0u8; 512]));
+    fpregs.copy_from_slice(&last_ucontext.fpregs);
+    curr_user_ctxt.fpregs = Box::into_raw(fpregs) as *mut u8;
+    curr_user_ctxt.fpregs_on_heap = 1; // indicates the fpregs is on heap
     Ok(())
 }
 
@@ -243,12 +256,30 @@ fn handle_signals_by_user(
         // signal handler. So we need to make sure the allocation is at least
         // 16-byte aligned.
         let ucontext = user_stack.alloc_aligned::<ucontext_t>(16)?;
+
         // TODO: set all fields in ucontext
         *ucontext = unsafe { std::mem::zeroed() };
         // Save the old sigmask
         ucontext.uc_sigmask = old_sigmask.to_c();
         // Save the user context
         ucontext.uc_mcontext.inner = *curr_user_ctxt;
+
+        // Save the floating point registers
+        if curr_user_ctxt.fpregs != ptr::null_mut() {
+            let fpregs =
+                unsafe { slice::from_raw_parts(curr_user_ctxt.fpregs, ucontext.fpregs.len()) };
+            ucontext.fpregs.copy_from_slice(fpregs);
+            // Clear the floating point registers, since we do not need to recover is when this syscall return
+            curr_user_ctxt.fpregs = ptr::null_mut();
+        } else {
+            // We need a correct fxsave structure in the buffer,
+            // because the app may modify part of it to update the
+            // floating point after the signal handler finished.
+            let mut fpregs: Aligned<A16, _> = Aligned([0u8; 512]);
+            unsafe { _fxsave(fpregs.as_mut_ptr()) };
+            ucontext.fpregs.copy_from_slice(fpregs.deref());
+        }
+
         ucontext as *mut ucontext_t
     };
     // 3. Set up the call return address on the stack before we "call" the signal handler

src/libos/src/syscall/mod.rs

@@ -7,6 +7,8 @@
 //! 3. Preprocess the system call and then call `dispatch_syscall` (in this file)
 //! 4. Call `do_*` to process the system call (in other modules)
 
+use aligned::{Aligned, A16};
+use core::arch::x86_64::_fxrstor;
 use std::any::Any;
 use std::convert::TryFrom;
 use std::ffi::{CStr, CString};
@@ -408,8 +410,8 @@ macro_rules! process_syscall_table_with_callback {
 
             // Occlum-specific system calls
             (Spawn = 360) => do_spawn(child_pid_ptr: *mut u32, path: *const i8, argv: *const *const i8, envp: *const *const i8, fdop_list: *const FdOp),
-            (HandleException = 361) => do_handle_exception(info: *mut sgx_exception_info_t, context: *mut CpuContext),
+            (HandleException = 361) => do_handle_exception(info: *mut sgx_exception_info_t, fpregs: *mut u8, context: *mut CpuContext),
-            (HandleInterrupt = 362) => do_handle_interrupt(info: *mut sgx_interrupt_info_t, context: *mut CpuContext),
+            (HandleInterrupt = 362) => do_handle_interrupt(info: *mut sgx_interrupt_info_t, fpregs: *mut u8, context: *mut CpuContext),
         }
     };
 }
@@ -622,10 +624,12 @@ fn do_syscall(user_context: &mut CpuContext) {
             syscall.args[0] = user_context as *mut _ as isize;
         } else if syscall_num == SyscallNum::HandleException {
             // syscall.args[0] == info
-            syscall.args[1] = user_context as *mut _ as isize;
+            // syscall.args[1] == fpregs
+            syscall.args[2] = user_context as *mut _ as isize;
         } else if syscall.num == SyscallNum::HandleInterrupt {
             // syscall.args[0] == info
-            syscall.args[1] = user_context as *mut _ as isize;
+            // syscall.args[1] == fpregs
+            syscall.args[2] = user_context as *mut _ as isize;
         } else if syscall.num == SyscallNum::Sigaltstack {
             // syscall.args[0] == new_ss
             // syscall.args[1] == old_ss
@@ -715,6 +719,20 @@ fn do_sysret(user_context: &mut CpuContext) -> ! {
         fn do_exit_task() -> !;
     }
     if current!().status() != ThreadStatus::Exited {
+        // Restore the floating point registers
+        // Todo: Is it correct to do fxstor in kernel?
+        let fpregs: *const u8 = user_context.fpregs;
+        if (fpregs != ptr::null()) {
+            unsafe { _fxrstor(fpregs) };
+
+            if user_context.fpregs_on_heap == 1 {
+                // Converting the raw pointer back into a Box with Box::from_raw for automatic cleanup
+                unsafe {
+                    let buf: &mut [u8] = core::slice::from_raw_parts_mut(user_context.fpregs, 512);
+                    let _ = Box::from_raw(buf);
+                }
+            }
+        }
         unsafe { __occlum_sysret(user_context) } // jump to user space
     } else {
         unsafe { do_exit_task() } // exit enclave
@@ -874,7 +892,7 @@ fn handle_unsupported() -> Result<isize> {
 ///
 /// Note. The definition of this struct must be kept in sync with the assembly
 /// code in `syscall_entry_x86-64.S`.
-#[derive(Clone, Copy, Debug, Default)]
+#[derive(Clone, Copy, Debug)]
 #[repr(C)]
 pub struct CpuContext {
     pub r8: u64,
@@ -895,6 +913,8 @@ pub struct CpuContext {
     pub rsp: u64,
     pub rip: u64,
     pub rflags: u64,
+    pub fpregs_on_heap: u64,
+    pub fpregs: *mut u8,
 }
 
 impl CpuContext {
@@ -918,6 +938,8 @@ impl CpuContext {
             rsp: src.rsp,
             rip: src.rip,
             rflags: src.rflags,
+            fpregs_on_heap: 0,
+            fpregs: ptr::null_mut(),
         }
     }
 }

src/libos/src/syscall/syscall_entry_x86-64.S

@@ -28,6 +28,8 @@ __occlum_syscall_linux_abi:
     // Save the target CPU state when `call __occlum_syscall` is returned in
     // a CpuContext struct. The registers are saved in the reverse order of 
     // the fields in CpuContext.
+    push $0         // default fpregs is NULL 
+    push $0         // default fpregs is allocated on stack
     pushfq
     push %rcx       // save %rip
     push %r11       // save %rsp

test/signal/main.c

@@ -230,6 +230,8 @@ int div_maybe_zero(int x, int y) {
     return x / y;
 }
 
+#define	fxsave(addr) __asm __volatile("fxsave %0" : "=m" (*(addr)))
+
 int test_handle_sigfpe() {
 #ifdef SGX_MODE_SIM
     printf("WARNING: Skip this test case as we do not support "
@@ -249,12 +251,21 @@ int test_handle_sigfpe() {
         THROW_ERROR("unexpected old sig handler");
     }
 
+    char x[512] __attribute__((aligned(16))) = {};
+    char y[512] __attribute__((aligned(16))) = {};
+
     // Trigger divide-by-zero exception
     int a = 1;
     int b = 0;
     // Use volatile to prevent compiler optimization
     volatile int c;
+    fxsave(x);
     c = div_maybe_zero(a, b);
+    fxsave(y);
+
+    if (memcmp(x, y, 512) != 0) {
+        THROW_ERROR("floating point registers are modified");
+    }
 
     printf("Signal handler successfully jumped over the divide-by-zero instruction\n");