]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Core/Dxe/SectionExtraction/CoreSectionExtraction.c
MdeModulePkg: Clean up source files
[mirror_edk2.git] / MdeModulePkg / Core / Dxe / SectionExtraction / CoreSectionExtraction.c
index cc129cc008f3583d1fd1a3c8e70fdf01cf119cc6..a1d5316c657db18908e310d845d7dee3997140d8 100644 (file)
@@ -27,7 +27,7 @@
   3) A support protocol is not found, and the data is not available to be read\r
      without it.  This results in EFI_PROTOCOL_ERROR.\r
 \r
-Copyright (c) 2006 - 2010, Intel Corporation. All rights reserved.<BR>\r
+Copyright (c) 2006 - 2018, 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
@@ -66,6 +66,11 @@ typedef struct {
   //\r
   UINTN                       EncapsulatedStreamHandle;\r
   EFI_GUID                    *EncapsulationGuid;\r
+  //\r
+  // If the section REQUIRES an extraction protocol, register for RPN\r
+  // when the required GUIDed extraction protocol becomes available.\r
+  //\r
+  EFI_EVENT                   Event;\r
 } CORE_SECTION_CHILD_NODE;\r
 \r
 #define CORE_SECTION_STREAM_SIGNATURE SIGNATURE_32('S','X','S','S')\r
@@ -91,7 +96,6 @@ typedef struct {
   CORE_SECTION_CHILD_NODE     *ChildNode;\r
   CORE_SECTION_STREAM_NODE    *ParentStream;\r
   VOID                        *Registration;\r
-  EFI_EVENT                   Event;\r
 } RPN_EVENT_CONTEXT;\r
 \r
 \r
@@ -269,7 +273,11 @@ IsValidSectionStream (
   SectionHeader = (EFI_COMMON_SECTION_HEADER *)SectionStream;\r
 \r
   while (TotalLength < SectionStreamLength) {\r
-    SectionLength = SECTION_SIZE (SectionHeader);\r
+    if (IS_SECTION2 (SectionHeader)) {\r
+      SectionLength = SECTION2_SIZE (SectionHeader);\r
+    } else {\r
+      SectionLength = SECTION_SIZE (SectionHeader);\r
+    }\r
     TotalLength += SectionLength;\r
 \r
     if (TotalLength == SectionStreamLength) {\r
@@ -432,7 +440,7 @@ OpenSectionStream (
   return OpenSectionStreamEx (\r
            SectionStreamLength,\r
            SectionStream,\r
-           TRUE,\r
+           FALSE,\r
            0,\r
            SectionStreamHandle\r
            );\r
@@ -475,7 +483,60 @@ ChildIsType (
     return TRUE;\r
   }\r
   GuidedSection = (EFI_GUID_DEFINED_SECTION * )(Stream->StreamBuffer + Child->OffsetInStream);\r
-  return CompareGuid (&GuidedSection->SectionDefinitionGuid, SectionDefinitionGuid);\r
+  if (IS_SECTION2 (GuidedSection)) {\r
+    return CompareGuid (&(((EFI_GUID_DEFINED_SECTION2 *) GuidedSection)->SectionDefinitionGuid), SectionDefinitionGuid);\r
+  } else {\r
+    return CompareGuid (&GuidedSection->SectionDefinitionGuid, SectionDefinitionGuid);\r
+  }\r
+}\r
+\r
+/**\r
+  Verify the Guided Section GUID by checking if there is the Guided Section GUID configuration table recorded the GUID itself.\r
+\r
+  @param GuidedSectionGuid          The Guided Section GUID.\r
+  @param GuidedSectionExtraction    A pointer to the pointer to the supported Guided Section Extraction Protocol\r
+                                    for the Guided Section.\r
+\r
+  @return TRUE      The GuidedSectionGuid could be identified, and the pointer to\r
+                    the Guided Section Extraction Protocol will be returned to *GuidedSectionExtraction.\r
+  @return FALSE     The GuidedSectionGuid could not be identified, or\r
+                    the Guided Section Extraction Protocol has not been installed yet.\r
+\r
+**/\r
+BOOLEAN\r
+VerifyGuidedSectionGuid (\r
+  IN  EFI_GUID                                  *GuidedSectionGuid,\r
+  OUT EFI_GUIDED_SECTION_EXTRACTION_PROTOCOL    **GuidedSectionExtraction\r
+  )\r
+{\r
+  EFI_GUID              *GuidRecorded;\r
+  VOID                  *Interface;\r
+  EFI_STATUS            Status;\r
+\r
+  Interface = NULL;\r
+\r
+  //\r
+  // Check if there is the Guided Section GUID configuration table recorded the GUID itself.\r
+  //\r
+  Status = EfiGetSystemConfigurationTable (GuidedSectionGuid, (VOID **) &GuidRecorded);\r
+  if (Status == EFI_SUCCESS) {\r
+    if (CompareGuid (GuidRecorded, GuidedSectionGuid)) {\r
+      //\r
+      // Found the recorded GuidedSectionGuid.\r
+      //\r
+      Status = CoreLocateProtocol (GuidedSectionGuid, NULL, (VOID **) &Interface);\r
+      if (!EFI_ERROR (Status) && Interface != NULL) {\r
+        //\r
+        // Found the supported Guided Section Extraction Porotocol for the Guided Section.\r
+        //\r
+        *GuidedSectionExtraction = (EFI_GUIDED_SECTION_EXTRACTION_PROTOCOL *) Interface;\r
+        return TRUE;\r
+      }\r
+      return FALSE;\r
+    }\r
+  }\r
+\r
+  return FALSE;\r
 }\r
 \r
 /**\r
@@ -500,17 +561,16 @@ NotifyGuidedExtraction (
   UINTN                                   NewStreamBufferSize;\r
   UINT32                                  AuthenticationStatus;\r
   RPN_EVENT_CONTEXT                       *Context;\r
-  \r
+\r
   Context = RpnContext;\r
-  \r
+\r
   GuidedHeader = (EFI_GUID_DEFINED_SECTION *) (Context->ParentStream->StreamBuffer + Context->ChildNode->OffsetInStream);\r
   ASSERT (GuidedHeader->CommonHeader.Type == EFI_SECTION_GUID_DEFINED);\r
-  \r
-  Status = gBS->LocateProtocol (Context->ChildNode->EncapsulationGuid, NULL, (VOID **)&GuidedExtraction);\r
-  if (EFI_ERROR (Status)) {\r
+\r
+  if (!VerifyGuidedSectionGuid (Context->ChildNode->EncapsulationGuid, &GuidedExtraction)) {\r
     return;\r
   }\r
-  \r
+\r
   Status = GuidedExtraction->ExtractSection (\r
                                GuidedExtraction,\r
                                GuidedHeader,\r
@@ -550,8 +610,9 @@ NotifyGuidedExtraction (
   //  Close the event when done.\r
   //\r
   gBS->CloseEvent (Event);\r
+  Context->ChildNode->Event = NULL;\r
   FreePool (Context);\r
-}  \r
+}\r
 \r
 /**\r
   Constructor for RPN event when a missing GUIDED_SECTION_EXTRACTION_PROTOCOL appears...\r
@@ -567,23 +628,23 @@ CreateGuidedExtractionRpnEvent (
   )\r
 {\r
   RPN_EVENT_CONTEXT *Context;\r
-  \r
+\r
   //\r
   // Allocate new event structure and context\r
   //\r
   Context = AllocatePool (sizeof (RPN_EVENT_CONTEXT));\r
   ASSERT (Context != NULL);\r
-  \r
+\r
   Context->ChildNode = ChildNode;\r
   Context->ParentStream = ParentStream;\r
\r
-  Context->Event = EfiCreateProtocolNotifyEvent (\r
-                    Context->ChildNode->EncapsulationGuid,\r
-                    TPL_NOTIFY,\r
-                    NotifyGuidedExtraction,\r
-                    Context,\r
-                    &Context->Registration\r
-                    );\r
+\r
+  Context->ChildNode->Event = EfiCreateProtocolNotifyEvent (\r
+                                Context->ChildNode->EncapsulationGuid,\r
+                                TPL_NOTIFY,\r
+                                NotifyGuidedExtraction,\r
+                                Context,\r
+                                &Context->Registration\r
+                                );\r
 }\r
 \r
 /**\r
@@ -611,7 +672,7 @@ EFI_STATUS
 CreateChildNode (\r
   IN     CORE_SECTION_STREAM_NODE              *Stream,\r
   IN     UINT32                                ChildOffset,\r
-     OUT CORE_SECTION_CHILD_NODE               **ChildNode\r
+  OUT    CORE_SECTION_CHILD_NODE               **ChildNode\r
   )\r
 {\r
   EFI_STATUS                                   Status;\r
@@ -625,7 +686,11 @@ CreateChildNode (
   UINT32                                       ScratchSize;\r
   UINTN                                        NewStreamBufferSize;\r
   UINT32                                       AuthenticationStatus;\r
-  UINT32                                       SectionLength;\r
+  VOID                                         *CompressionSource;\r
+  UINT32                                       CompressionSourceSize;\r
+  UINT32                                       UncompressedLength;\r
+  UINT8                                        CompressionType;\r
+  UINT16                                       GuidedSectionAttributes;\r
 \r
   CORE_SECTION_CHILD_NODE                      *Node;\r
 \r
@@ -634,7 +699,7 @@ CreateChildNode (
   //\r
   // Allocate a new node\r
   //\r
-  *ChildNode = AllocatePool (sizeof (CORE_SECTION_CHILD_NODE));\r
+  *ChildNode = AllocateZeroPool (sizeof (CORE_SECTION_CHILD_NODE));\r
   Node = *ChildNode;\r
   if (Node == NULL) {\r
     return EFI_OUT_OF_RESOURCES;\r
@@ -645,7 +710,11 @@ CreateChildNode (
   //\r
   Node->Signature = CORE_SECTION_CHILD_SIGNATURE;\r
   Node->Type = SectionHeader->Type;\r
-  Node->Size = SECTION_SIZE (SectionHeader);\r
+  if (IS_SECTION2 (SectionHeader)) {\r
+    Node->Size = SECTION2_SIZE (SectionHeader);\r
+  } else {\r
+    Node->Size = SECTION_SIZE (SectionHeader);\r
+  }\r
   Node->OffsetInStream = ChildOffset;\r
   Node->EncapsulatedStreamHandle = NULL_STREAM_HANDLE;\r
   Node->EncapsulationGuid = NULL;\r
@@ -658,27 +727,42 @@ CreateChildNode (
       //\r
       // Get the CompressionSectionHeader\r
       //\r
-      ASSERT (Node->Size >= sizeof (EFI_COMPRESSION_SECTION));\r
+      if (Node->Size < sizeof (EFI_COMPRESSION_SECTION)) {\r
+        CoreFreePool (Node);\r
+        return EFI_NOT_FOUND;\r
+      }\r
 \r
       CompressionHeader = (EFI_COMPRESSION_SECTION *) SectionHeader;\r
 \r
+      if (IS_SECTION2 (CompressionHeader)) {\r
+        CompressionSource = (VOID *) ((UINT8 *) CompressionHeader + sizeof (EFI_COMPRESSION_SECTION2));\r
+        CompressionSourceSize = (UINT32) (SECTION2_SIZE (CompressionHeader) - sizeof (EFI_COMPRESSION_SECTION2));\r
+        UncompressedLength = ((EFI_COMPRESSION_SECTION2 *) CompressionHeader)->UncompressedLength;\r
+        CompressionType = ((EFI_COMPRESSION_SECTION2 *) CompressionHeader)->CompressionType;\r
+      } else {\r
+        CompressionSource = (VOID *) ((UINT8 *) CompressionHeader + sizeof (EFI_COMPRESSION_SECTION));\r
+        CompressionSourceSize = (UINT32) (SECTION_SIZE (CompressionHeader) - sizeof (EFI_COMPRESSION_SECTION));\r
+        UncompressedLength = CompressionHeader->UncompressedLength;\r
+        CompressionType = CompressionHeader->CompressionType;\r
+      }\r
+\r
       //\r
       // Allocate space for the new stream\r
       //\r
-      if (CompressionHeader->UncompressedLength > 0) {\r
-        NewStreamBufferSize = CompressionHeader->UncompressedLength;\r
+      if (UncompressedLength > 0) {\r
+        NewStreamBufferSize = UncompressedLength;\r
         NewStreamBuffer = AllocatePool (NewStreamBufferSize);\r
         if (NewStreamBuffer == NULL) {\r
           CoreFreePool (Node);\r
           return EFI_OUT_OF_RESOURCES;\r
         }\r
 \r
-        if (CompressionHeader->CompressionType == EFI_NOT_COMPRESSED) {\r
+        if (CompressionType == EFI_NOT_COMPRESSED) {\r
           //\r
           // stream is not actually compressed, just encapsulated.  So just copy it.\r
           //\r
-          CopyMem (NewStreamBuffer, CompressionHeader + 1, NewStreamBufferSize);\r
-        } else if (CompressionHeader->CompressionType == EFI_STANDARD_COMPRESSION) {\r
+          CopyMem (NewStreamBuffer, CompressionSource, NewStreamBufferSize);\r
+        } else if (CompressionType == EFI_STANDARD_COMPRESSION) {\r
           //\r
           // Only support the EFI_SATNDARD_COMPRESSION algorithm.\r
           //\r
@@ -692,13 +776,19 @@ CreateChildNode (
 \r
           Status = Decompress->GetInfo (\r
                                  Decompress,\r
-                                 CompressionHeader + 1,\r
-                                 Node->Size - sizeof (EFI_COMPRESSION_SECTION),\r
+                                 CompressionSource,\r
+                                 CompressionSourceSize,\r
                                  (UINT32 *)&NewStreamBufferSize,\r
                                  &ScratchSize\r
                                  );\r
-          ASSERT_EFI_ERROR (Status);\r
-          ASSERT (NewStreamBufferSize == CompressionHeader->UncompressedLength);\r
+          if (EFI_ERROR (Status) || (NewStreamBufferSize != UncompressedLength)) {\r
+            CoreFreePool (Node);\r
+            CoreFreePool (NewStreamBuffer);\r
+            if (!EFI_ERROR (Status)) {\r
+              Status = EFI_BAD_BUFFER_SIZE;\r
+            }\r
+            return Status;\r
+          }\r
 \r
           ScratchBuffer = AllocatePool (ScratchSize);\r
           if (ScratchBuffer == NULL) {\r
@@ -709,15 +799,19 @@ CreateChildNode (
 \r
           Status = Decompress->Decompress (\r
                                  Decompress,\r
-                                 CompressionHeader + 1,\r
-                                 Node->Size - sizeof (EFI_COMPRESSION_SECTION),\r
+                                 CompressionSource,\r
+                                 CompressionSourceSize,\r
                                  NewStreamBuffer,\r
                                  (UINT32)NewStreamBufferSize,\r
                                  ScratchBuffer,\r
                                  ScratchSize\r
                                  );\r
-          ASSERT_EFI_ERROR (Status);\r
           CoreFreePool (ScratchBuffer);\r
+          if (EFI_ERROR (Status)) {\r
+            CoreFreePool (Node);\r
+            CoreFreePool (NewStreamBuffer);\r
+            return Status;\r
+          }\r
         }\r
       } else {\r
         NewStreamBuffer = NULL;\r
@@ -740,9 +834,14 @@ CreateChildNode (
 \r
     case EFI_SECTION_GUID_DEFINED:\r
       GuidedHeader = (EFI_GUID_DEFINED_SECTION *) SectionHeader;\r
-      Node->EncapsulationGuid = &GuidedHeader->SectionDefinitionGuid;\r
-      Status = CoreLocateProtocol (Node->EncapsulationGuid, NULL, (VOID **)&GuidedExtraction);\r
-      if (!EFI_ERROR (Status) && GuidedExtraction != NULL) {\r
+      if (IS_SECTION2 (GuidedHeader)) {\r
+        Node->EncapsulationGuid = &(((EFI_GUID_DEFINED_SECTION2 *) GuidedHeader)->SectionDefinitionGuid);\r
+        GuidedSectionAttributes = ((EFI_GUID_DEFINED_SECTION2 *) GuidedHeader)->Attributes;\r
+      } else {\r
+        Node->EncapsulationGuid = &GuidedHeader->SectionDefinitionGuid;\r
+        GuidedSectionAttributes = GuidedHeader->Attributes;\r
+      }\r
+      if (VerifyGuidedSectionGuid (Node->EncapsulationGuid, &GuidedExtraction)) {\r
         //\r
         // NewStreamBuffer is always allocated by ExtractSection... No caller\r
         // allocation here.\r
@@ -763,7 +862,7 @@ CreateChildNode (
         // Make sure we initialize the new stream with the correct\r
         // authentication status for both aggregate and local status fields.\r
         //\r
-        if ((GuidedHeader->Attributes & EFI_GUIDED_SECTION_AUTH_STATUS_VALID) != 0) {\r
+        if ((GuidedSectionAttributes & EFI_GUIDED_SECTION_AUTH_STATUS_VALID) != 0) {\r
           //\r
           // OR in the parent stream's aggregate status.\r
           //\r
@@ -792,10 +891,10 @@ CreateChildNode (
         //\r
         // There's no GUIDed section extraction protocol available.\r
         //\r
-        if ((GuidedHeader->Attributes & EFI_GUIDED_SECTION_PROCESSING_REQUIRED) != 0) {\r
+        if ((GuidedSectionAttributes & EFI_GUIDED_SECTION_PROCESSING_REQUIRED) != 0) {\r
           //\r
-          // If the section REQUIRES an extraction protocol, register for RPN \r
-          // when the required GUIDed extraction protocol becomes available. \r
+          // If the section REQUIRES an extraction protocol, register for RPN\r
+          // when the required GUIDed extraction protocol becomes available.\r
           //\r
           CreateGuidedExtractionRpnEvent (Stream, Node);\r
         } else {\r
@@ -804,14 +903,27 @@ CreateChildNode (
           //\r
           AuthenticationStatus = Stream->AuthenticationStatus;\r
 \r
-          SectionLength = SECTION_SIZE (GuidedHeader);\r
-          Status = OpenSectionStreamEx (\r
-                     SectionLength - GuidedHeader->DataOffset,\r
-                     (UINT8 *) GuidedHeader + GuidedHeader->DataOffset,\r
-                     TRUE,\r
-                     AuthenticationStatus,\r
-                     &Node->EncapsulatedStreamHandle\r
-                     );\r
+          if ((GuidedSectionAttributes & EFI_GUIDED_SECTION_AUTH_STATUS_VALID) == EFI_GUIDED_SECTION_AUTH_STATUS_VALID) {\r
+            AuthenticationStatus |= EFI_AUTH_STATUS_IMAGE_SIGNED | EFI_AUTH_STATUS_NOT_TESTED;\r
+          }\r
+\r
+          if (IS_SECTION2 (GuidedHeader)) {\r
+            Status = OpenSectionStreamEx (\r
+                       SECTION2_SIZE (GuidedHeader) - ((EFI_GUID_DEFINED_SECTION2 *) GuidedHeader)->DataOffset,\r
+                       (UINT8 *) GuidedHeader + ((EFI_GUID_DEFINED_SECTION2 *) GuidedHeader)->DataOffset,\r
+                       TRUE,\r
+                       AuthenticationStatus,\r
+                       &Node->EncapsulatedStreamHandle\r
+                       );\r
+          } else {\r
+            Status = OpenSectionStreamEx (\r
+                       SECTION_SIZE (GuidedHeader) - ((EFI_GUID_DEFINED_SECTION *) GuidedHeader)->DataOffset,\r
+                       (UINT8 *) GuidedHeader + ((EFI_GUID_DEFINED_SECTION *) GuidedHeader)->DataOffset,\r
+                       TRUE,\r
+                       AuthenticationStatus,\r
+                       &Node->EncapsulatedStreamHandle\r
+                       );\r
+          }\r
           if (EFI_ERROR (Status)) {\r
             CoreFreePool (Node);\r
             return Status;\r
@@ -1054,7 +1166,7 @@ FindStreamNode (
                                 non-null on input, then the buffer is caller\r
                                 allocated.  If Buffer is NULL, then the buffer\r
                                 is callee allocated.  In either case, the\r
-                                requried buffer size is returned in *BufferSize.\r
+                                required buffer size is returned in *BufferSize.\r
   @param  BufferSize            On input, indicates the size of *Buffer if\r
                                 *Buffer is non-null on input.  On output,\r
                                 indicates the required size (allocated size if\r
@@ -1075,6 +1187,7 @@ FindStreamNode (
                                 function returns anything other than\r
                                 EFI_SUCCESS, the value of *AuthenticationStatus\r
                                 is undefined.\r
+  @param  IsFfs3Fv              Indicates the FV format.\r
 \r
   @retval EFI_SUCCESS           Section was retrieved successfully\r
   @retval EFI_PROTOCOL_ERROR    A GUID defined section was encountered in the\r
@@ -1105,7 +1218,8 @@ GetSection (
   IN UINTN                                              SectionInstance,\r
   IN VOID                                               **Buffer,\r
   IN OUT UINTN                                          *BufferSize,\r
-  OUT UINT32                                            *AuthenticationStatus\r
+  OUT UINT32                                            *AuthenticationStatus,\r
+  IN BOOLEAN                                            IsFfs3Fv\r
   )\r
 {\r
   CORE_SECTION_STREAM_NODE                              *StreamNode;\r
@@ -1118,8 +1232,10 @@ GetSection (
   UINTN                                                 Instance;\r
   UINT8                                                 *CopyBuffer;\r
   UINTN                                                 SectionSize;\r
+  EFI_COMMON_SECTION_HEADER                             *Section;\r
 \r
 \r
+  ChildStreamNode = NULL;\r
   OldTpl = CoreRaiseTpl (TPL_NOTIFY);\r
   Instance = SectionInstance + 1;\r
 \r
@@ -1158,8 +1274,22 @@ GetSection (
     if (EFI_ERROR (Status)) {\r
       goto GetSection_Done;\r
     }\r
-    CopySize = ChildNode->Size - sizeof (EFI_COMMON_SECTION_HEADER);\r
-    CopyBuffer = ChildStreamNode->StreamBuffer + ChildNode->OffsetInStream + sizeof (EFI_COMMON_SECTION_HEADER);\r
+\r
+    Section = (EFI_COMMON_SECTION_HEADER *) (ChildStreamNode->StreamBuffer + ChildNode->OffsetInStream);\r
+\r
+    if (IS_SECTION2 (Section)) {\r
+      ASSERT (SECTION2_SIZE (Section) > 0x00FFFFFF);\r
+      if (!IsFfs3Fv) {\r
+        DEBUG ((DEBUG_ERROR, "It is a FFS3 formatted section in a non-FFS3 formatted FV.\n"));\r
+        Status = EFI_NOT_FOUND;\r
+        goto GetSection_Done;\r
+      }\r
+      CopySize = SECTION2_SIZE (Section) - sizeof (EFI_COMMON_SECTION_HEADER2);\r
+      CopyBuffer = (UINT8 *) Section + sizeof (EFI_COMMON_SECTION_HEADER2);\r
+    } else {\r
+      CopySize = SECTION_SIZE (Section) - sizeof (EFI_COMMON_SECTION_HEADER);\r
+      CopyBuffer = (UINT8 *) Section + sizeof (EFI_COMMON_SECTION_HEADER);\r
+    }\r
     *AuthenticationStatus = ExtractedAuthenticationStatus;\r
   }\r
 \r
@@ -1214,8 +1344,13 @@ FreeChildNode (
     // If it's an encapsulating section, we close the resulting section stream.\r
     // CloseSectionStream will free all memory associated with the stream.\r
     //\r
-    CloseSectionStream (ChildNode->EncapsulatedStreamHandle);\r
+    CloseSectionStream (ChildNode->EncapsulatedStreamHandle, TRUE);\r
   }\r
+\r
+  if (ChildNode->Event != NULL) {\r
+    gBS->CloseEvent (ChildNode->Event);\r
+  }\r
+\r
   //\r
   // Last, free the child node itself\r
   //\r
@@ -1227,6 +1362,8 @@ FreeChildNode (
   SEP member function.  Deletes an existing section stream\r
 \r
   @param  StreamHandleToClose    Indicates the stream to close\r
+  @param  FreeStreamBuffer       TRUE - Need to free stream buffer;\r
+                                 FALSE - No need to free stream buffer.\r
 \r
   @retval EFI_SUCCESS            The section stream is closed sucessfully.\r
   @retval EFI_OUT_OF_RESOURCES   Memory allocation failed.\r
@@ -1237,7 +1374,8 @@ FreeChildNode (
 EFI_STATUS\r
 EFIAPI\r
 CloseSectionStream (\r
-  IN  UINTN                                     StreamHandleToClose\r
+  IN  UINTN                                     StreamHandleToClose,\r
+  IN  BOOLEAN                                   FreeStreamBuffer\r
   )\r
 {\r
   CORE_SECTION_STREAM_NODE                      *StreamNode;\r
@@ -1262,7 +1400,9 @@ CloseSectionStream (
       ChildNode = CHILD_SECTION_NODE_FROM_LINK (Link);\r
       FreeChildNode (ChildNode);\r
     }\r
-    CoreFreePool (StreamNode->StreamBuffer);\r
+    if (FreeStreamBuffer) {\r
+      CoreFreePool (StreamNode->StreamBuffer);\r
+    }\r
     CoreFreePool (StreamNode);\r
     Status = EFI_SUCCESS;\r
   } else {\r
@@ -1411,7 +1551,9 @@ CustomGuidedSectionExtract (
     //\r
     AllocatedOutputBuffer = AllocatePool (OutputBufferSize);\r
     if (AllocatedOutputBuffer == NULL) {\r
-      FreePool (ScratchBuffer);\r
+      if (ScratchBuffer != NULL) {\r
+        FreePool (ScratchBuffer);\r
+      }\r
       return EFI_OUT_OF_RESOURCES;\r
     }\r
     *OutputBuffer = AllocatedOutputBuffer;\r