Clean up DEC files:

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

diff --git a/EdkCompatibilityPkg/Compatibility/SmmBaseOnSmmBase2Thunk/SmmBaseOnSmmBase2Thunk.c b/EdkCompatibilityPkg/Compatibility/SmmBaseOnSmmBase2Thunk/SmmBaseOnSmmBase2Thunk.c
index 8bbfd8049b6f7c8364f5cd527357999bf3c4ca35..187adff94175a1370c0ec684a1c77e3211ade8e6 100644 (file)
--- a/EdkCompatibilityPkg/Compatibility/SmmBaseOnSmmBase2Thunk/SmmBaseOnSmmBase2Thunk.c
+++ b/EdkCompatibilityPkg/Compatibility/SmmBaseOnSmmBase2Thunk/SmmBaseOnSmmBase2Thunk.c
@@ -8,8 +8,8 @@
   SMM BASE Protocol can be published immediately after SMM Base2 Protocol is installed to\r
   make SMM Base Protocol.InSmm() as early as possible.\r
 \r
-  Copyright (c) 2009 - 2010, Intel Corporation\r
-  All rights reserved. This program and the accompanying materials\r
+  Copyright (c) 2009 - 2010, Intel Corporation. All rights reserved.<BR>\r
+  This program and the accompanying materials\r
   are licensed and made available under the terms and conditions of the BSD License\r
   which accompanies this distribution.  The full text of the license may be found at\r
   http://opensource.org/licenses/bsd-license.php\r
@@ -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
@@ -75,6 +76,7 @@ SmmBaseHelperService (
   UINTN DataSize;\r
 \r
   mCommunicationData.FunctionData.Status = EFI_UNSUPPORTED;\r
+  mCommunicationData.FunctionData.SmmBaseImageHandle = mSmmBaseHandle;\r
 \r
   if ((mCommunicationData.FunctionData.Function != SmmBaseFunctionCommunicate) && IsInSmm()) {\r
     ///\r
@@ -132,7 +134,7 @@ SmmBaseRegister (
   IN      BOOLEAN                   LegacyIA32Binary\r
   )\r
 {\r
-  if (LegacyIA32Binary) {\r
+  if (mAtRuntime || LegacyIA32Binary) {\r
     return EFI_UNSUPPORTED;\r
   }\r
 \r
@@ -165,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
@@ -201,6 +207,10 @@ SmmBaseCommunicate (
   /// Note this is a runtime interface\r
   ///\r
 \r
+  if (CommunicationBuffer == NULL || BufferSize == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
   mCommunicationData.FunctionData.Function = SmmBaseFunctionCommunicate;\r
   mCommunicationData.FunctionData.Args.Communicate.ImageHandle = ImageHandle;\r
   mCommunicationData.FunctionData.Args.Communicate.CommunicationBuffer = CommunicationBuffer;\r
@@ -303,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
@@ -331,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
@@ -403,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
@@ -421,6 +457,8 @@ SmmBaseThunkMain (
   EFI_STATUS  Status;\r
   EFI_EVENT   Event;\r
 \r
+  mSmmBaseHandle = ImageHandle;\r
+\r
   //\r
   // Assume only one instance of SMM Base2 Protocol in the system\r
   // Locate SMM Base2 Protocol\r
@@ -432,7 +470,7 @@ SmmBaseThunkMain (
   // Assume only one instance of SMM Communication Protocol in the system\r
   // Locate SMM Communication Protocol\r
   //\r
-  gBS->LocateProtocol (&gEfiSmmCommunicationProtocolGuid, NULL, (VOID **) &mSmmCommunication);\r
+  Status = gBS->LocateProtocol (&gEfiSmmCommunicationProtocolGuid, NULL, (VOID **) &mSmmCommunication);\r
   ASSERT_EFI_ERROR (Status);\r
 \r
   //\r
@@ -442,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