From 33f15b44deee88a42041bc97b79ecf2c642783e6 Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Mon, 25 Dec 2023 16:08:53 +0300
Subject: [PATCH] Ring3: Added EFI_MEMORY_USER attribute.

---
 MdeModulePkg/Core/Dxe/Image/Image.c           |  2 +-
 MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c | 13 ++++++++----
 MdePkg/Include/Uefi/UefiSpec.h                |  7 ++++++-
 UefiCpuPkg/Library/CpuArchLib/CpuPageTable.c  | 21 +++++++++++++++++++
 4 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/Image/Image.c b/MdeModulePkg/Core/Dxe/Image/Image.c
index 9f0f7c0f64..5065c01fb3 100644
--- a/MdeModulePkg/Core/Dxe/Image/Image.c
+++ b/MdeModulePkg/Core/Dxe/Image/Image.c
@@ -828,7 +828,7 @@ Done:
   if (DstBufAlocated) {
     ZeroMem ((VOID *)(UINTN)BufferAddress, EFI_PAGES_TO_SIZE (Image->NumberOfPages));
     FreeAlignedPages ((VOID *)(UINTN)BufferAddress, Image->NumberOfPages);
-    Image->ImageBasePage             = 0;
+    Image->ImageBasePage = 0;
   }
 
   if (RelocationData != NULL) {
diff --git a/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c b/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
index 5b0fe44a23..8c173eaf6a 100644
--- a/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
+++ b/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
@@ -94,7 +94,8 @@ SetUefiImageMemoryAttributes (
 **/
 VOID
 SetUefiImageProtectionAttributes (
-  IN UEFI_IMAGE_RECORD  *ImageRecord
+  IN UEFI_IMAGE_RECORD  *ImageRecord,
+  IN BOOLEAN            IsUser
   )
 {
   UEFI_IMAGE_RECORD_SEGMENT  *ImageRecordSegment;
@@ -107,7 +108,7 @@ SetUefiImageProtectionAttributes (
     SetUefiImageMemoryAttributes (
       SectionAddress,
       ImageRecordSegment->Size,
-      ImageRecordSegment->Attributes
+      IsUser ? ImageRecordSegment->Attributes | (UINT32)EFI_MEMORY_USER : ImageRecordSegment->Attributes
       );
 
     SectionAddress += ImageRecordSegment->Size;
@@ -231,7 +232,11 @@ ProtectUefiImage (
     //
     // CPU ARCH present. Update memory attribute directly.
     //
-    SetUefiImageProtectionAttributes (ImageRecord);
+    if (AsciiStrStr (PdbPointer, "Ntfs") != NULL) {
+      SetUefiImageProtectionAttributes (ImageRecord, TRUE);
+    } else {
+      SetUefiImageProtectionAttributes (ImageRecord, FALSE);
+    }
   }
 
 Finish:
@@ -663,7 +668,7 @@ MemoryProtectionCpuArchProtocolNotify (
     //
     // CPU ARCH present. Update memory attribute directly.
     //
-    SetUefiImageProtectionAttributes (ImageRecord);
+    SetUefiImageProtectionAttributes (ImageRecord, FALSE);
   }
 
 Done:
diff --git a/MdePkg/Include/Uefi/UefiSpec.h b/MdePkg/Include/Uefi/UefiSpec.h
index bb395a2cb2..cfaa8363f1 100644
--- a/MdePkg/Include/Uefi/UefiSpec.h
+++ b/MdePkg/Include/Uefi/UefiSpec.h
@@ -114,6 +114,11 @@ typedef enum {
 // capable of being dynamically removed from the platform at runtime.
 //
 #define EFI_MEMORY_HOT_PLUGGABLE  0x0000000000100000
+//
+// If this flag is set, the memory region contains user code or data.
+// If this flag is clear, the memory region contains supervisor code or data.
+//
+#define EFI_MEMORY_USER  0x0000000000200000ULL
 
 //
 // Runtime memory attribute
@@ -140,7 +145,7 @@ typedef enum {
 //
 #define EFI_CACHE_ATTRIBUTE_MASK   (EFI_MEMORY_UC | EFI_MEMORY_WC | EFI_MEMORY_WT | EFI_MEMORY_WB | EFI_MEMORY_UCE | EFI_MEMORY_WP)
 #define EFI_MEMORY_ACCESS_MASK     (EFI_MEMORY_RP | EFI_MEMORY_XP | EFI_MEMORY_RO)
-#define EFI_MEMORY_ATTRIBUTE_MASK  (EFI_MEMORY_ACCESS_MASK | EFI_MEMORY_SP | EFI_MEMORY_CPU_CRYPTO)
+#define EFI_MEMORY_ATTRIBUTE_MASK  (EFI_MEMORY_ACCESS_MASK | EFI_MEMORY_SP | EFI_MEMORY_CPU_CRYPTO | EFI_MEMORY_USER)
 
 ///
 /// Memory descriptor version number.
diff --git a/UefiCpuPkg/Library/CpuArchLib/CpuPageTable.c b/UefiCpuPkg/Library/CpuArchLib/CpuPageTable.c
index 1092a30d28..3844030398 100644
--- a/UefiCpuPkg/Library/CpuArchLib/CpuPageTable.c
+++ b/UefiCpuPkg/Library/CpuArchLib/CpuPageTable.c
@@ -467,6 +467,27 @@ ConvertPageEntryAttribute (
     }
   }
 
+  if ((Attributes & EFI_MEMORY_USER) != 0) {
+    switch (PageAction) {
+      case PageActionAssign:
+      case PageActionSet:
+        NewPageEntry |= IA32_PG_U;
+        break;
+      case PageActionClear:
+        NewPageEntry &= ~(UINT64)IA32_PG_U;
+        break;
+    }
+  } else {
+    switch (PageAction) {
+      case PageActionAssign:
+        NewPageEntry &= ~(UINT64)IA32_PG_U;
+        break;
+      case PageActionSet:
+      case PageActionClear:
+        break;
+    }
+  }
+
   GetPagingDetails (&PagingContext->ContextData, NULL, &PageAttributes);
 
   if ((*PageAttributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_XD_ACTIVATED) != 0) {