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>

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