[PATCH 11/12] QuarkSocPkg/SmmChildDispatch: Add SmiHandlerProfile support.

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

Cc: Michael D Kinney <michael.d.kinney@intel.com>
Cc: Kelly Steele <kelly.steele@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jiewen Yao <jiewen.yao@intel.com>
---
 QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmCore.c         | 19 +++++++++++++++++--
 QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmDispatcher.inf |  3 ++-
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmCore.c b/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmCore.c
index c2f75f8..ea8b53c 100644
--- a/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmCore.c
+++ b/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmCore.c
@@ -2,7 +2,7 @@
 This driver is responsible for the registration of child drivers
 and the abstraction of the QNC SMI sources.
 
-Copyright (c) 2013-2016 Intel Corporation.
+Copyright (c) 2013-2017 Intel Corporation.
 
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
@@ -23,6 +23,8 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include "QNCSmm.h"
 #include "QNCSmmHelpers.h"
 
+#include <Library/SmiHandlerProfileLib.h>
+
 //
 // /////////////////////////////////////////////////////////////////////////////
 // MODULE / GLOBAL DATA
@@ -331,6 +333,7 @@ Returns:
   DATABASE_RECORD             *Record;
   QNC_SMM_QUALIFIED_PROTOCOL  *Qualified;
   INTN                        Index;
+  UINTN                       ContextSize;
 
   //
   // Check for invalid parameter
@@ -363,6 +366,7 @@ Returns:
   // Perform linked list housekeeping
   //
   Record->Signature = DATABASE_RECORD_SIGNATURE;
+  ContextSize = 0;
 
   switch (Qualified->Type) {
   //
@@ -383,6 +387,7 @@ Returns:
 
     InsertTailList (&mPrivateData.CallbackDataBase, &Record->Link);
     CopyMem (&Record->SrcDesc, &SX_SOURCE_DESC, sizeof (Record->SrcDesc));
+    ContextSize = sizeof(Record->ChildContext.Sx);
     //
     // use default clear source function
     //
@@ -425,6 +430,7 @@ Returns:
     InsertTailList (&mPrivateData.CallbackDataBase, &Record->Link);
     CopyMem (&Record->SrcDesc, &SW_SOURCE_DESC, sizeof (Record->SrcDesc));
     Record->BufferSize = sizeof (EFI_SMM_SW_REGISTER_CONTEXT);
+    ContextSize = sizeof(Record->ChildContext.Sw);
     //
     // use default clear source function
     //
@@ -434,6 +440,7 @@ Returns:
 
     InsertTailList (&mPrivateData.CallbackDataBase, &Record->Link);
     CopyMem (&Record->SrcDesc, &GPI_SOURCE_DESC, sizeof (Record->SrcDesc));
+    ContextSize = sizeof(Record->ChildContext.Gpi);
     //
     // use default clear source function
     //
@@ -450,6 +457,7 @@ Returns:
     InsertTailList (&mPrivateData.CallbackDataBase, &Record->Link);
     CopyMem (&Record->SrcDesc, &QNCN_SOURCE_DESCS[Record->ChildContext.QNCn.Type], sizeof (Record->SrcDesc));
     Record->ClearSource = QNCSmmQNCnClearSource;
+    ContextSize = sizeof(Record->ChildContext.QNCn);
     break;
 
   case PeriodicTimerType:
@@ -462,6 +470,7 @@ Returns:
     InsertTailList (&mPrivateData.CallbackDataBase, &Record->Link);
     Record->BufferSize = sizeof (EFI_SMM_PERIODIC_TIMER_CONTEXT);
     Record->ClearSource = QNCSmmPeriodicTimerClearSource;
+    ContextSize = sizeof(Record->ChildContext.PeriodicTimer);
     break;
 
   default:
@@ -488,6 +497,8 @@ Returns:
   //
   *DispatchHandle = (EFI_HANDLE) (&Record->Link);
 
+  SmiHandlerProfileRegisterHandler (Qualified->Guid, DispatchFunction, (UINTN)RETURN_ADDRESS (0), RegisterContext, ContextSize);
+
   return EFI_SUCCESS;
 
 Error:
@@ -522,6 +533,7 @@ Returns:
   DATABASE_RECORD *RecordToDelete;
   DATABASE_RECORD *RecordInDb;
   LIST_ENTRY      *LinkInDb;
+  QNC_SMM_QUALIFIED_PROTOCOL  *Qualified;
 
   if (DispatchHandle == NULL) {
     return EFI_INVALID_PARAMETER;
@@ -552,7 +564,10 @@ Returns:
   }
   if (SafeToDisable) {
     QNCSmmDisableSource( &RecordToDelete->SrcDesc );
-}
+  }
+
+  Qualified = QUALIFIED_PROTOCOL_FROM_GENERIC (This);
+  SmiHandlerProfileUnregisterHandler (Qualified->Guid, RecordToDelete->Callback);
 
   FreePool (RecordToDelete);
 
diff --git a/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmDispatcher.inf b/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmDispatcher.inf
index ed94825..23b806a 100644
--- a/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmDispatcher.inf
+++ b/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmDispatcher.inf
@@ -3,7 +3,7 @@
 #
 # This driver is responsible for the registration of child drivers
 #  and the abstraction of the ICH SMI sources.
-# Copyright (c) 2013-2015 Intel Corporation.
+# Copyright (c) 2013-2017 Intel Corporation.
 #
 # This program and the accompanying materials
 # are licensed and made available under the terms and conditions of the BSD License
@@ -66,6 +66,7 @@
   DevicePathLib
   S3IoLib
   QNCAccessLib
+  SmiHandlerProfileLib
 
 [Protocols]
   gEfiSmmCpuProtocolGuid                        # PROTOCOL ALWAYS_CONSUMED
-- 
2.7.4.windows.1