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 206be9b73684c456411745f4f603eb4fdff26ab6..3632ca8a7a4eb1e16db1a252067b1741f5e551b2 100644 (file)
--- a/ShellPkg/Library/UefiShellDebug1CommandsLib/Dblk.c
+++ b/ShellPkg/Library/UefiShellDebug1CommandsLib/Dblk.c
@@ -1,8 +1,8 @@
 /** @file\r
   Main file for Dblk shell Debug1 function.\r
 \r
-  Copyright (c) 2015, Hewlett-Packard Development Company, L.P.<BR>\r
-  Copyright (c) 2005 - 2011, Intel Corporation. All rights reserved.<BR>\r
+  (C) Copyright 2015 Hewlett-Packard Development Company, L.P.<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