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

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