From ada8ce2cb13f054cab6d4d2ee162d6b9e23043aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mario=20B=C4=83l=C4=83nic=C4=83?=
 <mariobalanica02@gmail.com>
Date: Thu, 5 Dec 2024 13:28:08 +0200
Subject: [PATCH] Add upstream EFI Memory Attribute Protocol state manager
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Mario Bălănică <mariobalanica02@gmail.com>
---
 .../Silicon/Rockchip/FvMainModules.fdf.inc    |  7 +-
 .../RK3588/Drivers/RK3588Dxe/RK3588Dxe.c      | 74 +------------------
 .../RK3588/Drivers/RK3588Dxe/RK3588Dxe.inf    |  1 -
 .../Silicon/Rockchip/Rockchip.dsc.inc         |  5 ++
 4 files changed, 12 insertions(+), 75 deletions(-)

diff --git a/edk2-rockchip/Silicon/Rockchip/FvMainModules.fdf.inc b/edk2-rockchip/Silicon/Rockchip/FvMainModules.fdf.inc
index 6c0a2645..e3b7b22e 100644
--- a/edk2-rockchip/Silicon/Rockchip/FvMainModules.fdf.inc
+++ b/edk2-rockchip/Silicon/Rockchip/FvMainModules.fdf.inc
@@ -1,7 +1,7 @@
 ## @file
 #
 #  Copyright (c) 2021-2022, Rockchip Limited. All rights reserved.
-#  Copyright (c) 2023, Mario Bălănică <mariobalanica02@gmail.com>
+#  Copyright (c) 2023-2024, Mario Bălănică <mariobalanica02@gmail.com>
 #
 #  SPDX-License-Identifier: BSD-2-Clause-Patent
 #
@@ -205,6 +205,11 @@
   INF Drivers/ASIX/Bus/Usb/UsbNetworking/Ax88772c/Ax88772c.inf
 !endif
 
+  #
+  # EFI Memory Attribute Protocol state manager (rhboot shim workaround)
+  #
+  INF EmbeddedPkg/Drivers/MemoryAttributeManagerDxe/MemoryAttributeManagerDxe.inf
+
   #
   # Bds
   #
diff --git a/edk2-rockchip/Silicon/Rockchip/RK3588/Drivers/RK3588Dxe/RK3588Dxe.c b/edk2-rockchip/Silicon/Rockchip/RK3588/Drivers/RK3588Dxe/RK3588Dxe.c
index 9c1bd075..c9a72f73 100644
--- a/edk2-rockchip/Silicon/Rockchip/RK3588/Drivers/RK3588Dxe/RK3588Dxe.c
+++ b/edk2-rockchip/Silicon/Rockchip/RK3588/Drivers/RK3588Dxe/RK3588Dxe.c
@@ -1,14 +1,13 @@
 /** @file
 *
 *  Copyright (c) 2021, Rockchip Limited. All rights reserved.
-*  Copyright (c) 2023, Mario Bălănică <mariobalanica02@gmail.com>
+*  Copyright (c) 2023-2024, Mario Bălănică <mariobalanica02@gmail.com>
 *
 *  SPDX-License-Identifier: BSD-2-Clause-Patent
 *
 **/
 
 #include <Uefi.h>
-#include <Protocol/MemoryAttribute.h>
 #include <Library/ArmLib.h>
 #include <Library/CacheMaintenanceLib.h>
 #include <Library/DebugLib.h>
@@ -408,69 +407,6 @@ RK3588NotifyReadyToBoot (
   gBS->CloseEvent (Event);
 }
 
-
-/**
-  This function uninstalls the recently added EFI_MEMORY_ATTRIBUTE_PROTOCOL
-  to workaround older versions of OS loaders/shims using it incorrectly and
-  throwing a Synchronous Exception.
-
-  See:
-    - https://github.com/microsoft/mu_silicon_arm_tiano/issues/124
-    - https://edk2.groups.io/g/devel/topic/99631663
-**/
-STATIC
-VOID
-EFIAPI
-UninstallMemoryAttributeProtocol (
-  VOID
-  )
-{
-  EFI_STATUS                      Status;
-  EFI_HANDLE                      *Handles;
-  UINTN                           HandleCount;
-  EFI_MEMORY_ATTRIBUTE_PROTOCOL   *MemoryAttributeProtocol;
-
-  Status = gBS->LocateHandleBuffer (
-                  ByProtocol,
-                  &gEfiMemoryAttributeProtocolGuid,
-                  NULL,
-                  &HandleCount,
-                  &Handles
-                  );
-  ASSERT_EFI_ERROR (Status);
-  ASSERT (HandleCount == 1);
-
-  Status = gBS->HandleProtocol (
-                  Handles[0],
-                  &gEfiMemoryAttributeProtocolGuid,
-                  (VOID **)&MemoryAttributeProtocol
-                  );
-  ASSERT_EFI_ERROR (Status);
-
-  Status = gBS->UninstallMultipleProtocolInterfaces (
-                  Handles[0],
-                  &gEfiMemoryAttributeProtocolGuid,
-                  MemoryAttributeProtocol,
-                  NULL
-                  );
-  ASSERT_EFI_ERROR (Status);
-
-  gBS->FreePool (Handles);
-}
-
-STATIC
-VOID
-EFIAPI
-RK3588NotifyEndOfDxe (
-  IN EFI_EVENT  Event,
-  IN VOID       *Context
-  )
-{
-  gBS->CloseEvent (Event);
-
-  UninstallMemoryAttributeProtocol ();
-}
-
 EFI_STATUS
 EFIAPI
 RK3588EntryPoint (
@@ -501,14 +437,6 @@ RK3588EntryPoint (
                                         &Event);
   ASSERT_EFI_ERROR (Status);
 
-  Status = gBS->CreateEventEx (EVT_NOTIFY_SIGNAL,
-                    TPL_CALLBACK,
-                    RK3588NotifyEndOfDxe,
-                    NULL,
-                    &gEfiEndOfDxeEventGroupGuid,
-                    &Event);
-  ASSERT_EFI_ERROR (Status);
-
   Status = RK3588InitPeripherals ();
   if (EFI_ERROR (Status)) {
     return Status;
diff --git a/edk2-rockchip/Silicon/Rockchip/RK3588/Drivers/RK3588Dxe/RK3588Dxe.inf b/edk2-rockchip/Silicon/Rockchip/RK3588/Drivers/RK3588Dxe/RK3588Dxe.inf
index 93d266fe..319bc41b 100644
--- a/edk2-rockchip/Silicon/Rockchip/RK3588/Drivers/RK3588Dxe/RK3588Dxe.inf
+++ b/edk2-rockchip/Silicon/Rockchip/RK3588/Drivers/RK3588Dxe/RK3588Dxe.inf
@@ -113,7 +113,6 @@
 
 [Guids]
   gRK3588DxeFormSetGuid
-  gEfiEndOfDxeEventGroupGuid
 
 [Depex]
   TRUE
diff --git a/edk2-rockchip/Silicon/Rockchip/Rockchip.dsc.inc b/edk2-rockchip/Silicon/Rockchip/Rockchip.dsc.inc
index ceea54aa..2c3ea6a5 100644
--- a/edk2-rockchip/Silicon/Rockchip/Rockchip.dsc.inc
+++ b/edk2-rockchip/Silicon/Rockchip/Rockchip.dsc.inc
@@ -734,6 +734,11 @@ FspiLib|Silicon/Rockchip/Library/FspiLib/FspiLib.inf
   Drivers/ASIX/Bus/Usb/UsbNetworking/Ax88772c/Ax88772c.inf
 !endif
 
+  #
+  # EFI Memory Attribute Protocol state manager (rhboot shim workaround)
+  #
+  EmbeddedPkg/Drivers/MemoryAttributeManagerDxe/MemoryAttributeManagerDxe.inf
+
   #
   # Bds
   #