From 3615418a84d3ab2ff51dd092dfe9a24c9f846537 Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Fri, 29 Nov 2024 14:33:17 +0300
Subject: [PATCH] Ring3: Added SetExceptionAddresses(), aligned exceptions'
 stacks.

---
 .../Include/Library/CpuExceptionHandlerLib.h  |  7 +++++
 UefiCpuPkg/Library/CpuArchLib/CpuMp.c         |  8 +++++-
 .../CpuExceptionHandlerLib/DxeException.c     | 26 ++++++++++++-------
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/MdeModulePkg/Include/Library/CpuExceptionHandlerLib.h b/MdeModulePkg/Include/Library/CpuExceptionHandlerLib.h
index 6864d8cb6c..36fe66ed32 100644
--- a/MdeModulePkg/Include/Library/CpuExceptionHandlerLib.h
+++ b/MdeModulePkg/Include/Library/CpuExceptionHandlerLib.h
@@ -113,4 +113,11 @@ GetExceptionAddresses (
   VOID
   );
 
+VOID
+EFIAPI
+SetExceptionAddresses (
+  IN VOID   *Buffer,
+  IN UINTN  BufferSize
+  );
+
 #endif
diff --git a/UefiCpuPkg/Library/CpuArchLib/CpuMp.c b/UefiCpuPkg/Library/CpuArchLib/CpuMp.c
index e85dce0d46..d4b0dc4e55 100644
--- a/UefiCpuPkg/Library/CpuArchLib/CpuMp.c
+++ b/UefiCpuPkg/Library/CpuArchLib/CpuMp.c
@@ -623,6 +623,7 @@ InitializeExceptionStackSwitchHandlers (
   EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;
   UINTN                           Index;
   EFI_STATUS                      Status;
+  MSR_IA32_APIC_BASE_REGISTER     ApicBaseMsr;
 
   Status = MpInitLibWhoAmI (&Index);
 
@@ -631,7 +632,8 @@ InitializeExceptionStackSwitchHandlers (
     return;
   }
 
-  SwitchStackData = (EXCEPTION_STACK_SWITCH_CONTEXT *)Buffer;
+  SwitchStackData    = (EXCEPTION_STACK_SWITCH_CONTEXT *)Buffer;
+  ApicBaseMsr.Uint64 = AsmReadMsr64 (MSR_IA32_APIC_BASE);
 
   //
   // This may be called twice for each Cpu. Only run InitializeSeparateExceptionStacks
@@ -639,6 +641,9 @@ InitializeExceptionStackSwitchHandlers (
   //
   if ((SwitchStackData[Index].Status == EFI_NOT_STARTED) || (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL)) {
     SwitchStackData[Index].Status = InitializeSeparateExceptionStacks (SwitchStackData[Index].Buffer, &SwitchStackData[Index].BufferSize);
+    if ((ApicBaseMsr.Bits.BSP != 0) && (SwitchStackData[Index].Status == EFI_SUCCESS)) {
+      SetExceptionAddresses (SwitchStackData[Index].Buffer, SwitchStackData[Index].BufferSize);
+    }
   }
 }
 
@@ -687,6 +692,7 @@ InitializeMpExceptionStackSwitchHandlers (
   for (Index = 0; Index < mNumberOfProcessors; ++Index) {
     if (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL) {
       ASSERT (SwitchStackData[Index].BufferSize != 0);
+      SwitchStackData[Index].BufferSize = ALIGN_VALUE (SwitchStackData[Index].BufferSize, EFI_PAGE_SIZE);
       BufferSize += SwitchStackData[Index].BufferSize;
     } else {
       ASSERT (SwitchStackData[Index].Status == EFI_SUCCESS);
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
index 91c2411e4c..9429df0355 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/DxeException.c
@@ -136,17 +136,13 @@ InitializeSeparateExceptionStacks (
   UINTN       LocalBufferSize;
   EFI_STATUS  Status;
 
-  mAddresses.ExceptionStackSize = CPU_STACK_SWITCH_EXCEPTION_NUMBER * CPU_KNOWN_GOOD_STACK_SIZE;
-
   if ((Buffer == NULL) && (BufferSize == NULL)) {
     SetMem (mBuffer, sizeof (mBuffer), 0);
     LocalBufferSize = sizeof (mBuffer);
     Status          = ArchSetupExceptionStack (mBuffer, &LocalBufferSize);
     ASSERT_EFI_ERROR (Status);
-    mAddresses.ExceptionStackBase = (UINTN)mBuffer;
     return Status;
   } else {
-    mAddresses.ExceptionStackBase = (UINTN)Buffer;
     return ArchSetupExceptionStack (Buffer, BufferSize);
   }
 }
@@ -209,11 +205,21 @@ GetExceptionAddresses (
 {
   mSwitchCr3Flag = 1;
 
-  mAddresses.ExceptionHandlerBase = (UINTN)&ExceptionHandlerBase;
-  mAddresses.ExceptionHandlerSize = (UINTN)&ExceptionHandlerEnd - mAddresses.ExceptionHandlerBase;
-  mAddresses.ExceptionDataBase    = (UINTN)&CorePageTable;
-
-  CorePageTable = AsmReadCr3 ();
-
   return &mAddresses;
 }
+
+VOID
+EFIAPI
+SetExceptionAddresses (
+  IN VOID   *Buffer,
+  IN UINTN  BufferSize
+  )
+{
+  mAddresses.ExceptionStackBase   = (UINTN)Buffer;
+  mAddresses.ExceptionStackSize   = BufferSize;
+  mAddresses.ExceptionHandlerBase = (UINTN)&ExceptionHandlerBase;
+  mAddresses.ExceptionHandlerSize = (UINTN)&ExceptionHandlerEnd - mAddresses.ExceptionHandlerBase;
+  mAddresses.ExceptionDataBase    = (UINTN)&CorePageTable;
+
+  CorePageTable = AsmReadCr3 ();
+}