From d6d858c4b7f0311f1c4d3b5d410ce87b1a261538 Mon Sep 17 00:00:00 2001
From: xli24 <xli24@6f19259b-4bc3-4df7-8a09-765794883524>
Date: Fri, 30 Jul 2010 01:49:51 +0000
Subject: Fix the risk of AP stack conflict.

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@10714 6f19259b-4bc3-4df7-8a09-765794883524
---
 .../X64/AsmInclude.inc                             |  1 +
 .../X64/MpFuncs.S                                  | 45 ++++++++++++----------
 .../X64/MpFuncs.asm                                | 36 +++++++++--------
 3 files changed, 44 insertions(+), 38 deletions(-)

(limited to 'EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64')

diff --git a/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/AsmInclude.inc b/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/AsmInclude.inc
index d0269760a6..b0e611108f 100644
--- a/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/AsmInclude.inc
+++ b/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/AsmInclude.inc
@@ -23,5 +23,6 @@ GdtrLocation                  equ        LockLocation + 20h
 IdtrLocation                  equ        LockLocation + 2Ah
 BufferStartLocation           equ        LockLocation + 34h
 Cr3OffsetLocation             equ        LockLocation + 38h
+ProcessorNumberLocation       equ        LockLocation + 3Ch
 
 ;-------------------------------------------------------------------------------
diff --git a/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/MpFuncs.S b/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/MpFuncs.S
index 8efd2d183e..5ec90bdaac 100644
--- a/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/MpFuncs.S
+++ b/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/MpFuncs.S
@@ -24,6 +24,7 @@
 .equ                   IdtrLocation,              RendezvousFunnelProcEnd - RendezvousFunnelProcStart + 0x2A
 .equ                   BufferStartLocation,       RendezvousFunnelProcEnd - RendezvousFunnelProcStart + 0x34
 .equ                   Cr3OffsetLocation,         RendezvousFunnelProcEnd - RendezvousFunnelProcStart + 0x38
+.equ                   ProcessorNumberLocation,   RendezvousFunnelProcEnd - RendezvousFunnelProcStart + 0x38
 
 #-------------------------------------------------------------------------------------
 
@@ -128,31 +129,33 @@ LongModeStart:
         .byte       0x66
         movw        %ax,%ss
 
-        movl        %esi, %edi
-        addl        $LockLocation, %edi
-        movb        $NotVacantFlag, %al
-TestLock:
-        xchgb       (%edi), %al
-        cmpb        $NotVacantFlag, %al
-        jz          TestLock
-
-ProgramStack:
+        #
+        # ProgramStack
+        #
+        movl        $0x1b, %ecx
+        rdmsr
+        andl        $0xfffff000, %eax
+        addl        $0x20, %eax
+        movl        (%eax), %ebx
+        shrl        $24, %ebx
+        
+        xorq        %rcx, %rcx
+        movl        %esi,%edi
+        addl        $ProcessorNumberLocation, %edi
+        movl        (%edi, %ebx, 4), %ecx
 
-        movl        %esi, %edi
-        addl        $StackSizeLocation, %edi 
+        movl        %esi,%edi
+        addl        $StackSizeLocation, %edi
         movq        (%edi), %rax
-        movl        %esi, %edi
-        addl        $StackStartAddressLocation, %edi
-        addq        (%edi), %rax
-        movq        %rax, %rsp
-        movq        %rax, (%edi) 
+        incq        %rcx
+        mulq        %rcx
 
-Releaselock:
+        movl        %esi,%edi
+        addl        $StackStartAddressLocation, %edi
+        movq        (%edi), %rbx
+        addq        %rbx, %rax
 
-        movb        $VacantFlag, %al
-        movl        %esi, %edi
-        addl        $LockLocation, %edi
-        xchgb       (%edi), %al
+        movq        %rax, %rsp
 
         #
         # Call C Function
diff --git a/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/MpFuncs.asm b/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/MpFuncs.asm
index 91bdb22451..5666322486 100644
--- a/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/MpFuncs.asm
+++ b/EdkCompatibilityPkg/Compatibility/MpServicesOnFrameworkMpServicesThunk/X64/MpFuncs.asm
@@ -112,31 +112,33 @@ LongModeStart::
         mov         es,  ax
         mov         ss,  ax
 
+        ;
+        ; ProgramStack
+        ;
+        mov         ecx, 1bh                          ; Read IA32_APIC_BASE MSR
+        rdmsr
+        and         eax, 0fffff000h
+        add         eax, 20h
+        mov         ebx, dword ptr [eax]
+        shr         ebx, 24
+        
+        xor         rcx, rcx
         mov         edi, esi
-        add         edi, LockLocation
-        mov         al,  NotVacantFlag
-TestLock::
-        xchg        byte ptr [edi], al
-        cmp         al, NotVacantFlag
-        jz          TestLock
-
-ProgramStack::
+        add         edi, ProcessorNumberLocation
+        mov         ecx, dword ptr [edi + 4 * ebx]    ; RCX = CpuNumber
 
         mov         edi, esi
         add         edi, StackSizeLocation
         mov         rax, qword ptr [edi]
+        inc         rcx
+        mul         rcx                               ; RAX = StackSize * (CpuNumber + 1)
+
         mov         edi, esi
         add         edi, StackStartAddressLocation
-        add         rax, qword ptr [edi]
-        mov         rsp, rax
-        mov         qword ptr [edi], rax
-
-Releaselock::
+        mov         rbx, qword ptr [edi]
+        add         rax, rbx                          ; RAX = StackStart + StackSize * (CpuNumber + 1)
 
-        mov         al,  VacantFlag
-        mov         edi, esi
-        add         edi, LockLocation
-        xchg        byte ptr [edi], al
+        mov         rsp, rax
 
         ;
         ; Call C Function
-- 
cgit v1.2.3