Add [depex] on gEfiMpServiceProtocolGuid so the number of CPUs is known

[mirror_edk2.git] / EdkCompatibilityPkg / Compatibility / SmmBaseHelper / SmmBaseHelper.c

diff --git a/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.c b/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.c
index a8f5d13b60885a14bc2772d97923852e8824c716..8a6214afd304bc380a47720e8143cc12cdedd7ae 100644 (file)
--- a/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.c
+++ b/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.c
@@ -30,6 +30,7 @@
 #include <Protocol/SmmCpu.h>\r
 #include <Protocol/LoadedImage.h>\r
 #include <Protocol/SmmCpuSaveState.h>\r
+#include <Protocol/MpService.h>\r
 \r
 ///\r
 /// Structure for tracking paired information of registered Framework SMI handler\r
@@ -61,6 +62,7 @@ EFI_SMM_CPU_PROTOCOL               *mSmmCpu;
 EFI_GUID                           mEfiSmmCpuIoGuid = EFI_SMM_CPU_IO_GUID;\r
 EFI_SMM_BASE_HELPER_READY_PROTOCOL *mSmmBaseHelperReady;\r
 EFI_SMM_SYSTEM_TABLE               *mFrameworkSmst;\r
+UINTN                              mNumberOfProcessors;\r
 \r
 LIST_ENTRY mCallbackInfoListHead = INITIALIZE_LIST_HEAD_VARIABLE (mCallbackInfoListHead);\r
 \r
@@ -156,7 +158,7 @@ ConstructFrameworkSmst (
   FrameworkSmst->Hdr.Revision = EFI_SMM_SYSTEM_TABLE_REVISION;\r
   CopyGuid (&FrameworkSmst->EfiSmmCpuIoGuid, &mEfiSmmCpuIoGuid);\r
 \r
-  FrameworkSmst->CpuSaveState = (EFI_SMM_CPU_SAVE_STATE *)AllocateZeroPool (gSmst->NumberOfCpus * sizeof (EFI_SMM_CPU_SAVE_STATE));\r
+  FrameworkSmst->CpuSaveState = (EFI_SMM_CPU_SAVE_STATE *)AllocateZeroPool (mNumberOfProcessors * sizeof (EFI_SMM_CPU_SAVE_STATE));\r
   ASSERT (FrameworkSmst->CpuSaveState != NULL);\r
 \r
   ///\r
@@ -437,7 +439,7 @@ CallbackThunk (
   /// and MP states in the Framework SMST.\r
   ///\r
 \r
-  for (CpuIndex = 0; CpuIndex < gSmst->NumberOfCpus; CpuIndex++) {\r
+  for (CpuIndex = 0; CpuIndex < mNumberOfProcessors; CpuIndex++) {\r
     State = (EFI_SMM_CPU_STATE *)gSmst->CpuSaveState[CpuIndex];\r
     SaveState = &mFrameworkSmst->CpuSaveState[CpuIndex].Ia32SaveState;\r
 \r
@@ -487,7 +489,7 @@ CallbackThunk (
   ///\r
   /// Save CPU Save States in case any of them was modified\r
   ///\r
-  for (CpuIndex = 0; CpuIndex < gSmst->NumberOfCpus; CpuIndex++) {\r
+  for (CpuIndex = 0; CpuIndex < mNumberOfProcessors; CpuIndex++) {\r
     for (Index = 0; Index < sizeof (mCpuSaveStateConvTable) / sizeof (CPU_SAVE_STATE_CONVERSION); Index++) {\r
       Status = mSmmCpu->WriteSaveState (\r
                           mSmmCpu,\r
@@ -676,7 +678,9 @@ SmmBaseHelperMain (
   )\r
 {\r
   EFI_STATUS  Status;\r
+  EFI_MP_SERVICES_PROTOCOL   *MpServices;\r
   EFI_HANDLE  Handle = NULL;\r
+  UINTN                      NumberOfEnabledProcessors;\r
 \r
   ///\r
   /// Locate SMM CPU Protocol which is used later to retrieve/update CPU Save States\r
@@ -684,6 +688,18 @@ SmmBaseHelperMain (
   Status = gSmst->SmmLocateProtocol (&gEfiSmmCpuProtocolGuid, NULL, (VOID **) &mSmmCpu);\r
   ASSERT_EFI_ERROR (Status);\r
 \r
+  //\r
+  // Get MP Services Protocol\r
+  //\r
+  Status = SystemTable->BootServices->LocateProtocol (&gEfiMpServiceProtocolGuid, NULL, (VOID **)&MpServices);\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  //\r
+  // Use MP Services Protocol to retrieve the number of processors and number of enabled processors\r
+  //\r
+  Status = MpServices->GetNumberOfProcessors (MpServices, &mNumberOfProcessors, &NumberOfEnabledProcessors);\r
+  ASSERT_EFI_ERROR (Status);\r
+  \r
   ///\r
   /// Interface structure of SMM BASE Helper Ready Protocol is allocated from UEFI pool\r
   /// instead of SMM pool so that SMM Base Thunk driver can access it in Non-SMM mode.\r