From 38e92b2072b1a5d05fae20532d360a106b37c18b Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Fri, 8 Mar 2024 13:04:47 +0300
Subject: [PATCH] Ring3: Refactored exception handling.

---
 MdeModulePkg/Core/Dxe/Image/Image.c           |  7 ++++++-
 .../Dxe/SysCall/X64/CoreBootServices.nasm     | 12 +++---------
 .../X64/ExceptionHandlerAsm.nasm              | 19 ++++++++++++-------
 3 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/Image/Image.c b/MdeModulePkg/Core/Dxe/Image/Image.c
index 170fbf08bc..8233ec96af 100644
--- a/MdeModulePkg/Core/Dxe/Image/Image.c
+++ b/MdeModulePkg/Core/Dxe/Image/Image.c
@@ -1691,13 +1691,18 @@ InitializeRing3 (
   DEBUG ((DEBUG_ERROR, "Core: gRing3CallStackTop = %p\n", gRing3CallStackTop));
 
   //
-  // Initialize MSR_IA32_STAR and MSR_IA32_LSTAR for SYSCALL and SYSRET.
+  // Initialize MSR_IA32_STAR, MSR_IA32_LSTAR and MSR_IA32_FMASK for SYSCALL and SYSRET.
   //
   Msr = (((((UINT64)RING3_CODE64_SEL - 16) | 3) << 16) | (UINT64)RING0_CODE64_SEL) << 32;
   AsmWriteMsr64 (MSR_IA32_STAR, Msr);
 
   Msr = (UINT64)(UINTN)CoreBootServices;
   AsmWriteMsr64 (MSR_IA32_LSTAR, Msr);
+  //
+  // Disable maskable interrupts at SYSCALL.
+  //
+  Msr = (UINT64)BIT9;
+  AsmWriteMsr64 (MSR_IA32_FMASK, Msr);
 
   return Status;
 }
diff --git a/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm b/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm
index 6f9099a5a8..16b182a764 100644
--- a/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm
+++ b/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm
@@ -126,7 +126,6 @@ copy:
 ;------------------------------------------------------------------------------
 global ASM_PFX(CoreBootServices)
 ASM_PFX(CoreBootServices):
-    cli
     ; Switch from User to Core data segment selectors.
     mov     ax, ss
     mov     ds, ax
@@ -160,10 +159,8 @@ ASM_PFX(CoreBootServices):
     mov     r8, [rbp + 8*6]
 
     sti
-
     call ASM_PFX(CallBootService)
     push    rax
-
     cli
 
     SetRing3DataSegmentSelectors
@@ -181,8 +178,6 @@ ASM_PFX(CoreBootServices):
     pop     rbp
     pop     rsp
 
-    sti
-
     ; SYSCALL saves RFLAGS into R11 and the RIP of the next instruction into RCX.
 o64 sysret
     ; SYSRET copies the value in RCX into RIP and loads RFLAGS from R11.
@@ -198,7 +193,10 @@ o64 sysret
 ;------------------------------------------------------------------------------
 global ASM_PFX(CallRing3)
 ASM_PFX(CallRing3):
+    pushfq
+    pop     r11
     cli
+
     ; Save input Arguments.
     push    rcx
 
@@ -207,8 +205,6 @@ ASM_PFX(CallRing3):
     ; Prepare SYSRET arguments.
     mov     rcx, [gRing3EntryPoint]
     pop     rdx
-    pushfq
-    pop     r11
 
     ; Save Core Stack pointers and switch to User Stack.
     mov     [ASM_PFX(CoreRsp)], rsp
@@ -216,8 +212,6 @@ ASM_PFX(CallRing3):
     mov     rsp, [ASM_PFX(gRing3CallStackTop)]
     mov     rbp, rsp
 
-    sti
-
     ; Pass control to user image
 o64 sysret
 
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm
index bd0ff04de8..a8c69026db 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm
@@ -207,6 +207,15 @@ HasErrorCode:
     mov     rax, gs
     push    rax
 
+    ; Check whether Ring3 process was interrupted.
+    and     rax, 3
+    cmp     rax, 3
+    jne     SkipHook
+    mov     rax, cr2
+    cmp     rax, 0xFFFFFFFFFFFFFFF8
+    jne     SkipHook
+    mov     rcx, 32
+SkipHook:
     mov     rax, ss
     mov     ds, rax
     mov     es, rax
@@ -433,14 +442,12 @@ CetDone:
     pop     r15
 
     ; Check whether Ring3 process was interrupted.
-    push    rax
-    mov     rax, ss
     push    rcx
     mov     rcx, ds
-    cmp     rax, rcx
-    jne     ReturnToRing3
+    and     rcx, 3
+    cmp     rcx, 3
     pop     rcx
-    pop     rax
+    je      ReturnToRing3
 
     mov     rsp, rbp
     pop     rbp
@@ -469,8 +476,6 @@ DoReturn:
 DoIret:
     iretq
 ReturnToRing3:
-    pop     rcx
-    pop     rax
     mov     rsp, rbp
     pop     rbp
     add     rsp, 16