From 048bcba1bcd082e4527dd1b17e5dfb6470285260 Mon Sep 17 00:00:00 2001
From: Ruiyu Ni <ruiyu.ni@intel.com>
Date: Tue, 1 Nov 2016 13:55:49 +0800
Subject: [PATCH] MdeModulePkg/BdsDxe: Check deferred images before booting to
 OS

The patch adds check of deferred images before booting to OS.

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Ruiyu Ni <ruiyu.ni@intel.com>
Reviewed-by: Liming Gao <liming.gao@intel.com>
Reviewed-by: Chao B Zhang <chao.b.zhang@intel.com>
Reviewed-by: Sunny Wang <sunnywang@hpe.com>
---
 MdeModulePkg/Universal/BdsDxe/Bds.h      |  4 +-
 MdeModulePkg/Universal/BdsDxe/BdsDxe.inf |  2 +
 MdeModulePkg/Universal/BdsDxe/BdsEntry.c | 90 ++++++++++++++++++++++++
 3 files changed, 95 insertions(+), 1 deletion(-)

diff --git a/MdeModulePkg/Universal/BdsDxe/Bds.h b/MdeModulePkg/Universal/BdsDxe/Bds.h
index d2439329d1..1f8a192424 100644
--- a/MdeModulePkg/Universal/BdsDxe/Bds.h
+++ b/MdeModulePkg/Universal/BdsDxe/Bds.h
@@ -1,7 +1,7 @@
 /** @file
   Head file for BDS Architectural Protocol implementation
 
-Copyright (c) 2004 - 2015, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2004 - 2016, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -20,10 +20,12 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Guid/ConnectConInEvent.h>
 #include <Guid/Performance.h>
 #include <Guid/StatusCodeDataTypeVariable.h>
+#include <Guid/EventGroup.h>
 
 #include <Protocol/Bds.h>
 #include <Protocol/LoadedImage.h>
 #include <Protocol/VariableLock.h>
+#include <Protocol/DeferredImageLoad.h>
 
 #include <Library/UefiDriverEntryPoint.h>
 #include <Library/DebugLib.h>
diff --git a/MdeModulePkg/Universal/BdsDxe/BdsDxe.inf b/MdeModulePkg/Universal/BdsDxe/BdsDxe.inf
index e03cfde83d..a00b44276d 100644
--- a/MdeModulePkg/Universal/BdsDxe/BdsDxe.inf
+++ b/MdeModulePkg/Universal/BdsDxe/BdsDxe.inf
@@ -79,11 +79,13 @@
   gConnectConInEventGuid                        ## SOMETIMES_CONSUMES ## Event
   gEdkiiStatusCodeDataTypeVariableGuid          ## SOMETIMES_CONSUMES ## GUID
   gPerformanceProtocolGuid                      ## SOMETIMES_PRODUCES ## Variable:L"PerfDataMemAddr" (The ACPI address of performance data)
+  gEfiEventReadyToBootGuid                      ## CONSUMES           ## Event
 
 [Protocols]
   gEfiBdsArchProtocolGuid                       ## PRODUCES
   gEfiSimpleTextInputExProtocolGuid             ## CONSUMES
   gEdkiiVariableLockProtocolGuid                ## CONSUMES
+  gEfiDeferredImageLoadProtocolGuid             ## CONSUMES
 
 [FeaturePcd]
   gEfiMdePkgTokenSpaceGuid.PcdUefiVariableDefaultLangDeprecate    ## CONSUMES
diff --git a/MdeModulePkg/Universal/BdsDxe/BdsEntry.c b/MdeModulePkg/Universal/BdsDxe/BdsEntry.c
index aacc4a6479..72adc5155c 100644
--- a/MdeModulePkg/Universal/BdsDxe/BdsEntry.c
+++ b/MdeModulePkg/Universal/BdsDxe/BdsEntry.c
@@ -86,6 +86,81 @@ BdsDxeOnConnectConInCallBack (
     DEBUG ((EFI_D_WARN, "[Bds] Connect ConIn failed - %r!!!\n", Status));
   }
 }
+/**
+  Notify function for event group EFI_EVENT_GROUP_READY_TO_BOOT. This is used to
+  check whether there is remaining deferred load images.
+
+  @param[in]  Event   The Event that is being processed.
+  @param[in]  Context The Event Context.
+
+**/
+VOID
+EFIAPI
+CheckDeferredLoadImageOnReadyToBoot (
+  IN EFI_EVENT        Event,
+  IN VOID             *Context
+  )
+{
+  EFI_STATUS                         Status;
+  EFI_DEFERRED_IMAGE_LOAD_PROTOCOL   *DeferredImage;
+  UINTN                              HandleCount;
+  EFI_HANDLE                         *Handles;
+  UINTN                              Index;
+  UINTN                              ImageIndex;
+  EFI_DEVICE_PATH_PROTOCOL           *ImageDevicePath;
+  VOID                               *Image;
+  UINTN                              ImageSize;
+  BOOLEAN                            BootOption;
+  CHAR16                             *DevicePathStr;
+
+  //
+  // Find all the deferred image load protocols.
+  //
+  HandleCount = 0;
+  Handles = NULL;
+  Status = gBS->LocateHandleBuffer (
+    ByProtocol,
+    &gEfiDeferredImageLoadProtocolGuid,
+    NULL,
+    &HandleCount,
+    &Handles
+  );
+  if (EFI_ERROR (Status)) {
+    return;
+  }
+
+  for (Index = 0; Index < HandleCount; Index++) {
+    Status = gBS->HandleProtocol (Handles[Index], &gEfiDeferredImageLoadProtocolGuid, (VOID **) &DeferredImage);
+    if (EFI_ERROR (Status)) {
+      continue;
+    }
+
+    for (ImageIndex = 0; ; ImageIndex++) {
+      //
+      // Load all the deferred images in this protocol instance.
+      //
+      Status = DeferredImage->GetImageInfo (
+        DeferredImage,
+        ImageIndex,
+        &ImageDevicePath,
+        (VOID **) &Image,
+        &ImageSize,
+        &BootOption
+      );
+      if (EFI_ERROR (Status)) {
+        break;
+      }
+      DevicePathStr = ConvertDevicePathToText (ImageDevicePath, FALSE, FALSE);
+      DEBUG ((DEBUG_LOAD, "[Bds] Image was deferred but not loaded: %s.\n", DevicePathStr));
+      if (DevicePathStr != NULL) {
+        FreePool (DevicePathStr);
+      }
+    }
+  }
+  if (Handles != NULL) {
+    FreePool (Handles);
+  }
+}
 
 /**
 
@@ -119,6 +194,21 @@ BdsInitialize (
                   );
   ASSERT_EFI_ERROR (Status);
 
+  DEBUG_CODE (
+    EFI_EVENT   Event;
+    //
+    // Register notify function to check deferred images on ReadyToBoot Event.
+    //
+    Status = gBS->CreateEventEx (
+                    EVT_NOTIFY_SIGNAL,
+                    TPL_CALLBACK,
+                    CheckDeferredLoadImageOnReadyToBoot,
+                    NULL,
+                    &gEfiEventReadyToBootGuid,
+                    &Event
+                    );
+    ASSERT_EFI_ERROR (Status);
+  );
   return Status;
 }
 
-- 
2.39.2