From: Ceping Sun <cepingx.sun@intel.com>

REF: https://bugzilla.tianocore.org/show_bug.cgi?id=4572

According to section 3.2 of the [GHCI] document, if the result of MapGPA
is "TDG.VP.VMCALL_RETRY", TDVF must retry mapping for pages in that region,
starting with the GPA specified in R11.

Reference:
[GHCI]: TDX Guest-Host-Communication Interface v1.0
https://cdrdv2.intel.com/v1/dl/getContent/726790

Cc: Erdem Aktas <erdemaktas@google.com>
Cc: James Bottomley <jejb@linux.ibm.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Min Xu <min.m.xu@intel.com>
Cc: Tom Lendacky <thomas.lendacky@amd.com>
Cc: Michael Roth <michael.roth@amd.com>
Cc: Gerd Hoffmann <kraxel@redhat.com>
Signed-off-by: Ceping Sun <cepingx.sun@intel.com>
---
 .../BaseMemEncryptTdxLib.inf                  |  1 +
 .../BaseMemEncryptTdxLib/MemoryEncryption.c   | 36 ++++++++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/OvmfPkg/Library/BaseMemEncryptTdxLib/BaseMemEncryptTdxLib.inf b/OvmfPkg/Library/BaseMemEncryptTdxLib/BaseMemEncryptTdxLib.inf
index 11768825f8ca..742b65a289ce 100644
--- a/OvmfPkg/Library/BaseMemEncryptTdxLib/BaseMemEncryptTdxLib.inf
+++ b/OvmfPkg/Library/BaseMemEncryptTdxLib/BaseMemEncryptTdxLib.inf
@@ -30,6 +30,7 @@
 [Sources]
   VirtualMemory.h
   MemoryEncryption.c
+  X64/TdVmCallMapGPA.nasm

 [LibraryClasses]
   BaseLib
diff --git a/OvmfPkg/Library/BaseMemEncryptTdxLib/MemoryEncryption.c b/OvmfPkg/Library/BaseMemEncryptTdxLib/MemoryEncryption.c
index b47f56b391a5..1f29f9194c30 100644
--- a/OvmfPkg/Library/BaseMemEncryptTdxLib/MemoryEncryption.c
+++ b/OvmfPkg/Library/BaseMemEncryptTdxLib/MemoryEncryption.c
@@ -38,6 +38,10 @@ typedef enum {

 STATIC PAGE_TABLE_POOL  *mPageTablePool = NULL;

+#define TDVMCALL_STATUS_RETRY  0x1
+
+#define MAX_RETRIES_PER_PAGE  3
+
 /**
   This function is used to help request the host VMM to map a GPA range as
   private or shared-memory mappings.
@@ -546,6 +550,13 @@ SetOrClearSharedBit (
   EFI_STATUS                    Status;
   EDKII_MEMORY_ACCEPT_PROTOCOL  *MemoryAcceptProtocol;

+  UINT64  MapGpaRetryaddr;

+  UINT32  RetryCount;
+  UINT64  EndAddress;
+
+  MapGpaRetryaddr = 0;
+  RetryCount      = 0;
+
   AddressEncMask = GetMemEncryptionAddressMask ();

   //
@@ -559,7 +570,30 @@ SetOrClearSharedBit (
     PhysicalAddress   &= ~AddressEncMask;
   }

-  TdStatus = TdVmCall (TDVMCALL_MAPGPA, PhysicalAddress, Length, 0, 0, NULL);
+  while (RetryCount < MAX_RETRIES_PER_PAGE) {
+    TdStatus = TdVmCallMapGPA (PhysicalAddress, Length, &MapGpaRetryaddr);

+    if (TdStatus != TDVMCALL_STATUS_RETRY) {
+      break;
+    }
+
+    DEBUG ((DEBUG_VERBOSE, "%a: TdVmcall(MAPGPA) Retry PhysicalAddress is %llx, MapGpaRetryaddr is %llx\n", __func__, PhysicalAddress, MapGpaRetryaddr));
+
+    EndAddress = PhysicalAddress + Length;
+    if ((MapGpaRetryaddr < PhysicalAddress) || (MapGpaRetryaddr > EndAddress)) {

+      DEBUG ((DEBUG_ERROR, "%a: TdVmcall(MAPGPA) failed Retry PhysicalAddress is %llx, MapGpaRetryaddr is %llx\n", __func__, PhysicalAddress, MapGpaRetryaddr));
+      break;
+    }
+
+    if (MapGpaRetryaddr == PhysicalAddress) {
+      RetryCount++;
+      continue;
+    }
+
+    PhysicalAddress = MapGpaRetryaddr;
+    Length          = EndAddress - PhysicalAddress;
+    RetryCount      = 0;
+  }
+
   if (TdStatus != 0) {
     DEBUG ((DEBUG_ERROR, "%a: TdVmcall(MAPGPA) failed with %llx\n", __func__, TdStatus));
     ASSERT (FALSE);
--
2.34.1