QuarkSocPkg/QncSmmDispatcher: Fix use after free issue

[mirror_edk2.git] / QuarkSocPkg / QuarkNorthCluster / Smm / DxeSmm / QncSmmDispatcher / QNCSmmCore.c

diff --git a/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmCore.c b/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmCore.c
index 47834064ffc121f9e86499792981fbd47904db9a..c2f75f86647ad4334ebfff2230dc68c559483132 100644 (file)
--- a/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmCore.c
+++ b/QuarkSocPkg/QuarkNorthCluster/Smm/DxeSmm/QncSmmDispatcher/QNCSmmCore.c
@@ -592,6 +592,7 @@ QNCSmmCoreDispatcher (
   BOOLEAN             ChildWasDispatched;\r
 \r
   DATABASE_RECORD     *RecordInDb;\r
+  DATABASE_RECORD     ActiveRecordInDb;\r
   LIST_ENTRY          *LinkInDb;\r
   DATABASE_RECORD     *RecordToExhaust;\r
   LIST_ENTRY          *LinkToExhaust;\r
@@ -613,6 +614,16 @@ QNCSmmCoreDispatcher (
   Status                = EFI_WARN_INTERRUPT_SOURCE_PENDING;\r
   ChildWasDispatched    = FALSE;\r
 \r
+  //\r
+  // Mark all child handlers as not processed\r
+  //\r
+  LinkInDb = GetFirstNode (&mPrivateData.CallbackDataBase);\r
+  while (!IsNull (&mPrivateData.CallbackDataBase, LinkInDb)) {\r
+    RecordInDb = DATABASE_RECORD_FROM_LINK (LinkInDb);\r
+    RecordInDb->Processed = FALSE;\r
+    LinkInDb = GetNextNode (&mPrivateData.CallbackDataBase, LinkInDb);\r
+  }\r
+\r
   //\r
   // Preserve Index registers\r
   //\r
@@ -634,6 +645,12 @@ QNCSmmCoreDispatcher (
 \r
       while ((!IsNull (&mPrivateData.CallbackDataBase, LinkInDb)) && (ResetListSearch == FALSE)) {\r
         RecordInDb = DATABASE_RECORD_FROM_LINK (LinkInDb);\r
+        //\r
+        // Make a copy of the record that contains an active SMI source,\r
+        // because un-register maybe invoked in callback function and\r
+        // RecordInDb maybe released\r
+        //\r
+        CopyMem (&ActiveRecordInDb, RecordInDb, sizeof (ActiveRecordInDb));\r
 \r
         //\r
         // look for the first active source\r
@@ -663,6 +680,13 @@ QNCSmmCoreDispatcher (
           //\r
           while (!IsNull (&mPrivateData.CallbackDataBase, LinkToExhaust)) {\r
             RecordToExhaust = DATABASE_RECORD_FROM_LINK (LinkToExhaust);\r
+            LinkToExhaust = GetNextNode (&mPrivateData.CallbackDataBase, LinkToExhaust);\r
+            if (RecordToExhaust->Processed) {\r
+              //\r
+              // Record has already been processed.  Continue with next child handler.\r
+              //\r
+              continue;\r
+            }\r
 \r
             if (CompareSources (&RecordToExhaust->SrcDesc, &ActiveSource)) {\r
               //\r
@@ -692,6 +716,11 @@ QNCSmmCoreDispatcher (
                 ContextsMatch = TRUE;\r
               }\r
 \r
+              //\r
+              // Mark this child handler so it will not be processed again\r
+              //\r
+              RecordToExhaust->Processed = TRUE;\r
+\r
               if (ContextsMatch) {\r
 \r
                 if (RecordToExhaust->BufferSize != 0) {\r
@@ -720,11 +749,13 @@ QNCSmmCoreDispatcher (
                   SxChildWasDispatched = TRUE;\r
                 }\r
               }\r
+              //\r
+              // Can not use RecordInDb after this point because Callback may have unregistered RecordInDb\r
+              // Restart processing of SMI handlers from the begining of the linked list because the\r
+              // state of the linked listed may have been modified due to unregister actions in the Callback.\r
+              //\r
+              LinkToExhaust = GetFirstNode (&mPrivateData.CallbackDataBase);\r
             }\r
-            //\r
-            // Get next record in DB\r
-            //\r
-            LinkToExhaust = GetNextNode (&mPrivateData.CallbackDataBase, &RecordToExhaust->Link);\r
           }\r
 \r
           if (RecordInDb->ClearSource == NULL) {\r
@@ -736,7 +767,7 @@ QNCSmmCoreDispatcher (
             //\r
             // This source requires special handling to clear\r
             //\r
-            RecordInDb->ClearSource (&ActiveSource);\r
+            ActiveRecordInDb.ClearSource (&ActiveSource);\r
           }\r
 \r
           if (ChildWasDispatched) {\r