[PATCH 2/2] UefiCpuPkg/PiSmmCpuDxeSmm: Add RSB stuffing before rsm instruction

[Hao Wu <hao.a.wu@intel.com>]

System Management Interrupt (SMI) handlers can leave the Return Stack
Buffer (RSB) in a state that application program or operating-system does
not expect.

In order to avoid RSB underflow on return from SMI, this commit will add
RSB stuffing logic before instruction 'rsm'.

After the stuffing, RSB entries will contain a trap like:

@SpecTrap:
    pause
    lfence
    jmp     @SpecTrap

to keep the speculative execution within control.

Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Hao Wu <hao.a.wu@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiEntry.nasm | 20 +++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmInit.nasm  | 21 ++++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmiEntry.nasm  | 20 +++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmInit.nasm   | 20 +++++++++++++++++++
 4 files changed, 81 insertions(+)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiEntry.nasm b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiEntry.nasm
index 509e7a0a66..e5875353a1 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiEntry.nasm
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiEntry.nasm
@@ -37,6 +37,8 @@
 %define PROTECT_MODE_DS 0x20
 %define TSS_SEGMENT 0x40
 
+%define RSB_STUFF_ENTRIES 0x20
+
 extern ASM_PFX(SmiRendezvous)
 extern ASM_PFX(FeaturePcdGet (PcdCpuSmmStackGuard))
 extern ASM_PFX(CpuSmmDebugEntry)
@@ -204,6 +206,24 @@ ASM_PFX(SmiHandler):
     wrmsr
 
 .7:
+    mov     eax, RSB_STUFF_ENTRIES / 2
+@Unroll1:
+    call    @Unroll2
+@SpecTrap1:
+    pause
+    lfence
+    jmp     @SpecTrap1
+@Unroll2:
+    call    @StuffLoop
+@SpecTrap2:
+    pause
+    lfence
+    jmp     @SpecTrap2
+@StuffLoop:
+    dec     eax
+    jnz     @Unroll1
+    add     esp, RSB_STUFF_ENTRIES * 4     ; Restore the stack pointer
+
     rsm
 
 ASM_PFX(gcSmiHandlerSize): DW $ - _SmiEntryPoint
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmInit.nasm b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmInit.nasm
index 5ff3cd2e73..fd559d25cd 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmInit.nasm
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmInit.nasm
@@ -33,6 +33,8 @@ global ASM_PFX(gcSmmInitTemplate)
 %define PROTECT_MODE_CS 0x8
 %define PROTECT_MODE_DS 0x20
 
+%define RSB_STUFF_ENTRIES 0x20
+
     SECTION .text
 
 ASM_PFX(gcSmiInitGdtr):
@@ -75,6 +77,25 @@ BITS 32
     mov     esp, strict dword 0         ; source operand will be patched
 ASM_PFX(gPatchSmmInitStack):
     call    ASM_PFX(SmmInitHandler)
+
+    mov     eax, RSB_STUFF_ENTRIES / 2
+@Unroll1:
+    call    @Unroll2
+@SpecTrap1:
+    pause
+    lfence
+    jmp     @SpecTrap1
+@Unroll2:
+    call    @StuffLoop
+@SpecTrap2:
+    pause
+    lfence
+    jmp     @SpecTrap2
+@StuffLoop:
+    dec     eax
+    jnz     @Unroll1
+    add     esp, RSB_STUFF_ENTRIES * 4     ; Restore the stack pointer
+
     rsm
 
 BITS 16
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmiEntry.nasm b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmiEntry.nasm
index 97c7b01d0d..b955fa1cf1 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmiEntry.nasm
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmiEntry.nasm
@@ -48,6 +48,8 @@
 %define TSS_SEGMENT 0x40
 %define GDT_SIZE 0x50
 
+%define RSB_STUFF_ENTRIES 0x20
+
 extern ASM_PFX(SmiRendezvous)
 extern ASM_PFX(gSmiHandlerIdtr)
 extern ASM_PFX(CpuSmmDebugEntry)
@@ -217,6 +219,24 @@ _SmiHandler:
     wrmsr
 
 .1:
+    mov     rax, RSB_STUFF_ENTRIES / 2
+@Unroll1:
+    call    @Unroll2
+@SpecTrap1:
+    pause
+    lfence
+    jmp     @SpecTrap1
+@Unroll2:
+    call    @StuffLoop
+@SpecTrap2:
+    pause
+    lfence
+    jmp     @SpecTrap2
+@StuffLoop:
+    dec     rax
+    jnz     @Unroll1
+    add     rsp, RSB_STUFF_ENTRIES * 8     ; Restore the stack pointer
+
     rsm
 
 ASM_PFX(gcSmiHandlerSize)    DW      $ - _SmiEntryPoint
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmInit.nasm b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmInit.nasm
index 0b0c3f28e5..bff14e809b 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmInit.nasm
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmInit.nasm
@@ -34,6 +34,8 @@ global ASM_PFX(gPatchSmmRelocationOriginalAddressPtr32)
 
 %define LONG_MODE_CS 0x38
 
+%define RSB_STUFF_ENTRIES 0x20
+
     DEFAULT REL
     SECTION .text
 
@@ -101,6 +103,24 @@ ASM_PFX(gPatchSmmInitStack):
     movdqa  xmm4, [rsp + 0x40]
     movdqa  xmm5, [rsp + 0x50]
 
+    mov     rax, RSB_STUFF_ENTRIES / 2
+@Unroll1:
+    call    @Unroll2
+@SpecTrap1:
+    pause
+    lfence
+    jmp     @SpecTrap1
+@Unroll2:
+    call    @StuffLoop
+@SpecTrap2:
+    pause
+    lfence
+    jmp     @SpecTrap2
+@StuffLoop:
+    dec     rax
+    jnz     @Unroll1
+    add     rsp, RSB_STUFF_ENTRIES * 8     ; Restore the stack pointer
+
     rsm
 
 BITS 16
-- 
2.12.0.windows.1