From 05aa27ef2369bd98380d56c57a677be66972ea67 Mon Sep 17 00:00:00 2001
From: Tim Crawford <tcrawford@system76.com>
Date: Tue, 3 Nov 2020 14:19:21 -0700
Subject: [PATCH] MdeModulePkg/BM: Update boot options on device change

Register an event to trigger when a block device is added or removed
that will update the list of boot options. Use a refreshguid to force
the form to display the changes if it already open.
---
 .../Library/BootManagerUiLib/BootManager.c    | 58 ++++++++++++++-----
 .../Library/BootManagerUiLib/BootManager.h    |  5 ++
 .../BootManagerUiLib/BootManagerUiLib.inf     |  1 +
 .../BootManagerUiLib/BootManagerVfr.Vfr       |  2 +
 4 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/MdeModulePkg/Library/BootManagerUiLib/BootManager.c b/MdeModulePkg/Library/BootManagerUiLib/BootManager.c
index c98e132528..e0e89ec4a3 100644
--- a/MdeModulePkg/Library/BootManagerUiLib/BootManager.c
+++ b/MdeModulePkg/Library/BootManagerUiLib/BootManager.c
@@ -8,6 +8,8 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 
 #include "BootManager.h"
 
+#include <Protocol/BlockIo.h>
+
 UINT16             mKeyInput;
 EFI_GUID           mBootManagerGuid = BOOT_MANAGER_FORMSET_GUID;
 //
@@ -27,6 +29,10 @@ UINT32    mBmSetupVerticalResolution     = 0;
 
 BOOLEAN   mBmModeInitialized             = FALSE;
 
+STATIC EFI_EVENT mBmRefreshEvent;
+STATIC CONST EFI_GUID mBmRefreshGuid = BOOT_MANAGER_REFRESH_GUID;
+STATIC VOID *mBlockIoRegistration; // Unused
+
 CHAR16             *mDeviceTypeStr[] = {
   L"Legacy BEV",
   L"Legacy Floppy",
@@ -488,19 +494,6 @@ UpdateBootManager (
 
   DeviceType = (UINT16) -1;
 
-  //
-  // for better user experience
-  // 1. User changes HD configuration (e.g.: unplug HDD), here we have a chance to remove the HDD boot option
-  // 2. User enables/disables UEFI PXE, here we have a chance to add/remove EFI Network boot option
-  //
-  EfiBootManagerRefreshAllBootOption ();
-
-  //
-  // BdsDxe doesn't group the legacy boot options for the same device type
-  // It's UI's choice.
-  //
-  GroupMultipleLegacyBootOption4SameType ();
-
   BootOption = EfiBootManagerGetLoadOptions (&BootOptionCount, LoadOptionTypeBoot);
 
   HiiHandle = gBootManagerPrivate.HiiHandle;
@@ -843,6 +836,26 @@ BootManagerCallback (
   return EFI_SUCCESS;
 }
 
+STATIC
+VOID
+EFIAPI
+RefreshBootOptions(
+  IN  EFI_EVENT   Event,
+  IN  VOID        *Context
+  )
+{
+  EFI_TPL OldTpl = gBS->RaiseTPL(TPL_CALLBACK);
+
+  EfiBootManagerRefreshAllBootOption();
+
+  // BdsDxe doesn't group the legacy boot options for the same device type. It's UI's choice.
+  GroupMultipleLegacyBootOption4SameType();
+
+  UpdateBootManager();
+
+  gBS->RestoreTPL (OldTpl);
+}
+
 /**
 
   Install Boot Manager Menu driver.
@@ -891,6 +904,23 @@ BootManagerUiLibConstructor (
 
   BmInitialBootModeInfo ();
 
+  Status = gBS->CreateEventEx(
+      EVT_NOTIFY_SIGNAL,
+      TPL_CALLBACK,
+      RefreshBootOptions,
+      NULL,
+      &mBmRefreshGuid,
+      &mBmRefreshEvent
+      );
+  ASSERT_EFI_ERROR(Status);
+
+  Status = gBS->RegisterProtocolNotify (
+      &gEfiBlockIoProtocolGuid,
+      mBmRefreshEvent,
+      &mBlockIoRegistration
+      );
+  ASSERT_EFI_ERROR (Status);
+
   return EFI_SUCCESS;
 }
 
@@ -911,6 +941,8 @@ BootManagerUiLibDestructor (
 {
   EFI_STATUS    Status;
 
+  gBS->CloseEvent(mBmRefreshEvent);
+
   Status = gBS->UninstallMultipleProtocolInterfaces (
                   gBootManagerPrivate.DriverHandle,
                   &gEfiDevicePathProtocolGuid,
diff --git a/MdeModulePkg/Library/BootManagerUiLib/BootManager.h b/MdeModulePkg/Library/BootManagerUiLib/BootManager.h
index 027995cbca..df37987f47 100644
--- a/MdeModulePkg/Library/BootManagerUiLib/BootManager.h
+++ b/MdeModulePkg/Library/BootManagerUiLib/BootManager.h
@@ -46,6 +46,11 @@ typedef struct {
   0x847bc3fe, 0xb974, 0x446d, {0x94, 0x49, 0x5a, 0xd5, 0x41, 0x2e, 0x99, 0x3b} \
   }
 
+#define BOOT_MANAGER_REFRESH_GUID \
+  { \
+    0x7648C827, 0xBE32, 0x44D2, { 0xA1, 0x33, 0xF1, 0x7F, 0x93, 0x76, 0xC7, 0x00 } \
+  }
+
 #define BOOT_MANAGER_FORM_ID     0x1030
 
 #define LABEL_BOOT_OPTION        0x00
diff --git a/MdeModulePkg/Library/BootManagerUiLib/BootManagerUiLib.inf b/MdeModulePkg/Library/BootManagerUiLib/BootManagerUiLib.inf
index 020cf75053..3f880f2720 100644
--- a/MdeModulePkg/Library/BootManagerUiLib/BootManagerUiLib.inf
+++ b/MdeModulePkg/Library/BootManagerUiLib/BootManagerUiLib.inf
@@ -51,6 +51,7 @@
   gEfiHiiConfigAccessProtocolGuid               ## CONSUMES
   gEfiDevicePathToTextProtocolGuid              ## CONSUMES
   gEdkiiFormBrowserEx2ProtocolGuid              ## CONSUMES
+  gEfiBlockIoProtocolGuid
 
 [FeaturePcd]
 
diff --git a/MdeModulePkg/Library/BootManagerUiLib/BootManagerVfr.Vfr b/MdeModulePkg/Library/BootManagerUiLib/BootManagerVfr.Vfr
index c49b6f50b1..77fd54ec35 100644
--- a/MdeModulePkg/Library/BootManagerUiLib/BootManagerVfr.Vfr
+++ b/MdeModulePkg/Library/BootManagerUiLib/BootManagerVfr.Vfr
@@ -7,6 +7,7 @@
 //
 //**/
 #define FORMSET_GUID  { 0x847bc3fe, 0xb974, 0x446d, 0x94, 0x49, 0x5a, 0xd5, 0x41, 0x2e, 0x99, 0x3b }
+#define BOOT_MANAGER_REFRESH_GUID { 0x7648C827, 0xBE32, 0x44D2, { 0xA1, 0x33, 0xF1, 0x7F, 0x93, 0x76, 0xC7, 0x00 }}
 
 #define BOOT_MANAGER_FORM_ID     0x1030
 
@@ -21,6 +22,7 @@ formset
 
   form formid = BOOT_MANAGER_FORM_ID,
     title  = STRING_TOKEN(STR_BM_BANNER);
+    refreshguid = BOOT_MANAGER_REFRESH_GUID;
 
     //
     //Add this invisable text in order to indicate enter Boot Manager form.