[PATCH v9 1/4] OvmfPkg: Introduce CocoDxe driver

["Dionna Glaze" <dionnaglaze@google.com>]

This driver is meant as a join point for all Confidential Compute
technologies to put shared behavior that doesn't belong anywhere else.

The first behavior added here is to accept all unaccepted memory at
ExitBootServices if the behavior is not disabled. This allows safe
upgrades for OS loaders to affirm their support for the unaccepted
memory type.

Cc: Gerd Hoffmann <kraxel@redhat.com>
Cc: James Bottomley <jejb@linux.ibm.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Tom Lendacky <thomas.lendacky@amd.com>
Cc: Ard Biesheuvel <ardb@kernel.org>
Cc: "Min M. Xu" <min.m.xu@intel.com>
Cc: Andrew Fish <afish@apple.com>
Cc: "Michael D. Kinney" <michael.d.kinney@intel.com>

Signed-off-by: Dionna Glaze <dionnaglaze@google.com>
---
 OvmfPkg/AmdSev/AmdSevX64.dsc     |   1 +
 OvmfPkg/AmdSev/AmdSevX64.fdf     |   1 +
 OvmfPkg/CocoDxe/CocoDxe.c        | 147 +++++++++++++++++++++++++++++++
 OvmfPkg/CocoDxe/CocoDxe.inf      |  45 ++++++++++
 OvmfPkg/IntelTdx/IntelTdxX64.dsc |   1 +
 OvmfPkg/IntelTdx/IntelTdxX64.fdf |   1 +
 OvmfPkg/OvmfPkgIa32X64.dsc       |   1 +
 OvmfPkg/OvmfPkgIa32X64.fdf       |   1 +
 OvmfPkg/OvmfPkgX64.dsc           |   1 +
 OvmfPkg/OvmfPkgX64.fdf           |   1 +
 10 files changed, 200 insertions(+)
 create mode 100644 OvmfPkg/CocoDxe/CocoDxe.c
 create mode 100644 OvmfPkg/CocoDxe/CocoDxe.inf

diff --git a/OvmfPkg/AmdSev/AmdSevX64.dsc b/OvmfPkg/AmdSev/AmdSevX64.dsc
index 36100f5fdc..5e5e9887bb 100644
--- a/OvmfPkg/AmdSev/AmdSevX64.dsc
+++ b/OvmfPkg/AmdSev/AmdSevX64.dsc
@@ -749,6 +749,7 @@
     <LibraryClasses>
     PciLib|MdePkg/Library/BasePciLibCf8/BasePciLibCf8.inf
   }
+  OvmfPkg/CocoDxe/CocoDxe.inf
   OvmfPkg/IoMmuDxe/IoMmuDxe.inf
 
   #
diff --git a/OvmfPkg/AmdSev/AmdSevX64.fdf b/OvmfPkg/AmdSev/AmdSevX64.fdf
index 5fb3b5d276..ae64693c28 100644
--- a/OvmfPkg/AmdSev/AmdSevX64.fdf
+++ b/OvmfPkg/AmdSev/AmdSevX64.fdf
@@ -302,6 +302,7 @@ INF  OvmfPkg/QemuRamfbDxe/QemuRamfbDxe.inf
 INF  OvmfPkg/VirtioGpuDxe/VirtioGpu.inf
 INF  OvmfPkg/PlatformDxe/Platform.inf
 INF  OvmfPkg/AmdSevDxe/AmdSevDxe.inf
+INF  OvmfPkg/CocoDxe/CocoDxe.inf
 INF  OvmfPkg/IoMmuDxe/IoMmuDxe.inf
 
 
diff --git a/OvmfPkg/CocoDxe/CocoDxe.c b/OvmfPkg/CocoDxe/CocoDxe.c
new file mode 100644
index 0000000000..da16af32a3
--- /dev/null
+++ b/OvmfPkg/CocoDxe/CocoDxe.c
@@ -0,0 +1,147 @@
+/** @file
+
+  Confidential Compute Dxe driver. This driver installs protocols that are
+  generic over confidential compute techonology.
+
+  Copyright (c) 2022, Google LLC. All rights reserved.<BR>
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include <Library/BaseLib.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/DebugLib.h>
+#include <Library/DxeServicesTableLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+#include <Library/MemEncryptSevLib.h>
+#include <Library/MemEncryptTdxLib.h>
+#include <Pi/PrePiDxeCis.h>
+#include <Protocol/MemoryAccept.h>
+
+STATIC BOOLEAN  mAcceptAllMemoryAtEBS = TRUE;
+
+STATIC EFI_EVENT  mAcceptAllMemoryEvent = NULL;
+
+STATIC
+EFI_STATUS
+AcceptAllMemory (
+  IN EDKII_MEMORY_ACCEPT_PROTOCOL  *AcceptMemory
+  )
+{
+  EFI_GCD_MEMORY_SPACE_DESCRIPTOR  *AllDescMap;
+  UINTN                            NumEntries;
+  UINTN                            Index;
+  EFI_STATUS                       Status;
+
+  DEBUG ((DEBUG_INFO, "Accepting all memory\n"));
+
+  /*
+   * Get a copy of the memory space map to iterate over while
+   * changing the map.
+   */
+  Status = gDS->GetMemorySpaceMap (&NumEntries, &AllDescMap);
+  if (EFI_ERROR (Status)) {
+    return Status;
+  }
+
+  for (Index = 0; Index < NumEntries; Index++) {
+    CONST EFI_GCD_MEMORY_SPACE_DESCRIPTOR  *Desc;
+
+    Desc = &AllDescMap[Index];
+    if (Desc->GcdMemoryType != EFI_GCD_MEMORY_TYPE_UNACCEPTED) {
+      continue;
+    }
+
+    Status = AcceptMemory->AcceptMemory (
+                             AcceptMemory,
+                             Desc->BaseAddress,
+                             Desc->Length
+                             );
+    if (EFI_ERROR (Status)) {
+      break;
+    }
+
+    Status = gDS->RemoveMemorySpace (Desc->BaseAddress, Desc->Length);
+    if (EFI_ERROR (Status)) {
+      break;
+    }
+
+    Status = gDS->AddMemorySpace (
+                    EfiGcdMemoryTypeSystemMemory,
+                    Desc->BaseAddress,
+                    Desc->Length,
+                    EFI_MEMORY_CPU_CRYPTO | EFI_MEMORY_XP | EFI_MEMORY_RO | EFI_MEMORY_RP
+                    );
+    if (EFI_ERROR (Status)) {
+      break;
+    }
+  }
+
+  gBS->FreePool (AllDescMap);
+  return Status;
+}
+
+VOID
+EFIAPI
+ResolveUnacceptedMemory (
+  IN EFI_EVENT  Event,
+  IN VOID       *Context
+  )
+{
+  EDKII_MEMORY_ACCEPT_PROTOCOL  *AcceptMemory;
+  EFI_STATUS                    Status;
+
+  if (!mAcceptAllMemoryAtEBS) {
+    return;
+  }
+
+  Status = gBS->LocateProtocol (
+                  &gEdkiiMemoryAcceptProtocolGuid,
+                  NULL,
+                  (VOID **)&AcceptMemory
+                  );
+  if (Status == EFI_NOT_FOUND) {
+    return;
+  }
+
+  ASSERT_EFI_ERROR (Status);
+
+  Status = AcceptAllMemory (AcceptMemory);
+  ASSERT_EFI_ERROR (Status);
+}
+
+EFI_STATUS
+EFIAPI
+CocoDxeEntryPoint (
+  IN EFI_HANDLE        ImageHandle,
+  IN EFI_SYSTEM_TABLE  *SystemTable
+  )
+{
+  EFI_STATUS  Status;
+
+  //
+  // Do nothing when confidential compute technologies that require memory
+  // acceptance are not enabled.
+  //
+  if (!MemEncryptSevSnpIsEnabled () &&
+      !MemEncryptTdxIsEnabled ())
+  {
+    return EFI_UNSUPPORTED;
+  }
+
+  Status = gBS->CreateEventEx (
+                  EVT_NOTIFY_SIGNAL,
+                  TPL_CALLBACK,
+                  ResolveUnacceptedMemory,
+                  NULL,
+                  &gEfiEventBeforeExitBootServicesGuid,
+                  &mAcceptAllMemoryEvent
+                  );
+
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "AllowUnacceptedMemory event creation for EventBeforeExitBootServices failed.\n"));
+  }
+
+  return EFI_SUCCESS;
+}
diff --git a/OvmfPkg/CocoDxe/CocoDxe.inf b/OvmfPkg/CocoDxe/CocoDxe.inf
new file mode 100644
index 0000000000..8d4452e94d
--- /dev/null
+++ b/OvmfPkg/CocoDxe/CocoDxe.inf
@@ -0,0 +1,45 @@
+#/** @file
+#
+#  Driver installs shared protocols needed for confidential compute
+#  technologies.
+#
+#  Copyright (c) 2022, Google LLC. All rights reserved.<BR>
+#
+#  SPDX-License-Identifier: BSD-2-Clause-Patent
+#
+#**/
+
+[Defines]
+  INF_VERSION                    = 1.25
+  BASE_NAME                      = CocoDxe
+  FILE_GUID                      = 08162f1e-5147-4d3e-b5a9-fa48c9808419
+  MODULE_TYPE                    = DXE_DRIVER
+  VERSION_STRING                 = 1.0
+  ENTRY_POINT                    = CocoDxeEntryPoint
+
+[Sources]
+  CocoDxe.c
+
+[Packages]
+  MdeModulePkg/MdeModulePkg.dec
+  MdePkg/MdePkg.dec
+  OvmfPkg/OvmfPkg.dec
+
+[LibraryClasses]
+  BaseLib
+  BaseMemoryLib
+  DebugLib
+  DxeServicesTableLib
+  MemEncryptSevLib
+  MemEncryptTdxLib
+  MemoryAllocationLib
+  UefiDriverEntryPoint
+
+[Depex]
+  TRUE
+
+[Guids]
+  gEfiEventBeforeExitBootServicesGuid
+
+[Protocols]
+  gEdkiiMemoryAcceptProtocolGuid
diff --git a/OvmfPkg/IntelTdx/IntelTdxX64.dsc b/OvmfPkg/IntelTdx/IntelTdxX64.dsc
index 81511e3556..c3e64d97c0 100644
--- a/OvmfPkg/IntelTdx/IntelTdxX64.dsc
+++ b/OvmfPkg/IntelTdx/IntelTdxX64.dsc
@@ -754,6 +754,7 @@
   OvmfPkg/IoMmuDxe/IoMmuDxe.inf
 
   OvmfPkg/TdxDxe/TdxDxe.inf
+  OvmfPkg/CocoDxe/CocoDxe.inf
 
   #
   # Variable driver stack (non-SMM)
diff --git a/OvmfPkg/IntelTdx/IntelTdxX64.fdf b/OvmfPkg/IntelTdx/IntelTdxX64.fdf
index a57bbcee89..f5765b50eb 100644
--- a/OvmfPkg/IntelTdx/IntelTdxX64.fdf
+++ b/OvmfPkg/IntelTdx/IntelTdxX64.fdf
@@ -269,6 +269,7 @@ INF  ShellPkg/Application/Shell/Shell.inf
 INF MdeModulePkg/Logo/LogoDxe.inf
 
 INF OvmfPkg/TdxDxe/TdxDxe.inf
+INF OvmfPkg/CocoDxe/CocoDxe.inf
 
 #
 # Usb Support
diff --git a/OvmfPkg/OvmfPkgIa32X64.dsc b/OvmfPkg/OvmfPkgIa32X64.dsc
index a9d422bd91..8e4d31bcea 100644
--- a/OvmfPkg/OvmfPkgIa32X64.dsc
+++ b/OvmfPkg/OvmfPkgIa32X64.dsc
@@ -966,6 +966,7 @@
     <LibraryClasses>
     PciLib|MdePkg/Library/BasePciLibCf8/BasePciLibCf8.inf
   }
+  OvmfPkg/CocoDxe/CocoDxe.inf
   OvmfPkg/IoMmuDxe/IoMmuDxe.inf
 
 !if $(SMM_REQUIRE) == TRUE
diff --git a/OvmfPkg/OvmfPkgIa32X64.fdf b/OvmfPkg/OvmfPkgIa32X64.fdf
index 4c5bd0dbc3..7d75140fe3 100644
--- a/OvmfPkg/OvmfPkgIa32X64.fdf
+++ b/OvmfPkg/OvmfPkgIa32X64.fdf
@@ -343,6 +343,7 @@ INF  OvmfPkg/QemuRamfbDxe/QemuRamfbDxe.inf
 INF  OvmfPkg/VirtioGpuDxe/VirtioGpu.inf
 INF  OvmfPkg/PlatformDxe/Platform.inf
 INF  OvmfPkg/AmdSevDxe/AmdSevDxe.inf
+INF  OvmfPkg/CocoDxe/CocoDxe.inf
 INF  OvmfPkg/IoMmuDxe/IoMmuDxe.inf
 
 !if $(SMM_REQUIRE) == TRUE
diff --git a/OvmfPkg/OvmfPkgX64.dsc b/OvmfPkg/OvmfPkgX64.dsc
index 8401d73900..a728610c86 100644
--- a/OvmfPkg/OvmfPkgX64.dsc
+++ b/OvmfPkg/OvmfPkgX64.dsc
@@ -1037,6 +1037,7 @@
   OvmfPkg/IoMmuDxe/IoMmuDxe.inf
 
   OvmfPkg/TdxDxe/TdxDxe.inf
+  OvmfPkg/CocoDxe/CocoDxe.inf
 
 !if $(SMM_REQUIRE) == TRUE
   OvmfPkg/SmmAccess/SmmAccess2Dxe.inf
diff --git a/OvmfPkg/OvmfPkgX64.fdf b/OvmfPkg/OvmfPkgX64.fdf
index 8c02dfe11e..9e0aee225c 100644
--- a/OvmfPkg/OvmfPkgX64.fdf
+++ b/OvmfPkg/OvmfPkgX64.fdf
@@ -370,6 +370,7 @@ INF  OvmfPkg/QemuRamfbDxe/QemuRamfbDxe.inf
 INF  OvmfPkg/VirtioGpuDxe/VirtioGpu.inf
 INF  OvmfPkg/PlatformDxe/Platform.inf
 INF  OvmfPkg/AmdSevDxe/AmdSevDxe.inf
+INF  OvmfPkg/CocoDxe/CocoDxe.inf
 INF  OvmfPkg/IoMmuDxe/IoMmuDxe.inf
 
 !if $(SMM_REQUIRE) == TRUE
-- 
2.39.0.314.g84b9a713c41-goog