Move ACPI OS identification to its own driver

To support multiple handlers.

Also add Linux kernel image detection while I'm at it.

Signed-off-by: Mario Bălănică <[email protected]>

edk2-rockchip/Silicon/Rockchip/Drivers/ExitBootServicesHookDxe/ExitBootServicesHook.h

@@ -0,0 +1,50 @@
+/** @file
+ *
+ *  Copyright (c) 2024-2025, Mario Bălănică <[email protected]>
+ *
+ *  SPDX-License-Identifier: BSD-2-Clause-Patent
+ *
+ **/
+
+#ifndef __EXIT_BOOT_SERVICES_HOOK_H__
+#define __EXIT_BOOT_SERVICES_HOOK_H__
+
+#include <Library/BaseLib.h>
+#include <Library/DebugLib.h>
+#include <Library/UefiLib.h>
+#include <Protocol/ExitBootServicesOsNotify.h>
+
+typedef struct {
+  UINT32                           Signature;
+  LIST_ENTRY                       Link;
+  EXIT_BOOT_SERVICES_OS_HANDLER    Handler;
+} EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY;
+#define EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY_SIGNATURE  SIGNATURE_32('E', 'b', 'S', 'h')
+#define EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY_FROM_LINK(a)  \
+  CR (a, EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY, Link, EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY_SIGNATURE)
+
+typedef struct {
+  UINT32                                   Signature;
+  EXIT_BOOT_SERVICES_OS_NOTIFY_PROTOCOL    Notify;
+  LIST_ENTRY                               Handlers;
+} EXIT_BOOT_SERVICES_OS_NOTIFY_INSTANCE;
+#define EXIT_BOOT_SERVICES_OS_NOTIFY_INSTANCE_SIGNATURE  SIGNATURE_32('E', 'b', 'S', 'n')
+#define EXIT_BOOT_SERVICES_OS_NOTIFY_INSTANCE_FROM_THIS(a) \
+  CR (a, EXIT_BOOT_SERVICES_OS_NOTIFY_INSTANCE, Notify, EXIT_BOOT_SERVICES_OS_NOTIFY_INSTANCE_SIGNATURE)
+
+EFI_PHYSICAL_ADDRESS
+FindPeImageBase (
+  IN EFI_PHYSICAL_ADDRESS  Base
+  );
+
+EXIT_BOOT_SERVICES_OS_TYPE
+IdentifyOsType (
+  IN EFI_PHYSICAL_ADDRESS  OsLoaderAddress
+  );
+
+CHAR8 *
+OsTypeToString (
+  IN EXIT_BOOT_SERVICES_OS_TYPE  OsType
+  );
+
+#endif

edk2-rockchip/Silicon/Rockchip/Drivers/ExitBootServicesHookDxe/ExitBootServicesHookDxe.c

@@ -0,0 +1,165 @@
+/** @file
+ *
+ *  Copyright (c) 2024-2025, Mario Bălănică <[email protected]>
+ *
+ *  SPDX-License-Identifier: BSD-2-Clause-Patent
+ *
+ **/
+
+#include <Library/MemoryAllocationLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+
+#include "ExitBootServicesHook.h"
+
+STATIC EFI_EXIT_BOOT_SERVICES  mOriginalExitBootServices;
+
+STATIC
+EFI_STATUS
+EFIAPI
+RegisterHandler (
+  IN EXIT_BOOT_SERVICES_OS_NOTIFY_PROTOCOL  *This,
+  IN EXIT_BOOT_SERVICES_OS_HANDLER          Handler
+  )
+{
+  EXIT_BOOT_SERVICES_OS_NOTIFY_INSTANCE  *Instance;
+  LIST_ENTRY                             *Link;
+  EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY    *Entry;
+
+  if (Handler == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  Instance = EXIT_BOOT_SERVICES_OS_NOTIFY_INSTANCE_FROM_THIS (This);
+
+  for ( Link = GetFirstNode (&Instance->Handlers)
+        ; !IsNull (&Instance->Handlers, Link)
+        ; Link = GetNextNode (&Instance->Handlers, Link)
+        )
+  {
+    Entry = EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY_FROM_LINK (Link);
+    if (Entry->Handler == Handler) {
+      return EFI_ALREADY_STARTED;
+    }
+  }
+
+  ASSERT (IsNull (&Instance->Handlers, Link));
+  Entry = AllocatePool (sizeof (*Entry));
+  if (Entry == NULL) {
+    return EFI_OUT_OF_RESOURCES;
+  }
+
+  Entry->Signature = EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY_SIGNATURE;
+  Entry->Handler   = Handler;
+  InsertTailList (&Instance->Handlers, &Entry->Link);
+  return EFI_SUCCESS;
+}
+
+STATIC
+EFI_STATUS
+EFIAPI
+UnregisterHandler (
+  IN EXIT_BOOT_SERVICES_OS_NOTIFY_PROTOCOL  *This,
+  IN EXIT_BOOT_SERVICES_OS_HANDLER          Handler
+  )
+{
+  EXIT_BOOT_SERVICES_OS_NOTIFY_INSTANCE  *Instance;
+  LIST_ENTRY                             *Link;
+  EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY    *Entry;
+
+  if (Handler == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  Instance = EXIT_BOOT_SERVICES_OS_NOTIFY_INSTANCE_FROM_THIS (This);
+
+  for ( Link = GetFirstNode (&Instance->Handlers)
+        ; !IsNull (&Instance->Handlers, Link)
+        ; Link = GetNextNode (&Instance->Handlers, Link)
+        )
+  {
+    Entry = EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY_FROM_LINK (Link);
+    if (Entry->Handler == Handler) {
+      RemoveEntryList (&Entry->Link);
+      FreePool (Entry);
+      return EFI_SUCCESS;
+    }
+  }
+
+  return EFI_INVALID_PARAMETER;
+}
+
+STATIC EXIT_BOOT_SERVICES_OS_NOTIFY_INSTANCE  mNotifyInstance = {
+  EXIT_BOOT_SERVICES_OS_NOTIFY_INSTANCE_SIGNATURE,
+  {
+    RegisterHandler,
+    UnregisterHandler
+  },
+  INITIALIZE_LIST_HEAD_VARIABLE (mNotifyInstance.Handlers)
+};
+
+STATIC
+EFI_STATUS
+EFIAPI
+ExitBootServicesHook (
+  IN EFI_HANDLE  ImageHandle,
+  IN UINTN       MapKey
+  )
+{
+  EXIT_BOOT_SERVICES_OS_CONTEXT        Context;
+  LIST_ENTRY                           *Link;
+  EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY  *Entry;
+
+  Context.ReturnAddress = (EFI_PHYSICAL_ADDRESS)RETURN_ADDRESS (0);
+  ASSERT (Context.ReturnAddress != 0);
+
+  Context.OsLoaderAddress = FindPeImageBase (Context.ReturnAddress);
+  Context.OsType          = IdentifyOsType (Context.OsLoaderAddress);
+
+  DEBUG ((
+    DEBUG_INFO,
+    "ExitBootServices: Booting %a OS at 0x%lx\n",
+    OsTypeToString (Context.OsType),
+    Context.OsLoaderAddress
+    ));
+
+  for ( Link = GetFirstNode (&mNotifyInstance.Handlers)
+        ; !IsNull (&mNotifyInstance.Handlers, Link)
+        ; Link = GetNextNode (&mNotifyInstance.Handlers, Link)
+        )
+  {
+    Entry = EXIT_BOOT_SERVICES_OS_HANDLER_ENTRY_FROM_LINK (Link);
+    Entry->Handler (&Context);
+  }
+
+  gBS->ExitBootServices = mOriginalExitBootServices;
+
+  return gBS->ExitBootServices (ImageHandle, MapKey);
+}
+
+EFI_STATUS
+EFIAPI
+ExitBootServicesHookDxeInitialize (
+  IN EFI_HANDLE        ImageHandle,
+  IN EFI_SYSTEM_TABLE  *SystemTable
+  )
+{
+  EFI_STATUS  Status;
+  EFI_HANDLE  Handle;
+
+  Handle = NULL;
+  Status = gBS->InstallMultipleProtocolInterfaces (
+                  &Handle,
+                  &gExitBootServicesOsNotifyProtocolGuid,
+                  &mNotifyInstance.Notify,
+                  NULL
+                  );
+  if (EFI_ERROR (Status)) {
+    ASSERT_EFI_ERROR (Status);
+    return Status;
+  }
+
+  mOriginalExitBootServices = gBS->ExitBootServices;
+  gBS->ExitBootServices     = ExitBootServicesHook;
+
+  return EFI_SUCCESS;
+}

edk2-rockchip/Silicon/Rockchip/Drivers/ExitBootServicesHookDxe/ExitBootServicesHookDxe.inf

@@ -0,0 +1,38 @@
+#/** @file
+#
+#  Copyright (c) 2024-2025, Mario Bălănică <[email protected]>
+#
+#  SPDX-License-Identifier: BSD-2-Clause-Patent
+#
+#**/
+
+[Defines]
+  INF_VERSION                    = 0x00010005
+  BASE_NAME                      = ExitBootServicesHookDxe
+  FILE_GUID                      = 6fd64e41-870c-44cd-b8b1-6752255fb399
+  MODULE_TYPE                    = DXE_DRIVER
+  VERSION_STRING                 = 1.0
+  ENTRY_POINT                    = ExitBootServicesHookDxeInitialize
+
+[Sources]
+  ExitBootServicesHookDxe.c
+  OsIdentification.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+  Silicon/Rockchip/RockchipPkg.dec
+
+[LibraryClasses]
+  BaseLib
+  DebugLib
+  MemoryAllocationLib
+  PeCoffGetEntryPointLib
+  UefiBootServicesTableLib
+  UefiDriverEntryPoint
+  UefiLib
+
+[Protocols]
+  gExitBootServicesOsNotifyProtocolGuid     ## PRODUCES
+
+[Depex]
+  TRUE

edk2-rockchip/Silicon/Rockchip/Drivers/ExitBootServicesHookDxe/OsIdentification.c

@@ -0,0 +1,130 @@
+/** @file
+ *
+ *  Copyright (c) 2024-2025, Mario Bălănică <[email protected]>
+ *
+ *  SPDX-License-Identifier: BSD-2-Clause-Patent
+ *
+ **/
+
+#include <IndustryStandard/PeImage.h>
+#include <Library/PeCoffGetEntryPointLib.h>
+
+#include "ExitBootServicesHook.h"
+
+STATIC CHAR8  *mOsTypeStrings[] = {
+  [ExitBootServicesOsUnknown] = "Unknown",
+  [ExitBootServicesOsWindows] = "Windows",
+  [ExitBootServicesOsLinux]   = "Linux",
+};
+STATIC_ASSERT (ARRAY_SIZE (mOsTypeStrings) == ExitBootServicesOsMax);
+
+#define LINUX_ARM64_MAGIC  0x644d5241
+#define LINUX_PE_MAGIC     0x818223cd
+
+STATIC
+BOOLEAN
+IsPeImageVmlinuz (
+  IN VOID  *PeImage
+  )
+{
+  UINT8  *Buf = PeImage;
+
+  switch (*(UINT32 *)(Buf + 0x38)) {
+    case LINUX_ARM64_MAGIC:
+    case LINUX_PE_MAGIC:
+      return TRUE;
+  }
+
+  return FALSE;
+}
+
+STATIC CHAR8  mWinLoadNameStr[] = "winload";
+#define PDB_NAME_MAX_LENGTH  256
+
+STATIC
+BOOLEAN
+IsPeImageWinLoader (
+  IN VOID  *PeImage
+  )
+{
+  CHAR8  *PdbStr;
+  UINTN  WinLoadNameStrLen;
+  UINTN  Index;
+
+  PdbStr = (CHAR8 *)PeCoffLoaderGetPdbPointer (PeImage);
+  if (PdbStr == NULL) {
+    return FALSE;
+  }
+
+  WinLoadNameStrLen = sizeof (mWinLoadNameStr) - sizeof (CHAR8);
+
+  for (Index = 0; Index < PDB_NAME_MAX_LENGTH && PdbStr[Index] != '\0'; Index++) {
+    if (AsciiStrnCmp (PdbStr + Index, mWinLoadNameStr, WinLoadNameStrLen) == 0) {
+      return TRUE;
+    }
+  }
+
+  return FALSE;
+}
+
+EFI_PHYSICAL_ADDRESS
+FindPeImageBase (
+  IN EFI_PHYSICAL_ADDRESS  Base
+  )
+{
+  EFI_IMAGE_DOS_HEADER                 *DosHdr;
+  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
+
+  Base &= ~(EFI_PAGE_SIZE - 1);
+
+  while (Base != 0) {
+    DosHdr = (EFI_IMAGE_DOS_HEADER *)Base;
+    if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
+      Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)(Base + DosHdr->e_lfanew);
+      if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
+        break;
+      }
+    }
+
+    Base -= EFI_PAGE_SIZE;
+  }
+
+  return Base;
+}
+
+EXIT_BOOT_SERVICES_OS_TYPE
+IdentifyOsType (
+  IN EFI_PHYSICAL_ADDRESS  OsLoaderAddress
+  )
+{
+  VOID  *OsLoaderImage;
+
+  if (OsLoaderAddress == 0) {
+    return ExitBootServicesOsUnknown;
+  }
+
+  OsLoaderImage = (VOID *)OsLoaderAddress;
+
+  if (IsPeImageVmlinuz (OsLoaderImage)) {
+    return ExitBootServicesOsLinux;
+  }
+
+  if (IsPeImageWinLoader (OsLoaderImage)) {
+    return ExitBootServicesOsWindows;
+  }
+
+  return ExitBootServicesOsUnknown;
+}
+
+CHAR8 *
+OsTypeToString (
+  IN EXIT_BOOT_SERVICES_OS_TYPE  OsType
+  )
+{
+  if ((OsType < ExitBootServicesOsUnknown) || (OsType >= ExitBootServicesOsMax)) {
+    ASSERT (FALSE);
+    return mOsTypeStrings[ExitBootServicesOsUnknown];
+  }
+
+  return mOsTypeStrings[OsType];
+}

edk2-rockchip/Silicon/Rockchip/FvMainModules.fdf.inc

@@ -30,6 +30,11 @@
   INF SecurityPkg/RandomNumberGenerator/RngDxe/RngDxe.inf
   INF SecurityPkg/Hash2DxeCrypto/Hash2DxeCrypto.inf
 
+  #
+  # ExitBootServices hook manager
+  #
+  INF Silicon/Rockchip/Drivers/ExitBootServicesHookDxe/ExitBootServicesHookDxe.inf
+
   #
   # Status LED support
   #