[PATCH V2 5/6] OvmfPkg: Enable APs to accept memory for TDVF

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

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

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

TDVF APs once did nothing but spin around to wait for the Wakeup command.
This patch enables APs to handle the AcceptPages command. Once APs find
the AcceptPages command, it set its stack and jump to the function of
ApAcceptMemoryResourceRange (which will be introduced in the following
patch).

Cc: Erdem Aktas <erdemaktas@google.com>
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>
Signed-off-by: Min Xu <min.m.xu@intel.com>
---
 OvmfPkg/Include/TdxCommondefs.inc         |  4 +-
 OvmfPkg/IntelTdx/Sec/X64/IntelTdxAPs.nasm | 61 +++++++++++++++++++++++
 2 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/OvmfPkg/Include/TdxCommondefs.inc b/OvmfPkg/Include/TdxCommondefs.inc
index 970eac96592a..a29d2fad4233 100644
--- a/OvmfPkg/Include/TdxCommondefs.inc
+++ b/OvmfPkg/Include/TdxCommondefs.inc
@@ -15,8 +15,8 @@ FirmwareArgsOffset                        equ       800h
 WakeupArgsRelocatedMailBox                equ       800h
 AcceptPageArgsPhysicalStart               equ       800h
 AcceptPageArgsPhysicalEnd                 equ       808h
-AcceptPageArgsChunkSize                   equ       810h
-AcceptPageArgsPageSize                    equ       818h
+AcceptPageArgsTopStackAddress             equ       810h
+AcceptPageArgsApStackSize                 equ       818h
 CpuArrivalOffset                          equ       900h
 CpusExitingOffset                         equ       0a00h
 TalliesOffset                             equ       0a08h
diff --git a/OvmfPkg/IntelTdx/Sec/X64/IntelTdxAPs.nasm b/OvmfPkg/IntelTdx/Sec/X64/IntelTdxAPs.nasm
index 034ac0ee9421..4a984ecc1058 100644
--- a/OvmfPkg/IntelTdx/Sec/X64/IntelTdxAPs.nasm
+++ b/OvmfPkg/IntelTdx/Sec/X64/IntelTdxAPs.nasm
@@ -40,9 +40,70 @@ do_wait_loop:
     cmp     eax, MpProtectedModeWakeupCommandWakeup
     je      .do_wakeup
 
+    cmp     eax, MpProtectedModeWakeupCommandAcceptPages
+    je      .do_accept_pages
+
     ; Don't support this command, so ignore
     jmp     .check_command
 
+.do_accept_pages:
+    ;
+    ; Read the top stack address from arguments
+    mov     rsi, [rsp + AcceptPageArgsTopStackAddress]
+
+    ;
+    ; Calculate the top stack address of the AP.
+    ; ApStackAddr = BaseStackAddr + (vCpuIndex) * ApStackSize
+    xor     rdx, rdx
+    xor     rbx, rbx
+    xor     rax, rax
+    mov     eax, [rsp + AcceptPageArgsApStackSize]
+    mov     ebx, r9d    ; vCpuIndex
+    mul     ebx
+    add     rsi, rax    ; now rsi is ApStackAddr
+
+.start_accept_pages:
+    ;
+    ; Read the function address which will be called
+    mov     rax, [rsp + WakeupVectorOffset]
+
+    ;
+    ; vCPU index as the first argument
+    mov     ecx, r9d
+    mov     rdx, [rsp + AcceptPageArgsPhysicalStart]
+    mov     r8, [rsp + AcceptPageArgsPhysicalEnd]
+
+    ; save the Mailbox address to rbx
+    mov     rbx, rsp
+
+    ;
+    ; set AP Stack
+    mov     rsp, rsi
+    nop
+
+    ; save rax (the Mailbox address)
+    push    rbx
+
+    call    rax
+
+    ; recove rsp
+    pop     rbx
+    mov     rsp, rbx
+    ;
+    ; recover r8, r9
+    mov     rax, 1
+    tdcall
+
+    mov     eax, 0FFFFFFFFh
+    lock xadd dword [rsp + CpusExitingOffset], eax
+    dec     eax
+
+.check_exiting_cnt:
+    cmp     eax, 0
+    je      do_wait_loop
+    mov     eax, dword[rsp + CpusExitingOffset]
+    jmp     .check_exiting_cnt
+
 .do_wakeup:
     ;
     ; BSP sets these variables before unblocking APs
-- 
2.29.2.windows.2