[RFC PATCH 13/14] OvmfPkg/AmdSev: Don't overwrite MH stack

["Tobin Feldman-Fitzthum" <tobin@linux.ibm.com>]

When restoring pages, the Migration Handler shoudl avoid overwriting
its own stack.

Signed-off-by: Tobin Feldman-Fitzthum <tobin@linux.ibm.com>
---
 .../ConfidentialMigrationDxe.inf              |   2 +
 OvmfPkg/AmdSev/ConfidentialMigration/MpLib.h  | 235 ++++++++++++++++++
 .../ConfidentialMigrationDxe.c                |  30 ++-
 3 files changed, 266 insertions(+), 1 deletion(-)
 create mode 100644 OvmfPkg/AmdSev/ConfidentialMigration/MpLib.h

diff --git a/OvmfPkg/AmdSev/ConfidentialMigration/ConfidentialMigrationDxe.inf b/OvmfPkg/AmdSev/ConfidentialMigration/ConfidentialMigrationDxe.inf
index 8dadfd1d13..2816952863 100644
--- a/OvmfPkg/AmdSev/ConfidentialMigration/ConfidentialMigrationDxe.inf
+++ b/OvmfPkg/AmdSev/ConfidentialMigration/ConfidentialMigrationDxe.inf
@@ -16,6 +16,7 @@
 [Sources]
   ConfidentialMigrationDxe.c
   VirtualMemory.h
+  MpLib.h
 
 [Packages]
   MdePkg/MdePkg.dec
@@ -36,6 +37,7 @@
   gUefiOvmfPkgTokenSpaceGuid.PcdIsConfidentialMigrationTarget
   gUefiOvmfPkgTokenSpaceGuid.PcdStartConfidentialMigrationHandler
   gUefiOvmfPkgTokenSpaceGuid.PcdConfidentialMigrationMailboxBase
+  gUefiCpuPkgTokenSpaceGuid.PcdCpuApStackSize
 
 [Depex]
   gEfiMpServiceProtocolGuid
diff --git a/OvmfPkg/AmdSev/ConfidentialMigration/MpLib.h b/OvmfPkg/AmdSev/ConfidentialMigration/MpLib.h
new file mode 100644
index 0000000000..5007e25243
--- /dev/null
+++ b/OvmfPkg/AmdSev/ConfidentialMigration/MpLib.h
@@ -0,0 +1,235 @@
+/** @file
+  Common header file for MP Initialize Library.
+  -- adapted from UefiCpuPkg/Library/MpInitLib/MpLib.h
+  Copyright (c) 2016 - 2020, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2020, AMD Inc. All rights reserved.<BR>
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+**/
+
+#ifndef _MP_LIB_H_
+#define _MP_LIB_H_
+
+#include <PiPei.h>
+
+#include <Library/MpInitLib.h>
+#include <Library/MemoryAllocationLib.h>
+#include <Library/CpuLib.h>
+#include <Library/HobLib.h>
+#include <Library/SynchronizationLib.h>
+#include <Library/MtrrLib.h>
+
+#define CPU_INIT_MP_LIB_HOB_GUID \
+  { \
+    0x58eb6a19, 0x3699, 0x4c68, { 0xa8, 0x36, 0xda, 0xcd, 0x8e, 0xdc, 0xad, 0x4a } \
+  }
+
+
+//
+// CPU exchange information for switch BSP
+//
+typedef struct {
+  UINT8             State;        // offset 0
+  UINTN             StackPointer; // offset 4 / 8
+  IA32_DESCRIPTOR   Gdtr;         // offset 8 / 16
+  IA32_DESCRIPTOR   Idtr;         // offset 14 / 26
+} CPU_EXCHANGE_ROLE_INFO;
+
+//
+// AP initialization state during APs wakeup
+//
+typedef enum {
+  ApInitConfig   = 1,
+  ApInitReconfig = 2,
+  ApInitDone     = 3
+} AP_INIT_STATE;
+
+//
+// AP state
+//
+// The state transitions for an AP when it process a procedure are:
+//  Idle ----> Ready ----> Busy ----> Idle
+//       [BSP]       [AP]       [AP]
+//
+typedef enum {
+  CpuStateIdle,
+  CpuStateReady,
+  CpuStateBusy,
+  CpuStateFinished,
+  CpuStateDisabled
+} CPU_STATE;
+
+//
+// CPU volatile registers around INIT-SIPI-SIPI
+//
+typedef struct {
+  UINTN                          Cr0;
+  UINTN                          Cr3;
+  UINTN                          Cr4;
+  UINTN                          Dr0;
+  UINTN                          Dr1;
+  UINTN                          Dr2;
+  UINTN                          Dr3;
+  UINTN                          Dr6;
+  UINTN                          Dr7;
+  IA32_DESCRIPTOR                Gdtr;
+  IA32_DESCRIPTOR                Idtr;
+  UINT16                         Tr;
+} CPU_VOLATILE_REGISTERS;
+
+//
+// AP related data
+//
+typedef struct {
+  SPIN_LOCK                      ApLock;
+  volatile UINT32                *StartupApSignal;
+  volatile UINTN                 ApFunction;
+  volatile UINTN                 ApFunctionArgument;
+  BOOLEAN                        CpuHealthy;
+  volatile CPU_STATE             State;
+  CPU_VOLATILE_REGISTERS         VolatileRegisters;
+  BOOLEAN                        Waiting;
+  BOOLEAN                        *Finished;
+  UINT64                         ExpectedTime;
+  UINT64                         CurrentTime;
+  UINT64                         TotalTime;
+  EFI_EVENT                      WaitEvent;
+  UINT32                         ProcessorSignature;
+  UINT8                          PlatformId;
+  UINT64                         MicrocodeEntryAddr;
+} CPU_AP_DATA;
+
+//
+// Basic CPU information saved in Guided HOB.
+// Because the contents will be shard between PEI and DXE,
+// we need to make sure the each fields offset same in different
+// architecture.
+//
+#pragma pack (1)
+typedef struct {
+  UINT32                         InitialApicId;
+  UINT32                         ApicId;
+  UINT32                         Health;
+  UINT64                         ApTopOfStack;
+} CPU_INFO_IN_HOB;
+#pragma pack ()
+
+//
+// AP reset code information including code address and size,
+// this structure will be shared be C code and assembly code.
+// It is natural aligned by design.
+//
+typedef struct {
+  UINT8             *RendezvousFunnelAddress;
+  UINTN             ModeEntryOffset;
+  UINTN             RendezvousFunnelSize;
+  UINT8             *RelocateApLoopFuncAddress;
+  UINTN             RelocateApLoopFuncSize;
+  UINTN             ModeTransitionOffset;
+} MP_ASSEMBLY_ADDRESS_MAP;
+
+typedef struct _CPU_MP_DATA  CPU_MP_DATA;
+
+#pragma pack(1)
+
+//
+// MP CPU exchange information for AP reset code
+// This structure is required to be packed because fixed field offsets
+// into this structure are used in assembly code in this module
+//
+typedef struct {
+  UINTN                 Lock;
+  UINTN                 StackStart;
+  UINTN                 StackSize;
+  UINTN                 CFunction;
+  IA32_DESCRIPTOR       GdtrProfile;
+  IA32_DESCRIPTOR       IdtrProfile;
+  UINTN                 BufferStart;
+  UINTN                 ModeOffset;
+  UINTN                 ApIndex;
+  UINTN                 CodeSegment;
+  UINTN                 DataSegment;
+  UINTN                 EnableExecuteDisable;
+  UINTN                 Cr3;
+  UINTN                 InitFlag;
+  CPU_INFO_IN_HOB       *CpuInfo;
+  UINTN                 NumApsExecuting;
+  CPU_MP_DATA           *CpuMpData;
+  UINTN                 InitializeFloatingPointUnitsAddress;
+  UINT32                ModeTransitionMemory;
+  UINT16                ModeTransitionSegment;
+  UINT32                ModeHighMemory;
+  UINT16                ModeHighSegment;
+  //
+  // Enable5LevelPaging indicates whether 5-level paging is enabled in long mode.
+  //
+  BOOLEAN               Enable5LevelPaging;
+} MP_CPU_EXCHANGE_INFO;
+
+#pragma pack()
+
+//
+// CPU MP Data save in memory
+//
+struct _CPU_MP_DATA {
+  UINT64                         CpuInfoInHob;
+  UINT32                         CpuCount;
+  UINT32                         BspNumber;
+  //
+  // The above fields data will be passed from PEI to DXE
+  // Please make sure the fields offset same in the different
+  // architecture.
+  //
+  SPIN_LOCK                      MpLock;
+  UINTN                          Buffer;
+  UINTN                          CpuApStackSize;
+  MP_ASSEMBLY_ADDRESS_MAP        AddressMap;
+  UINTN                          WakeupBuffer;
+  UINTN                          WakeupBufferHigh;
+  UINTN                          BackupBuffer;
+  UINTN                          BackupBufferSize;
+
+  volatile UINT32                FinishedCount;
+  UINT32                         RunningCount;
+  BOOLEAN                        SingleThread;
+  EFI_AP_PROCEDURE               Procedure;
+  VOID                           *ProcArguments;
+  BOOLEAN                        *Finished;
+  UINT64                         ExpectedTime;
+  UINT64                         CurrentTime;
+  UINT64                         TotalTime;
+  EFI_EVENT                      WaitEvent;
+  UINTN                          **FailedCpuList;
+
+  AP_INIT_STATE                  InitFlag;
+  BOOLEAN                        SwitchBspFlag;
+  UINTN                          NewBspNumber;
+  CPU_EXCHANGE_ROLE_INFO         BSPInfo;
+  CPU_EXCHANGE_ROLE_INFO         APInfo;
+  MTRR_SETTINGS                  MtrrTable;
+  UINT8                          ApLoopMode;
+  UINT8                          ApTargetCState;
+  UINT16                         PmCodeSegment;
+  CPU_AP_DATA                    *CpuData;
+  volatile MP_CPU_EXCHANGE_INFO  *MpCpuExchangeInfo;
+
+  UINT32                         CurrentTimerCount;
+  UINTN                          DivideValue;
+  UINT8                          Vector;
+  BOOLEAN                        PeriodicMode;
+  BOOLEAN                        TimerInterruptState;
+  UINT64                         MicrocodePatchAddress;
+  UINT64                         MicrocodePatchRegionSize;
+
+  //
+  // Whether need to use Init-Sipi-Sipi to wake up the APs.
+  // Two cases need to set this value to TRUE. One is in HLT
+  // loop mode, the other is resume from S3 which loop mode
+  // will be hardcode change to HLT mode by PiSmmCpuDxeSmm
+  // driver.
+  //
+  BOOLEAN                        WakeUpByInitSipiSipi;
+};
+
+extern EFI_GUID mCpuInitMpLibHobGuid;
+
+#endif
diff --git a/OvmfPkg/AmdSev/ConfidentialMigration/ConfidentialMigrationDxe.c b/OvmfPkg/AmdSev/ConfidentialMigration/ConfidentialMigrationDxe.c
index f609e16f8d..42b99be552 100644
--- a/OvmfPkg/AmdSev/ConfidentialMigration/ConfidentialMigrationDxe.c
+++ b/OvmfPkg/AmdSev/ConfidentialMigration/ConfidentialMigrationDxe.c
@@ -12,6 +12,8 @@
 #include <Library/BaseMemoryLib.h>
 
 #include "VirtualMemory.h"
+#include "MpLib.h"
+
 //
 // Functions implemented by the migration handler
 //
@@ -114,6 +116,7 @@ PrepareMigrationHandlerPageTables (
   mMigrationHelperPageTables = (UINT64)Start | AddressEncMask;
 }
 
+
 VOID
 SwitchToMigrationHelperPageTables(VOID)
 {
@@ -121,6 +124,25 @@ SwitchToMigrationHelperPageTables(VOID)
 }
 
 
+UINT64
+GetMHTopOfStack()
+{
+  EFI_HOB_GUID_TYPE       *GuidHob;
+  VOID                    *DataInHob;
+  CPU_MP_DATA             *CpuMpData;
+  CPU_INFO_IN_HOB         *CpuInfoInHob;
+
+  GuidHob = GetFirstGuidHob (&mCpuInitMpLibHobGuid);
+  ASSERT(GuidHob != NULL);
+
+  DataInHob = GET_GUID_HOB_DATA (GuidHob);
+  CpuMpData = (CPU_MP_DATA *) (*(UINTN *) DataInHob);
+  CpuInfoInHob = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;
+
+  return CpuInfoInHob[MigrationHandlerCpuIndex].ApTopOfStack;
+
+}
+
 
 VOID
 EFIAPI
@@ -132,6 +154,8 @@ MigrationHandlerMain (
   UINT64                       mailbox_end;
   UINT64                       pagetable_start;
   UINT64                       pagetable_end;
+  UINT64                       stack_start;
+  UINT64                       stack_end;
   UINT64                       params_base;
   MH_COMMAND_PARAMETERS        *params;
   VOID                         *page_va;
@@ -153,6 +177,9 @@ MigrationHandlerMain (
   pagetable_start = mMigrationHelperPageTables;
   pagetable_end = pagetable_start + 11 * EFI_PAGE_SIZE;
 
+  stack_end = GetMHTopOfStack();
+  stack_start = stack_end - PcdGet32(PcdCpuApStackSize);
+
   DisableInterrupts();
   params->go = 0;
 
@@ -177,7 +204,8 @@ MigrationHandlerMain (
       // Don't import a page that covers the mailbox or pagetables.
       //
       if ((params->gpa >= mailbox_start && params->gpa < mailbox_end) ||
-          (params->gpa >= pagetable_start && params->gpa < pagetable_end)) {
+          (params->gpa >= pagetable_start && params->gpa < pagetable_end) ||
+          (params->gpa >= stack_start && params->gpa < stack_end)) {
       }
       else {
         CopyMem((VOID *)params->gpa, page_va, 4096);
-- 
2.20.1