From fc48db0da0523576bb6959b59a268feea0f0f750 Mon Sep 17 00:00:00 2001
From: ydong10 <ydong10@6f19259b-4bc3-4df7-8a09-765794883524>
Date: Fri, 5 Nov 2010 06:28:33 +0000
Subject: [PATCH] Refine code to make code run more safely.

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@11006 6f19259b-4bc3-4df7-8a09-765794883524
---
 PerformancePkg/Dp_App/DpTrace.c | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/PerformancePkg/Dp_App/DpTrace.c b/PerformancePkg/Dp_App/DpTrace.c
index 593c269f72..ca71b62517 100644
--- a/PerformancePkg/Dp_App/DpTrace.c
+++ b/PerformancePkg/Dp_App/DpTrace.c
@@ -50,7 +50,7 @@ GatherStatistics(
   MEASUREMENT_RECORD        Measurement;
   UINT64                    Duration;
   UINTN                     LogEntryKey;
-  UINTN                     TIndex;
+  INTN                      TIndex;
 
   LogEntryKey = 0;
   while ((LogEntryKey = GetPerformanceMeasurement (
@@ -143,11 +143,14 @@ DumpAllTrace(
   // Get Handle information
   //
   Size = 0;
-  HandleBuffer = NULL;
+  HandleBuffer = &TempHandle;
   Status  = gBS->LocateHandle (AllHandles, NULL, NULL, &Size, &TempHandle);
   if (Status == EFI_BUFFER_TOO_SMALL) {
     HandleBuffer = AllocatePool (Size);
     ASSERT (HandleBuffer != NULL);
+    if (HandleBuffer == NULL) {
+      return;
+    }
     Status  = gBS->LocateHandle (AllHandles, NULL, NULL, &Size, HandleBuffer);
   }
   if (EFI_ERROR (Status)) {
@@ -222,7 +225,9 @@ DumpAllTrace(
       );
     }
   }
-  FreePool (HandleBuffer);
+  if (HandleBuffer != &TempHandle) {
+    FreePool (HandleBuffer);
+  }
 }
 
 /** Gather and print Raw Trace Records.
@@ -457,11 +462,14 @@ ProcessHandles(
               (StringPtr == NULL) ? ALit_UNKNOWN: StringPtr);
 
   Size = 0;
-  HandleBuffer = NULL;
+  HandleBuffer = &TempHandle;
   Status  = gBS->LocateHandle (AllHandles, NULL, NULL, &Size, &TempHandle);
   if (Status == EFI_BUFFER_TOO_SMALL) {
     HandleBuffer = AllocatePool (Size);
     ASSERT (HandleBuffer != NULL);
+    if (HandleBuffer == NULL) {
+      return;
+    }
     Status  = gBS->LocateHandle (AllHandles, NULL, NULL, &Size, HandleBuffer);
   }
   if (EFI_ERROR (Status)) {
@@ -520,7 +528,9 @@ ProcessHandles(
       }
     }
   }
-  FreePool (HandleBuffer);
+  if (HandleBuffer != &TempHandle) {
+    FreePool (HandleBuffer);
+  }
   return Status;
 }
 
-- 
2.39.2