From 2ec8f0c6407f062441b205b900038933865c7b3c Mon Sep 17 00:00:00 2001
From: Zhiguang Liu <zhiguang.liu@intel.com>
Date: Fri, 1 Mar 2024 11:01:33 +0800
Subject: StandaloneMmPkg: Disallow unregister MMI handler in other MMI handler

In last patch, we add code support to unregister MMI handler inside
itself. However, the code doesn't support unregister MMI handler
insider other MMI handler. While this is not a must-have usage.
So add check to disallow unregister MMI handler in other MMI handler.

Cc: Liming Gao <gaoliming@byosoft.com.cn>
Cc: Jiaxin Wu <jiaxin.wu@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Ard Biesheuvel <ardb+tianocore@kernel.org>
Cc: Sami Mujawar <sami.mujawar@arm.com>
Cc: Ray Ni <ray.ni@intel.com>
Signed-off-by: Zhiguang Liu <zhiguang.liu@intel.com>
Message-Id: <20240301030133.628-5-zhiguang.liu@intel.com>
Reviewed-by: Ray Ni <ray.ni@intel.com>
Reviewed-by: Laszlo Ersek <lersek@redhat.com>
---
 StandaloneMmPkg/Core/Mmi.c | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

(limited to 'StandaloneMmPkg')

diff --git a/StandaloneMmPkg/Core/Mmi.c b/StandaloneMmPkg/Core/Mmi.c
index c1a1d76e85..9e52072bf7 100644
--- a/StandaloneMmPkg/Core/Mmi.c
+++ b/StandaloneMmPkg/Core/Mmi.c
@@ -36,8 +36,9 @@ typedef struct {
   MMI_ENTRY                     *MmiEntry;
 } MMI_HANDLER;
 
-LIST_ENTRY  mRootMmiHandlerList = INITIALIZE_LIST_HEAD_VARIABLE (mRootMmiHandlerList);
-LIST_ENTRY  mMmiEntryList       = INITIALIZE_LIST_HEAD_VARIABLE (mMmiEntryList);
+LIST_ENTRY   mRootMmiHandlerList = INITIALIZE_LIST_HEAD_VARIABLE (mRootMmiHandlerList);
+LIST_ENTRY   mMmiEntryList       = INITIALIZE_LIST_HEAD_VARIABLE (mMmiEntryList);
+MMI_HANDLER  *mCurrentMmiHandler = NULL;
 
 /**
   Finds the MMI entry for the requested handler type.
@@ -161,13 +162,19 @@ MmiManage (
     // get next node before handler is executed, since LIST_ENTRY that
     // Link points to may be freed if unregister MMI handler.
     //
-    Link   = Link->ForwardLink;
-    Status = MmiHandler->Handler (
-                           (EFI_HANDLE)MmiHandler,
-                           Context,
-                           CommBuffer,
-                           CommBufferSize
-                           );
+    Link = Link->ForwardLink;
+    //
+    // Assign gCurrentMmiHandle before calling the MMI handler and
+    // set to NULL when it returns.
+    //
+    mCurrentMmiHandler = MmiHandler;
+    Status             = MmiHandler->Handler (
+                                       (EFI_HANDLE)MmiHandler,
+                                       Context,
+                                       CommBuffer,
+                                       CommBufferSize
+                                       );
+    mCurrentMmiHandler = NULL;
 
     switch (Status) {
       case EFI_INTERRUPT_PENDING:
@@ -314,6 +321,13 @@ MmiHandlerUnRegister (
     return EFI_INVALID_PARAMETER;
   }
 
+  //
+  // Do not allow to unregister MMI Handler inside other MMI Handler
+  //
+  if ((mCurrentMmiHandler != NULL) && (mCurrentMmiHandler != MmiHandler)) {
+    return EFI_INVALID_PARAMETER;
+  }
+
   MmiEntry = MmiHandler->MmiEntry;
 
   RemoveEntryList (&MmiHandler->Link);
-- 
cgit v1.2.3