From e2701217a1fa4f5edef14a870e8ed75474d45f04 Mon Sep 17 00:00:00 2001
From: lgao4 <lgao4@6f19259b-4bc3-4df7-8a09-765794883524>
Date: Fri, 19 Oct 2007 09:11:42 +0000
Subject: [PATCH] Update ExtractGuidedSectionLib instance to cover the same
 handler is registered multi times, the later will override the pre setting.

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@4183 6f19259b-4bc3-4df7-8a09-765794883524
---
 .../DxeExtractGuidedSectionLib.c              | 20 +++++++++++++++++++
 .../PeiExtractGuidedSectionLib.c              | 20 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/MdePkg/Library/DxeExtractGuidedSectionLib/DxeExtractGuidedSectionLib.c b/MdePkg/Library/DxeExtractGuidedSectionLib/DxeExtractGuidedSectionLib.c
index 164691389b..acffa264fe 100644
--- a/MdePkg/Library/DxeExtractGuidedSectionLib/DxeExtractGuidedSectionLib.c
+++ b/MdePkg/Library/DxeExtractGuidedSectionLib/DxeExtractGuidedSectionLib.c
@@ -114,12 +114,32 @@ ExtractGuidedSectionRegisterHandlers (
   IN        EXTRACT_GUIDED_SECTION_DECODE_HANDLER    DecodeHandler
   )
 {
+  UINT32 Index;
   //
   // Check input paramter.
   //
   if (SectionGuid == NULL) {
     return RETURN_INVALID_PARAMETER;
   }
+
+  //
+  // Search the match registered GetInfo handler for the input guided section.
+  //
+  for (Index = 0; Index < mNumberOfExtractHandler; Index ++) {
+    if (CompareGuid (&mExtractHandlerGuidTable[Index], SectionGuid)) {
+      break;
+    }
+  }
+
+  //
+  // If the guided handler has been registered before, only update its handler.
+  //
+  if (Index < mNumberOfExtractHandler) {
+    mExtractDecodeHandlerTable [Index] = DecodeHandler;
+    mExtractGetInfoHandlerTable [Index] = GetInfoHandler;
+    return RETURN_SUCCESS;
+  }
+  
   //
   // Check the global table is enough to contain new Handler.
   //
diff --git a/MdePkg/Library/PeiExtractGuidedSectionLib/PeiExtractGuidedSectionLib.c b/MdePkg/Library/PeiExtractGuidedSectionLib/PeiExtractGuidedSectionLib.c
index bbe6b50872..ce42a5dcd2 100644
--- a/MdePkg/Library/PeiExtractGuidedSectionLib/PeiExtractGuidedSectionLib.c
+++ b/MdePkg/Library/PeiExtractGuidedSectionLib/PeiExtractGuidedSectionLib.c
@@ -151,6 +151,7 @@ ExtractGuidedSectionRegisterHandlers (
   )
 {
   EFI_STATUS Status;
+  UINT32     Index;
   PEI_EXTRACT_GUIDED_SECTION_HANDLER_INFO *HandlerInfo;
 
   //
@@ -167,6 +168,25 @@ ExtractGuidedSectionRegisterHandlers (
   if (EFI_ERROR (Status)) {
     return Status;
   }
+
+  //
+  // Search the match registered GetInfo handler for the input guided section.
+  //
+  for (Index = 0; Index < HandlerInfo->NumberOfExtractHandler; Index ++) {
+    if (CompareGuid (&(HandlerInfo->ExtractHandlerGuidTable[Index]), SectionGuid)) {
+      break;
+    }
+  }
+
+  //
+  // If the guided handler has been registered before, only update its handler.
+  //
+  if (Index < HandlerInfo->NumberOfExtractHandler) {
+    HandlerInfo->ExtractDecodeHandlerTable [Index] = DecodeHandler;
+    HandlerInfo->ExtractGetInfoHandlerTable [Index] = GetInfoHandler;
+    return RETURN_SUCCESS;
+  }
+
   //
   // Check the global table is enough to contain new Handler.
   //
-- 
2.39.2