[PATCH 3/3] MdeModulePkg/SmmCore: Add Context in SmiHandlerProfileUnregister.

[Jiewen Yao <jiewen.yao@intel.com>]

The reason is that we observe that a platform may use same Handler
for different context.

In order to support Unregister such handler, we have to input
context information as well.

Cc: Jeff Fan <jeff.fan@intel.com>
Cc: Feng Tian <feng.tian@intel.com>
Cc: Star Zeng <star.zeng@intel.com>
Cc: Bret Barkelew <Bret.Barkelew@microsoft.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jiewen Yao <jiewen.yao@intel.com>
---
 MdeModulePkg/Core/PiSmmCore/PiSmmCore.h         |   8 +-
 MdeModulePkg/Core/PiSmmCore/SmiHandlerProfile.c | 103 ++++++++++++++++----
 2 files changed, 89 insertions(+), 22 deletions(-)

diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.h b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.h
index 45b9d97..c12805a 100644
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.h
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.h
@@ -1139,6 +1139,10 @@ SmiHandlerProfileRegisterHandler (
                          For the SmmChildDispatch protocol, the HandlerGuid
                          must be the GUID of SmmChildDispatch protocol.
   @param Handler         The SMI handler.
+  @param Context         The context of the SMI handler.
+                         If it is NOT NULL, it will be used to check what is registered.
+  @param ContextSize     The size of the context in bytes.
+                         If Context is NOT NULL, it will be used to check what is registered.
 
   @retval EFI_SUCCESS           The original record is removed.
   @retval EFI_NOT_FOUND         There is no record for the HandlerGuid and handler.
@@ -1148,7 +1152,9 @@ EFIAPI
 SmiHandlerProfileUnregisterHandler (
   IN SMI_HANDLER_PROFILE_PROTOCOL   *This,
   IN EFI_GUID                       *HandlerGuid,
-  IN EFI_SMM_HANDLER_ENTRY_POINT2   Handler
+  IN EFI_SMM_HANDLER_ENTRY_POINT2   Handler,
+  IN VOID                           *Context, OPTIONAL
+  IN UINTN                          ContextSize OPTIONAL
   );
 
 extern UINTN                    mFullSmramRangeCount;
diff --git a/MdeModulePkg/Core/PiSmmCore/SmiHandlerProfile.c b/MdeModulePkg/Core/PiSmmCore/SmiHandlerProfile.c
index f85c0f0..1e36039 100644
--- a/MdeModulePkg/Core/PiSmmCore/SmiHandlerProfile.c
+++ b/MdeModulePkg/Core/PiSmmCore/SmiHandlerProfile.c
@@ -1089,6 +1089,40 @@ SmmCoreFindHardwareSmiEntry (
 }
 
 /**
+  Convert EFI_SMM_USB_REGISTER_CONTEXT to SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT.
+
+  @param UsbContext                   A pointer to EFI_SMM_USB_REGISTER_CONTEXT
+  @param UsbContextSize               The size of EFI_SMM_USB_REGISTER_CONTEXT in bytes
+  @param SmiHandlerUsbContextSize     The size of SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT in bytes
+
+  @return SmiHandlerUsbContext  A pointer to SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT
+**/
+SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT *
+ConvertSmiHandlerUsbContext (
+  IN EFI_SMM_USB_REGISTER_CONTEXT   *UsbContext,
+  IN UINTN                          UsbContextSize,
+  OUT UINTN                         *SmiHandlerUsbContextSize
+  )
+{
+  UINTN                                     DevicePathSize;
+  SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT  *SmiHandlerUsbContext;
+
+  ASSERT (UsbContextSize == sizeof(EFI_SMM_USB_REGISTER_CONTEXT));
+
+  DevicePathSize = GetDevicePathSize (UsbContext->Device);
+  SmiHandlerUsbContext = AllocatePool (sizeof (SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT) + DevicePathSize);
+  if (SmiHandlerUsbContext == NULL) {
+    *SmiHandlerUsbContextSize = 0;
+    return NULL;
+  }
+  SmiHandlerUsbContext->Type = UsbContext->Type;
+  SmiHandlerUsbContext->DevicePathSize = (UINT32)DevicePathSize;
+  CopyMem (SmiHandlerUsbContext + 1, UsbContext->Device, DevicePathSize);
+  *SmiHandlerUsbContextSize = sizeof (SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT) + DevicePathSize;
+  return SmiHandlerUsbContext;
+}
+
+/**
   This function is called by SmmChildDispatcher module to report
   a new SMI handler is registered, to SmmCore.
 
@@ -1123,6 +1157,11 @@ SmiHandlerProfileRegisterHandler (
   SMI_ENTRY    *SmiEntry;
   LIST_ENTRY   *List;
 
+  if (((ContextSize == 0) && (Context != NULL)) ||
+      ((ContextSize != 0) && (Context == NULL))) {
+    return EFI_INVALID_PARAMETER;
+  }
+
   SmiHandler = AllocateZeroPool (sizeof (SMI_HANDLER));
   if (SmiHandler == NULL) {
     return EFI_OUT_OF_RESOURCES;
@@ -1131,33 +1170,24 @@ SmiHandlerProfileRegisterHandler (
   SmiHandler->Signature = SMI_HANDLER_SIGNATURE;
   SmiHandler->Handler = Handler;
   SmiHandler->CallerAddr = (UINTN)CallerAddress;
-  if (ContextSize != 0 && Context != NULL) {
+  SmiHandler->Context = Context;
+  SmiHandler->ContextSize = ContextSize;
+
+  if (Context != NULL) {
     if (CompareGuid (HandlerGuid, &gEfiSmmUsbDispatch2ProtocolGuid)) {
-      EFI_SMM_USB_REGISTER_CONTEXT              *UsbContext;
-      UINTN                                     DevicePathSize;
-      SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT  *SmiHandlerUsbContext;
-
-      ASSERT (ContextSize == sizeof(EFI_SMM_USB_REGISTER_CONTEXT));
-
-      UsbContext = (EFI_SMM_USB_REGISTER_CONTEXT *)Context;
-      DevicePathSize = GetDevicePathSize (UsbContext->Device);
-      SmiHandlerUsbContext = AllocatePool (sizeof (SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT) + DevicePathSize);
-      if (SmiHandlerUsbContext != NULL) {
-        SmiHandlerUsbContext->Type = UsbContext->Type;
-        SmiHandlerUsbContext->DevicePathSize = (UINT32)DevicePathSize;
-        CopyMem (SmiHandlerUsbContext + 1, UsbContext->Device, DevicePathSize);
-        SmiHandler->Context = SmiHandlerUsbContext;
-      }
+      SmiHandler->Context = ConvertSmiHandlerUsbContext (Context, ContextSize, &SmiHandler->ContextSize);
     } else {
       SmiHandler->Context = AllocateCopyPool (ContextSize, Context);
     }
   }
-  if (SmiHandler->Context != NULL) {
-    SmiHandler->ContextSize = ContextSize;
+  if (SmiHandler->Context == NULL) {
+    SmiHandler->ContextSize = 0;
   }
 
   SmiEntry = SmmCoreFindHardwareSmiEntry (HandlerGuid, TRUE);
   if (SmiEntry == NULL) {
+    FreePool (SmiHandler->Context);
+    FreePool (SmiHandler);
     return EFI_OUT_OF_RESOURCES;
   }
 
@@ -1178,6 +1208,10 @@ SmiHandlerProfileRegisterHandler (
                          For the SmmChildDispatch protocol, the HandlerGuid
                          must be the GUID of SmmChildDispatch protocol.
   @param Handler         The SMI handler.
+  @param Context         The context of the SMI handler.
+                         If it is NOT NULL, it will be used to check what is registered.
+  @param ContextSize     The size of the context in bytes.
+                         If Context is NOT NULL, it will be used to check what is registered.
 
   @retval EFI_SUCCESS           The original record is removed.
   @retval EFI_NOT_FOUND         There is no record for the HandlerGuid and handler.
@@ -1187,7 +1221,9 @@ EFIAPI
 SmiHandlerProfileUnregisterHandler (
   IN SMI_HANDLER_PROFILE_PROTOCOL   *This,
   IN EFI_GUID                       *HandlerGuid,
-  IN EFI_SMM_HANDLER_ENTRY_POINT2   Handler
+  IN EFI_SMM_HANDLER_ENTRY_POINT2   Handler,
+  IN VOID                           *Context, OPTIONAL
+  IN UINTN                          ContextSize OPTIONAL
   )
 {
   LIST_ENTRY   *Link;
@@ -1195,21 +1231,46 @@ SmiHandlerProfileUnregisterHandler (
   SMI_HANDLER  *SmiHandler;
   SMI_ENTRY    *SmiEntry;
   SMI_HANDLER  *TargetSmiHandler;
+  VOID         *SearchContext;
+  UINTN        SearchContextSize;
+
+  if (((ContextSize == 0) && (Context != NULL)) ||
+      ((ContextSize != 0) && (Context == NULL))) {
+    return EFI_INVALID_PARAMETER;
+  }
 
   SmiEntry = SmmCoreFindHardwareSmiEntry (HandlerGuid, FALSE);
   if (SmiEntry == NULL) {
     return EFI_NOT_FOUND;
   }
 
+  SearchContext = Context;
+  SearchContextSize = ContextSize;
+  if (Context != NULL) {
+    if (CompareGuid (HandlerGuid, &gEfiSmmUsbDispatch2ProtocolGuid)) {
+      SearchContext = ConvertSmiHandlerUsbContext (Context, ContextSize, &SearchContextSize);
+    }
+  }
+
   TargetSmiHandler = NULL;
   Head = &SmiEntry->SmiHandlers;
   for (Link = Head->ForwardLink; Link != Head; Link = Link->ForwardLink) {
     SmiHandler = CR (Link, SMI_HANDLER, Link, SMI_HANDLER_SIGNATURE);
     if (SmiHandler->Handler == Handler) {
-      TargetSmiHandler = SmiHandler;
-      break;
+      if ((SearchContext == NULL) ||
+          ((SearchContextSize == SmiHandler->ContextSize) && (CompareMem (SearchContext, SmiHandler->Context, SearchContextSize) == 0))) {
+        TargetSmiHandler = SmiHandler;
+        break;
+      }
+    }
+  }
+
+  if (SearchContext != NULL) {
+    if (CompareGuid (HandlerGuid, &gEfiSmmUsbDispatch2ProtocolGuid)) {
+      FreePool (SearchContext);
     }
   }
+
   if (TargetSmiHandler == NULL) {
     return EFI_NOT_FOUND;
   }
-- 
2.7.4.windows.1