1 files changed, 4 insertions, 0 deletions
diff --git a/Silicon/Intel/KabylakeSiliconPkg/Pch/PchSmiDispatcher/Smm/IoTrap.c b/Silicon/Intel/KabylakeSiliconPkg/Pch/PchSmiDispatcher/Smm/IoTrap.c
index 5c3ea07646..ffb2105719 100644
--- a/Silicon/Intel/KabylakeSiliconPkg/Pch/PchSmiDispatcher/Smm/IoTrap.c
+++ b/Silicon/Intel/KabylakeSiliconPkg/Pch/PchSmiDispatcher/Smm/IoTrap.c
@@ -457,6 +457,8 @@ IoTrapRegister (
   DEBUG ((DEBUG_INFO, "RegisterContext->Address:%x! \n", RegisterContext->Address));
   DEBUG ((DEBUG_INFO, "RegisterContext->Length:%x! \n", RegisterContext->Length));
 
+  SmiHandlerProfileRegisterHandler (&gEfiSmmIoTrapDispatch2ProtocolGuid, DispatchFunction, (UINTN)RETURN_ADDRESS (0), RegisterContext, sizeof(*RegisterContext));
+
   return EFI_SUCCESS;
 }
 
@@ -584,6 +586,8 @@ IoTrapUnRegister (
     }
   }
 
+  SmiHandlerProfileUnregisterHandler (&gEfiSmmIoTrapDispatch2ProtocolGuid, RecordToDelete->Callback, &RecordToDelete->Context, sizeof(RecordToDelete->Context));
+
   RemoveEntryList (&RecordToDelete->Link);
   Status = gSmst->SmmFreePool (RecordToDelete);
   ASSERT_EFI_ERROR (Status);