From f3ce095edbc2ebd36c365a1862ab23f9b4626337 Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Tue, 5 Mar 2024 10:04:50 +0300
Subject: [PATCH] Ring3: Fixed interrrupts handling.

---
 .../Dxe/SysCall/X64/CoreBootServices.nasm     | 11 +++++++
 .../X64/ExceptionHandlerAsm.nasm              | 29 ++++++++++++++++---
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm b/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm
index 79ba899087..7a89cded04 100644
--- a/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm
+++ b/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm
@@ -113,6 +113,7 @@ copy:
 ;------------------------------------------------------------------------------
 global ASM_PFX(CoreBootServices)
 ASM_PFX(CoreBootServices):
+    cli
     ; Save User data segment selector temporarily in R11.
     mov     r11, ds
 
@@ -148,8 +149,12 @@ ASM_PFX(CoreBootServices):
     mov     rdx, rbp
     mov     r8, [rbp + 8*6]
 
+    sti
+
     call ASM_PFX(CallBootService)
 
+    cli
+
     ; Step over Arguments [1..3].
     add     rsp, 8*3
 
@@ -170,6 +175,8 @@ o16 mov     gs, r11
     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.
@@ -185,6 +192,7 @@ o64 sysret
 ;------------------------------------------------------------------------------
 global ASM_PFX(CallRing3)
 ASM_PFX(CallRing3):
+    cli
     ; Save input Arguments.
     push    r12
     mov     r12, rcx
@@ -217,6 +225,8 @@ ASM_PFX(CallRing3):
     mov     rsp, [ASM_PFX(gRing3CallStackTop)]
     mov     rbp, rsp
 
+    sti
+
     ; Pass control to user image
 o64 sysret
 
@@ -224,6 +234,7 @@ coreReturnAddress:
     mov     rsp, [ASM_PFX(CoreRsp)]
     mov     rbp, [ASM_PFX(CoreRbp)]
     mov     rax, rdx
+    sti
     ret
 
 SECTION .data
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm
index 4863686f14..bd0ff04de8 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm
@@ -207,6 +207,12 @@ HasErrorCode:
     mov     rax, gs
     push    rax
 
+    mov     rax, ss
+    mov     ds, rax
+    mov     es, rax
+    mov     fs, rax
+    mov     gs, rax
+
     mov     [rbp + 8], rcx               ; save vector number
 
 ;; UINT64  Rip;
@@ -356,7 +362,6 @@ DrFinish:
     incsspq rax                 ; SSP should be 0xFC0 now
 CetDone:
 %endif
-
     cli
 ;; UINT64  ExceptionData;
     add     rsp, 8
@@ -398,10 +403,9 @@ CetDone:
 
 ;; UINT64  Gs, Fs, Es, Ds, Cs, Ss;
     pop     rax
-    ; mov     gs, rax ; not for gs
+    mov     gs, rax
     pop     rax
-    ; mov     fs, rax ; not for fs
-    ; (X64 will not use fs and gs, so we do not restore it)
+    mov     fs, rax
     pop     rax
     mov     es, rax
     pop     rax
@@ -428,6 +432,16 @@ CetDone:
     pop     r14
     pop     r15
 
+    ; Check whether Ring3 process was interrupted.
+    push    rax
+    mov     rax, ss
+    push    rcx
+    mov     rcx, ds
+    cmp     rax, rcx
+    jne     ReturnToRing3
+    pop     rcx
+    pop     rax
+
     mov     rsp, rbp
     pop     rbp
     add     rsp, 16
@@ -454,6 +468,13 @@ DoReturn:
     retfq
 DoIret:
     iretq
+ReturnToRing3:
+    pop     rcx
+    pop     rax
+    mov     rsp, rbp
+    pop     rbp
+    add     rsp, 16
+    iretq
 
 ;-------------------------------------------------------------------------------------
 ;  GetTemplateAddressMap (&AddressMap);