[RFC PATCH v2 42/44] UefiCpuPkg/MpInitLib: Prepare SEV-ES guest APs for OS use

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

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

BZ: https://bugzilla.tianocore.org/show_bug.cgi?id=2198

Before UEFI transfers control to the OS, it must park the AP. This is
done using the AsmRelocateApLoop function to transition into 32-bit
non-paging mode. For an SEV-ES guest, a few additional things must be
done:
  - AsmRelocateApLoop must be updated to support SEV-ES. This means
    performing a VMGEXIT AP Reset Hold instead of an MWAIT or HLT loop.
  - Since the AP must transition to real mode, a small routine is copied
    to the WakeupBuffer area. Since the WakeupBuffer will be used by
    the AP during OS booting, it must be placed in reserved memory.
    Additionally, the AP stack must be located where it can be accessed
    in real mode.
  - Once the AP is in real mode it will transfer control to the
    destination specified by the OS in the SEV-ES AP Jump Table. The
    SEV-ES AP Jump Table address is communicated to the OS using the BSP
    GHCB MSR.

Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Signed-off-by: Tom Lendacky <thomas.lendacky@amd.com>
---
 UefiCpuPkg/Library/MpInitLib/MpLib.h          |   8 +-
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c       |  54 ++++++-
 UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm | 134 ++++++++++++++++--
 3 files changed, 176 insertions(+), 20 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index 5966510d4a1b..2d38a0e85a40 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -277,7 +277,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 {
@@ -331,8 +332,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
   );
 
 /**
diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index 16603ef3f20e..cf53b5026aa4 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;
@@ -86,6 +85,13 @@ GetWakeupBuffer (
 {
   EFI_STATUS              Status;
   EFI_PHYSICAL_ADDRESS    StartAddress;
+  EFI_MEMORY_TYPE         MemoryType;
+
+  if (PcdGetBool (PcdSevEsActive)) {
+    MemoryType = EfiReservedMemoryType;
+  } else {
+    MemoryType = EfiBootServicesData;
+  }
 
   //
   // Try to allocate buffer below 1M for waking vector.
@@ -98,7 +104,7 @@ GetWakeupBuffer (
   StartAddress = 0x88000;
   Status = gBS->AllocatePages (
                   AllocateMaxAddress,
-                  EfiBootServicesData,
+                  MemoryType,
                   EFI_SIZE_TO_PAGES (WakeupBufferSize),
                   &StartAddress
                   );
@@ -331,17 +337,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
@@ -895,9 +910,34 @@ MpFinalize (
   IN CPU_MP_DATA   *CpuMpData
   )
 {
-  //
-  // DXE phase will do this transition, but just return EFI_SUCCESS for now.
-  //
+  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/X64/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
index bbc7432740ff..3cb0cd5bb306 100644
--- a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
+++ b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
@@ -465,6 +465,10 @@ BITS 16
     ;     - IP for Real Mode (two bytes)
     ;     - CS for Real Mode (two bytes)
     ;
+    ; This label is also used with AsmRelocateApLoop. During MP finalization,
+    ; the code from PM16Mode to SwitchToRealProcEnd is copied to the start of
+    ; the WakeupBuffer, allowing a parked AP to be booted by an OS.
+    ;
 PM16Mode:
     mov        eax, cr0                    ; Read CR0
     btr        eax, 0                      ; Set PE=0
@@ -487,32 +491,95 @@ 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
+    ;   - Must be done while in 64-bit long mode so that writes to
+    ;     the GHCB memory will be unencrypted.
+    ;   - No NAE events can be generated once this is set otherwise
+    ;     the AP_RESET_HOLD 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_RESET_HOLD
+    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
@@ -525,6 +592,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
@@ -538,10 +607,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