From 0fe43214e33e4276741872a783528d896d82e785 Mon Sep 17 00:00:00 2001
From: qhuang8 <qhuang8@6f19259b-4bc3-4df7-8a09-765794883524>
Date: Fri, 17 Apr 2009 07:38:57 +0000
Subject: Save segment registers on stack in case the thunk code assembly calls
 CF9 soft reset and the x64 registers get cleared.

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@8123 6f19259b-4bc3-4df7-8a09-765794883524
---
 MdePkg/Library/BaseLib/X64/Thunk16.S   | 34 ++++++++++++++++++++++------------
 MdePkg/Library/BaseLib/X64/Thunk16.asm | 30 ++++++++++++++++++++----------
 2 files changed, 42 insertions(+), 22 deletions(-)

diff --git a/MdePkg/Library/BaseLib/X64/Thunk16.S b/MdePkg/Library/BaseLib/X64/Thunk16.S
index b829685b43..f9770b258c 100644
--- a/MdePkg/Library/BaseLib/X64/Thunk16.S
+++ b/MdePkg/Library/BaseLib/X64/Thunk16.S
@@ -136,10 +136,13 @@ SavedCr4:   .space  4
 SavedCr0:    .space      4
     movq    %rax, %cr0
     .byte 0x66,0xea                     # jmp far cs:L_64Bit
-L_64Eip:      .space      4
+L_64Eip:     .space      4
 SavedCs:     .space      2
 L_64BitCode: 
-    movq     %r8, %rsp
+    .byte   0x90
+    .byte   0x67,0xbc                  # mov esp, imm32
+SavedSp:    .space 4                   # restore stack
+    nop
     ret
 
 _EntryPoint: .long      ASM_PFX(ToUserCode) - ASM_PFX(m16Start)
@@ -227,7 +230,6 @@ ASM_PFX(_32Data):
 #   IN OUT  VOID                      *Transition
 #   );
 #------------------------------------------------------------------------------
-# MISMATCH: "InternalAsmThunk16  PROC    USES    rbp rbx rsi rdi"
 
 .globl ASM_PFX(InternalAsmThunk16)
 ASM_PFX(InternalAsmThunk16):
@@ -236,9 +238,13 @@ ASM_PFX(InternalAsmThunk16):
     pushq   %rsi
     pushq   %rdi
     
-    movl    %ds, %r10d                  # r9 ~ r11 are not accessible in 16-bit
-    movl    %es, %r11d                  # so use them for saving seg registers
-    movl    %ss, %r9d
+    movq    %ds, %rbx
+    pushq   %rbx      # Save ds segment register on the stack
+    movq    %es, %rbx
+    pushq   %rbx      # Save es segment register on the stack
+    movq    %ss, %rbx
+    pushq   %rbx      # Save ss segment register on the stack
+
     .byte   0x0f, 0xa0                  #push   fs
     .byte   0x0f, 0xa8                  #push   gs
     movq    %rcx, %rsi
@@ -259,7 +265,7 @@ ASM_PFX(InternalAsmThunk16):
     lea     (_BackFromUserCode - ASM_PFX(m16Start))(%rdx), %ax
     stosl                               # [edi] <- return address of user code
     sgdt    (SavedGdt - SavedCr4)(%rcx) 
-    sidt    0x38(%rsp)
+    sidt    0x50(%rsp)
     movq    %cr0, %rax
     movl    %eax, (SavedCr0 - SavedCr4)(%rcx)
     andl    $0x7ffffffe,%eax            # clear PE, PG bits
@@ -277,18 +283,22 @@ ASM_PFX(InternalAsmThunk16):
     pushq   %r8
     movl    %cs, %r8d
     movw    %r8w, (SavedCs - SavedCr4)(%rcx)
-    movq     %rsp, %r8
+    movl    %esp, (SavedSp - SavedCr4)(%rcx)
     .byte   0xff, 0x69                  #  jmp (_EntryPoint - SavedCr4)(%rcx)
     .byte   _EntryPoint - SavedCr4
 L_RetFromRealMode: 
     popfq
-    lidt    0x38(%rsp)
+    lidt    0x50(%rsp)
     lea     -IA32_REGS_SIZE(%rbp), %eax
     .byte 0x0f, 0xa9                    # pop gs
     .byte 0x0f, 0xa1                    # pop fs
-    movl    %r9d, %ss
-    movl    %r11d, %es
-    movl    %r10d, %ds
+    
+    popq     %rbx
+    movq     %rbx, %ss
+    popq     %rbx
+    movq     %rbx, %es
+    popq     %rbx
+    movq     %rbx, %ds
     
     popq    %rdi
     popq    %rsi
diff --git a/MdePkg/Library/BaseLib/X64/Thunk16.asm b/MdePkg/Library/BaseLib/X64/Thunk16.asm
index d93af44597..cef63d98e1 100644
--- a/MdePkg/Library/BaseLib/X64/Thunk16.asm
+++ b/MdePkg/Library/BaseLib/X64/Thunk16.asm
@@ -140,7 +140,10 @@ SavedCr0    DD      ?
 @64Eip      DD      ?
 SavedCs     DW      ?
 @64BitCode:
-    mov     rsp, r8                     ; restore stack
+    db      090h 
+    db      067h, 0bch                 ; mov esp, imm32
+SavedSp     DD   ?                     ; restore stack
+    nop
     ret
 _BackFromUserCode   ENDP
 
@@ -230,9 +233,13 @@ GDT_SIZE = $ - _NullSeg
 ;   );
 ;------------------------------------------------------------------------------
 InternalAsmThunk16  PROC    USES    rbp rbx rsi rdi
-    mov     r10d, ds                    ; r9 ~ r11 are not accessible in 16-bit
-    mov     r11d, es                    ; so use them for saving seg registers
-    mov     r9d, ss
+    mov     rbx, ds
+    push    rbx          ; Save ds segment register on the stack
+    mov     rbx, es
+    push    rbx          ; Save es segment register on the stack
+    mov     rbx, ss
+    push    rbx          ; Save ss segment register on the stack
+    
     push    fs
     push    gs
     mov     rsi, rcx
@@ -252,7 +259,7 @@ InternalAsmThunk16  PROC    USES    rbp rbx rsi rdi
     lea     ax, [rdx + (_BackFromUserCode - m16Start)]  ; offset address
     stosd                               ; [edi] <- return address of user code
     sgdt    fword ptr [rcx + (SavedGdt - SavedCr4)]
-    sidt    fword ptr [rsp + 38h]       ; save IDT stack in argument space
+    sidt    fword ptr [rsp + 50h]       ; save IDT stack in argument space
     mov     rax, cr0
     mov     [rcx + (SavedCr0 - SavedCr4)], eax
     and     eax, 7ffffffeh              ; clear PE, PG bits
@@ -270,17 +277,20 @@ InternalAsmThunk16  PROC    USES    rbp rbx rsi rdi
     push    r8
     mov     r8d, cs
     mov     [rcx + (SavedCs - SavedCr4)], r8w
-    mov     r8, rsp
+    mov     [rcx + (SavedSp - SavedCr4)], esp
     jmp     fword ptr [rcx + (_EntryPoint - SavedCr4)]
 @RetFromRealMode:
     popfq
-    lidt    fword ptr [rsp + 38h]       ; restore protected mode IDTR
+    lidt    fword ptr [rsp + 50h]       ; restore protected mode IDTR
     lea     eax, [rbp - sizeof (IA32_REGS)]
     pop     gs
     pop     fs
-    mov     ss, r9d
-    mov     es, r11d
-    mov     ds, r10d
+    pop     rbx
+    mov     ss, rbx
+    pop     rbx
+    mov     es, rbx
+    pop     rbx
+    mov     ds, rbx
     ret
 InternalAsmThunk16  ENDP
 
-- 
cgit v1.2.3