[RFC PATCH 14/19] OvmfPkg/MemEncryptSevLib: Add support to validate RAM in PEI phase

["Brijesh Singh" <brijesh.singh@amd.com>]

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

MemEncryptSevSnpValidateSystemRam() is used for validating the system
RAM. During the validation process, we must avoid double validation
cases. The double validation can lead to security issues. Extend the
MemEncryptSevSnpValidateSystemRam() to use the interval search tree to
keep track of the validated range; if the requested range is already
validated, then do nothing. If the requested range overlaps with the
previous validation, then validate only non-overlapped ranges.

Cc: James Bottomley <jejb@linux.ibm.com>
Cc: Min Xu <min.m.xu@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Tom Lendacky <thomas.lendacky@amd.com>
Cc: Jordan Justen <jordan.l.justen@intel.com>
Cc: Ard Biesheuvel <ardb+tianocore@kernel.org>
Cc: Laszlo Ersek <lersek@redhat.com>
Signed-off-by: Brijesh Singh <brijesh.singh@amd.com>
---
 OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf       |   8 ++
 OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLibInternal.c |  20 ++++
 OvmfPkg/Library/BaseMemEncryptSevLib/X64/PeiSnpSystemRamValidate.c | 105 +++++++++++++++++
 OvmfPkg/Library/BaseMemEncryptSevLib/X64/SnpPageStateTrack.c       | 119 ++++++++++++++++++++
 OvmfPkg/Library/BaseMemEncryptSevLib/X64/SnpPageStateTrack.h       |  36 ++++++
 5 files changed, 288 insertions(+)

diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf b/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
index 03a78c32df..cb9dd2bb21 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
@@ -32,10 +32,15 @@
 [Sources]
   PeiDxeMemEncryptSevLibInternal.c
   PeiMemEncryptSevLibInternal.c
+  SnpPageStateChange.h
 
 [Sources.X64]
   X64/MemEncryptSevLib.c
   X64/PeiDxeVirtualMemory.c
+  X64/PeiSnpSystemRamValidate.c
+  X64/SnpPageStateTrack.c
+  X64/SnpPageStateChangeInternal.c
+  X64/SnpPageStateTrack.h
   X64/VirtualMemory.c
   X64/VirtualMemory.h
 
@@ -49,9 +54,12 @@
   DebugLib
   MemoryAllocationLib
   PcdLib
+  VmgExitLib
 
 [FeaturePcd]
   gUefiOvmfPkgTokenSpaceGuid.PcdSmmSmramRequire
 
 [FixedPcd]
+  gUefiOvmfPkgTokenSpaceGuid.PcdOvmfSnpCpuidBase
+  gUefiOvmfPkgTokenSpaceGuid.PcdOvmfDecompressionScratchEnd
   gUefiCpuPkgTokenSpaceGuid.PcdSevEsWorkAreaBase
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLibInternal.c b/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLibInternal.c
index b561f211f5..9863722e9d 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLibInternal.c
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLibInternal.c
@@ -17,6 +17,8 @@
 #include <Register/Cpuid.h>
 #include <Uefi/UefiBaseType.h>
 
+#include "SnpPageStateChange.h"
+
 STATIC BOOLEAN mSevStatus = FALSE;
 STATIC BOOLEAN mSevEsStatus = FALSE;
 STATIC BOOLEAN mSevSnpStatus = FALSE;
@@ -184,3 +186,21 @@ MemEncryptSevGetEncryptionMask (
 
   return mSevEncryptionMask;
 }
+
+/**
+  If SEV-SNP is active then set the page state of the specified virtual
+  address range. This should be called in SEC and PEI phases only.
+
+  @param[in]  BaseAddress             Base address
+  @param[in]  NumPages                Number of pages starting from the base address
+
+**/
+VOID
+EFIAPI
+MemEncryptSevSnpValidateSystemRam (
+  IN PHYSICAL_ADDRESS                   BaseAddress,
+  IN UINTN                              NumPages
+  )
+{
+  SevSnpValidateSystemRam (BaseAddress, NumPages);
+}
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/PeiSnpSystemRamValidate.c b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/PeiSnpSystemRamValidate.c
new file mode 100644
index 0000000000..ce8a05bb1f
--- /dev/null
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/PeiSnpSystemRamValidate.c
@@ -0,0 +1,105 @@
+/** @file
+
+  SEV-SNP Page Validation functions.
+
+  Copyright (c) 2020 - 2021, AMD Incorporated. All rights reserved.<BR>
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include <Uefi/UefiBaseType.h>
+#include <Library/BaseLib.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/MemEncryptSevLib.h>
+#include <Library/DebugLib.h>
+
+#include "../SnpPageStateChange.h"
+#include "SnpPageStateTrack.h"
+
+STATIC SNP_VALIDATED_RANGE     *mRootNode;
+
+STATIC
+SNP_VALIDATED_RANGE *
+SetPageStateChangeInitialize (
+  VOID
+  )
+{
+  UINTN                      StartAddress, EndAddress;
+  SNP_VALIDATED_RANGE        *RootNode;
+
+  //
+  // The memory range from PcdOvmfSnpCpuidBase to PcdOvmfDecompressionScratchEnd is
+  // prevalidated before we enter into the Pei phase. The pre-validation breakdown
+  // looks like this:
+  //
+  //    SnpCpuidBase                                      (VMM)
+  //    SnpSecretBase                                     (VMM)
+  //    SnpLaunchValidatedStart - SnpLaunchValidatedEnd   (VMM)
+  //    SnpLaunchValidatedEnd - DecompressionScratchEnd   (SecMain)
+  //
+  // Add the range in system ram region tracker interval tree. The interval tree will
+  // used to check whether there is an overlap with the pre-validated region. We will
+  // skip validating the pre-validated region.
+  //
+  StartAddress = (UINTN) PcdGet32 (PcdOvmfSnpCpuidBase);
+  EndAddress = (UINTN) PcdGet32 (PcdOvmfDecompressionScratchEnd);
+
+  RootNode = AddRangeToIntervalTree (NULL, StartAddress, EndAddress);
+  if (RootNode == NULL) {
+    DEBUG ((DEBUG_ERROR, "Failed to add range to interval tree\n"));
+    ASSERT (FALSE);
+  }
+
+  return RootNode;
+}
+
+VOID
+SevSnpValidateSystemRam (
+  IN UINTN      BaseAddress,
+  IN UINTN      NumPages
+  )
+{
+  UINTN                   EndAddress;
+  SNP_VALIDATED_RANGE     *Range;
+
+  EndAddress = BaseAddress + EFI_PAGES_TO_SIZE (NumPages);
+
+  //
+  // If the Root is NULL then its the first call. Lets initialize the List before
+  // we process the request.
+  //
+  if (mRootNode == NULL) {
+    mRootNode = SetPageStateChangeInitialize ();
+  }
+
+  //
+  // Check if the range is already validated
+  //
+  EndAddress = BaseAddress + EFI_PAGES_TO_SIZE(NumPages);
+  Range = FindOverlapRange (mRootNode, BaseAddress, EndAddress);
+
+  //
+  // Range is not validated
+  if (Range == NULL) {
+    SetPageStateInternal (BaseAddress, NumPages, SevSnpPagePrivate, TRUE);
+    AddRangeToIntervalTree (mRootNode, BaseAddress, EndAddress);
+    return;
+  }
+
+  //
+  // The input range overlaps with the pre-validated range. Calculate the non-overlap
+  // region and validate them.
+  //
+  if (BaseAddress < Range->StartAddress) {
+    NumPages = EFI_SIZE_TO_PAGES (Range->StartAddress - BaseAddress);
+    SetPageStateInternal (BaseAddress, NumPages, SevSnpPagePrivate, TRUE);
+    AddRangeToIntervalTree (mRootNode, BaseAddress, Range->StartAddress);
+  }
+
+  if (EndAddress > Range->EndAddress) {
+    NumPages = EFI_SIZE_TO_PAGES (EndAddress - Range->EndAddress);
+    SetPageStateInternal (Range->EndAddress, NumPages, SevSnpPagePrivate, TRUE);
+    AddRangeToIntervalTree (mRootNode, Range->StartAddress, EndAddress);
+  }
+}
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/SnpPageStateTrack.c b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/SnpPageStateTrack.c
new file mode 100644
index 0000000000..91b4fc8db4
--- /dev/null
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/SnpPageStateTrack.c
@@ -0,0 +1,119 @@
+/** @file
+
+  Provides a simple interval search tree implementation that will be used
+  by the SnpValidateSystemRam() to keep track of the memory range validated
+  during the SEC/PEI phases.
+
+  Copyright (c) 2020 - 2021, AMD Incorporated. All rights reserved.<BR>
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include <Library/MemoryAllocationLib.h>
+
+#include "SnpPageStateTrack.h"
+
+STATIC
+SNP_VALIDATED_RANGE *
+AllocNewNode (
+  IN  UINTN     StartAddress,
+  IN  UINTN     EndAddress
+  )
+{
+  SNP_VALIDATED_RANGE   *Node;
+
+  Node = AllocatePool (sizeof (SNP_VALIDATED_RANGE));
+  if (Node == NULL) {
+    return NULL;
+  }
+
+  Node->StartAddress  = StartAddress;
+  Node->EndAddress = EndAddress;
+  Node->MaxAddress = Node->EndAddress;
+  Node->Left = Node->Right = NULL;
+
+  return Node;
+}
+
+STATIC
+BOOLEAN
+RangeIsOverlap (
+  IN  SNP_VALIDATED_RANGE     *Node,
+  IN  UINTN                   StartAddress,
+  IN  UINTN                   EndAddress
+  )
+{
+  if (Node->StartAddress < EndAddress && StartAddress < Node->EndAddress) {
+    return TRUE;
+  }
+
+  return FALSE;
+}
+
+
+/**
+ Function to find the overlapping range within the interval tree. If range is not
+ found then NULL is returned.
+
+ */
+SNP_VALIDATED_RANGE *
+FindOverlapRange (
+  IN  SNP_VALIDATED_RANGE   *RootNode,
+  IN  UINTN                 StartAddress,
+  IN  UINTN                 EndAddress
+  )
+{
+  // Tree is empty or no overlap found
+  if (RootNode == NULL) {
+    return NULL;
+  }
+
+  // Check with the range exist in the root node
+  if (RangeIsOverlap(RootNode, StartAddress, EndAddress)) {
+    return RootNode;
+  }
+
+  //
+  // If the left child of root is present and the max of the left child is
+  // greater than or equal to a given range then requested range will overlap
+  // with left subtree
+  //
+  if (RootNode->Left != NULL && (RootNode->Left->MaxAddress >= StartAddress)) {
+    return FindOverlapRange (RootNode->Left, StartAddress, EndAddress);
+  }
+
+  // The range can only overlap with the right subtree
+  return FindOverlapRange (RootNode->Right, StartAddress, EndAddress);
+}
+
+/**
+ Function to insert the validated range in the interval search tree.
+
+ */
+SNP_VALIDATED_RANGE *
+AddRangeToIntervalTree (
+  IN  SNP_VALIDATED_RANGE   *RootNode,
+  IN  UINTN                 StartAddress,
+  IN  UINTN                 EndAddress
+  )
+{
+  // Tree is empty or we reached to the leaf
+  if (RootNode == NULL) {
+    return AllocNewNode (StartAddress, EndAddress);
+  }
+
+  // If the StartAddress is smaller then the BaseAddress then go to the left in the tree.
+  if (StartAddress < RootNode->StartAddress) {
+    RootNode->Left = AddRangeToIntervalTree (RootNode->Left, StartAddress, EndAddress);
+  } else {
+    RootNode->Right = AddRangeToIntervalTree (RootNode->Right, StartAddress, EndAddress);
+  }
+
+  // Update the max value of the ancestor if needed
+  if (RootNode->MaxAddress < EndAddress) {
+    RootNode->MaxAddress = EndAddress;
+  }
+
+  return RootNode;
+}
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/SnpPageStateTrack.h b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/SnpPageStateTrack.h
new file mode 100644
index 0000000000..106c3411f0
--- /dev/null
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/SnpPageStateTrack.h
@@ -0,0 +1,36 @@
+/** @file
+
+
+  Copyright (c) 2020 - 2021, AMD Incorporated. All rights reserved.<BR>
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#ifndef SNP_PAGE_STATE_TRACK_INTERNAL_H_
+#define SNP_PAGE_STATE_TRACK_INTERNAL_H_
+
+#include <Uefi/UefiBaseType.h>
+
+typedef struct SNP_VALIDATED_RANGE {
+  UINT64    StartAddress, EndAddress;
+  UINT64    MaxAddress;
+
+  struct SNP_VALIDATED_RANGE  *Left, *Right;
+} SNP_VALIDATED_RANGE;
+
+SNP_VALIDATED_RANGE *
+FindOverlapRange (
+  IN  SNP_VALIDATED_RANGE   *RootNode,
+  IN  UINTN                 StartAddress,
+  IN  UINTN                 EndAddress
+  );
+
+SNP_VALIDATED_RANGE *
+AddRangeToIntervalTree (
+  IN  SNP_VALIDATED_RANGE   *RootNode,
+  IN  UINTN                 StartAddress,
+  IN  UINTN                 EndAddress
+  );
+
+#endif
-- 
2.17.1