UefiCpuPkg/MpInitLib: Fix SNP AP creation when using known APIC IDs

A typical initial AP boot up will choose a CpuNumber based on the ApIndex
value that it gets back after a locked increment of the ApIndex value.
The ApIndex to APIC ID relationship is random, which is not an issue when
a broadcast INIT-SIPI is performed.

With SNP and a hypervisor that supports retrieval of the known APIC IDs,
the broadcast INIT-SIPI method is replaced by waking each individual vCPU.
In this situation, a specific VMSA is associated with a specific APIC ID.
However, random assignment of an ApIndex can break this association. This
isn't typically an issue, because the AP bring-up finishes with the AP
issuing a HLT instruction, which is intercepted by the hypervisor and the
AP won't run again until the next INIT-SIPI. However, when HLT isn't
intercepted by the hypervisor (Qemu '-overcommit cpu-pm=on' parameter),
then the HLT does not exit to the hypervisor. On the next INIT-SIPI, it
can happen that a VMRUN is executed with a different VMSA address than
was originally used, and if that VMSA is still in a VMRUN on another AP,
then the executing VMRUN will fail, crashing the guest.

To fix this issue, add a CPU exchange info field, SevSnpKnownInitApicId,
that indicates the APs are starting with an already known initial APIC ID
and set the initial APIC ID and APIC ID in the CPU_INFO_IN_HOB HOB.
During AP boot, the SevSnpKnownInitApicId field will result in the
CpuNumber being set to the index with a matching APIC ID (similar to AP
booting when the InitFlag != ApInitConfig).

Signed-off-by: Tom Lendacky <thomas.lendacky@amd.com>

UefiCpuPkg/Library/MpInitLib/AmdSev.c

@@ -293,6 +293,8 @@ FillExchangeInfoDataSevEs (
       );
     ExchangeInfo->ExtTopoAvail = !!ExtTopoEbx.Bits.LogicalProcessors;
   }
+
+  ExchangeInfo->SevSnpKnownInitApicId = FALSE;
 }
 
 /**

UefiCpuPkg/Library/MpInitLib/MpEqu.inc

@@ -97,6 +97,7 @@ struc MP_CPU_EXCHANGE_INFO
   .SevSnpIsEnabled               CTYPE_BOOLEAN 1
   .GhcbBase:                     CTYPE_UINTN 1
   .ExtTopoAvail:                 CTYPE_BOOLEAN 1
+  .SevSnpKnownInitApicId:        CTYPE_BOOLEAN 1
 endstruc
 
 MP_CPU_EXCHANGE_INFO_OFFSET equ (Flat32Start - RendezvousFunnelProcStart)

UefiCpuPkg/Library/MpInitLib/MpLib.h

@@ -239,6 +239,7 @@ typedef struct {
   BOOLEAN            SevSnpIsEnabled;
   UINTN              GhcbBase;
   BOOLEAN            ExtTopoAvail;
+  BOOLEAN            SevSnpKnownInitApicId;
 } MP_CPU_EXCHANGE_INFO;
 
 #pragma pack()

UefiCpuPkg/Library/MpInitLib/X64/AmdSev.c

@@ -271,18 +271,27 @@ SevSnpCreateAP (
   IN INTN         ProcessorNumber
   )
 {
-  CPU_INFO_IN_HOB    *CpuInfoInHob;
+  CPU_INFO_IN_HOB                *CpuInfoInHob;
-  CPU_AP_DATA        *CpuData;
+  CPU_AP_DATA                    *CpuData;
-  UINTN              Index;
+  UINTN                          Index;
-  UINTN              MaxIndex;
+  UINTN                          MaxIndex;
-  UINT32             ApicId;
+  UINT32                         ApicId;
-  EFI_HOB_GUID_TYPE  *GuidHob;
+  EFI_HOB_GUID_TYPE              *GuidHob;
-  GHCB_APIC_IDS      *GhcbApicIds;
+  GHCB_APIC_IDS                  *GhcbApicIds;
+  volatile MP_CPU_EXCHANGE_INFO  *ExchangeInfo;
 
   ASSERT (CpuMpData->MpCpuExchangeInfo->BufferStart < 0x100000);
 
+  ExchangeInfo = CpuMpData->MpCpuExchangeInfo;
   CpuInfoInHob = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;
 
+  //
+  // Set to FALSE by default. This is only set to TRUE when the InitFlag
+  // is equal to ApInitConfig and the GHCB APIC ID List NAE event has
+  // been called.
+  //
+  ExchangeInfo->SevSnpKnownInitApicId = FALSE;
+
   if (ProcessorNumber < 0) {
     if (CpuMpData->InitFlag == ApInitConfig) {
       //
@@ -294,6 +303,14 @@ SevSnpCreateAP (
       GuidHob     = GetFirstGuidHob (&gGhcbApicIdsGuid);
       GhcbApicIds = (GHCB_APIC_IDS *)(*(UINTN *)GET_GUID_HOB_DATA (GuidHob));
       MaxIndex    = MIN (GhcbApicIds->NumEntries, PcdGet32 (PcdCpuMaxLogicalProcessorNumber));
+
+      //
+      // Set to TRUE so that ApicId and SEV-SNP SaveArea stay in sync. When
+      // the InitFlag is ApInitConfig, the random order of AP initialization
+      // can end up with an Index / ApicId mismatch (see X64/AmdSev.nasm).
+      //
+      ExchangeInfo->SevSnpKnownInitApicId = TRUE;
+      DEBUG ((DEBUG_INFO, "SEV-SNP: Using known initial APIC IDs\n"));
     } else {
       //
       // APs have been previously started.
@@ -308,6 +325,13 @@ SevSnpCreateAP (
         if (CpuMpData->InitFlag == ApInitConfig) {
           ApicId = GhcbApicIds->ApicIds[Index];
 
+          //
+          // Set the ApicId values so that the proper AP data structure
+          // can be found during boot.
+          //
+          CpuInfoInHob[Index].InitialApicId = ApicId;
+          CpuInfoInHob[Index].ApicId        = ApicId;
+
           //
           // For the first boot, use the BSP register information.
           //

UefiCpuPkg/Library/MpInitLib/X64/AmdSev.nasm

@@ -15,6 +15,57 @@
 
 %define SIZE_4KB    0x1000
 
+;
+; This function will ensure that CpuNumber and ApicId are in sync when using
+; the ApicIds retrieved via the GHCB APIC ID List NAE event to start the APs
+; when the InitFlag is ApInitConfig. If this is not done, the CpuNumber to
+; ApicId relationship may not hold, which would result in the ApicId to VSMA
+; relationship getting out of sync after the first AP boot.
+;
+SevSnpGetInitCpuNumber:
+    ;
+    ; If not an SNP guest, leave EBX (CpuNumber) as is
+    ;
+    lea        edi, [esi + MP_CPU_EXCHANGE_INFO_FIELD (SevSnpIsEnabled)]
+    cmp        byte [edi], 1        ; SevSnpIsEnabled
+    jne        SevSnpGetCpuNumberDone
+
+    ;
+    ; If not starting the AP with a specific ApicId, leave EBX (CpuNumber) as is
+    ;
+    lea        edi, [esi + MP_CPU_EXCHANGE_INFO_FIELD (SevSnpKnownInitApicId)]
+    cmp        byte [edi], 1        ; SevSnpKnownInitApicId
+    jne        SevSnpGetCpuNumberDone
+
+    ;
+    ; Use existing code to retrieve the ApicId. SevEsGetApicId will return to
+    ; the SevSnpGetInitApicId label if SevSnpKnownInitApicId is set.
+    ;
+    jmp        SevEsGetApicId
+
+SevSnpGetInitApicId:
+    ;
+    ; EDX holds the ApicId, get processor number for this AP
+    ;
+    xor        ebx, ebx
+    lea        eax, [esi + MP_CPU_EXCHANGE_INFO_FIELD (CpuInfo)]
+    mov        rdi, [eax]
+
+SevSnpGetNextProcNumber:
+    cmp        dword [rdi + CPU_INFO_IN_HOB.InitialApicId], edx         ; APIC ID match?
+    jz         SevSnpGetCpuNumberDone
+    add        rdi, CPU_INFO_IN_HOB_size
+    inc        ebx
+    jmp        SevSnpGetNextProcNumber
+
+SevSnpGetCpuNumberDone:
+    ;
+    ; If SevSnpKnownInitApicId is set, EBX now holds the CpuNumber for this
+    ; ApicId, which matches how it was started in SevSnpCreateAP(). Otherwise,
+    ; EBX is unchanged and holds the CpuNumber based on the startup order.
+    ;
+    OneTimeCallRet    SevSnpGetInitCpuNumber
+
 RegisterGhcbGpa:
     ;
     ; Register GHCB GPA when SEV-SNP is enabled
@@ -193,7 +244,14 @@ RestoreGhcb:
 
     mov        rdx, rbx
 
-    ; x2APIC ID or APIC ID is in EDX
+    ;
+    ; x2APIC ID or APIC ID is in EDX. If SevSnpKnownInitApicId is set, then
+    ; return to SevSnpGetInitApicId, otherwise return to GetProcessorNumber.
+    ;
+    lea        edi, [esi + MP_CPU_EXCHANGE_INFO_FIELD (SevSnpKnownInitApicId)]
+    cmp        byte [edi], 1        ; SevSnpKnownInitApicId
+    je         SevSnpGetInitApicId
+
     jmp        GetProcessorNumber
 
 SevEsGetApicIdExit:

UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm

@@ -173,6 +173,10 @@ LongModeStart:
     lock xadd  dword [edi], ebx                 ; EBX = ApIndex++
     inc        ebx                              ; EBX is CpuNumber
 
+    ; If running under AMD SEV-SNP and starting with a known ApicId,
+    ; adjust EBX to be the actual CpuNumber
+    OneTimeCall SevSnpGetInitCpuNumber
+
     ; program stack
     mov        edi, esi
     add        edi, MP_CPU_EXCHANGE_INFO_FIELD (StackSize)