Free the buffer allocated by GetSectionFromAnyFv() when exit, and add Error Status...

Signed-off-by: Star Zeng <star.zeng@intel.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>
git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@13933 6f19259b-4bc3-4df7-8a09-765794883524

diff --git a/EdkCompatibilityPkg/Compatibility/BootScriptSaveOnS3SaveStateThunk/ScriptSave.c b/EdkCompatibilityPkg/Compatibility/BootScriptSaveOnS3SaveStateThunk/ScriptSave.c
index eb8de8ce6d7e7c9f352822fb7ec48166d540a1f7..4a63977115bc78ed0dd7c8919ffd775720a94add 100644 (file)
--- a/EdkCompatibilityPkg/Compatibility/BootScriptSaveOnS3SaveStateThunk/ScriptSave.c
+++ b/EdkCompatibilityPkg/Compatibility/BootScriptSaveOnS3SaveStateThunk/ScriptSave.c
@@ -2,7 +2,7 @@
   Implementation for S3 Boot Script Save thunk driver.\r
   This thunk driver consumes PI S3SaveState protocol to produce framework S3BootScriptSave Protocol \r
   \r
-  Copyright (c) 2010 - 2011, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2010 - 2012, 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
@@ -830,7 +830,7 @@ InitializeScriptSaveOnS3SaveState (
     // This is the first-time loaded by DXE core. reload itself to NVS mem\r
     //\r
     //\r
-    // A workarouond: Here we install a dummy handle\r
+    // A workaround: Here we install a dummy handle\r
     //\r
     NewImageHandle = NULL;\r
     Status = gBS->InstallProtocolInterface (\r
@@ -839,6 +839,7 @@ InitializeScriptSaveOnS3SaveState (
                     EFI_NATIVE_INTERFACE,\r
                     NULL\r
                     );\r
+    ASSERT_EFI_ERROR (Status);\r
 \r
     Status = GetSectionFromAnyFv  (\r
                &gEfiCallerIdGuid,\r
@@ -847,15 +848,14 @@ InitializeScriptSaveOnS3SaveState (
                (VOID **) &Buffer,\r
                &BufferSize\r
                );\r
+    ASSERT_EFI_ERROR (Status);\r
     ImageContext.Handle    = Buffer;\r
     ImageContext.ImageRead = PeCoffLoaderImageReadFromMemory;\r
     //\r
     // Get information about the image being loaded\r
     //\r
     Status = PeCoffLoaderGetImageInfo (&ImageContext);\r
-    if (EFI_ERROR (Status)) {\r
-      return Status;\r
-    }\r
+    ASSERT_EFI_ERROR (Status);\r
 \r
     MemoryAddress = SIZE_4GB - 1;\r
     PageNumber    = EFI_SIZE_TO_PAGES (BufferSize + ImageContext.SectionAlignment);\r
@@ -865,9 +865,7 @@ InitializeScriptSaveOnS3SaveState (
                      PageNumber,\r
                      &MemoryAddress\r
                      );\r
-    if (EFI_ERROR (Status)) {\r
-      return EFI_OUT_OF_RESOURCES;\r
-    }\r
+    ASSERT_EFI_ERROR (Status);\r
     ImageContext.ImageAddress = (PHYSICAL_ADDRESS)(UINTN)MemoryAddress;\r
     //\r
     // Align buffer on section boundry\r
@@ -878,30 +876,26 @@ InitializeScriptSaveOnS3SaveState (
     // Load the image to our new buffer\r
     //\r
     Status = PeCoffLoaderLoadImage (&ImageContext);\r
-    if (EFI_ERROR (Status)) {\r
-      gBS->FreePages (MemoryAddress, PageNumber);\r
-      return Status;\r
-    }\r
+    ASSERT_EFI_ERROR (Status);\r
 \r
     //\r
     // Relocate the image in our new buffer\r
     //\r
     Status = PeCoffLoaderRelocateImage (&ImageContext);\r
+    ASSERT_EFI_ERROR (Status);\r
+\r
+    //\r
+    // Free the buffer allocated by ReadSection since the image has been relocated in the new buffer\r
+    //\r
+    gBS->FreePool (Buffer);\r
 \r
-    if (EFI_ERROR (Status)) {\r
-      PeCoffLoaderUnloadImage (&ImageContext);\r
-      gBS->FreePages (MemoryAddress, PageNumber);\r
-      return Status;\r
-    }\r
     //\r
     // Flush the instruction cache so the image data is written before we execute it\r
     //\r
     InvalidateInstructionCacheRange ((VOID *)(UINTN)ImageContext.ImageAddress, (UINTN)ImageContext.ImageSize);\r
     Status = ((EFI_IMAGE_ENTRY_POINT)(UINTN)(ImageContext.EntryPoint)) (NewImageHandle, SystemTable);\r
-    if (EFI_ERROR (Status)) {\r
-      gBS->FreePages (MemoryAddress, PageNumber);\r
-      return Status;\r
-    }\r
+    ASSERT_EFI_ERROR (Status);\r
+\r
     //\r
     // Additional step for BootScriptThunk integrity\r
     //\r
@@ -910,9 +904,7 @@ InitializeScriptSaveOnS3SaveState (
     // Allocate BootScriptThunkData\r
     //\r
     BootScriptThunkData = AllocatePool (sizeof (BOOT_SCRIPT_THUNK_DATA));\r
-    if (BootScriptThunkData == NULL) {\r
-      return EFI_OUT_OF_RESOURCES;\r
-    }\r
+    ASSERT (BootScriptThunkData != NULL);\r
 \r
     BootScriptThunkData->BootScriptThunkBase   = ImageContext.ImageAddress;\r
     BootScriptThunkData->BootScriptThunkLength = ImageContext.ImageSize;\r

diff --git a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c
index e671164a83e17bc16d7c44b57d27108a49778c00..d5698d4682646114047b3de0c19bb981927583fc 100644 (file)
--- a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c
+++ b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c
@@ -235,7 +235,7 @@ BootScriptExecutorEntryPoint (
       // This is the first-time loaded by DXE core. reload itself to NVS mem\r
       //\r
       //\r
-      // A workarouond: Here we install a dummy handle\r
+      // A workaround: Here we install a dummy handle\r
       //\r
       NewImageHandle = NULL;\r
       Status = gBS->InstallProtocolInterface (\r
@@ -244,6 +244,7 @@ BootScriptExecutorEntryPoint (
                   EFI_NATIVE_INTERFACE,\r
                   NULL\r
                   );\r
+      ASSERT_EFI_ERROR (Status);\r
 \r
       Status = GetSectionFromAnyFv  (\r
                  &gEfiCallerIdGuid,\r
@@ -252,15 +253,14 @@ BootScriptExecutorEntryPoint (
                  (VOID **) &Buffer,\r
                  &BufferSize\r
                  );\r
+      ASSERT_EFI_ERROR (Status);\r
       ImageContext.Handle    = Buffer;\r
       ImageContext.ImageRead = PeCoffLoaderImageReadFromMemory;\r
       //\r
       // Get information about the image being loaded\r
       //\r
       Status = PeCoffLoaderGetImageInfo (&ImageContext);\r
-      if (EFI_ERROR (Status)) {\r
-        return Status;\r
-      }\r
+      ASSERT_EFI_ERROR (Status);\r
       Pages = EFI_SIZE_TO_PAGES(BufferSize + ImageContext.SectionAlignment);\r
       FfsBuffer = 0xFFFFFFFF;\r
       Status = gBS->AllocatePages (\r
@@ -269,9 +269,7 @@ BootScriptExecutorEntryPoint (
                     Pages,\r
                     &FfsBuffer\r
                     );\r
-      if (EFI_ERROR (Status)) {\r
-        return EFI_OUT_OF_RESOURCES;\r
-      }\r
+      ASSERT_EFI_ERROR (Status);\r
       ImageContext.ImageAddress = (PHYSICAL_ADDRESS)(UINTN)FfsBuffer;\r
       //\r
       // Align buffer on section boundry\r
@@ -282,30 +280,26 @@ BootScriptExecutorEntryPoint (
       // Load the image to our new buffer\r
       //\r
       Status = PeCoffLoaderLoadImage (&ImageContext);\r
-      if (EFI_ERROR (Status)) {\r
-        gBS->FreePages (FfsBuffer, Pages);\r
-        return Status;\r
-      }\r
+      ASSERT_EFI_ERROR (Status);\r
 \r
       //\r
       // Relocate the image in our new buffer\r
       //\r
       Status = PeCoffLoaderRelocateImage (&ImageContext);\r
+      ASSERT_EFI_ERROR (Status);\r
+\r
+      //\r
+      // Free the buffer allocated by ReadSection since the image has been relocated in the new buffer\r
+      //\r
+      gBS->FreePool (Buffer);\r
 \r
-      if (EFI_ERROR (Status)) {\r
-        PeCoffLoaderUnloadImage (&ImageContext);\r
-        gBS->FreePages (FfsBuffer, Pages);\r
-        return Status;\r
-      }\r
       //\r
       // Flush the instruction cache so the image data is written before we execute it\r
       //\r
       InvalidateInstructionCacheRange ((VOID *)(UINTN)ImageContext.ImageAddress, (UINTN)ImageContext.ImageSize);\r
       Status = ((EFI_IMAGE_ENTRY_POINT)(UINTN)(ImageContext.EntryPoint)) (NewImageHandle, SystemTable);\r
-      if (EFI_ERROR (Status)) {\r
-        gBS->FreePages (FfsBuffer, Pages);\r
-        return Status;\r
-      }\r
+      ASSERT_EFI_ERROR (Status);\r
+\r
       //\r
       // Additional step for BootScript integrity\r
       // Save BootScriptExecutor image\r
@@ -334,9 +328,7 @@ BootScriptExecutorEntryPoint (
                       Pages,\r
                       &BootScriptExecutorBuffer\r
                       );\r
-      if (EFI_ERROR (Status)) {\r
-        return EFI_OUT_OF_RESOURCES;\r
-      }\r
+      ASSERT_EFI_ERROR (Status);\r
 \r
       EfiBootScriptExecutorVariable = (BOOT_SCRIPT_EXECUTOR_VARIABLE *)(UINTN)BootScriptExecutorBuffer;\r
       EfiBootScriptExecutorVariable->BootScriptExecutorEntrypoint = (UINTN) S3BootScriptExecutorEntryFunction ;\r