From d1fa366ba25b8f2ac909384c64961e92a5f00f4f Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Fri, 17 May 2024 13:50:08 +0300
Subject: [PATCH] Ring3: Initialized DxeRing3 with Supervisor privileges.

---
 MdeModulePkg/Core/Dxe/DxeMain.h               | 14 +++++++
 MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c | 38 +++++++++++++++++++
 .../Core/Dxe/SysCall/Initialization.c         |  6 +++
 3 files changed, 58 insertions(+)

diff --git a/MdeModulePkg/Core/Dxe/DxeMain.h b/MdeModulePkg/Core/Dxe/DxeMain.h
index b73721928c..2a78eb063d 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain.h
+++ b/MdeModulePkg/Core/Dxe/DxeMain.h
@@ -2635,6 +2635,20 @@ UnprotectUefiImage (
   IN EFI_DEVICE_PATH_PROTOCOL   *LoadedImageDevicePath
   );
 
+/**
+  Change UEFI image owner: Supervisor / Privileged or User / Unprivileged.
+
+  @param[in]  LoadedImage              The loaded image protocol
+  @param[in]  LoadedImageDevicePath    The loaded image device path protocol
+  @param[in]  IsUser                   Whether UEFI image record is User Image.
+**/
+VOID
+ChangeUefiImageRing (
+  IN EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage,
+  IN EFI_DEVICE_PATH_PROTOCOL   *LoadedImageDevicePath,
+  IN BOOLEAN                    IsUser
+  );
+
 /**
   ExitBootServices Callback function for memory protection.
 **/
diff --git a/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c b/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
index 4be5f37fbf..1652a59c68 100644
--- a/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
+++ b/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
@@ -294,6 +294,44 @@ UnprotectUefiImage (
   }
 }
 
+/**
+  Change UEFI image owner: Supervisor / Privileged or User / Unprivileged.
+
+  @param[in]  LoadedImage              The loaded image protocol
+  @param[in]  LoadedImageDevicePath    The loaded image device path protocol
+  @param[in]  IsUser                   Whether UEFI image record is User Image.
+**/
+VOID
+ChangeUefiImageRing (
+  IN EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage,
+  IN EFI_DEVICE_PATH_PROTOCOL   *LoadedImageDevicePath,
+  IN BOOLEAN                    IsUser
+  )
+{
+  UEFI_IMAGE_RECORD          *ImageRecord;
+  LIST_ENTRY                 *ImageRecordLink;
+
+  for (ImageRecordLink = mProtectedImageRecordList.ForwardLink;
+        ImageRecordLink != &mProtectedImageRecordList;
+        ImageRecordLink = ImageRecordLink->ForwardLink)
+    {
+    ImageRecord = CR (
+                    ImageRecordLink,
+                    UEFI_IMAGE_RECORD,
+                    Link,
+                    UEFI_IMAGE_RECORD_SIGNATURE
+                    );
+
+    if (ImageRecord->StartAddress == (EFI_PHYSICAL_ADDRESS)(UINTN)LoadedImage->ImageBase) {
+      ASSERT (gCpu != NULL);
+
+      SetUefiImageProtectionAttributes (ImageRecord, IsUser);
+
+      return;
+    }
+  }
+}
+
 /**
   Return the EFI memory permission attribute associated with memory
   type 'MemoryType' under the configured DXE memory protection policy.
diff --git a/MdeModulePkg/Core/Dxe/SysCall/Initialization.c b/MdeModulePkg/Core/Dxe/SysCall/Initialization.c
index 3bf1d4d74c..0b108d34fa 100644
--- a/MdeModulePkg/Core/Dxe/SysCall/Initialization.c
+++ b/MdeModulePkg/Core/Dxe/SysCall/Initialization.c
@@ -50,9 +50,15 @@ InitializeRing3 (
   gRing3Data = (RING3_DATA *)(UINTN)Physical;
 
   CopyMem ((VOID *)gRing3Data, (VOID *)Image->Info.SystemTable, sizeof (EFI_SYSTEM_TABLE));
+  //
+  // Initialize DxeRing3 with Supervisor privileges.
+  //
+  ChangeUefiImageRing (&Image->Info, Image->LoadedImageDevicePath, FALSE);
 
   Status = Image->EntryPoint (ImageHandle, (EFI_SYSTEM_TABLE *)gRing3Data);
 
+  ChangeUefiImageRing (&Image->Info, Image->LoadedImageDevicePath, TRUE);
+
   gRing3EntryPoint = gRing3Data->EntryPoint;
 
   gRing3Data->SystemTable.BootServices    = gRing3Data->BootServices;