1 files changed, 23 insertions, 6 deletions
diff --git a/OvmfPkg/CpuHotplugSmm/CpuHotplug.c b/OvmfPkg/CpuHotplugSmm/CpuHotplug.c
index bf68fcd429..ee1497b931 100644
--- a/OvmfPkg/CpuHotplugSmm/CpuHotplug.c
+++ b/OvmfPkg/CpuHotplugSmm/CpuHotplug.c
@@ -45,13 +45,16 @@ STATIC CPU_HOT_PLUG_DATA *mCpuHotPlugData;
 // don't want to allocate SMRAM at OS runtime, and potentially fail (or
 // fragment the SMRAM map).
 //
-// These arrays provide room for ("possible CPU count" minus one) APIC IDs
-// each, as we don't expect every possible CPU to appear, or disappear, in a
-// single MMI. The numbers of used (populated) elements in the arrays are
+// The first array stores APIC IDs for hot-plug events, the second and the
+// third store APIC IDs and QEMU CPU Selectors (both indexed similarly) for
+// hot-unplug events. All of these provide room for "possible CPU count" minus
+// one elements as we don't expect every possible CPU to appear, or disappear,
+// in a single MMI. The numbers of used (populated) elements in the arrays are
 // determined on every MMI separately.
 //
 STATIC APIC_ID *mPluggedApicIds;
 STATIC APIC_ID *mToUnplugApicIds;
+STATIC UINT32  *mToUnplugSelectors;
 //
 // Address of the non-SMRAM reserved memory page that contains the Post-SMM Pen
 // for hot-added CPUs.
@@ -289,6 +292,7 @@ CpuHotplugMmi (
              mPluggedApicIds,
              &PluggedCount,
              mToUnplugApicIds,
+             mToUnplugSelectors,
              &ToUnplugCount
              );
   if (EFI_ERROR (Status)) {
@@ -333,7 +337,9 @@ CpuHotplugEntry (
   )
 {
   EFI_STATUS Status;
+  UINTN      Len;
   UINTN      Size;
+  UINTN      SizeSel;
 
   //
   // This module should only be included when SMM support is required.
@@ -387,8 +393,9 @@ CpuHotplugEntry (
   //
   // Allocate the data structures that depend on the possible CPU count.
   //
-  if (RETURN_ERROR (SafeUintnSub (mCpuHotPlugData->ArrayLength, 1, &Size)) ||
-      RETURN_ERROR (SafeUintnMult (sizeof (APIC_ID), Size, &Size))) {
+  if (RETURN_ERROR (SafeUintnSub (mCpuHotPlugData->ArrayLength, 1, &Len)) ||
+      RETURN_ERROR (SafeUintnMult (sizeof (APIC_ID), Len, &Size)) ||
+      RETURN_ERROR (SafeUintnMult (sizeof (UINT32), Len, &SizeSel))) {
     Status = EFI_ABORTED;
     DEBUG ((DEBUG_ERROR, "%a: invalid CPU_HOT_PLUG_DATA\n", __FUNCTION__));
     goto Fatal;
@@ -405,6 +412,12 @@ CpuHotplugEntry (
     DEBUG ((DEBUG_ERROR, "%a: MmAllocatePool(): %r\n", __FUNCTION__, Status));
     goto ReleasePluggedApicIds;
   }
+  Status = gMmst->MmAllocatePool (EfiRuntimeServicesData, SizeSel,
+                    (VOID **)&mToUnplugSelectors);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%a: MmAllocatePool(): %r\n", __FUNCTION__, Status));
+    goto ReleaseToUnplugApicIds;
+  }
 
   //
   // Allocate the Post-SMM Pen for hot-added CPUs.
@@ -412,7 +425,7 @@ CpuHotplugEntry (
   Status = SmbaseAllocatePostSmmPen (&mPostSmmPenAddress,
              SystemTable->BootServices);
   if (EFI_ERROR (Status)) {
-    goto ReleaseToUnplugApicIds;
+    goto ReleaseToUnplugSelectors;
   }
 
   //
@@ -472,6 +485,10 @@ ReleasePostSmmPen:
   SmbaseReleasePostSmmPen (mPostSmmPenAddress, SystemTable->BootServices);
   mPostSmmPenAddress = 0;
 
+ReleaseToUnplugSelectors:
+  gMmst->MmFreePool (mToUnplugSelectors);
+  mToUnplugSelectors = NULL;
+
 ReleaseToUnplugApicIds:
   gMmst->MmFreePool (mToUnplugApicIds);
   mToUnplugApicIds = NULL;