From d19bee4d2d438eb557ce7cfc1f6f02f606f78180 Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Mon, 19 Feb 2024 15:00:10 +0300
Subject: [PATCH] Ring3: Some refactoring.

---
 .../Core/Dxe/DxeRing3/Ring3UefiBootServices.c |  2 +-
 MdeModulePkg/Core/Dxe/Image/Image.c           | 11 +-------
 MdeModulePkg/Core/Dxe/SysCall/BootServices.c  | 25 +++++++++++--------
 .../Core/Dxe/SysCall/SupportedProtocols.c     | 19 +++++++++++++-
 .../Dxe/SysCall/X64/CoreBootServices.nasm     |  1 -
 5 files changed, 35 insertions(+), 23 deletions(-)

diff --git a/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c b/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c
index 84b30af06c..fb04ac7f32 100644
--- a/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c
+++ b/MdeModulePkg/Core/Dxe/DxeRing3/Ring3UefiBootServices.c
@@ -397,7 +397,7 @@ Ring3OpenProtocol (
              Attributes
              );
   if (EFI_ERROR (Status)) {
-    DEBUG ((DEBUG_ERROR, "Ring3: Failed to open protocol %g - %r\n", Protocol, Status));
+    // DEBUG ((DEBUG_ERROR, "Ring3: Failed to open protocol %g - %r\n", Protocol, Status));
     return Status;
   }
 
diff --git a/MdeModulePkg/Core/Dxe/Image/Image.c b/MdeModulePkg/Core/Dxe/Image/Image.c
index f9d1520705..229a7673e5 100644
--- a/MdeModulePkg/Core/Dxe/Image/Image.c
+++ b/MdeModulePkg/Core/Dxe/Image/Image.c
@@ -1698,7 +1698,7 @@ InitializeRing3 (
   //
   // Initialize MSR_IA32_STAR and MSR_IA32_LSTAR for SYSCALL and SYSRET.
   //
-  Msr = ((((UINT64)RING3_CODE64_SEL - 16) << 16) | (UINT64)RING0_CODE64_SEL) << 32;
+  Msr = (((((UINT64)RING3_CODE64_SEL - 16) | 3) << 16) | (UINT64)RING0_CODE64_SEL) << 32;
   AsmWriteMsr64 (MSR_IA32_STAR, Msr);
 
   Msr = (UINT64)(UINTN)CoreBootServices;
@@ -1837,15 +1837,6 @@ CoreStartImage (
       gCpu->GetMemoryAttributes (gCpu, (EFI_PHYSICAL_ADDRESS)Image->EntryPoint, &Attributes);
       ASSERT ((Attributes & EFI_MEMORY_USER) != 0);
 
-      //
-      // Necessary fix for ProcessLibraryConstructorList() -> DxeCcProbeLibConstructor()
-      //
-      SetUefiImageMemoryAttributes (
-        FixedPcdGet32 (PcdOvmfWorkAreaBase),
-        FixedPcdGet32 (PcdOvmfWorkAreaSize),
-        EFI_MEMORY_XP | EFI_MEMORY_USER
-        );
-
       Image->Status = GoToRing3 (
                         2,
                         (VOID *)Image->EntryPoint,
diff --git a/MdeModulePkg/Core/Dxe/SysCall/BootServices.c b/MdeModulePkg/Core/Dxe/SysCall/BootServices.c
index 51a24a71b4..c72242c3a1 100644
--- a/MdeModulePkg/Core/Dxe/SysCall/BootServices.c
+++ b/MdeModulePkg/Core/Dxe/SysCall/BootServices.c
@@ -155,10 +155,12 @@ CallBootService (
                       );
 
       DisableSMAP ();
-      Interface = AllocateRing3Copy (Interface, MemoryCoreSize, MemoryCoreSize);
-      if (Interface == NULL) {
-        EnableSMAP ();
-        return EFI_OUT_OF_RESOURCES;
+      if (Interface != NULL) {
+        Interface = AllocateRing3Copy (Interface, MemoryCoreSize, MemoryCoreSize);
+        if (Interface == NULL) {
+          EnableSMAP ();
+          return EFI_OUT_OF_RESOURCES;
+        }      
       }
 
       *(VOID **)CoreRbp->Argument3 = Interface;
@@ -197,13 +199,15 @@ CallBootService (
                       );
 
       DisableSMAP ();
-      Interface = AllocateRing3Copy (Interface, MemoryCoreSize, MemoryCoreSize);
-      if (Interface == NULL) {
-        EnableSMAP ();
-        return EFI_OUT_OF_RESOURCES;
-      }
+      if (Interface != NULL) {
+        Interface = AllocateRing3Copy (Interface, MemoryCoreSize, MemoryCoreSize);
+        if (Interface == NULL) {
+          EnableSMAP ();
+          return EFI_OUT_OF_RESOURCES;
+        }
 
-      FixInterface (CoreProtocol, Interface);
+        FixInterface (CoreProtocol, Interface);
+      }
 
       *(VOID **)CoreRbp->Argument3 = Interface;
       EnableSMAP ();
@@ -285,6 +289,7 @@ CallBootService (
 
       return Status;
     default:
+      DEBUG ((DEBUG_ERROR, "Ring0: Unknown syscall type.\n"));
       break;
   }
 
diff --git a/MdeModulePkg/Core/Dxe/SysCall/SupportedProtocols.c b/MdeModulePkg/Core/Dxe/SysCall/SupportedProtocols.c
index 4557536998..bf7f7b6200 100644
--- a/MdeModulePkg/Core/Dxe/SysCall/SupportedProtocols.c
+++ b/MdeModulePkg/Core/Dxe/SysCall/SupportedProtocols.c
@@ -44,7 +44,24 @@ GoToRing3 (
   VA_END (Marker);
   EnableSMAP ();
 
-  return CallRing3 (Input);
+  //
+  // Necessary fix for ProcessLibraryConstructorList() -> DxeCcProbeLibConstructor()
+  //
+  SetUefiImageMemoryAttributes (
+    FixedPcdGet32 (PcdOvmfWorkAreaBase),
+    FixedPcdGet32 (PcdOvmfWorkAreaSize),
+    EFI_MEMORY_XP | EFI_MEMORY_USER
+    );
+
+  Status = CallRing3 (Input);
+
+  SetUefiImageMemoryAttributes (
+    FixedPcdGet32 (PcdOvmfWorkAreaBase),
+    FixedPcdGet32 (PcdOvmfWorkAreaSize),
+    EFI_MEMORY_XP
+    );
+
+  return Status;
 }
 
 EFI_STATUS
diff --git a/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm b/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm
index 4f3071bb41..0987cf9320 100644
--- a/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm
+++ b/MdeModulePkg/Core/Dxe/SysCall/X64/CoreBootServices.nasm
@@ -197,7 +197,6 @@ ASM_PFX(CallRing3):
     add     rax, 8
 
     ; Set Data selectors
-    or      rax, 3H ; RPL = 3
     mov     ds, ax
     mov     es, ax
     mov     fs, ax