6 files changed, 122 insertions, 0 deletions
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiEntry.nasm b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiEntry.nasm
index 509e7a0a66..6bbc339c53 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiEntry.nasm
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmiEntry.nasm
@@ -18,6 +18,8 @@
 ;
 ;-------------------------------------------------------------------------------
 
+%include "StuffRsb.inc"
+
 %define MSR_IA32_MISC_ENABLE 0x1A0
 %define MSR_EFER      0xc0000080
 %define MSR_EFER_XD   0x800
@@ -204,6 +206,7 @@ ASM_PFX(SmiHandler):
     wrmsr
 
 .7:
+    StuffRsb32
     rsm
 
 ASM_PFX(gcSmiHandlerSize): DW $ - _SmiEntryPoint
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmInit.nasm b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmInit.nasm
index 5ff3cd2e73..322b1ab556 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmInit.nasm
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/SmmInit.nasm
@@ -18,6 +18,8 @@
 ;
 ;-------------------------------------------------------------------------------
 
+%include "StuffRsb.inc"
+
 extern ASM_PFX(SmmInitHandler)
 extern ASM_PFX(mRebasedFlag)
 extern ASM_PFX(mSmmRelocationOriginalAddress)
@@ -75,6 +77,7 @@ BITS 32
     mov     esp, strict dword 0         ; source operand will be patched
 ASM_PFX(gPatchSmmInitStack):
     call    ASM_PFX(SmmInitHandler)
+    StuffRsb32
     rsm
 
 BITS 16
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/StuffRsb.inc b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/StuffRsb.inc
new file mode 100644
index 0000000000..14267c3fde
--- /dev/null
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/StuffRsb.inc
@@ -0,0 +1,55 @@
+;------------------------------------------------------------------------------
+;
+; Copyright (c) 2018, Intel Corporation. All rights reserved.<BR>
+; This program and the accompanying materials
+; are licensed and made available under the terms and conditions of the BSD License
+; which accompanies this distribution.  The full text of the license may be found at
+; http://opensource.org/licenses/bsd-license.php.
+;
+; THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+; WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+;
+; Abstract:
+;
+;   This file provides macro definitions for stuffing the Return Stack Buffer (RSB).
+;
+;------------------------------------------------------------------------------
+
+%define RSB_STUFF_ENTRIES 0x20
+
+;
+; parameters:
+; @param 1: register to use as counter (e.g. IA32:eax, X64:rax)
+; @param 2: stack pointer to restore   (IA32:esp, X64:rsp)
+; @param 3: the size of a stack frame  (IA32:4, X64:8)
+;
+%macro StuffRsb 3
+      mov     %1, RSB_STUFF_ENTRIES / 2
+  %%Unroll1:
+      call    %%Unroll2
+  %%SpecTrap1:
+      pause
+      lfence
+      jmp     %%SpecTrap1
+  %%Unroll2:
+      call    %%StuffLoop
+  %%SpecTrap2:
+      pause
+      lfence
+      jmp     %%SpecTrap2
+  %%StuffLoop:
+      dec     %1
+      jnz     %%Unroll1
+      add     %2, RSB_STUFF_ENTRIES * %3 ; Restore the stack pointer
+%endmacro
+
+;
+; RSB stuffing macros for IA32 and X64
+;
+%macro StuffRsb32 0
+      StuffRsb     eax, esp, 4
+%endmacro
+
+%macro StuffRsb64 0
+      StuffRsb     rax, rsp, 8
+%endmacro
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmiEntry.nasm b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmiEntry.nasm
index 97c7b01d0d..315d0f8670 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmiEntry.nasm
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmiEntry.nasm
@@ -18,6 +18,8 @@
 ;
 ;-------------------------------------------------------------------------------
 
+%include "StuffRsb.inc"
+
 ;
 ; Variables referrenced by C code
 ;
@@ -217,6 +219,7 @@ _SmiHandler:
     wrmsr
 
 .1:
+    StuffRsb64
     rsm
 
 ASM_PFX(gcSmiHandlerSize)    DW      $ - _SmiEntryPoint
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmInit.nasm b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmInit.nasm
index 0b0c3f28e5..24357d5870 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmInit.nasm
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmInit.nasm
@@ -18,6 +18,8 @@
 ;
 ;-------------------------------------------------------------------------------
 
+%include "StuffRsb.inc"
+
 extern ASM_PFX(SmmInitHandler)
 extern ASM_PFX(mRebasedFlag)
 extern ASM_PFX(mSmmRelocationOriginalAddress)
@@ -101,6 +103,7 @@ ASM_PFX(gPatchSmmInitStack):
     movdqa  xmm4, [rsp + 0x40]
     movdqa  xmm5, [rsp + 0x50]
 
+    StuffRsb64
     rsm
 
 BITS 16
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/StuffRsb.inc b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/StuffRsb.inc
new file mode 100644
index 0000000000..14267c3fde
--- /dev/null
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/StuffRsb.inc
@@ -0,0 +1,55 @@
+;------------------------------------------------------------------------------
+;
+; Copyright (c) 2018, Intel Corporation. All rights reserved.<BR>
+; This program and the accompanying materials
+; are licensed and made available under the terms and conditions of the BSD License
+; which accompanies this distribution.  The full text of the license may be found at
+; http://opensource.org/licenses/bsd-license.php.
+;
+; THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+; WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+;
+; Abstract:
+;
+;   This file provides macro definitions for stuffing the Return Stack Buffer (RSB).
+;
+;------------------------------------------------------------------------------
+
+%define RSB_STUFF_ENTRIES 0x20
+
+;
+; parameters:
+; @param 1: register to use as counter (e.g. IA32:eax, X64:rax)
+; @param 2: stack pointer to restore   (IA32:esp, X64:rsp)
+; @param 3: the size of a stack frame  (IA32:4, X64:8)
+;
+%macro StuffRsb 3
+      mov     %1, RSB_STUFF_ENTRIES / 2
+  %%Unroll1:
+      call    %%Unroll2
+  %%SpecTrap1:
+      pause
+      lfence
+      jmp     %%SpecTrap1
+  %%Unroll2:
+      call    %%StuffLoop
+  %%SpecTrap2:
+      pause
+      lfence
+      jmp     %%SpecTrap2
+  %%StuffLoop:
+      dec     %1
+      jnz     %%Unroll1
+      add     %2, RSB_STUFF_ENTRIES * %3 ; Restore the stack pointer
+%endmacro
+
+;
+; RSB stuffing macros for IA32 and X64
+;
+%macro StuffRsb32 0
+      StuffRsb     eax, esp, 4
+%endmacro
+
+%macro StuffRsb64 0
+      StuffRsb     rax, rsp, 8
+%endmacro