[PATCH 13/16] OvmfPkg/CpuHotplugSmm: complete root MMI handler for CPU hotplug

["Laszlo Ersek" <lersek@redhat.com>]

With the help of the Post-SMM Pen and the SMBASE relocation functions
added in the previous patches, we can now complete the root MMI handler
for CPU hotplug.

In the driver's entry point function:

- allocate the pen (in a reserved page in normal RAM),

- install the default ("first") SMI handler for hot-added CPUs (which
  includes priming the exchange area between the MM Monarch and the
  hot-added CPUs, i.e., shutting the APIC ID gate).

In the root MMI handler, for each hot-added CPU:

- record the APIC ID of the new CPU in CPU_HOT_PLUG_DATA,

- relocate the SMBASE of the new CPU,

- inform PiSmmCpuDxeSmm by calling
  EFI_SMM_CPU_SERVICE_PROTOCOL.AddProcessor().

Cc: Ard Biesheuvel <ard.biesheuvel@linaro.org>
Cc: Igor Mammedov <imammedo@redhat.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Jordan Justen <jordan.l.justen@intel.com>
Cc: Michael Kinney <michael.d.kinney@intel.com>
Cc: Philippe Mathieu-Daudé <philmd@redhat.com>
Ref: https://bugzilla.tianocore.org/show_bug.cgi?id=1512
Signed-off-by: Laszlo Ersek <lersek@redhat.com>
---

Notes:
    (1) At this stage, CPU hotplug with SMM works.
    
    Test CPU topology (libvirt domain XML snippet):
    
    >   <vcpus>
    >     <vcpu id='0' enabled='yes' hotpluggable='no'/>
    >     <vcpu id='1' enabled='no' hotpluggable='yes'/>
    >     <vcpu id='2' enabled='yes' hotpluggable='yes'/>
    >     <vcpu id='3' enabled='no' hotpluggable='yes'/>
    >   </vcpus>
    
    The firmware logs the following during normal boot:
    
    > SmbaseAllocatePostSmmPen: Post-SMM Pen at 0x9F000
    
    CPU hotplug command on the host side:
    
    > virsh setvcpu ovmf.fedora.q35 1 --enable --live
    
    Firmware log in response:
    
    > QemuCpuhpCollectApicIds: CurrentSelector=1: insert
    > QemuCpuhpCollectApicIds: ApicId=0x00000001
    > QemuCpuhpCollectApicIds: CurrentSelector=2 PendingSelector=1:
    >                          wrap-around
    > QemuCpuhpCollectApicIds: PluggedCount=1 ToUnplugCount=0
    > CpuHotplugMmi: hot-added APIC ID 0x00000001, SMBASE 0x7D082000,
    >                EFI_SMM_CPU_SERVICE_PROTOCOL assigned number 2
    
    Query the CPU register state on the host side (without onlining the
    hot-added CPU in the OS just yet):
    
    > virsh qemu-monitor-command ovmf.fedora.q35 --hmp 'info registers -a'
    
    Output, confirming penned status (pen at 0x9F000):
    
    > CPU#1
    > EAX=00009f00 EBX=00000000 ECX=00000000 EDX=00000600
    > ESI=00000000 EDI=00000000 EBP=00000000 ESP=00000000
    > EIP=0000000c EFL=00000002 [-------] CPL=0 II=0 A20=1 SMM=0 HLT=1
    > ES =0000 00000000 0000ffff 00009300
    > CS =9f00 0009f000 0000ffff 00009b00
    > SS =0000 00000000 0000ffff 00009300
    > DS =9f00 0009f000 0000ffff 00009300
    > FS =0000 00000000 0000ffff 00009300
    > GS =0000 00000000 0000ffff 00009300
    > LDT=0000 00000000 0000ffff 00008200
    > TR =0000 00000000 0000ffff 00008b00
    > GDT=     00000000 0000ffff
    > IDT=     00000000 0000ffff
    > CR0=60000010 CR2=00000000 CR3=00000000 CR4=00000000
    > DR0=0000000000000000 DR1=0000000000000000 DR2=0000000000000000
    > DR3=0000000000000000
    > DR6=00000000ffff0ff0 DR7=0000000000000400
    > EFER=0000000000000000
    > FCW=037f FSW=0000 [ST=0] FTW=00 MXCSR=00001f80
    > FPR0=0000000000000000 0000 FPR1=0000000000000000 0000
    > FPR2=0000000000000000 0000 FPR3=0000000000000000 0000
    > FPR4=0000000000000000 0000 FPR5=0000000000000000 0000
    > FPR6=0000000000000000 0000 FPR7=0000000000000000 0000
    > XMM00=00000000000000000000000000000000
    > XMM01=00000000000000000000000000000000
    > XMM02=00000000000000000000000000000000
    > XMM03=00000000000000000000000000000000
    > XMM04=00000000000000000000000000000000
    > XMM05=00000000000000000000000000000000
    > XMM06=00000000000000000000000000000000
    > XMM07=00000000000000000000000000000000
    
    Additionally, the dmesg in the Linux guest contains:
    
    > CPU2 has been hot-added
    
    We can now boot the hot-addad CPU in the Linux guest:
    
    > echo 1 > /sys/devices/system/cpu/cpu2/online
    
    Dmesg in response:
    
    > smpboot: Booting Node 0 Processor 2 APIC 0x1
    > kvm-clock: cpu 2, msr 6e01081, secondary cpu clock
    > KVM setup async PF for cpu 2
    > kvm-stealtime: cpu 2, msr 7b52a040
    > Will online and init hotplugged CPU: 2
    
    Still in the guest, call UEFI variable services on each CPU (including
    the hot-added one), per
    <https://github.com/tianocore/tianocore.github.io/wiki/Testing-SMM-with-QEMU,-KVM-and-libvirt#uefi-variable-access-test>:
    
    > taskset -c 0 efibootmgr
    > taskset -c 1 efibootmgr
    > taskset -c 2 efibootmgr
    
    No delays, no instability.
    
    Having booted the hot-added CPU in the Linux guest, the host-side
    command
    
    > virsh qemu-monitor-command ovmf.fedora.q35 --hmp 'info registers -a'
    
    now shows:
    
    > CPU#1
    > RAX=ffffffffb39e4b50 RBX=0000000000000002 RCX=0000000000000000
    > RDX=0000000000000002
    > RSI=0000000000000002 RDI=ffff932cbb51c520 RBP=0000000000000002
    > RSP=ffffbc5bc0083eb0
    > R8 =000000cd42e4dffb R9 =0000005a22502599 R10=0000005a22502599
    > R11=0000005a22502599
    > R12=ffff932cbae25dc0 R13=0000000000000000 R14=0000000000000000
    > R15=ffff932cbae25dc0
    > RIP=ffffffffb39e4f2e RFL=00000246 [---Z-P-] CPL=0 II=0 A20=1 SMM=0
    > HLT=1
    > ES =0000 0000000000000000 ffffffff 00c00000
    > CS =0010 0000000000000000 ffffffff 00a09b00 DPL=0 CS64 [-RA]
    > SS =0018 0000000000000000 ffffffff 00c09300 DPL=0 DS   [-WA]
    > DS =0000 0000000000000000 ffffffff 00c00000
    > FS =0000 0000000000000000 ffffffff 00c00000
    > GS =0000 ffff932cbb500000 ffffffff 00c00000
    > LDT=0000 0000000000000000 000fffff 00000000
    > TR =0040 fffffe0000069000 0000206f 00008b00 DPL=0 TSS64-busy
    > GDT=     fffffe0000067000 0000007f
    > IDT=     fffffe0000000000 00000fff
    > CR0=80050033 CR2=00007f54893cbf20 CR3=0000000078418001 CR4=001606e0
    > DR0=0000000000000000 DR1=0000000000000000 DR2=0000000000000000
    > DR3=0000000000000000
    > DR6=00000000ffff0ff0 DR7=0000000000000400
    > EFER=0000000000000d01
    > FCW=037f FSW=0000 [ST=0] FTW=00 MXCSR=00001fa0
    > FPR0=0000000000000000 0000 FPR1=0000000000000000 0000
    > FPR2=0000000000000000 0000 FPR3=0000000000000000 0000
    > FPR4=0000000000000000 0000 FPR5=0000000000000000 0000
    > FPR6=0000000000000000 0000 FPR7=0000000000000000 0000
    > XMM00=000000000000000040ed2be000000000
    > XMM01=0000000000000000404ddf1a9fbe76c9
    > XMM02=00000000000000000000000000000000
    > XMM03=00000000000000000000000000000000
    > XMM04=00000000000000003ff0000000000000
    > XMM05=00000000000000000000000000000000
    > XMM06=00000000000000004078002d0ac487b7
    > XMM07=0000000000000000404ddf1a9fbe76c9
    > XMM08=72223d444955412022746f6f72223d44
    > XMM09=00000000000000000000000000000000
    > XMM10=00000000000000000000000000000000
    > XMM11=00000000000000000000000000000000
    > XMM12=00000000000000000000000000000000
    > XMM13=00000000000000000000000000000000
    > XMM14=00000000000000000000000000000000
    > XMM15=00000000000000000000000000000000
    
    Hotplug another CPU:
    
    > virsh setvcpu ovmf.fedora.q35 3 --enable --live
    
    Firmware log in response:
    
    > QemuCpuhpCollectApicIds: CurrentSelector=3: insert
    > QemuCpuhpCollectApicIds: ApicId=0x00000003
    > QemuCpuhpCollectApicIds: PluggedCount=1 ToUnplugCount=0
    > CpuHotplugMmi: hot-added APIC ID 0x00000003, SMBASE 0x7D084000,
    >                EFI_SMM_CPU_SERVICE_PROTOCOL assigned number 3
    
    Online the hot-added CPU in the Linux guest:
    
    > echo 1 > /sys/devices/system/cpu/cpu3/online
    
    Guest dmesg in response:
    
    > smpboot: Booting Node 0 Processor 3 APIC 0x3
    > kvm-clock: cpu 3, msr 6e010c1, secondary cpu clock
    > KVM setup async PF for cpu 3
    > kvm-stealtime: cpu 3, msr 7b5aa040
    > Will online and init hotplugged CPU: 3
    
    (2) Alternative order of actions:
    
    Hotplug both CPUs first, on the host:
    
    > virsh setvcpu ovmf.fedora.q35 1 --enable --live
    > virsh setvcpu ovmf.fedora.q35 3 --enable --live
    
    *then* online both hot-added CPUs in the Linux guest:
    
    > echo 1 > /sys/devices/system/cpu/cpu2/online
    > echo 1 > /sys/devices/system/cpu/cpu3/online
    
    The only difference here is that the broadcast SMI for the second
    hotplug finds the first hotplugged CPU still in the pen (i.e., not
    onlined by Linux). There is no difference in observable behavior (beyond
    the matching log messages and register dumps).
    
    (3) What doesn't work yet:
    
    - CPU hotplug w/ SMM in SEV guests.
    
      See <https://bugzilla.tianocore.org/show_bug.cgi?id=1512#c14>.
    
      (NB. I'm uncertain if CPU hotplug works in SEV guests without
      SMM_REQUIRE to begin with.)
    
    - S3 resume w/ SMM after hot-adding a CPU.
    
      The remaining patches in the series take care of that. (Only in the
      absence of SEV; SEV+S3 looks broken regardless of hotplug.)
    
    - CPU hot-unplug w/ SMM.
    
      Needs separate investigation. EFI_SMM_REMOVE_PROCESSOR in
      "UefiCpuPkg/Include/Protocol/SmmCpuService.h" is documented as
      follows: "After this API is called, the removed processor must not
      respond to SMIs in the coherence domain". This could raise challenging
      ordering questions between QEMU, the firmware, and the OS.

 OvmfPkg/CpuHotplugSmm/CpuHotplug.c | 97 +++++++++++++++++++-
 1 file changed, 95 insertions(+), 2 deletions(-)

diff --git a/OvmfPkg/CpuHotplugSmm/CpuHotplug.c b/OvmfPkg/CpuHotplugSmm/CpuHotplug.c
index 42e023cb85c0..20e6bec04f41 100644
--- a/OvmfPkg/CpuHotplugSmm/CpuHotplug.c
+++ b/OvmfPkg/CpuHotplugSmm/CpuHotplug.c
@@ -1,76 +1,82 @@
 /** @file
   Root SMI handler for VCPU hotplug SMIs.
 
   Copyright (c) 2020, Red Hat, Inc.
 
   SPDX-License-Identifier: BSD-2-Clause-Patent
 **/
 
 #include <CpuHotPlugData.h>                  // CPU_HOT_PLUG_DATA
 #include <IndustryStandard/Q35MchIch9.h>     // ICH9_APM_CNT
 #include <IndustryStandard/QemuCpuHotplug.h> // QEMU_CPUHP_CMD_GET_PENDING
 #include <Library/BaseLib.h>                 // CpuDeadLoop()
 #include <Library/DebugLib.h>                // ASSERT()
 #include <Library/MmServicesTableLib.h>      // gMmst
 #include <Library/PcdLib.h>                  // PcdGetBool()
 #include <Library/SafeIntLib.h>              // SafeUintnSub()
 #include <Protocol/MmCpuIo.h>                // EFI_MM_CPU_IO_PROTOCOL
 #include <Protocol/SmmCpuService.h>          // EFI_SMM_CPU_SERVICE_PROTOCOL
 #include <Uefi/UefiBaseType.h>               // EFI_STATUS
 
 #include "ApicId.h"                          // APIC_ID
 #include "QemuCpuhp.h"                       // QemuCpuhpWriteCpuSelector()
+#include "Smbase.h"                          // SmbaseAllocatePostSmmPen()
 
 //
 // We use this protocol for accessing IO Ports.
 //
 STATIC EFI_MM_CPU_IO_PROTOCOL *mMmCpuIo;
 //
 // The following protocol is used to report the addition or removal of a CPU to
 // the SMM CPU driver (PiSmmCpuDxeSmm).
 //
 STATIC EFI_SMM_CPU_SERVICE_PROTOCOL *mMmCpuService;
 //
 // This structure is a communication side-channel between the
 // EFI_SMM_CPU_SERVICE_PROTOCOL consumer (i.e., this driver) and provider
 // (i.e., PiSmmCpuDxeSmm).
 //
 STATIC CPU_HOT_PLUG_DATA *mCpuHotPlugData;
 //
 // SMRAM arrays for fetching the APIC IDs of processors with pending events (of
 // known event types), for the time of just one MMI.
 //
 // The lifetimes of these arrays match that of this driver only because we
 // 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
 // determined on every MMI separately.
 //
 STATIC APIC_ID *mPluggedApicIds;
 STATIC APIC_ID *mToUnplugApicIds;
 //
+// Address of the non-SMRAM reserved memory page that contains the Post-SMM Pen
+// for hot-added CPUs.
+//
+STATIC UINT32 mPostSmmPenAddress;
+//
 // Represents the registration of the CPU Hotplug MMI handler.
 //
 STATIC EFI_HANDLE mDispatchHandle;
 
 
 /**
   CPU Hotplug MMI handler function.
 
   This is a root MMI handler.
 
   @param[in] DispatchHandle      The unique handle assigned to this handler by
                                  EFI_MM_SYSTEM_TABLE.MmiHandlerRegister().
 
   @param[in] Context             Context passed in by
                                  EFI_MM_SYSTEM_TABLE.MmiManage(). Due to
                                  CpuHotplugMmi() being a root MMI handler,
                                  Context is ASSERT()ed to be NULL.
 
   @param[in,out] CommBuffer      Ignored, due to CpuHotplugMmi() being a root
                                  MMI handler.
 
   @param[in,out] CommBufferSize  Ignored, due to CpuHotplugMmi() being a root
@@ -97,44 +103,46 @@ STATIC EFI_HANDLE mDispatchHandle;
 
   @retval EFI_WARN_INTERRUPT_SOURCE_PENDING   The MMI source is still pending,
                                               and other handlers should still
                                               be called.
 
   @retval EFI_INTERRUPT_PENDING               The MMI source could not be
                                               quiesced.
 **/
 STATIC
 EFI_STATUS
 EFIAPI
 CpuHotplugMmi (
   IN EFI_HANDLE DispatchHandle,
   IN CONST VOID *Context        OPTIONAL,
   IN OUT VOID   *CommBuffer     OPTIONAL,
   IN OUT UINTN  *CommBufferSize OPTIONAL
   )
 {
   EFI_STATUS Status;
   UINT8      ApmControl;
   UINT32     PluggedCount;
   UINT32     ToUnplugCount;
+  UINT32     PluggedIdx;
+  UINT32     NewSlot;
 
   //
   // Assert that we are entering this function due to our root MMI handler
   // registration.
   //
   ASSERT (DispatchHandle == mDispatchHandle);
   //
   // When MmiManage() is invoked to process root MMI handlers, the caller (the
   // MM Core) is expected to pass in a NULL Context. MmiManage() then passes
   // the same NULL Context to individual handlers.
   //
   ASSERT (Context == NULL);
   //
   // Read the MMI command value from the APM Control Port, to see if this is an
   // MMI we should care about.
   //
   Status = mMmCpuIo->Io.Read (mMmCpuIo, MM_IO_UINT8, ICH9_APM_CNT, 1,
                           &ApmControl);
   if (EFI_ERROR (Status)) {
     DEBUG ((DEBUG_ERROR, "%a: failed to read ICH9_APM_CNT: %r\n", __FUNCTION__,
       Status));
     //
@@ -152,49 +160,116 @@ CpuHotplugMmi (
 
   //
   // Collect the CPUs with pending events.
   //
   Status = QemuCpuhpCollectApicIds (
              mMmCpuIo,
              mCpuHotPlugData->ArrayLength,     // PossibleCpuCount
              mCpuHotPlugData->ArrayLength - 1, // ApicIdCount
              mPluggedApicIds,
              &PluggedCount,
              mToUnplugApicIds,
              &ToUnplugCount
              );
   if (EFI_ERROR (Status)) {
     goto Fatal;
   }
   if (ToUnplugCount > 0) {
     DEBUG ((DEBUG_ERROR, "%a: hot-unplug is not supported yet\n",
       __FUNCTION__));
     goto Fatal;
   }
 
+  //
+  // Process hot-added CPUs.
+  //
+  // The Post-SMM Pen need not be reinstalled multiple times within a single
+  // root MMI handling. Even reinstalling once per root MMI is only prudence;
+  // in theory installing the pen in the driver's entry point function should
+  // suffice.
+  //
+  SmbaseReinstallPostSmmPen (mPostSmmPenAddress);
+
+  PluggedIdx = 0;
+  NewSlot = 0;
+  while (PluggedIdx < PluggedCount) {
+    APIC_ID NewApicId;
+    UINTN   NewProcessorNumberByProtocol;
+
+    NewApicId = mPluggedApicIds[PluggedIdx];
+    //
+    // Find the first empty slot in CPU_HOT_PLUG_DATA.
+    //
+    while (NewSlot < mCpuHotPlugData->ArrayLength &&
+           mCpuHotPlugData->ApicId[NewSlot] != MAX_UINT64) {
+      NewSlot++;
+    }
+    if (NewSlot == mCpuHotPlugData->ArrayLength) {
+      DEBUG ((DEBUG_ERROR, "%a: no room for APIC ID " FMT_APIC_ID "\n",
+        __FUNCTION__, NewApicId));
+      goto Fatal;
+    }
+
+    //
+    // Store the APIC ID of the new processor to the slot.
+    //
+    mCpuHotPlugData->ApicId[NewSlot] = NewApicId;
+
+    //
+    // Relocate the SMBASE of the new CPU.
+    //
+    Status = SmbaseRelocate (NewApicId, mCpuHotPlugData->SmBase[NewSlot],
+               mPostSmmPenAddress);
+    if (EFI_ERROR (Status)) {
+      goto RevokeNewSlot;
+    }
+
+    //
+    // Add the new CPU with EFI_SMM_CPU_SERVICE_PROTOCOL.
+    //
+    Status = mMmCpuService->AddProcessor (mMmCpuService, NewApicId,
+                              &NewProcessorNumberByProtocol);
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "%a: AddProcessor(" FMT_APIC_ID "): %r\n",
+        __FUNCTION__, NewApicId, Status));
+      goto RevokeNewSlot;
+    }
+
+    DEBUG ((DEBUG_INFO, "%a: hot-added APIC ID " FMT_APIC_ID ", SMBASE 0x%Lx, "
+      "EFI_SMM_CPU_SERVICE_PROTOCOL assigned number %Lu\n", __FUNCTION__,
+      NewApicId, (UINT64)mCpuHotPlugData->SmBase[NewSlot],
+      (UINT64)NewProcessorNumberByProtocol));
+
+    NewSlot++;
+    PluggedIdx++;
+  }
+
   //
   // We've handled this MMI.
   //
   return EFI_SUCCESS;
 
+RevokeNewSlot:
+  mCpuHotPlugData->ApicId[NewSlot] = MAX_UINT64;
+
 Fatal:
   ASSERT (FALSE);
   CpuDeadLoop ();
   //
   // We couldn't handle this MMI.
   //
   return EFI_INTERRUPT_PENDING;
 }
 
 
 //
 // Entry point function of this driver.
 //
 EFI_STATUS
 EFIAPI
 CpuHotplugEntry (
   IN EFI_HANDLE       ImageHandle,
   IN EFI_SYSTEM_TABLE *SystemTable
   )
 {
   EFI_STATUS Status;
   UINTN      Size;
@@ -251,83 +326,101 @@ 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))) {
     Status = EFI_ABORTED;
     DEBUG ((DEBUG_ERROR, "%a: invalid CPU_HOT_PLUG_DATA\n", __FUNCTION__));
     goto Fatal;
   }
   Status = gMmst->MmAllocatePool (EfiRuntimeServicesData, Size,
                     (VOID **)&mPluggedApicIds);
   if (EFI_ERROR (Status)) {
     DEBUG ((DEBUG_ERROR, "%a: MmAllocatePool(): %r\n", __FUNCTION__, Status));
     goto Fatal;
   }
   Status = gMmst->MmAllocatePool (EfiRuntimeServicesData, Size,
                     (VOID **)&mToUnplugApicIds);
   if (EFI_ERROR (Status)) {
     DEBUG ((DEBUG_ERROR, "%a: MmAllocatePool(): %r\n", __FUNCTION__, Status));
     goto ReleasePluggedApicIds;
   }
 
+  //
+  // Allocate the Post-SMM Pen for hot-added CPUs.
+  //
+  Status = SmbaseAllocatePostSmmPen (&mPostSmmPenAddress,
+             SystemTable->BootServices);
+  if (EFI_ERROR (Status)) {
+    goto ReleaseToUnplugApicIds;
+  }
+
   //
   // Sanity-check the CPU hotplug interface.
   //
   // Both of the following features are part of QEMU 5.0, introduced primarily
   // in commit range 3e08b2b9cb64..3a61c8db9d25:
   //
   // (a) the QEMU_CPUHP_CMD_GET_ARCH_ID command of the modern CPU hotplug
   //     interface,
   //
   // (b) the "SMRAM at default SMBASE" feature.
   //
   // From these, (b) is restricted to 5.0+ machine type versions, while (a)
   // does not depend on machine type version. Because we ensured the stricter
   // condition (b) through PcdQ35SmramAtDefaultSmbase above, the (a)
   // QEMU_CPUHP_CMD_GET_ARCH_ID command must now be available too. While we
   // can't verify the presence of precisely that command, we can still verify
   // (sanity-check) that the modern interface is active, at least.
   //
   // Consult the "Typical usecases | Detecting and enabling modern CPU hotplug
   // interface" section in QEMU's "docs/specs/acpi_cpu_hotplug.txt", on the
   // following.
   //
   QemuCpuhpWriteCpuSelector (mMmCpuIo, 0);
   QemuCpuhpWriteCpuSelector (mMmCpuIo, 0);
   QemuCpuhpWriteCommand (mMmCpuIo, QEMU_CPUHP_CMD_GET_PENDING);
   if (QemuCpuhpReadCommandData2 (mMmCpuIo) != 0) {
     Status = EFI_NOT_FOUND;
     DEBUG ((DEBUG_ERROR, "%a: modern CPU hotplug interface: %r\n",
       __FUNCTION__, Status));
-    goto ReleaseToUnplugApicIds;
+    goto ReleasePostSmmPen;
   }
 
   //
   // Register the handler for the CPU Hotplug MMI.
   //
   Status = gMmst->MmiHandlerRegister (
                     CpuHotplugMmi,
                     NULL,            // HandlerType: root MMI handler
                     &mDispatchHandle
                     );
   if (EFI_ERROR (Status)) {
     DEBUG ((DEBUG_ERROR, "%a: MmiHandlerRegister(): %r\n", __FUNCTION__,
       Status));
-    goto ReleaseToUnplugApicIds;
+    goto ReleasePostSmmPen;
   }
 
+  //
+  // Install the handler for the hot-added CPUs' first SMI.
+  //
+  SmbaseInstallFirstSmiHandler ();
+
   return EFI_SUCCESS;
 
+ReleasePostSmmPen:
+  SmbaseReleasePostSmmPen (mPostSmmPenAddress, SystemTable->BootServices);
+  mPostSmmPenAddress = 0;
+
 ReleaseToUnplugApicIds:
   gMmst->MmFreePool (mToUnplugApicIds);
   mToUnplugApicIds = NULL;
 
 ReleasePluggedApicIds:
   gMmst->MmFreePool (mPluggedApicIds);
   mPluggedApicIds = NULL;
 
 Fatal:
   ASSERT (FALSE);
   CpuDeadLoop ();
   return Status;
 }
-- 
2.19.1.3.g30247aa5d201