[PATCH 6/6] OvmfPkg/ResetVector: Update ResetVector to support Tdx

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

RFC: https://bugzilla.tianocore.org/show_bug.cgi?id=3429

In Tdx all CPUs "reset" to run on 32-bit protected mode with flat
descriptor (paging disabled). But in Non-Td guest the initial state of
CPUs is 16-bit real mode. To resolve this conflict, BITS 16/32 is used
in the very beginning of ResetVector. It will check the 32-bit protected
mode or 16-bit real mode, then jump to the corresponding entry point.
This is done in OvmfPkg/ResetVector/Ia16/ResetVectorVtf0.asm.

InitTdx.asm is called to record the Tdx signature ('TDXG') and other tdx
information in a TDX_WORK_AREA which can be used by the other routines in
ResetVector.

Main32 in UefiCpuPkg/ResetVector/Vtf0/Main.asm is the entry point for
Tdx guest. It call ReloadFlat32 to load the GDT and set the CR0, then
jump to flat32. After that InitTdx is called to do the above Tdx initialization.

Then Tdx jumps to 64-bit long mode by doing following tasks:
1. SetCr3ForPageTables64
   For OVMF, some initial page tables is built at:
     PcdOvmfSecPageTablesBase - (PcdOvmfSecPageTablesBase + 0x6000)
   This page table supports the 4-level page table.
   But Tdx support 4-level and 5-level page table based on the CPU GPA width.
   48bit is 4-level paging, 52-bit is 5-level paging.
   If 5-level page table is supported (GPAW is 52), then a top level
   page directory pointers (1 * 256TB entry) is generated in the
   TdxPageTable.
2. Set Cr4
   Enable PAE.
3. Adjust Cr3
   If GPAW is 48, then Cr3 is PT_ADDR (0). If GPAW is 52, then Cr3 is
   TDX_PT_ADDR (0).

Tdx MailBox [0x10, 0x800] is reserved for OS. So we initialize piece of this
area ([0x10, 0x20]) to record the Tdx flag ('TDXG') and other Tdx info so that
they can be used in the following flow.

After all above is successfully done, Tdx jump to SecEntry.

Cc: Brijesh Singh <brijesh.singh@amd.com>
Cc: Erdem Aktas <erdemaktas@google.com>
Cc: James Bottomley <jejb@linux.ibm.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Tom Lendacky <thomas.lendacky@amd.com>
Signed-off-by: Min Xu <min.m.xu@intel.com>
---
 OvmfPkg/ResetVector/Ia16/ResetVectorVtf0.asm | 21 ++++++++
 OvmfPkg/ResetVector/Ia32/Flat32ToFlat64.asm  | 47 ++++++++++++++++
 OvmfPkg/ResetVector/Ia32/InitTdx.asm         | 57 ++++++++++++++++++++
 OvmfPkg/ResetVector/Ia32/PageTables64.asm    | 41 ++++++++++++++
 OvmfPkg/ResetVector/ResetVector.nasmb        | 17 ++++++
 5 files changed, 183 insertions(+)
 create mode 100644 OvmfPkg/ResetVector/Ia32/InitTdx.asm

diff --git a/OvmfPkg/ResetVector/Ia16/ResetVectorVtf0.asm b/OvmfPkg/ResetVector/Ia16/ResetVectorVtf0.asm
index ac86ce69ebe8..a390ed81d021 100644
--- a/OvmfPkg/ResetVector/Ia16/ResetVectorVtf0.asm
+++ b/OvmfPkg/ResetVector/Ia16/ResetVectorVtf0.asm
@@ -155,10 +155,31 @@ resetVector:
 ;
 ; This is where the processor will begin execution
 ;
+; In IA32 we follow the standard reset vector flow. While in X64, Td guest
+; may be supported. Td guest requires the startup mode to be 32-bit
+; protected mode but the legacy VM startup mode is 16-bit real mode.
+; To make NASM generate such shared entry code that behaves correctly in
+; both 16-bit and 32-bit mode, more BITS directives are added.
+;
+%ifdef ARCH_IA32
+
     nop
     nop
     jmp     EarlyBspInitReal16
 
+%else
+
+    smsw    ax
+    test    al, 1
+    jz      .Real
+BITS 32
+    jmp     Main32
+BITS 16
+.Real:
+    jmp     EarlyBspInitReal16
+
+%endif
+
 ALIGN   16
 
 fourGigabytes:
diff --git a/OvmfPkg/ResetVector/Ia32/Flat32ToFlat64.asm b/OvmfPkg/ResetVector/Ia32/Flat32ToFlat64.asm
index c6d0d898bcd1..2206ca719593 100644
--- a/OvmfPkg/ResetVector/Ia32/Flat32ToFlat64.asm
+++ b/OvmfPkg/ResetVector/Ia32/Flat32ToFlat64.asm
@@ -17,6 +17,9 @@ Transition32FlatTo64Flat:
 
     OneTimeCall SetCr3ForPageTables64
 
+    cmp     dword[TDX_WORK_AREA], 0x47584454 ; 'TDXG'
+    jz      TdxTransition32FlatTo64Flat
+
     mov     eax, cr4
     bts     eax, 5                      ; enable PAE
     mov     cr4, eax
@@ -65,10 +68,54 @@ EnablePaging:
     bts     eax, 31                     ; set PG
     mov     cr0, eax                    ; enable paging
 
+    jmp     _jumpTo64Bit
+
+;
+; Tdx Transition from 32Flat to 64Flat
+;
+TdxTransition32FlatTo64Flat:
+
+    mov     eax, cr4
+    bts     eax, 5                      ; enable PAE
+
+    ;
+    ; byte[TDX_WORK_AREA_PAGELEVEL5] holds the indicator whether 52bit is supported.
+    ; if it is the case, need to set LA57 and use 5-level paging
+    ;
+    cmp     byte[TDX_WORK_AREA_PAGELEVEL5], 0
+    jz      .set_cr4
+    bts     eax, 12
+.set_cr4:
+    mov     cr4, eax
+    mov     ebx, cr3
+
+    ;
+    ; if la57 is not set, we are ok
+    ; if using 5-level paging, adjust top-level page directory
+    ;
+    bt      eax, 12
+    jnc     .set_cr3
+    mov     ebx, TDX_PT_ADDR (0)
+.set_cr3:
+    mov     cr3, ebx
+
+    mov     eax, cr0
+    bts     eax, 31                     ; set PG
+    mov     cr0, eax                    ; enable paging
+
+_jumpTo64Bit:
     jmp     LINEAR_CODE64_SEL:ADDR_OF(jumpTo64BitAndLandHere)
+
 BITS    64
 jumpTo64BitAndLandHere:
 
+    ;
+    ; For Td guest we are done and jump to the end
+    ;
+    mov     eax, TDX_WORK_AREA
+    cmp     dword [eax], 0x47584454 ; 'TDXG'
+    jz      GoodCompare
+
     ;
     ; Check if the second step of the SEV-ES mitigation is to be performed.
     ;
diff --git a/OvmfPkg/ResetVector/Ia32/InitTdx.asm b/OvmfPkg/ResetVector/Ia32/InitTdx.asm
new file mode 100644
index 000000000000..de8273da6a0c
--- /dev/null
+++ b/OvmfPkg/ResetVector/Ia32/InitTdx.asm
@@ -0,0 +1,57 @@
+;------------------------------------------------------------------------------
+; @file
+;   Initialize TDX_WORK_AREA to record the Tdx flag ('TDXG') and other Tdx info
+;   so that the following codes can use these information.
+;
+; Copyright (c) 2021, Intel Corporation. All rights reserved.<BR>
+; SPDX-License-Identifier: BSD-2-Clause-Patent
+;
+;------------------------------------------------------------------------------
+
+BITS 32
+
+;
+; Modified:  EBP
+;
+InitTdx:
+    ;
+    ; In Td guest, BSP/AP shares the same entry point
+    ; BSP builds up the page table, while APs shouldn't do the same task.
+    ; Instead, APs just leverage the page table which is built by BSP.
+    ; APs will wait until the page table is ready.
+    ; In Td guest, vCPU 0 is treated as the BSP, the others are APs.
+    ; ESI indicates the vCPU ID.
+    ;
+    cmp     esi, 0
+    je      tdBspEntry
+
+apWait:
+    cmp     byte[TDX_WORK_AREA_PGTBL_READY], 0
+    je      apWait
+    jmp     doneTdxInit
+
+tdBspEntry:
+    ;
+    ; It is of Tdx Guest
+    ; Save the Tdx info in TDX_WORK_AREA so that the following code can use
+    ; these information.
+    ;
+    mov     dword [TDX_WORK_AREA], 0x47584454 ; 'TDXG'
+
+    ;
+    ; EBP[6:0] CPU supported GPA width
+    ;
+    and     ebp, 0x3f
+    cmp     ebp, 52
+    jl      NotPageLevel5
+    mov     byte[TDX_WORK_AREA_PAGELEVEL5], 1
+
+NotPageLevel5:
+    ;
+    ; ECX[31:0] TDINITVP - Untrusted Configuration
+    ;
+    mov     DWORD[TDX_WORK_AREA_INITVP], ecx
+    mov     DWORD[TDX_WORK_AREA_INFO], ebp
+
+doneTdxInit:
+    OneTimeCallRet InitTdx
diff --git a/OvmfPkg/ResetVector/Ia32/PageTables64.asm b/OvmfPkg/ResetVector/Ia32/PageTables64.asm
index 5fae8986d9da..508df6cf5967 100644
--- a/OvmfPkg/ResetVector/Ia32/PageTables64.asm
+++ b/OvmfPkg/ResetVector/Ia32/PageTables64.asm
@@ -218,6 +218,24 @@ SevEsDisabled:
 ;
 SetCr3ForPageTables64:
 
+    ;
+    ; Check Td guest
+    ;
+    cmp     dword[TDX_WORK_AREA], 0x47584454 ; 'TDXG'
+    jnz     CheckSev
+
+    xor     edx, edx
+
+    ;
+    ; In Td guest, BSP builds the page table and set the flag of
+    ; TDX_WORK_AREA_PGTBL_READY. APs check this flag and then set
+    ; cr3 directly.
+    ;
+    cmp     byte[TDX_WORK_AREA_PGTBL_READY], 1
+    jz      SetCr3
+    jmp     SevNotActive
+
+CheckSev:
     OneTimeCall   CheckSevFeatures
     xor     edx, edx
     test    eax, eax
@@ -277,6 +295,29 @@ pageTableEntriesLoop:
     mov     [(ecx * 8 + PT_ADDR (0x2000 - 8)) + 4], edx
     loop    pageTableEntriesLoop
 
+    ;
+    ; If it is Td guest, TdxExtraPageTable should be initialized as well
+    ;
+    cmp     dword[TDX_WORK_AREA], 0x47584454 ; 'TDXG'
+    jnz     IsSevEs
+
+    xor     eax, eax
+    mov     ecx, 0x400
+tdClearTdxPageTablesMemoryLoop:
+    mov     dword [ecx * 4 + TDX_PT_ADDR (0) - 4], eax
+    loop    tdClearTdxPageTablesMemoryLoop
+
+    xor     edx, edx
+    ;
+    ; Top level Page Directory Pointers (1 * 256TB entry)
+    ;
+    mov     dword[TDX_PT_ADDR (0)], PT_ADDR (0) + PAGE_PDP_ATTR
+    mov     dword[TDX_PT_ADDR (4)], edx
+
+    mov     byte[TDX_WORK_AREA_PGTBL_READY], 1
+    jmp     SetCr3
+
+IsSevEs:
     OneTimeCall   IsSevEsEnabled
     test    eax, eax
     jz      SetCr3
diff --git a/OvmfPkg/ResetVector/ResetVector.nasmb b/OvmfPkg/ResetVector/ResetVector.nasmb
index b653fe87abd6..47ea23095c0a 100644
--- a/OvmfPkg/ResetVector/ResetVector.nasmb
+++ b/OvmfPkg/ResetVector/ResetVector.nasmb
@@ -106,6 +106,21 @@
   %define TDX_EXTRA_PAGE_TABLE_BASE FixedPcdGet32 (PcdOvmfSecGhcbPageTableBase)
   %define TDX_EXTRA_PAGE_TABLE_SIZE FixedPcdGet32 (PcdOvmfSecGhcbPageTableSize)
 
+  ;
+  ; TdMailboxBase [0x10, 0x800] is reserved for OS.
+  ; Td guest initialize piece of this area (TdMailboxBase [0x10,0x20]) to
+  ; record the Td guest info so that this information can be used in the
+  ; following ResetVector flow.
+  ;
+  %define TD_MAILBOX_WORKAREA_OFFSET    0x10
+  %define TDX_WORK_AREA                 (TDX_MAILBOX_MEMORY_BASE + TD_MAILBOX_WORKAREA_OFFSET)
+  %define TDX_WORK_AREA_PAGELEVEL5      (TDX_WORK_AREA + 4)
+  %define TDX_WORK_AREA_PGTBL_READY     (TDX_WORK_AREA + 5)
+  %define TDX_WORK_AREA_INITVP          (TDX_WORK_AREA + 8)
+  %define TDX_WORK_AREA_INFO            (TDX_WORK_AREA + 8 + 4)
+
+  %define TDX_PT_ADDR(Offset)  (TDX_EXTRA_PAGE_TABLE_BASE + (Offset))
+
   %define PT_ADDR(Offset) (FixedPcdGet32 (PcdOvmfSecPageTablesBase) + (Offset))
 
   %define GHCB_PT_ADDR (FixedPcdGet32 (PcdOvmfSecGhcbPageTableBase))
@@ -117,6 +132,8 @@
   %define SEV_ES_VC_TOP_OF_STACK (FixedPcdGet32 (PcdOvmfSecPeiTempRamBase) + FixedPcdGet32 (PcdOvmfSecPeiTempRamSize))
 
   %include "X64/TdxMetadata.asm"
+  %include "Ia32/InitTdx.asm"
+  %include "Ia32/ReloadFlat32.asm"
 
   %include "Ia32/Flat32ToFlat64.asm"
   %include "Ia32/PageTables64.asm"
-- 
2.29.2.windows.2