From 27af6f9d0ac93cb3a2ed59bdb3a34745e099b163 Mon Sep 17 00:00:00 2001
From: mdkinney <mdkinney@6f19259b-4bc3-4df7-8a09-765794883524>
Date: Tue, 16 Feb 2010 04:07:39 +0000
Subject: [PATCH] Use Memory Allocation Library instance for modules of type
 DXE_SMM_DRIVER

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@10015 6f19259b-4bc3-4df7-8a09-765794883524
---
 .../SmmBaseHelper/SmmBaseHelper.c             | 88 ++++++++-----------
 .../SmmBaseHelper/SmmBaseHelper.inf           |  1 +
 2 files changed, 39 insertions(+), 50 deletions(-)

diff --git a/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.c b/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.c
index 347e83d01f..a8f5d13b60 100644
--- a/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.c
+++ b/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.c
@@ -24,6 +24,7 @@
 #include <Library/PeCoffLib.h>
 #include <Library/DevicePathLib.h>
 #include <Library/CacheMaintenanceLib.h>
+#include <Library/MemoryAllocationLib.h>
 #include <Guid/SmmBaseThunkCommunication.h>
 #include <Protocol/SmmBaseHelperReady.h>
 #include <Protocol/SmmCpu.h>
@@ -134,15 +135,10 @@ ConstructFrameworkSmst (
   VOID
   )
 {
-  EFI_STATUS            Status;
   EFI_SMM_SYSTEM_TABLE  *FrameworkSmst;
 
-  Status = gSmst->SmmAllocatePool (
-                    EfiRuntimeServicesData,
-                    sizeof (EFI_SMM_SYSTEM_TABLE),
-                    (VOID **)&FrameworkSmst
-                    );
-  ASSERT_EFI_ERROR (Status);
+  FrameworkSmst = (EFI_SMM_SYSTEM_TABLE  *)AllocatePool (sizeof (EFI_SMM_SYSTEM_TABLE));
+  ASSERT (FrameworkSmst != NULL);
 
   ///
   /// Copy same things from PI SMST to Framework SMST
@@ -160,13 +156,8 @@ ConstructFrameworkSmst (
   FrameworkSmst->Hdr.Revision = EFI_SMM_SYSTEM_TABLE_REVISION;
   CopyGuid (&FrameworkSmst->EfiSmmCpuIoGuid, &mEfiSmmCpuIoGuid);
 
-  Status = gSmst->SmmAllocatePool (
-                    EfiRuntimeServicesData,
-                    gSmst->NumberOfCpus * sizeof (EFI_SMM_CPU_SAVE_STATE),
-                    (VOID **)&FrameworkSmst->CpuSaveState
-                    );
-  ASSERT_EFI_ERROR (Status);
-  ZeroMem (FrameworkSmst->CpuSaveState, gSmst->NumberOfCpus * sizeof (EFI_SMM_CPU_SAVE_STATE));
+  FrameworkSmst->CpuSaveState = (EFI_SMM_CPU_SAVE_STATE *)AllocateZeroPool (gSmst->NumberOfCpus * sizeof (EFI_SMM_CPU_SAVE_STATE));
+  ASSERT (FrameworkSmst->CpuSaveState != NULL);
 
   ///
   /// Do not support floating point state now
@@ -335,8 +326,8 @@ LoadImage (
   }
 
 Error:
-  gSmst->SmmFreePages (Buffer, PageCount);
-  return Status;
+  FreePages ((VOID *)(UINTN)Buffer, PageCount);
+  return EFI_SUCCESS;
 }
 
 /** 
@@ -534,10 +525,9 @@ CallbackThunk (
 **/
 VOID
 RegisterCallback (
-  IN OUT SMMBASE_FUNCTION_DATA *FunctionData
+  IN OUT SMMBASE_FUNCTION_DATA  *FunctionData
   )
 {
-  EFI_STATUS     Status;
   CALLBACK_INFO  *Buffer;
 
   ///
@@ -547,36 +537,35 @@ RegisterCallback (
   ///
   /// Allocate buffer for callback thunk information
   ///
-  Status = gSmst->SmmAllocatePool (
-                    EfiRuntimeServicesCode,
-                    sizeof (CALLBACK_INFO),
-                    (VOID **)&Buffer
-                    );
-  if (!EFI_ERROR (Status)) {
-    ///
-    /// Fill SmmImageHandle and CallbackAddress into the thunk
-    ///
-    Buffer->SmmImageHandle = FunctionData->Args.RegisterCallback.SmmImageHandle;
-    Buffer->CallbackAddress = FunctionData->Args.RegisterCallback.CallbackAddress;
-
-    ///
-    /// Register the thunk code as a root SMI handler
-    ///
-    Status = gSmst->SmiHandlerRegister (
-                      CallbackThunk,
-                      NULL,
-                      &Buffer->DispatchHandle
-                      );
-    if (!EFI_ERROR (Status)) {
-      ///
-      /// Save this callback info
-      ///
-      InsertTailList (&mCallbackInfoListHead, &Buffer->Link);
-    } else {
-      gSmst->SmmFreePool (Buffer);
-    }
+  Buffer = (CALLBACK_INFO *)AllocatePool (sizeof (CALLBACK_INFO));
+  if (Buffer == NULL) {
+    FunctionData->Status = EFI_OUT_OF_RESOURCES;
+    return;
   }
-  FunctionData->Status = Status;
+
+  ///
+  /// Fill SmmImageHandle and CallbackAddress into the thunk
+  ///
+  Buffer->SmmImageHandle = FunctionData->Args.RegisterCallback.SmmImageHandle;
+  Buffer->CallbackAddress = FunctionData->Args.RegisterCallback.CallbackAddress;
+
+  ///
+  /// Register the thunk code as a root SMI handler
+  ///
+  FunctionData->Status = gSmst->SmiHandlerRegister (
+                                  CallbackThunk,
+                                  NULL,
+                                  &Buffer->DispatchHandle
+                                  );
+  if (EFI_ERROR (FunctionData->Status)) {
+    FreePool (Buffer);
+    return;
+  }
+
+  ///
+  /// Save this callback info
+  ///
+  InsertTailList (&mCallbackInfoListHead, &Buffer->Link);
 }
 
 
@@ -607,9 +596,8 @@ HelperFreePool (
   IN OUT SMMBASE_FUNCTION_DATA *FunctionData
   )
 {
-  FunctionData->Status = gSmst->SmmFreePool (
-                                  FunctionData->Args.FreePool.Buffer
-                                  );
+  FreePool (FunctionData->Args.FreePool.Buffer);
+  FunctionData->Status = EFI_SUCCESS;
 }
 
 /**
diff --git a/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.inf b/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.inf
index cd425adfa0..89ac7fb50c 100644
--- a/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.inf
+++ b/EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.inf
@@ -44,6 +44,7 @@
   PeCoffLib
   DevicePathLib
   CacheMaintenanceLib
+  MemoryAllocationLib
 
 [Guids]
   gEfiSmmBaseThunkCommunicationGuid