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

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

> v2
> a. Fix an issue which will cause setting capability failure if size is smaller
>    than a page.

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.

More detailed information, please refer to
    https://bugzilla.tianocore.org/show_bug.cgi?id=753

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

diff --git a/UefiCpuPkg/CpuDxe/CpuPageTable.c b/UefiCpuPkg/CpuDxe/CpuPageTable.c
index d312eb66f8..4a7827ebc9 100644
--- a/UefiCpuPkg/CpuDxe/CpuPageTable.c
+++ b/UefiCpuPkg/CpuDxe/CpuPageTable.c
@@ -809,7 +809,9 @@ RefreshGcdMemoryAttributesFromPaging (
   PageLength    = 0;
 
   for (Index = 0; Index < NumberOfDescriptors; Index++) {
-    if (MemorySpaceMap[Index].GcdMemoryType == EfiGcdMemoryTypeNonExistent) {
+    if (MemorySpaceMap[Index].GcdMemoryType == EfiGcdMemoryTypeNonExistent
+        || (MemorySpaceMap[Index].BaseAddress & EFI_PAGE_MASK) != 0
+        || (MemorySpaceMap[Index].Length & EFI_PAGE_MASK) != 0) {
       continue;
     }
 
@@ -829,6 +831,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 +857,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 +864,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