[PATCH] UefiCpuPkg/CpuDxe: Fix multiple entries of RT_CODE in memory map

[Jian J Wang <jian.j.wang@intel.com>]

More than one entry of RT_CODE memory might cause boot problem for some
old OSs. This patch will fix this issue to keep OS compatibility as much
as possible.

Cc: Eric Dong <eric.dong@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Jian J Wang <jian.j.wang@intel.com>
---
 UefiCpuPkg/CpuDxe/CpuPageTable.c | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/UefiCpuPkg/CpuDxe/CpuPageTable.c b/UefiCpuPkg/CpuDxe/CpuPageTable.c
index d312eb66f8..0802464b9d 100644
--- a/UefiCpuPkg/CpuDxe/CpuPageTable.c
+++ b/UefiCpuPkg/CpuDxe/CpuPageTable.c
@@ -829,6 +829,15 @@ RefreshGcdMemoryAttributesFromPaging (
     // Sync real page attributes to GCD
     BaseAddress       = MemorySpaceMap[Index].BaseAddress;
     MemorySpaceLength = MemorySpaceMap[Index].Length;
+    Capabilities      = MemorySpaceMap[Index].Capabilities |
+                        EFI_MEMORY_PAGETYPE_MASK;
+    Status = gDS->SetMemorySpaceCapabilities (
+                    BaseAddress,
+                    MemorySpaceLength,
+                    Capabilities
+                    );
+    ASSERT_EFI_ERROR (Status);
+
     while (MemorySpaceLength > 0) {
       if (PageLength == 0) {
         PageEntry = GetPageTableEntry (&PagingContext, BaseAddress, &PageAttribute);
@@ -846,7 +855,6 @@ RefreshGcdMemoryAttributesFromPaging (
         if (Attributes != (MemorySpaceMap[Index].Attributes & EFI_MEMORY_PAGETYPE_MASK)) {
           DoUpdate = TRUE;
           Attributes |= (MemorySpaceMap[Index].Attributes & ~EFI_MEMORY_PAGETYPE_MASK);
-          Capabilities = Attributes | MemorySpaceMap[Index].Capabilities;
         } else {
           DoUpdate = FALSE;
         }
@@ -854,8 +862,8 @@ RefreshGcdMemoryAttributesFromPaging (
 
       Length = MIN (PageLength, MemorySpaceLength);
       if (DoUpdate) {
-        gDS->SetMemorySpaceCapabilities (BaseAddress, Length, Capabilities);
-        gDS->SetMemorySpaceAttributes (BaseAddress, Length, Attributes);
+        Status = gDS->SetMemorySpaceAttributes (BaseAddress, Length, Attributes);
+        ASSERT_EFI_ERROR (Status);
         DEBUG ((DEBUG_INFO, "Update memory space attribute: [%02d] %016lx - %016lx (%08lx -> %08lx)\r\n",
                              Index, BaseAddress, BaseAddress + Length - 1,
                              MemorySpaceMap[Index].Attributes, Attributes));
-- 
2.14.1.windows.1