[Pal/Linux-SGX] Honor red zone when handling async exits

Before .Lhandle_exception used the stack directly and thereby messed
with the red zone of the interrupted code.

Also ensure that the stack is aligned before calling
_DkExceptionHandler.

Pal/src/host/Linux-SGX/db_exception.c

@@ -157,37 +157,6 @@ void _DkExceptionRealHandler (int event, PAL_NUM arg, struct pal_frame * frame,
     _DkGenericSignalHandle(event, arg, frame, context);
 }
 
-void restore_sgx_context (sgx_context_t * uc)
-{
-    /* prepare the return address */
-    uc->rsp -= 8;
-    *(uint64_t *) uc->rsp = uc->rip;
-
-    /* now pop the stack */
-    __asm__ volatile (
-                  "mov %0, %%rsp\n"
-                  "pop %%rax\n"
-                  "pop %%rcx\n"
-                  "pop %%rdx\n"
-                  "pop %%rbx\n"
-                  "add $8, %%rsp\n" /* don't pop RSP yet */
-                  "pop %%rbp\n"
-                  "pop %%rsi\n"
-                  "pop %%rdi\n"
-                  "pop %%r8\n"
-                  "pop %%r9\n"
-                  "pop %%r10\n"
-                  "pop %%r11\n"
-                  "pop %%r12\n"
-                  "pop %%r13\n"
-                  "pop %%r14\n"
-                  "pop %%r15\n"
-                  "popfq\n"
-                  "mov -104(%%rsp), %%rsp\n"
-                  "ret\n"
-                  :: "r"(uc) : "memory");
-}
-
 /*
  * return value:
  *  true:  #UD is handled.
@@ -355,7 +324,6 @@ void _DkRaiseFailure (int error)
 void _DkExceptionReturn (void * event)
 {
     PAL_EVENT * e = event;
-    sgx_context_t uc;
     PAL_CONTEXT * ctx = e->context;
 
     if (!ctx) {
@@ -372,26 +340,30 @@ void _DkExceptionReturn (void * event)
                       "retq\r\n" ::: "memory");
     }
 
-    uc.rax = ctx->rax;
-    uc.rbx = ctx->rbx;
-    uc.rcx = ctx->rcx;
-    uc.rdx = ctx->rdx;
-    uc.rsp = ctx->rsp;
-    uc.rbp = ctx->rbp;
-    uc.rsi = ctx->rsi;
-    uc.rdi = ctx->rdi;
-    uc.r8  = ctx->r8;
-    uc.r9  = ctx->r9;
-    uc.r10 = ctx->r10;
-    uc.r11 = ctx->r11;
-    uc.r12 = ctx->r12;
-    uc.r13 = ctx->r13;
-    uc.r14 = ctx->r14;
-    uc.r15 = ctx->r15;
-    uc.rflags = ctx->efl;
-    uc.rip = ctx->rip;
-
-    restore_sgx_context(&uc);
+    // Allocate sgx_context_t just below the "normal" stack (honoring the red
+    // zone) and then copy the content of ctx there. This is needed by
+    // restore_sgx_context.
+    sgx_context_t * uc = (void *)ctx->rsp - sizeof(sgx_context_t) - RED_ZONE_SIZE;
+    uc->rax = ctx->rax;
+    uc->rbx = ctx->rbx;
+    uc->rcx = ctx->rcx;
+    uc->rdx = ctx->rdx;
+    uc->rsp = ctx->rsp;
+    uc->rbp = ctx->rbp;
+    uc->rsi = ctx->rsi;
+    uc->rdi = ctx->rdi;
+    uc->r8  = ctx->r8;
+    uc->r9  = ctx->r9;
+    uc->r10 = ctx->r10;
+    uc->r11 = ctx->r11;
+    uc->r12 = ctx->r12;
+    uc->r13 = ctx->r13;
+    uc->r14 = ctx->r14;
+    uc->r15 = ctx->r15;
+    uc->rflags = ctx->efl;
+    uc->rip = ctx->rip;
+
+    restore_sgx_context(uc);
 }
 
 void _DkHandleExternalEvent (PAL_NUM event, sgx_context_t * uc)

Pal/src/host/Linux-SGX/enclave_entry.S

@@ -165,7 +165,7 @@ enclave_entry:
 	je 1f
 	movq %rax, %rsi
 1:
-	subq $SGX_CONTEXT_SIZE, %rsi
+	subq $(SGX_CONTEXT_SIZE + RED_ZONE_SIZE), %rsi
 
 	# we have exitinfo in RDI, swap with the one on GPR
 	# and dump into the context
@@ -209,9 +209,13 @@ enclave_entry:
 	movq SGX_GPR_RIP(%rbx), %rdi
 	movq %rdi, SGX_CONTEXT_RIP(%rsi)
 
-	movq %rsi, SGX_GPR_RSP(%rbx)
+	# Pass pointer to sgx_context_t to _DkExceptionHandler
 	movq %rsi, SGX_GPR_RSI(%rbx)
 
+	# Align the stack for _DkExceptionHandler
+	andq $STACK_ALIGN, %rsi
+	movq %rsi, SGX_GPR_RSP(%rbx)
+
 	# new RIP is the exception handler
 	leaq _DkExceptionHandler(%rip), %rdi
 	movq %rdi, SGX_GPR_RIP(%rbx)
@@ -435,3 +439,53 @@ wrfsbase:
 
 	.cfi_endproc
 	.size wrfsbase, .-wrfsbase
+
+/*
+ * Restore an sgx_context_t as generated by .Lhandle_exception. Execution will
+ * continue as specified by the rip in the context.
+ *
+ * It is required that:
+ *
+ *     %rdi == *(%rdi + SGX_CONTEXT_RSP) - (SGX_CONTEXT_SIZE + RED_ZONE_SIZE)
+ *
+ * This holds for the original sgx_context allocated by .Lhandle_exception.
+ * restore_sgx_context is a safe wrapper which checks this.
+ */
+	.global _restore_sgx_context
+	.type _restore_sgx_context, @function
+
+_restore_sgx_context:
+	movq SGX_CONTEXT_RAX(%rdi), %rax
+	movq SGX_CONTEXT_RCX(%rdi), %rcx
+	movq SGX_CONTEXT_RDX(%rdi), %rdx
+	movq SGX_CONTEXT_RBX(%rdi), %rbx
+	# For %rsp see below.
+	movq SGX_CONTEXT_RBP(%rdi), %rbp
+	movq SGX_CONTEXT_RSI(%rdi), %rsi
+	# For %rdi see below.
+	movq SGX_CONTEXT_R8(%rdi), %r8
+	movq SGX_CONTEXT_R9(%rdi), %r9
+	movq SGX_CONTEXT_R10(%rdi), %r10
+	movq SGX_CONTEXT_R11(%rdi), %r11
+	movq SGX_CONTEXT_R12(%rdi), %r12
+	movq SGX_CONTEXT_R13(%rdi), %r13
+	movq SGX_CONTEXT_R14(%rdi), %r14
+	movq SGX_CONTEXT_R15(%rdi), %r15
+
+	# We need to make sure that %rsp - RED_ZONE_SIZE never points above
+	# anything we still need. Otherwise .Lhandle_exception might mess with
+	# it. SGX_CONTEXT_RDI - SGX_CONTEXT_RFLAGS <= RED_ZONE_SIZE, see
+	# sgx_arch.h.
+	leaq SGX_CONTEXT_RFLAGS(%rdi), %rsp
+	popfq # remember to not touch any flags after here
+
+	movq SGX_CONTEXT_RDI(%rdi), %rdi
+	# Now %rdi is restored so we need to use the stack to access the
+	# context.
+
+	# Now pop %rip and fix stack pointer in one operation (to avoid
+	# problems with nesting, see comment above). SGX_CONTEXT_RIP is
+	# directly after SGX_CONTEXT_RFLAGS, see sgx_arch.h. Note that retq
+	# decreases %rsp by 8 for the poped %rip additionaly to the passed
+	# offset.
+	retq $(SGX_CONTEXT_SIZE + RED_ZONE_SIZE - SGX_CONTEXT_RIP - 8)

Pal/src/host/Linux-SGX/enclave_framework.c

@@ -1193,3 +1193,28 @@ out:
     DkStreamDelete(stream, 0);
     return ret;
 }
+
+/*
+ * Restore an sgx_context_t as generated by .Lhandle_exception. Execution will
+ * continue as specified by the rip in the context.
+ *
+ * It is required that:
+ *
+ *     ctx == ctx->rsp - (sizeof(sgx_context_t) + RED_ZONE_SIZE)
+ *
+ * This means that the ctx is allocated directly below the "normal" stack
+ * (honoring its red zone). This is needed to properly restore the old state
+ * (see _restore_sgx_context for details).
+ *
+ * For the original sgx_context_t allocated by .Lhandle_exception this is true.
+ * This is a safe wrapper around _restore_sgx_context, which checks this
+ * preconditon.
+ */
+void restore_sgx_context(sgx_context_t *ctx) {
+    if (((uint64_t) ctx) != ctx->rsp - (sizeof(sgx_context_t) + RED_ZONE_SIZE)) {
+        SGX_DBG(DBG_E, "Invalid sgx_context_t pointer passed to restore_sgx_context!\n");
+        ocall_exit(1);
+    }
+
+    _restore_sgx_context(ctx);
+}

Pal/src/host/Linux-SGX/pal_linux_defs.h

@@ -26,4 +26,6 @@
 #define MAX_ARGS_SIZE       (10000000)
 #define MAX_ENV_SIZE        (10000000)
 
+#define RED_ZONE_SIZE       (128)
+
 #endif /* PAL_LINUX_DEFS_H */

Pal/src/host/Linux-SGX/sgx_api.h

@@ -48,4 +48,7 @@ uint32_t rdrand (void);
 uint64_t rdfsbase (void);
 void wrfsbase (uint64_t addr);
 
+void restore_sgx_context(sgx_context_t *ctx);
+void _restore_sgx_context(sgx_context_t *ctx);
+
 #endif /* SGX_API_H */

Pal/src/host/Linux-SGX/sgx_arch.h

@@ -126,6 +126,13 @@ typedef struct {
     uint64_t rip;
 } sgx_context_t;
 
+// Required by _restore_sgx_context, see enclave_entry.S.
+_Static_assert(offsetof(sgx_context_t, rip) - offsetof(sgx_context_t, rflags) ==
+               sizeof(((sgx_context_t) {0}).rflags),
+               "rip must be directly after rflags in sgx_context_t");
+_Static_assert(offsetof(sgx_context_t, rflags) - offsetof(sgx_context_t, rdi) <= RED_ZONE_SIZE,
+               "rdi needs to be within red zone distance from rflags");
+
 typedef struct {
     uint32_t vector:8;
     uint32_t type:3;