ShellPkg/Dp: Add null pointer check

[mirror_edk2.git] / ShellPkg / Library / UefiShellDebug1CommandsLib / Dblk.c

diff --git a/ShellPkg/Library/UefiShellDebug1CommandsLib/Dblk.c b/ShellPkg/Library/UefiShellDebug1CommandsLib/Dblk.c
index a1fe9d7bef2945d1bcb921b72b4de3b3390d0f76..3632ca8a7a4eb1e16db1a252067b1741f5e551b2 100644 (file)
--- a/ShellPkg/Library/UefiShellDebug1CommandsLib/Dblk.c
+++ b/ShellPkg/Library/UefiShellDebug1CommandsLib/Dblk.c
@@ -2,7 +2,7 @@
   Main file for Dblk shell Debug1 function.\r
 \r
   (C) Copyright 2015 Hewlett-Packard Development Company, L.P.<BR>\r
-  Copyright (c) 2005 - 2011, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2005 - 2017, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials\r
   are licensed and made available under the terms and conditions of the BSD License\r
   which accompanies this distribution.  The full text of the license may be found at\r
@@ -26,7 +26,6 @@
   @retval SHELL_SUCCESS   The display was successful.\r
 **/\r
 SHELL_STATUS\r
-EFIAPI\r
 DisplayTheBlocks(\r
   IN CONST EFI_DEVICE_PATH_PROTOCOL *DevPath,\r
   IN CONST UINT64                   Lba,\r
@@ -38,6 +37,7 @@ DisplayTheBlocks(
   EFI_STATUS                Status;\r
   SHELL_STATUS              ShellStatus;\r
   UINT8                     *Buffer;\r
+  UINT8                     *OriginalBuffer;\r
   UINTN                     BufferSize;\r
 \r
   ShellStatus = SHELL_SUCCESS;\r
@@ -53,11 +53,17 @@ DisplayTheBlocks(
   }\r
 \r
   BufferSize = BlockIo->Media->BlockSize * BlockCount;\r
+  if(BlockIo->Media->IoAlign == 0) {\r
+    BlockIo->Media->IoAlign = 1;\r
+  }\r
+\r
   if (BufferSize > 0) {\r
-    Buffer     = AllocateZeroPool(BufferSize);\r
+    OriginalBuffer = AllocateZeroPool(BufferSize + BlockIo->Media->IoAlign);\r
+    Buffer         = ALIGN_POINTER (OriginalBuffer,BlockIo->Media->IoAlign);\r
   } else {\r
     ShellPrintEx(-1,-1,L"  BlockSize: 0x%08x, BlockCount: 0x%08x\r\n", BlockIo->Media->BlockSize, BlockCount);\r
-    Buffer    = NULL;\r
+    OriginalBuffer = NULL;\r
+    Buffer         = NULL;\r
   }\r
 \r
   Status = BlockIo->ReadBlocks(BlockIo, BlockIo->Media->MediaId, Lba, BufferSize, Buffer);\r
@@ -79,8 +85,8 @@ DisplayTheBlocks(
     ShellStatus = SHELL_DEVICE_ERROR;\r
   }\r
 \r
-  if (Buffer != NULL) {\r
-    FreePool(Buffer);\r
+  if (OriginalBuffer != NULL) {\r
+    FreePool (OriginalBuffer);\r
   }\r
 \r
   gBS->CloseProtocol(BlockIoHandle, &gEfiBlockIoProtocolGuid, gImageHandle, NULL);\r