Security enhancement to SMM Base thunk drivers: Framework SMM drivers can't be loaded...

[mirror_edk2.git] / EdkCompatibilityPkg / Compatibility / SmmBaseOnSmmBase2Thunk / SmmBaseOnSmmBase2Thunk.c

diff --git a/EdkCompatibilityPkg/Compatibility/SmmBaseOnSmmBase2Thunk/SmmBaseOnSmmBase2Thunk.c b/EdkCompatibilityPkg/Compatibility/SmmBaseOnSmmBase2Thunk/SmmBaseOnSmmBase2Thunk.c
index e470da0af034e92f9378a21e34d28cb535f4e1ad..4834036a6bef612e7d8d7c71fb39a891fa473e4f 100644 (file)
--- a/EdkCompatibilityPkg/Compatibility/SmmBaseOnSmmBase2Thunk/SmmBaseOnSmmBase2Thunk.c
+++ b/EdkCompatibilityPkg/Compatibility/SmmBaseOnSmmBase2Thunk/SmmBaseOnSmmBase2Thunk.c
@@ -44,6 +44,7 @@ EFI_HANDLE                         mSmmBaseHandle = NULL;
 EFI_SMM_BASE2_PROTOCOL             *mSmmBase2 = NULL;\r
 EFI_SMM_COMMUNICATION_PROTOCOL     *mSmmCommunication = NULL;\r
 EFI_SMM_BASE_HELPER_READY_PROTOCOL *mSmmBaseHelperReady = NULL;\r
+BOOLEAN                            mAtRuntime = FALSE;\r
 \r
 /**\r
   Determine if in SMM mode.\r
@@ -133,7 +134,7 @@ SmmBaseRegister (
   IN      BOOLEAN                   LegacyIA32Binary\r
   )\r
 {\r
-  if (LegacyIA32Binary) {\r
+  if (mAtRuntime || LegacyIA32Binary) {\r
     return EFI_UNSUPPORTED;\r
   }\r
 \r
@@ -166,6 +167,10 @@ SmmBaseUnregister (
   IN      EFI_HANDLE                ImageHandle\r
   )\r
 {\r
+  if (mAtRuntime) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
   mCommunicationData.FunctionData.Function = SmmBaseFunctionUnregister;\r
   mCommunicationData.FunctionData.Args.UnRegister.ImageHandle = ImageHandle;\r
 \r
@@ -308,6 +313,10 @@ SmmBaseSmmAllocatePool (
   OUT     VOID                      **Buffer\r
   )\r
 {\r
+  if (mAtRuntime) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
   mCommunicationData.FunctionData.Function = SmmBaseFunctionAllocatePool;\r
   mCommunicationData.FunctionData.Args.AllocatePool.PoolType = PoolType;\r
   mCommunicationData.FunctionData.Args.AllocatePool.Size = Size;\r
@@ -336,6 +345,10 @@ SmmBaseSmmFreePool (
   IN      VOID                      *Buffer\r
   )\r
 {\r
+  if (mAtRuntime) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
   mCommunicationData.FunctionData.Function = SmmBaseFunctionFreePool;\r
   mCommunicationData.FunctionData.Args.FreePool.Buffer = Buffer;\r
 \r
@@ -408,6 +421,24 @@ EFI_SMM_BASE_PROTOCOL  mSmmBase = {
   SmmBaseGetSmstLocation\r
 };\r
 \r
+/**\r
+  Notification function on Exit Boot Services Event.\r
+\r
+  This function sets a flag indicating it is in Runtime phase.\r
+\r
+  @param  Event        Event whose notification function is being invoked\r
+  @param  Context      Pointer to the notification function's context\r
+**/\r
+VOID\r
+EFIAPI\r
+SmmBaseExitBootServicesEventNotify (\r
+  IN EFI_EVENT  Event,\r
+  IN VOID       *Context\r
+  )\r
+{\r
+  mAtRuntime = TRUE;\r
+}\r
+\r
 /**\r
   Entry Point for SMM Base Protocol on SMM Base2 Protocol Thunk driver.\r
 \r
@@ -449,6 +480,19 @@ SmmBaseThunkMain (
   Status = gBS->LocateProtocol (&gEfiSmmBaseHelperReadyProtocolGuid, NULL, (VOID **) &mSmmBaseHelperReady);\r
   ASSERT_EFI_ERROR (Status);\r
 \r
+  //\r
+  // Create event notification on Exit Boot Services event.\r
+  //\r
+  Status = gBS->CreateEventEx (\r
+                  EVT_NOTIFY_SIGNAL,\r
+                  TPL_NOTIFY,\r
+                  SmmBaseExitBootServicesEventNotify,\r
+                  NULL,\r
+                  &gEfiEventExitBootServicesGuid,\r
+                  &Event\r
+                  );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
   //\r
   // Create event on SetVirtualAddressMap() to convert mSmmCommunication from a physical address to a virtual address\r
   //\r