From c55cf3f7958969b59dee0513c1826a3694de50b0 Mon Sep 17 00:00:00 2001
From: Min M Xu <min.m.xu@intel.com>
Date: Tue, 20 Dec 2022 16:42:39 +0800
Subject: [PATCH] OvmfPkg: Enable APs to accept memory for TDVF

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>
Reviewed-by: Jiewen Yao <jiewen.yao@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 970eac9659..a29d2fad42 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 034ac0ee94..4a984ecc10 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