From ed2992b3f7c0ac4da2c04a82e127c62ee87a4b00 Mon Sep 17 00:00:00 2001
From: Dandan Bi <dandan.bi@intel.com>
Date: Mon, 29 Feb 2016 14:37:07 +0800
Subject: [PATCH] SecurityPkg/SecureBootConfigDxe: Handle allocation failure
 gracefully

The function AllocateCopyPool may return NULL, so need to do check
after calling it. This patch is to enhance the related logic.

Cc: Chao Zhang <chao.b.zhang@intel.com>
Cc: Qiu Shumin <shumin.qiu@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Dandan Bi <dandan.bi@intel.com>
Reviewed-by: Chao Zhang <chao.b.zhang@intel.com>
Reviewed-by: Eric Dong <eric.dong@intel.com>
---
 .../SecureBootConfigFileExplorer.c            | 30 +++++++++++++------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/SecurityPkg/VariableAuthenticated/SecureBootConfigDxe/SecureBootConfigFileExplorer.c b/SecurityPkg/VariableAuthenticated/SecureBootConfigDxe/SecureBootConfigFileExplorer.c
index 2adb85ce90..05d97dc43b 100644
--- a/SecurityPkg/VariableAuthenticated/SecureBootConfigDxe/SecureBootConfigFileExplorer.c
+++ b/SecurityPkg/VariableAuthenticated/SecureBootConfigDxe/SecureBootConfigFileExplorer.c
@@ -222,11 +222,13 @@ OpenFileByDevicePath(
 
 /**
   Extract filename from device path. The returned buffer is allocated using AllocateCopyPool.
-  The caller is responsible for freeing the allocated buffer using FreePool().
+  The caller is responsible for freeing the allocated buffer using FreePool(). If return NULL
+  means not enough memory resource.
 
   @param DevicePath       Device path.
 
-  @return                 A new allocated string that represents the file name.
+  @retval NULL            Not enough memory resourece for AllocateCopyPool.
+  @retval Other           A new allocated string that represents the file name.
 
 **/
 CHAR16 *
@@ -245,6 +247,7 @@ ExtractFileNameFromDevicePath (
   String = DevicePathToStr(DevicePath);
   MatchString = String;
   LastMatch   = String;
+  FileName    = NULL;
 
   while(MatchString != NULL){
     LastMatch   = MatchString + 1;
@@ -253,7 +256,9 @@ ExtractFileNameFromDevicePath (
 
   Length = StrLen(LastMatch);
   FileName = AllocateCopyPool ((Length + 1) * sizeof(CHAR16), LastMatch);
-  *(FileName + Length) = 0;
+  if (FileName != NULL) {
+    *(FileName + Length) = 0;
+  }
 
   FreePool(String);
 
@@ -280,14 +285,21 @@ UpdatePage(
   CHAR16                *FileName;
   EFI_STRING_ID         StringToken;
 
-  if (FilePath != NULL){
+  FileName = NULL;
+
+  if (FilePath != NULL) {
     FileName = ExtractFileNameFromDevicePath(FilePath);
-    StringToken = HiiSetString (gSecureBootPrivateData->HiiHandle, 0, FileName, NULL);
-  } else {
-    FileName = HiiGetString (gSecureBootPrivateData->HiiHandle, STRING_TOKEN (STR_NULL), NULL);
-    ASSERT (FileName != NULL);
-    StringToken =  HiiSetString (gSecureBootPrivateData->HiiHandle, 0, FileName, NULL);
   }
+  if (FileName == NULL) {
+    //
+    // FileName = NULL has two case:
+    // 1. FilePath == NULL, not select file.
+    // 2. FilePath != NULL, but ExtractFileNameFromDevicePath return NULL not enough memory resource.
+    // In these two case, no need to update the form, and exit the caller function.
+    //
+    return TRUE;
+  }
+  StringToken =  HiiSetString (gSecureBootPrivateData->HiiHandle, 0, FileName, NULL);
 
   gSecureBootPrivateData->FileContext->FileName = FileName;