From 1f8f6ddc81963e7a1ec526e87e41463e47feb888 Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Mon, 20 Jan 2025 16:28:26 +0300
Subject: [PATCH] CpuExceptionHandlerLib: Refactored out mSwitchCr3Flag.

---
 UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c  | 3 ---
 .../CpuExceptionHandlerLib/Ia32/ExceptionHandlerAsm.nasm  | 8 +++-----
 .../CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm   | 8 +++-----
 3 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
index 9429df0355..ce4f32a58f 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
@@ -34,7 +34,6 @@ EXCEPTION_ADDRESSES  mAddresses;
 
 extern UINTN         ExceptionHandlerBase;
 extern UINTN         ExceptionHandlerEnd;
-extern UINT8         mSwitchCr3Flag;
 extern UINTN         CorePageTable;
 
 /**
@@ -203,8 +202,6 @@ GetExceptionAddresses (
   VOID
   )
 {
-  mSwitchCr3Flag = 1;
-
   return &mAddresses;
 }
 
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ExceptionHandlerAsm.nasm b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ExceptionHandlerAsm.nasm
index ee2a061fd5..07fef59095 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ExceptionHandlerAsm.nasm
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ExceptionHandlerAsm.nasm
@@ -38,10 +38,6 @@ global ASM_PFX(UserPageTable)
 ASM_PFX(UserPageTable):
   resq 1
 
-global ASM_PFX(mSwitchCr3Flag)
-ASM_PFX(mSwitchCr3Flag):
-  db 0x0
-
 ALIGN   4096
 Padding:
   db 0x0
@@ -105,7 +101,9 @@ HookAfterStubHeaderEnd:
 global ASM_PFX(CommonInterruptEntry)
 ASM_PFX(CommonInterruptEntry):
     cli
-    cmp     byte [ASM_PFX(mSwitchCr3Flag)], 0
+    ; Check whether User Space process was interrupted.
+    mov     eax, ds
+    and     eax, 3
     jz      NoCr3Switch
     mov     eax, cr3
     mov     [ASM_PFX(UserPageTable)], eax
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm
index 9ba3ae5172..278ff4f592 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ExceptionHandlerAsm.nasm
@@ -63,10 +63,6 @@ global ASM_PFX(UserPageTable)
 ASM_PFX(UserPageTable):
   resq 1
 
-global ASM_PFX(mSwitchCr3Flag)
-ASM_PFX(mSwitchCr3Flag):
-  db 0x0
-
 ALIGN   4096
 Padding:
   db 0x0
@@ -140,7 +136,9 @@ HookAfterStubHeaderEnd:
 global ASM_PFX(CommonInterruptEntry)
 ASM_PFX(CommonInterruptEntry):
     cli
-    cmp     byte [ASM_PFX(mSwitchCr3Flag)], 0
+    ; Check whether User Space process was interrupted.
+    mov     rax, ds
+    and     rax, 3
     jz      NoCr3Switch
     mov     rax, cr3
     mov     [ASM_PFX(UserPageTable)], rax