From 3e0587010e36d36d6374c8ca9956e9bfbc8da7a8 Mon Sep 17 00:00:00 2001
From: ydong10 <ydong10@6f19259b-4bc3-4df7-8a09-765794883524>
Date: Fri, 27 Apr 2012 02:48:46 +0000
Subject: Enhance the error handling for AllocatePool and AllocatePages
 function.

Signed-off-by: Eric Dong <eric.dong@intel.com>
Reviewed-by: Liming Gao <liming.gao@intel.com>

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@13225 6f19259b-4bc3-4df7-8a09-765794883524
---
 MdeModulePkg/Core/Dxe/DxeMain.h  | 2 +-
 MdeModulePkg/Core/Dxe/Mem/Page.c | 6 +++++-
 MdeModulePkg/Core/Dxe/Mem/Pool.c | 8 ++++++--
 3 files changed, 12 insertions(+), 4 deletions(-)

(limited to 'MdeModulePkg/Core/Dxe')

diff --git a/MdeModulePkg/Core/Dxe/DxeMain.h b/MdeModulePkg/Core/Dxe/DxeMain.h
index 2a8161554a..4ec895c0a2 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain.h
+++ b/MdeModulePkg/Core/Dxe/DxeMain.h
@@ -1247,7 +1247,7 @@ CoreGetMemoryMap (
   @param  Buffer                 The address to return a pointer to the allocated
                                  pool
 
-  @retval EFI_INVALID_PARAMETER  PoolType not valid
+  @retval EFI_INVALID_PARAMETER  PoolType not valid or Buffer is NULL
   @retval EFI_OUT_OF_RESOURCES   Size exceeds max pool size or allocation failed.
   @retval EFI_SUCCESS            Pool successfully allocated.
 
diff --git a/MdeModulePkg/Core/Dxe/Mem/Page.c b/MdeModulePkg/Core/Dxe/Mem/Page.c
index 04844e3305..b4a62b9f0c 100644
--- a/MdeModulePkg/Core/Dxe/Mem/Page.c
+++ b/MdeModulePkg/Core/Dxe/Mem/Page.c
@@ -1,7 +1,7 @@
 /** @file
   UEFI Memory page management functions.
 
-Copyright (c) 2007 - 2011, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2007 - 2012, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -1103,6 +1103,10 @@ CoreAllocatePages (
     return EFI_INVALID_PARAMETER;
   }
 
+  if (Memory == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
   Alignment = EFI_DEFAULT_PAGE_ALLOCATION_ALIGNMENT;
 
   if  (MemoryType == EfiACPIReclaimMemory   ||
diff --git a/MdeModulePkg/Core/Dxe/Mem/Pool.c b/MdeModulePkg/Core/Dxe/Mem/Pool.c
index 2808905eab..a435b9760c 100644
--- a/MdeModulePkg/Core/Dxe/Mem/Pool.c
+++ b/MdeModulePkg/Core/Dxe/Mem/Pool.c
@@ -1,7 +1,7 @@
 /** @file
   UEFI Memory pool management functions.
 
-Copyright (c) 2006 - 2010, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2006 - 2012, Intel Corporation. All rights reserved.<BR>
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -167,7 +167,7 @@ LookupPoolHead (
   @param  Buffer                 The address to return a pointer to the allocated
                                  pool
 
-  @retval EFI_INVALID_PARAMETER  PoolType not valid
+  @retval EFI_INVALID_PARAMETER  PoolType not valid or Buffer is NULL. 
   @retval EFI_OUT_OF_RESOURCES   Size exceeds max pool size or allocation failed.
   @retval EFI_SUCCESS            Pool successfully allocated.
 
@@ -190,6 +190,10 @@ CoreAllocatePool (
     return EFI_INVALID_PARAMETER;
   }
 
+  if (Buffer == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
+
   *Buffer = NULL;
 
   //
-- 
cgit v1.2.3