[PATCH 3/4] MdeModulePkg/BdsDxe: Check deferred images before booting to OS

[Ruiyu Ni <ruiyu.ni@intel.com>]

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>
Cc: Liming Gao <liming.gao@intel.com>
Cc: Chao B Zhang <chao.b.zhang@intel.com>
---
 MdeModulePkg/Universal/BdsDxe/Bds.h      |  4 +-
 MdeModulePkg/Universal/BdsDxe/BdsDxe.inf |  2 +
 MdeModulePkg/Universal/BdsDxe/BdsEntry.c | 89 ++++++++++++++++++++++++++++++++
 3 files changed, 94 insertions(+), 1 deletion(-)

diff --git a/MdeModulePkg/Universal/BdsDxe/Bds.h b/MdeModulePkg/Universal/BdsDxe/Bds.h
index d243932..1f8a192 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 e03cfde..a00b442 100644
--- a/MdeModulePkg/Universal/BdsDxe/BdsDxe.inf
+++ b/MdeModulePkg/Universal/BdsDxe/BdsDxe.inf
@@ -79,11 +79,13 @@ [Guids]
   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 aacc4a6..ffa884a 100644
--- a/MdeModulePkg/Universal/BdsDxe/BdsEntry.c
+++ b/MdeModulePkg/Universal/BdsDxe/BdsEntry.c
@@ -86,7 +86,83 @@ 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;
+
+  //
+  // 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;
+      }
+      DEBUG_CODE (
+        CHAR16    *DevicePathStr;
+        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);
+  }
+}
 /**
 
   Install Boot Device Selection Protocol
@@ -108,6 +184,7 @@ BdsInitialize (
 {
   EFI_STATUS  Status;
   EFI_HANDLE  Handle;
+  EFI_EVENT   Event;
   //
   // Install protocol interface
   //
@@ -119,6 +196,18 @@ BdsInitialize (
                   );
   ASSERT_EFI_ERROR (Status);
 
+  //
+  // 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.9.0.windows.1