[RFC PATCH 28/28] UefiCpuPkg/MpInitLib: Introduce an MP finalization routine to support SEV-ES

["Lendacky, Thomas" <thomas.lendacky@amd.com>]

From: Tom Lendacky <thomas.lendacky@amd.com>

Introduce a finalization routine to the MP library.  This routine is
used at the end of UEFI before transferring control to the OS and allows
for SEV-ES related AP state and information to be communicated to the OS.
The APs will be parked using VMGEXIT AP Reset Hold and the GHCB will be
modified to communicate a reserved page of memory that will be used by the
OS to direct the "initial" AP boot in the OS.

Signed-off-by: Tom Lendacky <thomas.lendacky@amd.com>
---
 MdePkg/Include/Protocol/Cpu.h                 |   6 +
 UefiCpuPkg/CpuDxe/CpuDxe.h                    |  14 ++
 UefiCpuPkg/Include/Library/MpInitLib.h        |  17 +++
 UefiCpuPkg/Library/MpInitLib/MpLib.h          |  18 ++-
 MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c       |   5 +
 OvmfPkg/PlatformPei/AmdSev.c                  |   2 +-
 UefiCpuPkg/CpuDxe/CpuDxe.c                    |  10 ++
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c       |  57 +++++++-
 UefiCpuPkg/Library/MpInitLib/MpLib.c          |  25 ++++
 UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm | 130 ++++++++++++++++--
 10 files changed, 265 insertions(+), 19 deletions(-)

diff --git a/MdePkg/Include/Protocol/Cpu.h b/MdePkg/Include/Protocol/Cpu.h
index e392f4cd9a13..79a701bc0d93 100644
--- a/MdePkg/Include/Protocol/Cpu.h
+++ b/MdePkg/Include/Protocol/Cpu.h
@@ -257,6 +257,11 @@ EFI_STATUS
   IN  UINT64                            Attributes
   );
 
+typedef
+EFI_STATUS
+(EFIAPI *EFI_CPU_FINALIZE)(
+  IN EFI_CPU_ARCH_PROTOCOL              *This
+  );
 
 ///
 /// The EFI_CPU_ARCH_PROTOCOL is used to abstract processor-specific functions from the DXE
@@ -273,6 +278,7 @@ struct _EFI_CPU_ARCH_PROTOCOL {
   EFI_CPU_REGISTER_INTERRUPT_HANDLER  RegisterInterruptHandler;
   EFI_CPU_GET_TIMER_VALUE             GetTimerValue;
   EFI_CPU_SET_MEMORY_ATTRIBUTES       SetMemoryAttributes;
+  EFI_CPU_FINALIZE                    Finalize;
   ///
   /// The number of timers that are available in a processor. The value in this
   /// field is a constant that must not be modified after the CPU Architectural
diff --git a/UefiCpuPkg/CpuDxe/CpuDxe.h b/UefiCpuPkg/CpuDxe/CpuDxe.h
index b029be430b4c..c5b9ada72ac9 100644
--- a/UefiCpuPkg/CpuDxe/CpuDxe.h
+++ b/UefiCpuPkg/CpuDxe/CpuDxe.h
@@ -232,6 +232,20 @@ CpuSetMemoryAttributes (
   IN UINT64                     Attributes
   );
 
+/**
+  Set ...
+
+  @param  This                   Protocol instance structure
+
+  @retval EFI_SUCCESS            If ...
+
+**/
+EFI_STATUS
+EFIAPI
+CpuFinalize (
+  IN EFI_CPU_ARCH_PROTOCOL      *This
+  );
+
 /**
   Initialize Global Descriptor Table.
 
diff --git a/UefiCpuPkg/Include/Library/MpInitLib.h b/UefiCpuPkg/Include/Library/MpInitLib.h
index fa8252937313..2095fb758664 100644
--- a/UefiCpuPkg/Include/Library/MpInitLib.h
+++ b/UefiCpuPkg/Include/Library/MpInitLib.h
@@ -344,4 +344,21 @@ MpInitLibWhoAmI (
   OUT UINTN                    *ProcessorNumber
   );
 
+/**
+  MP Exit ...
+
+  This service ...
+
+  This service must be invoked before ...
+
+  @retval  EFI_SUCCESS           MP initialization succeeds.
+  @retval  Others                MP initialization fails.
+
+**/
+EFI_STATUS
+EFIAPI
+MpLibFinalize (
+  VOID
+  );
+
 #endif
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index f2ba1a508715..a2ba6de0278f 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -273,7 +273,8 @@ struct _CPU_MP_DATA {
   UINT64                         GhcbBase;
 };
 
-#define AP_RESET_STACK_SIZE 64
+#define AP_SAFE_STACK_SIZE  128
+#define AP_RESET_STACK_SIZE AP_SAFE_STACK_SIZE
 
 typedef union {
   struct {
@@ -327,8 +328,11 @@ VOID
   IN BOOLEAN                 MwaitSupport,
   IN UINTN                   ApTargetCState,
   IN UINTN                   PmCodeSegment,
+  IN UINTN                   Pm16CodeSegment,
   IN UINTN                   TopOfApStack,
-  IN UINTN                   NumberToFinish
+  IN UINTN                   NumberToFinish,
+  IN UINTN                   SevEsAPJumpTable,
+  IN UINTN                   WakeupBuffer
   );
 
 /**
@@ -645,5 +649,15 @@ EnableDebugAgent (
   VOID
   );
 
+/**
+  MP finalize ...
+
+  @param[in] CpuMpData  The pointer to CPU MP Data structure will be saved.
+**/
+EFI_STATUS
+MpFinalize (
+  IN CPU_MP_DATA   *CpuMpData
+  );
+
 #endif
 
diff --git a/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c b/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
index 514d1aa75ada..13c962247243 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
+++ b/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
@@ -785,6 +785,11 @@ CoreExitBootServices (
   //
   gCpu->DisableInterrupt (gCpu);
 
+  //
+  // Finalize CPU
+  //
+  gCpu->Finalize (gCpu);
+
   //
   // Clear the non-runtime values of the EFI System Table
   //
diff --git a/OvmfPkg/PlatformPei/AmdSev.c b/OvmfPkg/PlatformPei/AmdSev.c
index fc396a6f229d..f0a18f026460 100644
--- a/OvmfPkg/PlatformPei/AmdSev.c
+++ b/OvmfPkg/PlatformPei/AmdSev.c
@@ -65,7 +65,7 @@ AmdSevEsInitialize (
   BuildMemoryAllocationHob (
     GhcbBasePa,
     EFI_PAGES_TO_SIZE (GhcbPageCount),
-    EfiBootServicesData
+    EfiReservedMemoryType
     );
 
   SetMem (GhcbBase, GhcbPageCount * SIZE_4KB, 0);
diff --git a/UefiCpuPkg/CpuDxe/CpuDxe.c b/UefiCpuPkg/CpuDxe/CpuDxe.c
index 7d7270e10b4a..7003f74e7d87 100644
--- a/UefiCpuPkg/CpuDxe/CpuDxe.c
+++ b/UefiCpuPkg/CpuDxe/CpuDxe.c
@@ -92,6 +92,7 @@ EFI_CPU_ARCH_PROTOCOL  gCpu = {
   CpuRegisterInterruptHandler,
   CpuGetTimerValue,
   CpuSetMemoryAttributes,
+  CpuFinalize,
   1,                          // NumberOfTimers
   4                           // DmaBufferAlignment
 };
@@ -499,6 +500,15 @@ CpuSetMemoryAttributes (
   return AssignMemoryPageAttributes (NULL, BaseAddress, Length, MemoryAttributes, NULL);
 }
 
+EFI_STATUS
+EFIAPI
+CpuFinalize (
+  IN EFI_CPU_ARCH_PROTOCOL     *This
+  )
+{
+  return MpLibFinalize ();
+}
+
 /**
   Initializes the valid bits mask and valid address mask for MTRRs.
 
diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index 127f64eb87e1..6e1bdbeed259 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -18,7 +18,6 @@
 #include <Protocol/Timer.h>
 
 #define  AP_CHECK_INTERVAL     (EFI_TIMER_PERIOD_MILLISECONDS (100))
-#define  AP_SAFE_STACK_SIZE    128
 
 CPU_MP_DATA      *mCpuMpData = NULL;
 EFI_EVENT        mCheckAllApsEvent = NULL;
@@ -98,7 +97,7 @@ GetWakeupBuffer (
   StartAddress = 0x88000;
   Status = gBS->AllocatePages (
                   AllocateMaxAddress,
-                  EfiBootServicesData,
+                  EfiReservedMemoryType,
                   EFI_SIZE_TO_PAGES (WakeupBufferSize),
                   &StartAddress
                   );
@@ -328,17 +327,26 @@ RelocateApLoop (
   BOOLEAN                MwaitSupport;
   ASM_RELOCATE_AP_LOOP   AsmRelocateApLoopFunc;
   UINTN                  ProcessorNumber;
+  UINTN                  StackStart;
 
   MpInitLibWhoAmI (&ProcessorNumber);
   CpuMpData    = GetCpuMpData ();
   MwaitSupport = IsMwaitSupport ();
+  if (CpuMpData->SevEsActive) {
+    StackStart = CpuMpData->SevEsAPResetStackStart;
+  } else {
+    StackStart = mReservedTopOfApStack;
+  }
   AsmRelocateApLoopFunc = (ASM_RELOCATE_AP_LOOP) (UINTN) mReservedApLoopFunc;
   AsmRelocateApLoopFunc (
     MwaitSupport,
     CpuMpData->ApTargetCState,
     CpuMpData->PmCodeSegment,
-    mReservedTopOfApStack - ProcessorNumber * AP_SAFE_STACK_SIZE,
-    (UINTN) &mNumberToFinish
+    CpuMpData->Pm16CodeSegment,
+    StackStart - ProcessorNumber * AP_SAFE_STACK_SIZE,
+    (UINTN) &mNumberToFinish,
+    CpuMpData->SevEsAPBuffer,
+    CpuMpData->WakeupBuffer
     );
   //
   // It should never reach here
@@ -880,3 +888,44 @@ MpInitLibEnableDisableAP (
 
   return Status;
 }
+
+/**
+  MP finalize ...
+
+  @param[in] CpuMpData  The pointer to CPU MP Data structure will be saved.
+**/
+EFI_STATUS
+MpFinalize (
+  IN CPU_MP_DATA   *CpuMpData
+  )
+{
+  if (CpuMpData->SevEsActive) {
+    //
+    // Perform SEV-ES specific finalization
+    //
+    if (CpuMpData->WakeupBuffer == (UINTN) -1) {
+      //
+      // No APs parked in UEFI, clear the GHCB
+      //
+      AsmWriteMsr64 (MSR_SEV_ES_GHCB, 0);
+    } else {
+      //
+      // Re-use reserved memory area below 1MB from WakeupBuffer
+      //
+      CopyMem (
+        (VOID *) CpuMpData->WakeupBuffer,
+        (VOID *) CpuMpData->AddressMap.RendezvousFunnelAddress +
+                   CpuMpData->AddressMap.SwitchToRealPM16ModeOffset,
+        CpuMpData->AddressMap.SwitchToRealPM16ModeSize
+        );
+
+      //
+      // Point the GHCB at the AP jump table to communicate the address to
+      // the booting system.
+      //
+      AsmWriteMsr64 (MSR_SEV_ES_GHCB, (CpuMpData->SevEsAPBuffer) | 0x03);
+    }
+  }
+
+  return EFI_SUCCESS;
+}
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index 0939019d7b8c..bc800a69527e 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -1707,6 +1707,31 @@ CheckAllAPs (
   return EFI_NOT_READY;
 }
 
+/**
+  MP Exit ...
+
+  This service ...
+
+  This service must be invoked before ...
+
+  @retval  EFI_SUCCESS           MP initialization succeeds.
+  @retval  Others                MP initialization fails.
+
+**/
+EFI_STATUS
+EFIAPI
+MpLibFinalize (
+  VOID
+  )
+{
+  CPU_MP_DATA              *CpuMpData;
+
+  CpuMpData = GetCpuMpData ();
+  MpFinalize (CpuMpData);
+
+  return EFI_SUCCESS;
+}
+
 /**
   MP Initialize Library initialization.
 
diff --git a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
index 286fa297791c..8936963913c4 100644
--- a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
+++ b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
@@ -446,7 +446,7 @@ CompatMode:
 
 BITS 16
     ;
-    ; At entry to this label
+    ; At entry to this label (used also by AsmRelocateApLoop):
     ;   - RDX will have its reset value
     ;   - On the top of the stack
     ;     - Alignment data (two bytes) to be discarded
@@ -475,32 +475,93 @@ PM16Mode:
 SwitchToRealProcEnd:
 
 ;-------------------------------------------------------------------------------------
-;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, PmCodeSegment, TopOfApStack, CountTofinish);
+;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, PmCodeSegment, Pm16CodeSegment, TopOfApStack, CountTofinish, SevEsAPJumpTable, WakeupBuffer);
 ;-------------------------------------------------------------------------------------
 global ASM_PFX(AsmRelocateApLoop)
 ASM_PFX(AsmRelocateApLoop):
 AsmRelocateApLoopStart:
 BITS 64
+    cmp        qword [rsp + 56], 0
+    je         NoSevEs
+
+    ;
+    ; Perform some SEV-ES related setup before leaving 64-bit mode
+    ;
+    push       rcx
+    push       rdx
+
+    ;
+    ; Get the RDX reset value using CPUID
+    ;
+    mov        rax, 1
+    cpuid
+    mov        rsi, rax          ; Save off the reset value for RDX
+
+    ;
+    ; Prepare the GHCB for the AP_HLT_LOOP VMGEXIT call
+    ;   - No NAE events can be generated once this is set otherwise
+    ;     the AP_HLT_LOOP SW_EXITCODE will be overwritten.
+    ;
+    mov        rcx, 0xc0010130
+    rdmsr                        ; Retrieve current GHCB address
+    shl        rdx, 32
+    or         rdx, rax
+
+    mov        rdi, rdx
+    xor        rax, rax
+    mov        rcx, 0x800
+    shr        rcx, 3
+    rep stosq                    ; Clear the GHCB
+
+    mov        rax, 0x80000004   ; VMGEXIT AP_HLT_LOOP
+    mov        [rdx + 0x390], rax
+
+    pop        rdx
+    pop        rcx
+
+NoSevEs:
     cli                          ; Disable interrupt before switching to 32-bit mode
-    mov        rax, [rsp + 40]   ; CountTofinish
+    mov        rax, [rsp + 48]   ; CountTofinish
     lock dec   dword [rax]       ; (*CountTofinish)--
-    mov        rsp, r9
-    push       rcx
-    push       rdx
 
-    lea        rsi, [PmEntry]    ; rsi <- The start address of transition code
+    mov        rax, [rsp + 56]   ; SevEsAPJumpTable
+    mov        rbx, [rsp + 64]   ; WakeupBuffer
+    mov        rsp, [rsp + 40]   ; TopOfApStack
+
+    push       rax               ; Save SevEsAPJumpTable
+    push       rbx               ; Save WakeupBuffer
+    push       r9                ; Save Pm16CodeSegment
+    push       rcx               ; Save MwaitSupport
+    push       rdx               ; Save ApTargetCState
+
+    lea        rax, [PmEntry]    ; rax <- The start address of transition code
 
     push       r8
-    push       rsi
-    DB         0x48
-    retf
+    push       rax
+
+    ;
+    ; Clear R8 - R15, for reset, before going into 32-bit mode
+    ;
+    xor        r8, r8
+    xor        r9, r9
+    xor        r10, r10
+    xor        r11, r11
+    xor        r12, r12
+    xor        r13, r13
+    xor        r14, r14
+    xor        r15, r15
+
+    ;
+    ; Far return into 32-bit mode
+    ;
+o64 retf
+
 BITS 32
 PmEntry:
     mov        eax, cr0
     btr        eax, 31           ; Clear CR0.PG
     mov        cr0, eax          ; Disable paging and caches
 
-    mov        ebx, edx          ; Save EntryPoint to rbx, for rdmsr will overwrite rdx
     mov        ecx, 0xc0000080
     rdmsr
     and        ah, ~ 1           ; Clear LME
@@ -513,6 +574,8 @@ PmEntry:
     add        esp, 4
     pop        ecx,
     add        esp, 4
+
+MwaitCheck:
     cmp        cl, 1              ; Check mwait-monitor support
     jnz        HltLoop
     mov        ebx, edx           ; Save C-State to ebx
@@ -526,10 +589,53 @@ MwaitLoop:
     shl        eax, 4
     mwait
     jmp        MwaitLoop
+
 HltLoop:
+    pop        edx                ; PM16CodeSegment
+    add        esp, 4
+    pop        ebx                ; WakeupBuffer
+    add        esp, 4
+    pop        eax                ; SevEsAPJumpTable
+    add        esp, 4
+    cmp        eax, 0             ; Check for SEV-ES
+    je         DoHlt
+
+    cli
+    ;
+    ; SEV-ES is active, use VMGEXIT (GHCB information already
+    ; set by caller)
+    ;
+    ; VMGEXIT is rep vmmcall
+    ;
+    db         0xf3
+    db         0x0f
+    db         0x01
+    db         0xd9
+
+    ;
+    ; Back from VMGEXIT AP_HLT_LOOP
+    ;   Push the FLAGS/CS/IP values to use
+    ;
+    push       word 0x0002        ; EFLAGS
+    xor        ecx, ecx
+    mov        cx, [eax + 2]      ; CS
+    push       cx
+    mov        cx, [eax]          ; IP
+    push       cx
+    push       word 0x0000        ; For alignment, will be discarded
+
+    push       edx
+    push       ebx
+
+    mov        edx, esi           ; Restore RDX reset value
+
+    retf
+
+DoHlt:
     cli
     hlt
-    jmp        HltLoop
+    jmp        DoHlt
+
 BITS 64
 AsmRelocateApLoopEnd:
 
-- 
2.17.1