From 18c2a7621d5e06fdaf38002d5cbcc064f576cf64 Mon Sep 17 00:00:00 2001
From: Hao Wu <hao.a.wu@intel.com>
Date: Tue, 11 Oct 2016 10:21:01 +0800
Subject: [PATCH] BaseTools/GenPage: Avoid possible NULL pointer dereference

Cc: Liming Gao <liming.gao@intel.com>
Cc: Yonghong Zhu <yonghong.zhu@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Hao Wu <hao.a.wu@intel.com>
Reviewed-by: Liming Gao <liming.gao@intel.com>
---
 BaseTools/Source/C/GenPage/GenPage.c | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/BaseTools/Source/C/GenPage/GenPage.c b/BaseTools/Source/C/GenPage/GenPage.c
index ef69823fda..e22485c4ae 100644
--- a/BaseTools/Source/C/GenPage/GenPage.c
+++ b/BaseTools/Source/C/GenPage/GenPage.c
@@ -15,7 +15,7 @@
                           Directory-Ptr Directory {512}
                         ) {4}
 
-Copyright (c) 2006 - 2014, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2006 - 2016, 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        
@@ -146,6 +146,10 @@ Return:
   X64_PAGE_TABLE_ENTRY_2M                       *PageDirectoryEntry2MB;
 
   PageTable = (void *)malloc (EFI_PAGE_NUMBER * EFI_SIZE_OF_PAGE);
+  if (PageTable == NULL) {
+    Error (NULL, 0, 4001, "Resource", "memory cannot be allocated!");
+    return NULL;
+  }
   memset (PageTable, 0, (EFI_PAGE_NUMBER * EFI_SIZE_OF_PAGE));
   PageTablePtr = PageTable;
 
@@ -417,6 +421,10 @@ main (
   // Create X64 page table
   //
   BaseMemory = CreateIdentityMappingPageTables ();
+  if (BaseMemory == NULL) {
+    Error (NULL, 0, 4001, "Resource", "memory cannot be allocated!");
+    return STATUS_ERROR;
+  }
 
   //
   // Add page table to binary file