From 799350d3de31db6b4545322fac5c25af1b8460ca Mon Sep 17 00:00:00 2001
From: zhubojun <zhubojun.zbj@antgroup.com>
Date: Thu, 23 Mar 2023 23:44:56 -0400
Subject: [PATCH] [libos] Retain red zone area and FLAGS unmodified across
 syscall

1. The previous implementation of __occlum_syret usse the red zone to
save the register temporarily, which may overwrite the red zone area.
This change avoids using the red zone when handling syscall.

2. Save and restore the FLAGS register before and after the operations
on PKRU to keep the FLAGS unmodified.
---
 .../src/entry/context_switch/switch_x86-64.S  | 35 +++++++++++++------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/src/libos/src/entry/context_switch/switch_x86-64.S b/src/libos/src/entry/context_switch/switch_x86-64.S
index 04dec47d9..e4d07b313 100644
--- a/src/libos/src/entry/context_switch/switch_x86-64.S
+++ b/src/libos/src/entry/context_switch/switch_x86-64.S
@@ -65,6 +65,10 @@ __switch_to_user:
     pop %rcx
     pop %rsp
 
+    // Skip the red zone aera (128 bytes according to x86_64 System V ABI)
+    // Use `lea` here since it does not modify FLAGS
+    lea -128(%rsp), %rsp
+
     // Store RFLAGS since `cmp` operation may overwrite it
     pushfq
     push %rax
@@ -76,31 +80,34 @@ __switch_to_user:
     pop %rax
     popfq
 
+    lea 128(%rsp), %rsp
     jmp *%gs:(TD_USER_RIP)
     // This should never happen
     ud2
 
 update_pkru_at_switch_to_user:
-    pop %rax
-    popfq
-
-    sub $0x20, %rsp
-    mov %rax, (%rsp)
-    mov %rdx, 0x8(%rsp)
-    mov %rcx, 0x10(%rsp)
+    // Prepare temporary stack area for saving the %rdx, %rdx and return address
+    lea -3*8(%rsp), %rsp
+    mov %rdx, (%rsp)
+    mov %rcx, 8(%rsp)
     mov %gs:(TD_USER_RIP), %rcx
-    mov %rcx, 0x18(%rsp)
+    mov %rcx, 2*8(%rsp)
 
     xor %ecx, %ecx
     xor %edx, %edx
     mov $PKRU_USER, %eax
     wrpkru
 
-    pop %rax
     pop %rdx
     pop %rcx
-    ret
+    // Skip the return address
+    lea 8(%rsp), %rsp
 
+    pop %rax
+    popfq
+    lea 128(%rsp), %rsp     // Restore the %rsp
+
+    jmp *(-128-3*8)(%rsp)
     // This should never happen
     ud2
 
@@ -134,6 +141,11 @@ update_pkru_at_switch_to_user:
     .global __syscall_entry_linux_pku_abi
     .type __syscall_entry_linux_pku_abi, @function
 __syscall_entry_linux_pku_abi:
+    // Skip the red zone aera (128 bytes according to x86_64 System V ABI)
+    // Use `lea` here since it does not modify FLAGS
+    lea -128(%rsp), %rsp
+
+    pushfq
     pushq %rcx
     pushq %rdx
     pushq %rax
@@ -146,6 +158,9 @@ __syscall_entry_linux_pku_abi:
     popq %rax
     popq %rdx
     popq %rcx
+    popfq
+
+    lea 128(%rsp), %rsp
 
     .global __syscall_entry_linux_abi
     .type __syscall_entry_linux_abi, @function