ShellPkg/map: Fix out-of-bound read when "map fsn"
[mirror_edk2.git] / ShellPkg / Library / UefiShellLevel2CommandsLib / Map.c
index 20eb528..9166ca2 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   Main file for map shell level 2 command.\r
 \r
-  Copyright (c) 2009 - 2015, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2009 - 2018, Intel Corporation. All rights reserved.<BR>\r
   (C) Copyright 2013-2015 Hewlett-Packard Development Company, L.P.<BR>\r
   (C) Copyright 2016 Hewlett Packard Enterprise Development LP<BR>\r
   \r
@@ -220,19 +220,25 @@ MappingListHasType(
   IN CONST BOOLEAN    Consist\r
   )\r
 {\r
-  CHAR16 *NewSpecific;\r
-  RETURN_STATUS  Status;\r
+  CHAR16              *NewSpecific;\r
+  RETURN_STATUS       Status;\r
+  UINTN               Length;\r
   \r
   //\r
   // specific has priority\r
   //\r
   if (Specific != NULL) {\r
-    NewSpecific = AllocateCopyPool(StrSize(Specific) + sizeof(CHAR16), Specific);\r
+    Length      = StrLen (Specific);\r
+    //\r
+    // Allocate enough buffer for Specific and potential ":"\r
+    //\r
+    NewSpecific = AllocatePool ((Length + 2) * sizeof(CHAR16));\r
     if (NewSpecific == NULL){\r
       return FALSE;\r
     }\r
-    if (NewSpecific[StrLen(NewSpecific)-1] != L':') {\r
-      Status = StrnCatS(NewSpecific, (StrSize(Specific) + sizeof(CHAR16))/sizeof(CHAR16), L":", StrLen(L":"));\r
+    StrCpyS (NewSpecific, Length + 2, Specific);\r
+    if (Specific[Length - 1] != L':') {\r
+      Status = StrnCatS(NewSpecific, Length + 2, L":", StrLen(L":"));\r
       if (EFI_ERROR (Status)) {\r
         FreePool(NewSpecific);\r
         return FALSE;\r
@@ -979,6 +985,57 @@ STATIC CONST SHELL_PARAM_ITEM MapParamList[] = {
   {NULL, TypeMax}\r
   };\r
 \r
+/**\r
+  The routine issues dummy read for every physical block device to cause\r
+  the BlockIo re-installed if media change happened.\r
+**/\r
+VOID\r
+ProbeForMediaChange (\r
+  VOID\r
+  )\r
+{\r
+  EFI_STATUS                            Status;\r
+  UINTN                                 HandleCount;\r
+  EFI_HANDLE                            *Handles;\r
+  EFI_BLOCK_IO_PROTOCOL                 *BlockIo;\r
+  UINTN                                 Index;\r
+\r
+  gBS->LocateHandleBuffer (\r
+         ByProtocol,\r
+         &gEfiBlockIoProtocolGuid,\r
+         NULL,\r
+         &HandleCount,\r
+         &Handles\r
+         );\r
+  //\r
+  // Probe for media change for every physical block io\r
+  //\r
+  for (Index = 0; Index < HandleCount; Index++) {\r
+    Status = gBS->HandleProtocol (\r
+                    Handles[Index],\r
+                    &gEfiBlockIoProtocolGuid,\r
+                    (VOID **) &BlockIo\r
+                    );\r
+    if (!EFI_ERROR (Status)) {\r
+      if (!BlockIo->Media->LogicalPartition) {\r
+        //\r
+        // Per spec:\r
+        //   The function (ReadBlocks) must return EFI_NO_MEDIA or\r
+        //   EFI_MEDIA_CHANGED even if LBA, BufferSize, or Buffer are invalid so the caller can probe\r
+        //   for changes in media state.\r
+        //\r
+        BlockIo->ReadBlocks (\r
+                   BlockIo,\r
+                   BlockIo->Media->MediaId,\r
+                   0,\r
+                   0,\r
+                   NULL\r
+                   );\r
+      }\r
+    }\r
+  }\r
+}\r
+\r
 /**\r
   Function for 'map' command.\r
 \r
@@ -1087,6 +1144,7 @@ ShellCommandRunMap (
                || ShellCommandLineGetFlag(Package, L"-u")\r
                || ShellCommandLineGetFlag(Package, L"-t")\r
               ){\r
+        ProbeForMediaChange ();\r
         if ( ShellCommandLineGetFlag(Package, L"-r")) {\r
           //\r
           // Do the reset\r