From a75ae00c63ad2859351e9682026462048f1cf83e Mon Sep 17 00:00:00 2001
From: Jethro Beekman <jethro@fortanix.com>
Date: Tue, 29 Jan 2019 17:10:22 +0530
Subject: [PATCH] SGX target: improve panic & exit handling

---
 src/libstd/sys/sgx/abi/entry.S          | 20 ++++++++---------
 src/libstd/sys/sgx/abi/panic.rs         | 30 ++++++++++++++++---------
 src/libstd/sys/sgx/abi/usercalls/mod.rs |  2 +-
 src/libstd/sys/sgx/mod.rs               |  2 +-
 4 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/src/libstd/sys/sgx/abi/entry.S b/src/libstd/sys/sgx/abi/entry.S
index ac7f95d4eae8..9b46c2180d9a 100644
--- a/src/libstd/sys/sgx/abi/entry.S
+++ b/src/libstd/sys/sgx/abi/entry.S
@@ -66,7 +66,7 @@ IMAGE_BASE:
     globvar EH_FRM_HDR_SIZE 8
 
 .Lreentry_panic_msg:
-    .asciz "Re-entered panicked enclave!"
+    .asciz "Re-entered aborted enclave!"
 .Lreentry_panic_msg_end:
 
 .Lusercall_panic_msg:
@@ -80,7 +80,7 @@ IMAGE_BASE:
     .org .+48 /*  reserved bits */
 
 .data
-.Lpanicked:
+.Laborted:
     .byte 0
 
 /*  TCS local storage section */
@@ -134,6 +134,9 @@ sgx_entry:
     jz .Lskip_debug_init
     mov %r10,%gs:tcsls_debug_panic_buf_ptr
 .Lskip_debug_init:
+/*  check for abort */
+    bt $0,.Laborted(%rip)
+    jc .Lreentry_panic
 /*  check if returning from usercall */
     mov %gs:tcsls_last_rsp,%r11
     test %r11,%r11
@@ -164,9 +167,6 @@ sgx_entry:
     mov %r14,%r8
     mov %r15,%r9
 .Lskip_init:
-/*  check for panic */
-    bt $0,.Lpanicked(%rip)
-    jc .Lreentry_panic
 /*  call into main entry point */
     load_tcsls_flag_secondary_bool cx /* RCX = entry() argument: secondary: bool */
     call entry /* RDI, RSI, RDX, R8, R9 passed in from userspace */
@@ -237,18 +237,18 @@ sgx_entry:
     stmxcsr (%rsp)
 .endm
 
-.global panic_exit
-panic_exit:
+.global usercall_exit
+usercall_exit:
 /* save registers in DEBUG mode, so that debugger can reconstruct the stack */
     testb $0xff,DEBUG(%rip)
     jz .Lskip_save_registers
     push_callee_saved_registers
     movq %rsp,%gs:tcsls_panic_last_rsp
 .Lskip_save_registers:
-/* set panicked bit */
-    movb $1,.Lpanicked(%rip)
+/* set aborted bit */
+    movb $1,.Laborted(%rip)
 /* call usercall exit(true) */
-    mov $1,%esi   /*  RSI = usercall() argument: panic = true */
+    /* NOP: mov %rsi,%rsi */ /*  RSI = usercall() argument: panic */
     xor %rdx,%rdx /*  RDX cleared */
     movq $usercall_nr_exit,%rdi /*  RDI = usercall exit */
     jmp .Lexit
diff --git a/src/libstd/sys/sgx/abi/panic.rs b/src/libstd/sys/sgx/abi/panic.rs
index 5ace7fb3368b..d23fa9a9ec6f 100644
--- a/src/libstd/sys/sgx/abi/panic.rs
+++ b/src/libstd/sys/sgx/abi/panic.rs
@@ -1,12 +1,18 @@
+use super::usercalls::alloc::UserRef;
+use cmp;
 use io::{self, Write};
-use slice::from_raw_parts_mut;
+use mem;
 
 extern "C" {
     fn take_debug_panic_buf_ptr() -> *mut u8;
     static DEBUG: u8;
 }
 
-pub(crate) struct SgxPanicOutput(Option<&'static mut [u8]>);
+pub(crate) struct SgxPanicOutput(Option<&'static mut UserRef<[u8]>>);
+
+fn empty_user_slice() -> &'static mut UserRef<[u8]> {
+    unsafe { UserRef::from_raw_parts_mut(1 as *mut u8, 0) }
+}
 
 impl SgxPanicOutput {
     pub(crate) fn new() -> Option<Self> {
@@ -17,32 +23,36 @@ impl SgxPanicOutput {
         }
     }
 
-    fn init(&mut self) -> &mut &'static mut [u8] {
+    fn init(&mut self) -> &mut &'static mut UserRef<[u8]> {
         self.0.get_or_insert_with(|| unsafe {
             let ptr = take_debug_panic_buf_ptr();
             if ptr.is_null() {
-                &mut []
+                empty_user_slice()
             } else {
-                from_raw_parts_mut(ptr, 1024)
+                UserRef::from_raw_parts_mut(ptr, 1024)
             }
         })
     }
 }
 
 impl Write for SgxPanicOutput {
-    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        self.init().write(buf)
+    fn write(&mut self, src: &[u8]) -> io::Result<usize> {
+        let dst = mem::replace(self.init(), empty_user_slice());
+        let written = cmp::min(src.len(), dst.len());
+        dst[..written].copy_from_enclave(&src[..written]);
+        self.0 = Some(&mut dst[written..]);
+        Ok(written)
     }
 
     fn flush(&mut self) -> io::Result<()> {
-        self.init().flush()
+        Ok(())
     }
 }
 
 #[no_mangle]
 pub extern "C" fn panic_msg(msg: &str) -> ! {
     let _ = SgxPanicOutput::new().map(|mut out| out.write(msg.as_bytes()));
-    unsafe { panic_exit(); }
+    unsafe { usercall_exit(true); }
 }
 
-extern "C" { pub fn panic_exit() -> !; }
+extern "C" { pub fn usercall_exit(panic: bool) -> !; }
diff --git a/src/libstd/sys/sgx/abi/usercalls/mod.rs b/src/libstd/sys/sgx/abi/usercalls/mod.rs
index 58903761ebe4..4e889c172ef3 100644
--- a/src/libstd/sys/sgx/abi/usercalls/mod.rs
+++ b/src/libstd/sys/sgx/abi/usercalls/mod.rs
@@ -119,7 +119,7 @@ pub unsafe fn launch_thread() -> IoResult<()> {
 /// Usercall `exit`. See the ABI documentation for more information.
 #[unstable(feature = "sgx_platform", issue = "56975")]
 pub fn exit(panic: bool) -> ! {
-    unsafe { raw::exit(panic) }
+    unsafe { super::panic::usercall_exit(panic) }
 }
 
 /// Usercall `wait`. See the ABI documentation for more information.
diff --git a/src/libstd/sys/sgx/mod.rs b/src/libstd/sys/sgx/mod.rs
index 7f8550490a19..f2593c35bed1 100644
--- a/src/libstd/sys/sgx/mod.rs
+++ b/src/libstd/sys/sgx/mod.rs
@@ -125,7 +125,7 @@ pub unsafe fn strlen(mut s: *const c_char) -> usize {
 }
 
 pub unsafe fn abort_internal() -> ! {
-    abi::panic::panic_exit()
+    abi::panic::usercall_exit(true)
 }
 
 pub fn hashmap_random_keys() -> (u64, u64) {