[PATCH V1 3/3] OvmfPkg/BaseMemEncryptTdxLib: Refactor error handle of SetOrClearSharedBit

["Min Xu" <min.m.xu@intel.com>]

From: Min M Xu <min.m.xu@intel.com>

The previous implementation of SetOrClearSharedBit doesn't handle the
error correctly. In this patch SetOrClearSharedBit is changed to return
error code so that the caller can handle it.

Cc: Erdem Aktas <erdemaktas@google.com>
Cc: James Bottomley <jejb@linux.ibm.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Gerd Hoffmann <kraxel@redhat.com>
Cc: Tom Lendacky <thomas.lendacky@amd.com>
Cc: Michael Roth <michael.roth@amd.com>
Signed-off-by: Min Xu <min.m.xu@intel.com>
---
 .../BaseMemEncryptTdxLib/MemoryEncryption.c   | 48 +++++++++++++++----
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/OvmfPkg/Library/BaseMemEncryptTdxLib/MemoryEncryption.c b/OvmfPkg/Library/BaseMemEncryptTdxLib/MemoryEncryption.c
index 503f626d75c6..5b13042512ad 100644
--- a/OvmfPkg/Library/BaseMemEncryptTdxLib/MemoryEncryption.c
+++ b/OvmfPkg/Library/BaseMemEncryptTdxLib/MemoryEncryption.c
@@ -510,8 +510,11 @@ Split1GPageTo2M (
   @param[in]      PagetablePoint        Page table entry pointer (PTE).
   @param[in]      Mode                  Set or Clear shared bit
 
+  @retval         EFI_SUCCESS           Successfully set or clear the memory shared bit
+  @retval         Others                Other error as indicated
 **/
-STATIC VOID
+STATIC
+EFI_STATUS
 SetOrClearSharedBit (
   IN   OUT     UINT64              *PageTablePointer,
   IN           TDX_PAGETABLE_MODE  Mode,
@@ -520,7 +523,8 @@ SetOrClearSharedBit (
   )
 {
   UINT64                        AddressEncMask;
-  UINT64                        Status;
+  UINT64                        TdStatus;
+  EFI_STATUS                    Status;
   EDKII_MEMORY_ACCEPT_PROTOCOL  *MemoryAcceptProtocol;
 
   AddressEncMask = GetMemEncryptionAddressMask ();
@@ -536,16 +540,30 @@ SetOrClearSharedBit (
     PhysicalAddress   &= ~AddressEncMask;
   }
 
-  Status = TdVmCall (TDVMCALL_MAPGPA, PhysicalAddress, Length, 0, 0, NULL);
+  TdStatus = TdVmCall (TDVMCALL_MAPGPA, PhysicalAddress, Length, 0, 0, NULL);
+  if (TdStatus != 0) {
+    DEBUG ((DEBUG_ERROR, "%a: TdVmcall(MAPGPA) failed with %llx\n", __FUNCTION__, TdStatus));
+    ASSERT (FALSE);
+    return EFI_DEVICE_ERROR;
+  }
 
   //
   // If changing shared to private, must accept-page again
   //
   if (Mode == ClearSharedBit) {
     Status = gBS->LocateProtocol (&gEdkiiMemoryAcceptProtocolGuid, NULL, (VOID **)&MemoryAcceptProtocol);
-    ASSERT (!EFI_ERROR (Status));
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "%a: Failed to locate MemoryAcceptProtocol with %r\n", __FUNCTION__, Status));
+      ASSERT (FALSE);
+      return Status;
+    }
+
     Status = MemoryAcceptProtocol->AcceptMemory (MemoryAcceptProtocol, PhysicalAddress, Length);
-    ASSERT (!EFI_ERROR (Status));
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "%a: Failed to AcceptMemory with %r\n", __FUNCTION__, Status));
+      ASSERT (FALSE);
+      return Status;
+    }
   }
 
   DEBUG ((
@@ -558,6 +576,8 @@ SetOrClearSharedBit (
     Mode,
     Status
     ));
+
+  return EFI_SUCCESS;
 }
 
 /**
@@ -747,7 +767,11 @@ SetMemorySharedOrPrivate (
       // If we have at least 1GB to go, we can just update this entry
       //
       if (!(PhysicalAddress & (BIT30 - 1)) && (Length >= BIT30)) {
-        SetOrClearSharedBit (&PageDirectory1GEntry->Uint64, Mode, PhysicalAddress, BIT30);
+        Status = SetOrClearSharedBit (&PageDirectory1GEntry->Uint64, Mode, PhysicalAddress, BIT30);
+        if (EFI_ERROR (Status)) {
+          goto Done;
+        }
+
         DEBUG ((
           DEBUG_VERBOSE,
           "%a:%a: updated 1GB entry for Physical=0x%Lx\n",
@@ -809,7 +833,11 @@ SetMemorySharedOrPrivate (
         // If we have at least 2MB left to go, we can just update this entry
         //
         if (!(PhysicalAddress & (BIT21-1)) && (Length >= BIT21)) {
-          SetOrClearSharedBit (&PageDirectory2MEntry->Uint64, Mode, PhysicalAddress, BIT21);
+          Status = SetOrClearSharedBit (&PageDirectory2MEntry->Uint64, Mode, PhysicalAddress, BIT21);
+          if (EFI_ERROR (Status)) {
+            goto Done;
+          }
+
           PhysicalAddress += BIT21;
           Length          -= BIT21;
         } else {
@@ -856,7 +884,11 @@ SetMemorySharedOrPrivate (
           goto Done;
         }
 
-        SetOrClearSharedBit (&PageTableEntry->Uint64, Mode, PhysicalAddress, EFI_PAGE_SIZE);
+        Status = SetOrClearSharedBit (&PageTableEntry->Uint64, Mode, PhysicalAddress, EFI_PAGE_SIZE);
+        if (EFI_ERROR (Status)) {
+          goto Done;
+        }
+
         PhysicalAddress += EFI_PAGE_SIZE;
         Length          -= EFI_PAGE_SIZE;
       }
-- 
2.29.2.windows.2