From 17b28722008eab745ce186b72cd325944cbe6bf0 Mon Sep 17 00:00:00 2001
From: Zhiguang Liu <zhiguang.liu@intel.com>
Date: Fri, 1 Mar 2024 11:01:31 +0800
Subject: [PATCH] MdeModulePkg/SMM: Disallow unregister SMI handler in other
 SMI handler

In last patch, we add code support to unregister SMI handler inside
itself. However, the code doesn't support unregister SMI handler
insider other SMI handler. While this is not a must-have usage.
So add check to disallow unregister SMI handler in other SMI 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>
Signed-off-by: Zhiguang Liu <zhiguang.liu@intel.com>
Message-Id: <20240301030133.628-3-zhiguang.liu@intel.com>
Reviewed-by: Ray Ni <ray.ni@intel.com>
Reviewed-by: Laszlo Ersek <lersek@redhat.com>
---
 MdeModulePkg/Core/PiSmmCore/Smi.c | 29 +++++++++++++++++++++--------
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/MdeModulePkg/Core/PiSmmCore/Smi.c b/MdeModulePkg/Core/PiSmmCore/Smi.c
index 3489c130fd..b3a81ac877 100644
--- a/MdeModulePkg/Core/PiSmmCore/Smi.c
+++ b/MdeModulePkg/Core/PiSmmCore/Smi.c
@@ -8,7 +8,8 @@
 
 #include "PiSmmCore.h"
 
-LIST_ENTRY  mSmiEntryList = INITIALIZE_LIST_HEAD_VARIABLE (mSmiEntryList);
+SMI_HANDLER  *mCurrentSmiHandler = NULL;
+LIST_ENTRY   mSmiEntryList       = INITIALIZE_LIST_HEAD_VARIABLE (mSmiEntryList);
 
 SMI_ENTRY  mRootSmiEntry = {
   SMI_ENTRY_SIGNATURE,
@@ -142,13 +143,18 @@ SmiManage (
     // Link points to may be freed if unregister SMI handler.
     //
     Link = Link->ForwardLink;
-
-    Status = SmiHandler->Handler (
-                           (EFI_HANDLE)SmiHandler,
-                           Context,
-                           CommBuffer,
-                           CommBufferSize
-                           );
+    //
+    // Assign gCurrentSmiHandle before calling the SMI handler and
+    // set to NULL when it returns.
+    //
+    mCurrentSmiHandler = SmiHandler;
+    Status             = SmiHandler->Handler (
+                                       (EFI_HANDLE)SmiHandler,
+                                       Context,
+                                       CommBuffer,
+                                       CommBufferSize
+                                       );
+    mCurrentSmiHandler = NULL;
 
     switch (Status) {
       case EFI_INTERRUPT_PENDING:
@@ -328,6 +334,13 @@ SmiHandlerUnRegister (
     return EFI_INVALID_PARAMETER;
   }
 
+  //
+  // Do not allow to unregister SMI Handler inside other SMI Handler
+  //
+  if ((mCurrentSmiHandler != NULL) && (mCurrentSmiHandler != SmiHandler)) {
+    return EFI_INVALID_PARAMETER;
+  }
+
   SmiEntry = SmiHandler->SmiEntry;
 
   RemoveEntryList (&SmiHandler->Link);