Verify memory allocations were successful.

[mirror_edk2.git] / ShellPkg / Library / UefiShellLevel2CommandsLib / Map.c

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c
index c4c8c87ff409d32bf56a3aba32c682cc1c02d58a..f696523f82f32682a3c9dad33bb97d4af49baf81 100644 (file)
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c
@@ -99,7 +99,7 @@ UpdateMapping (
   //\r
   // Find each handle with Simple File System\r
   //\r
-  HandleList = GetHandleListByPotocol(&gEfiSimpleFileSystemProtocolGuid);\r
+  HandleList = GetHandleListByProtocol(&gEfiSimpleFileSystemProtocolGuid);\r
   if (HandleList != NULL) {\r
     //\r
     // Do a count of the handles\r
@@ -503,6 +503,9 @@ PerformMappingDisplay(
     HandleBuffer);\r
   if (Status == EFI_BUFFER_TOO_SMALL) {\r
     HandleBuffer = AllocatePool(BufferSize);\r
+    if (HandleBuffer == NULL) {\r
+      return (SHELL_OUT_OF_RESOURCES);\r
+    }\r
     Status = gBS->LocateHandle(\r
       ByProtocol,\r
       &gEfiDevicePathProtocolGuid,\r
@@ -542,6 +545,9 @@ PerformMappingDisplay(
   if (Status == EFI_BUFFER_TOO_SMALL) {\r
     FreePool(HandleBuffer);\r
     HandleBuffer = AllocatePool(BufferSize);\r
+    if (HandleBuffer == NULL) {\r
+      return (SHELL_OUT_OF_RESOURCES);\r
+    }\r
     Status = gBS->LocateHandle(\r
       ByProtocol,\r
       &gEfiBlockIoProtocolGuid,\r
@@ -549,37 +555,37 @@ PerformMappingDisplay(
       &BufferSize,\r
       HandleBuffer);\r
   }\r
-  ASSERT_EFI_ERROR(Status);\r
-\r
-  //\r
-  // Get the map name(s) for each one.\r
-  //\r
-  for ( LoopVar = 0\r
-      ; LoopVar < BufferSize / sizeof(EFI_HANDLE)\r
-      ; LoopVar ++\r
-     ){\r
+  if (!EFI_ERROR(Status)) {\r
     //\r
-    // Skip any that were already done...\r
+    // Get the map name(s) for each one.\r
     //\r
-    if (gBS->OpenProtocol(\r
-      HandleBuffer[LoopVar],\r
-      &gEfiDevicePathProtocolGuid,\r
-      NULL,\r
-      gImageHandle,\r
-      NULL,\r
-      EFI_OPEN_PROTOCOL_TEST_PROTOCOL) == EFI_SUCCESS) {\r
-        continue;\r
+    for ( LoopVar = 0\r
+        ; LoopVar < BufferSize / sizeof(EFI_HANDLE)\r
+        ; LoopVar ++\r
+       ){\r
+      //\r
+      // Skip any that were already done...\r
+      //\r
+      if (gBS->OpenProtocol(\r
+        HandleBuffer[LoopVar],\r
+        &gEfiDevicePathProtocolGuid,\r
+        NULL,\r
+        gImageHandle,\r
+        NULL,\r
+        EFI_OPEN_PROTOCOL_TEST_PROTOCOL) == EFI_SUCCESS) {\r
+          continue;\r
+      }\r
+      PerformSingleMappingDisplay(\r
+        Verbose,\r
+        Consist,\r
+        Normal,\r
+        Test,\r
+        SFO,\r
+        Specific,\r
+        HandleBuffer[LoopVar]);\r
     }\r
-    PerformSingleMappingDisplay(\r
-      Verbose,\r
-      Consist,\r
-      Normal,\r
-      Test,\r
-      SFO,\r
-      Specific,\r
-      HandleBuffer[LoopVar]);\r
+    FreePool(HandleBuffer);\r
   }\r
-  FreePool(HandleBuffer);\r
   return (SHELL_SUCCESS);\r
 }\r
 \r