From 083e5f72ea9851fb0677f6e9917658d7cb8d0408 Mon Sep 17 00:00:00 2001
From: longpanda <support@iventoy.com>
Date: Thu, 25 Jun 2026 22:36:33 +0800
Subject: [PATCH] Optimization for Secure Boot process.

---
 .../Application/VtoyShim/VtoyShim.c           | 93 +++++++++++--------
 1 file changed, 53 insertions(+), 40 deletions(-)

diff --git a/EDK2/edk2_mod/edk2-edk2-stable201911/MdeModulePkg/Application/VtoyShim/VtoyShim.c b/EDK2/edk2_mod/edk2-edk2-stable201911/MdeModulePkg/Application/VtoyShim/VtoyShim.c
index c4f574c7..8618f539 100644
--- a/EDK2/edk2_mod/edk2-edk2-stable201911/MdeModulePkg/Application/VtoyShim/VtoyShim.c
+++ b/EDK2/edk2_mod/edk2-edk2-stable201911/MdeModulePkg/Application/VtoyShim/VtoyShim.c
@@ -46,6 +46,10 @@ STATIC SHIM_LOCK gShimLock;
 STATIC EFI_EXIT_BOOT_SERVICES gSysExitBootServices = NULL;
 STATIC EFI_GET_VARIABLE gSysGetVariable = NULL;
 
+STATIC VOID EFIAPI HookSystemService(VOID);
+STATIC VOID EFIAPI UnHookSystemService(VOID);
+
+
 STATIC VOID EFIAPI VtoyLog(CONST CHAR16 *Format, ...)
 {
     VA_LIST         Marker;
@@ -432,6 +436,7 @@ STATIC EFI_STATUS EFIAPI Security2PolicyAuth
         return EFI_SUCCESS;
     }
 
+
     /*
      * Step 1:
      * Use original UEFI firmware auth API.
@@ -623,9 +628,7 @@ STATIC BOOLEAN EFIAPI IsSetupMode(VOID)
 STATIC EFI_STATUS EFIAPI ShimEfiMain
 (
     IN EFI_HANDLE         ImageHandle,
-    IN EFI_SYSTEM_TABLE  *SystemTable,
-    IN BOOLEAN            IsSecureBoot,
-    IN BOOLEAN            IsSetup
+    IN EFI_SYSTEM_TABLE  *SystemTable
 )
 {
     EFI_STATUS Status;
@@ -633,17 +636,6 @@ STATIC EFI_STATUS EFIAPI ShimEfiMain
     shim_void_func_pf Func1 = NULL;
     shim_void_func_pf Func2 = NULL;
 
-    /* If secure boot is not enabled or in SetupMode, nothing needed, just launch Ventoy grub */
-    if (!IsSecureBoot || IsSetup)
-    {
-        Status = LaunchRealGrub(ImageHandle, REAL_GRUB_FILE);
-        if (EFI_ERROR(Status))
-        {
-            vErr(L"Failed to launch %s", REAL_GRUB_FILE);
-        }
-        return Status;
-    }
-
     /* We must be launched by shim */
     Status = gBS->LocateProtocol(&gShimLockGUID, NULL, (VOID**)&ShimLock);
     if (EFI_ERROR(Status) || !ShimLock)
@@ -691,6 +683,7 @@ STATIC EFI_STATUS EFIAPI ShimEfiMain
     Func1(); /* call shim unhook_system_services() */
     Func2(); /* call shim uninstall_shim_protocols() */
 
+    HookSystemService();
 
     /* Hook the system security policy */
     Status = HookSecurityPolicy();
@@ -715,24 +708,11 @@ END:
 
     UnInstallVtoyShimProtocol();
 
+    UnHookSystemService();
+
     return Status;
 }
 
-STATIC EFI_STATUS EFIAPI VtoyExitBootServices
-(
-    IN  EFI_HANDLE  ImageHandle,
-    IN  UINTN       MapKey
-)
-{
-    UnHookSecurityPolicy();
-    UnInstallVtoyShimProtocol();
-
-    gST->RuntimeServices->GetVariable = gSysGetVariable;
-    gBS->ExitBootServices = gSysExitBootServices;
-
-    return gSysExitBootServices(ImageHandle, MapKey);
-}
-
 EFI_STATUS EFIAPI VtoyGetVariable
 (
     IN     CHAR16                      *VariableName,
@@ -763,6 +743,43 @@ EFI_STATUS EFIAPI VtoyGetVariable
     return Status;
 }
 
+STATIC VOID EFIAPI UnHookSystemService(VOID)
+{
+    if (gSysExitBootServices)
+    {
+        gBS->ExitBootServices = gSysExitBootServices;
+        gSysExitBootServices = NULL;
+    }
+
+    if (gSysGetVariable)
+    {
+        gST->RuntimeServices->GetVariable = gSysGetVariable;
+        gSysGetVariable = NULL;
+    }
+}
+
+
+STATIC EFI_STATUS EFIAPI VtoyExitBootServices
+(
+    IN  EFI_HANDLE  ImageHandle,
+    IN  UINTN       MapKey
+)
+{
+    UnHookSecurityPolicy();
+    UnInstallVtoyShimProtocol();
+    UnHookSystemService();
+
+    return gSysExitBootServices(ImageHandle, MapKey);
+}
+
+STATIC VOID EFIAPI HookSystemService(VOID)
+{
+    gSysExitBootServices = gBS->ExitBootServices;
+    gBS->ExitBootServices = VtoyExitBootServices;
+
+    gSysGetVariable = gST->RuntimeServices->GetVariable;
+    gST->RuntimeServices->GetVariable = VtoyGetVariable;
+}
 
 EFI_STATUS EFIAPI VtoyShimEfiMain
 (
@@ -779,20 +796,16 @@ EFI_STATUS EFIAPI VtoyShimEfiMain
 
     if (!IsSecureBoot || IsSetup)
     {
-        Status = ShimEfiMain(ImageHandle, SystemTable, IsSecureBoot, IsSetup);
+        /* If secure boot is not enabled or in SetupMode, nothing needed, just launch Ventoy grub */
+        Status = LaunchRealGrub(ImageHandle, REAL_GRUB_FILE);
+        if (EFI_ERROR(Status))
+        {
+            vErr(L"Failed to launch %s", REAL_GRUB_FILE);
+        }
     }
     else
     {
-        gSysExitBootServices = gBS->ExitBootServices;
-        gBS->ExitBootServices = VtoyExitBootServices;
-
-        gSysGetVariable = gST->RuntimeServices->GetVariable;
-        gST->RuntimeServices->GetVariable = VtoyGetVariable;
-
-        Status = ShimEfiMain(ImageHandle, SystemTable, IsSecureBoot, IsSetup);
-
-        gBS->ExitBootServices = gSysExitBootServices;
-        gST->RuntimeServices->GetVariable = gSysGetVariable;
+        Status = ShimEfiMain(ImageHandle, SystemTable);
     }
 
     return Status;