[PATCH v5 4/7] 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 protocol 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        | 149 ++++++++++++++++++++
 OvmfPkg/CocoDxe/CocoDxe.inf      |  43 ++++++
 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(+)

diff --git a/OvmfPkg/AmdSev/AmdSevX64.dsc b/OvmfPkg/AmdSev/AmdSevX64.dsc
index 90e8a213ef..ad6b73ca4a 100644
--- a/OvmfPkg/AmdSev/AmdSevX64.dsc
+++ b/OvmfPkg/AmdSev/AmdSevX64.dsc
@@ -747,6 +747,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 4658e1d30e..3717ec9094 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..dc37c292f4
--- /dev/null
+++ b/OvmfPkg/CocoDxe/CocoDxe.c
@@ -0,0 +1,149 @@
+/** @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 <Protocol/ExitBootServicesCallback.h>
+#include <Protocol/MemoryAccept.h>
+
+STATIC EFI_HANDLE mCocoDxeHandle = NULL;
+
+STATIC
+EFI_STATUS
+AcceptAllUnacceptedMemory (
+  IN EFI_MEMORY_ACCEPT_PROTOCOL *AcceptMemory
+  )
+{
+  EFI_GCD_MEMORY_SPACE_DESCRIPTOR  *AllDescMap;
+  UINTN                            NumEntries;
+  UINTN                            Index;
+  EFI_STATUS                       Status;
+  BOOLEAN                          AcceptedAny;
+
+  DEBUG ((DEBUG_INFO, "Accepting all memory\n"));
+  AcceptedAny = FALSE;
+  /*
+   * 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 != EfiGcdMemoryTypeUnaccepted) {
+      continue;
+    }
+
+    Status = AcceptMemory->AcceptMemory (
+      AcceptMemory,
+      Desc->BaseAddress,
+      Desc->Length
+      );
+    if (EFI_ERROR(Status)) {
+      goto done;
+    }
+
+    Status = gDS->RemoveMemorySpace(Desc->BaseAddress, Desc->Length);
+    if (EFI_ERROR(Status)) {
+      goto done;
+    }
+
+    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)) {
+      goto done;
+    }
+
+    AcceptedAny = TRUE;
+  }
+
+  // If any memory is accepted, cause ExitBootServices to fail with
+  // EFI_INVALID_PARAMETER in order to force the caller to refresh
+  // their view of the MemoryMap.
+  if (AcceptedAny) {
+    Status = EFI_INVALID_PARAMETER;
+  }
+
+done:
+  gBS->FreePool (AllDescMap);
+  return Status;
+}
+
+EFI_STATUS
+EFIAPI
+ResolveUnacceptedMemory (
+  IN EDKII_EXIT_BOOT_SERVICES_CALLBACK_PROTOCOL  *This
+  )
+{
+  EFI_MEMORY_ACCEPT_PROTOCOL *AcceptMemory;
+  EFI_STATUS                 Status;
+
+  if (This->Disabled) {
+    return EFI_SUCCESS;
+  }
+
+  Status = gBS->LocateProtocol (&gEfiMemoryAcceptProtocolGuid, NULL,
+    (VOID **)&AcceptMemory);
+  if (Status == EFI_NOT_FOUND) {
+    return EFI_SUCCESS;
+  }
+  ASSERT_EFI_ERROR (Status);
+
+  return AcceptAllUnacceptedMemory(AcceptMemory);
+}
+
+STATIC EDKII_EXIT_BOOT_SERVICES_CALLBACK_PROTOCOL mExitBootServicesCallbackProcotol = {
+  ResolveUnacceptedMemory,
+  FALSE,
+};
+
+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->InstallProtocolInterface (&mCocoDxeHandle,
+                  &gEdkiiExitBootServicesCallbackProtocolGuid,
+                  EFI_NATIVE_INTERFACE,
+                  &mExitBootServicesCallbackProcotol
+                  );
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Install EdkiiExitBootServicesCallbackProtocol failed.\n"));
+  }
+
+  return EFI_SUCCESS;
+}
diff --git a/OvmfPkg/CocoDxe/CocoDxe.inf b/OvmfPkg/CocoDxe/CocoDxe.inf
new file mode 100644
index 0000000000..3ff2a6fade
--- /dev/null
+++ b/OvmfPkg/CocoDxe/CocoDxe.inf
@@ -0,0 +1,43 @@
+#/** @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
+
+[Protocols]
+  gEdkiiExitBootServicesCallbackProtocolGuid
+  gEfiMemoryAcceptProtocolGuid
diff --git a/OvmfPkg/IntelTdx/IntelTdxX64.dsc b/OvmfPkg/IntelTdx/IntelTdxX64.dsc
index c0c1a15b09..8136d50eb2 100644
--- a/OvmfPkg/IntelTdx/IntelTdxX64.dsc
+++ b/OvmfPkg/IntelTdx/IntelTdxX64.dsc
@@ -753,6 +753,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 6923eb8831..e612608c0c 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 af566b953f..2cfb3fbc6b 100644
--- a/OvmfPkg/OvmfPkgIa32X64.dsc
+++ b/OvmfPkg/OvmfPkgIa32X64.dsc
@@ -965,6 +965,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 80de4fa2c0..2ab7f3b95b 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 f39d9cd117..3ead476b61 100644
--- a/OvmfPkg/OvmfPkgX64.dsc
+++ b/OvmfPkg/OvmfPkgX64.dsc
@@ -1036,6 +1036,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 c0f5a1ef3c..5dd452f42b 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.38.0.rc1.362.ged0d419d3c-goog