[PATCH 2/3] MdeModulePkg/SmiHandlerProfile: Add Context support in Unregister

[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/Include/Guid/SmiHandlerProfile.h                          | 41 +++++++++++++++++++-
 MdeModulePkg/Library/SmmSmiHandlerProfileLib/SmmSmiHandlerProfileLib.c | 10 ++++-
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/MdeModulePkg/Include/Guid/SmiHandlerProfile.h b/MdeModulePkg/Include/Guid/SmiHandlerProfile.h
index b81631d..c5d29e8 100644
--- a/MdeModulePkg/Include/Guid/SmiHandlerProfile.h
+++ b/MdeModulePkg/Include/Guid/SmiHandlerProfile.h
@@ -150,6 +150,26 @@ extern EFI_GUID gSmiHandlerProfileGuid;
 
 typedef struct _SMI_HANDLER_PROFILE_PROTOCOL  SMI_HANDLER_PROFILE_PROTOCOL;
 
+/**
+  This function is called by SmmChildDispatcher module to report
+  a new SMI handler is registered, to SmmCore.
+
+  @param This            The protocol instance
+  @param HandlerGuid     The GUID to identify the type of the handler.
+                         For the SmmChildDispatch protocol, the HandlerGuid
+                         must be the GUID of SmmChildDispatch protocol.
+  @param Handler         The SMI handler.
+  @param CallerAddress   The address of the module who registers the SMI handler.
+  @param Context         The context of the SMI handler.
+                         For the SmmChildDispatch protocol, the Context
+                         must match the one defined for SmmChildDispatch protocol.
+  @param ContextSize     The size of the context in bytes.
+                         For the SmmChildDispatch protocol, the Context
+                         must match the one defined for SmmChildDispatch protocol.
+
+  @retval EFI_SUCCESS           The information is recorded.
+  @retval EFI_OUT_OF_RESOURCES  There is no enough resource to record the information.
+**/
 typedef
 EFI_STATUS
 (EFIAPI  *SMI_HANDLER_PROFILE_REGISTER_HANDLER) (
@@ -161,12 +181,31 @@ EFI_STATUS
   IN UINTN                          ContextSize OPTIONAL
   );
 
+/**
+  This function is called by SmmChildDispatcher module to report
+  an existing SMI handler is unregistered, to SmmCore.
+
+  @param This            The protocol instance
+  @param HandlerGuid     The GUID to identify the type of the handler.
+                         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.
+**/
 typedef
 EFI_STATUS
 (EFIAPI  *SMI_HANDLER_PROFILE_UNREGISTER_HANDLER) (
   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
   );
 
 struct _SMI_HANDLER_PROFILE_PROTOCOL {
diff --git a/MdeModulePkg/Library/SmmSmiHandlerProfileLib/SmmSmiHandlerProfileLib.c b/MdeModulePkg/Library/SmmSmiHandlerProfileLib/SmmSmiHandlerProfileLib.c
index 2edc71b..2911619 100644
--- a/MdeModulePkg/Library/SmmSmiHandlerProfileLib/SmmSmiHandlerProfileLib.c
+++ b/MdeModulePkg/Library/SmmSmiHandlerProfileLib/SmmSmiHandlerProfileLib.c
@@ -63,6 +63,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_UNSUPPORTED       The feature is unsupported.
@@ -72,11 +76,13 @@ EFI_STATUS
 EFIAPI
 SmiHandlerProfileUnregisterHandler (
   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
   )
 {
   if (mSmiHandlerProfile != NULL) {
-    return mSmiHandlerProfile->UnregisterHandler (mSmiHandlerProfile, HandlerGuid, Handler);
+    return mSmiHandlerProfile->UnregisterHandler (mSmiHandlerProfile, HandlerGuid, Handler, Context, ContextSize);
   }
   return EFI_UNSUPPORTED;
 }
-- 
2.7.4.windows.1