From 1a5a22e45075feaa1c02f3fe8259778401c45243 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         | 11 ++++++--
 .../CpuExceptionHandlerLib/DxeException.c     | 26 ++++++++++++-------
 3 files changed, 32 insertions(+), 12 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 e7575d9b80..1d13e044ee 100644
--- a/UefiCpuPkg/Library/CpuArchLib/CpuMp.c
+++ b/UefiCpuPkg/Library/CpuArchLib/CpuMp.c
@@ -622,9 +622,11 @@ InitializeExceptionStackSwitchHandlers (
 {
   EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;
   UINTN                           Index;
+  MSR_IA32_APIC_BASE_REGISTER     ApicBaseMsr;
 
   MpInitLibWhoAmI (&Index);
-  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
@@ -632,6 +634,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);
+    }
   }
 }
 
@@ -674,6 +679,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);
@@ -682,8 +688,9 @@ InitializeMpExceptionStackSwitchHandlers (
   }
 
   if (BufferSize != 0) {
-    Buffer = AllocateRuntimeZeroPool (BufferSize);
+    Buffer = AllocateRuntimePages (EFI_SIZE_TO_PAGES (BufferSize));
     ASSERT (Buffer != NULL);
+    SetMem (Buffer, BufferSize, 0);
     BufferSize = 0;
     for (Index = 0; Index < mNumberOfProcessors; ++Index) {
       if (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL) {
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 ();
+}