MdeModulePkg DxeCore: Enhance memory profile for memory leak detection

1. Implement include GetRecordingState/SetRecordingState/Record for
memory profile protocol.
2. Consume PcdMemoryProfilePropertyMask to support disable recording
at the start.
3. Consume PcdMemoryProfileDriverPath to control which drivers need
memory profile data.

Cc: Jiewen Yao <jiewen.yao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Star Zeng <star.zeng@intel.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>

5 files changed, 657 insertions, 233 deletions

diff --git a/MdeModulePkg/Core/Dxe/DxeMain.h b/MdeModulePkg/Core/Dxe/DxeMain.h
index e6b9114d2e..743221f675 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain.h
+++ b/MdeModulePkg/Core/Dxe/DxeMain.h
@@ -2780,11 +2780,13 @@ MemoryProfileInstallProtocol (
   @param DriverEntry    Image info.

   @param FileType       Image file type.

 

-  @retval TRUE          Register success.

-  @retval FALSE         Register fail.

+  @return EFI_SUCCESS           Register successfully.

+  @return EFI_UNSUPPORTED       Memory profile unsupported,

+                                or memory profile for the image is not required.

+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.

 

 **/

-BOOLEAN

+EFI_STATUS

 RegisterMemoryProfileImage (

   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,

   IN EFI_FV_FILETYPE            FileType

@@ -2795,11 +2797,13 @@ RegisterMemoryProfileImage (
 

   @param DriverEntry    Image info.

 

-  @retval TRUE          Unregister success.

-  @retval FALSE         Unregister fail.

+  @return EFI_SUCCESS           Unregister successfully.

+  @return EFI_UNSUPPORTED       Memory profile unsupported,

+                                or memory profile for the image is not required.

+  @return EFI_NOT_FOUND         The image is not found.

 

 **/

-BOOLEAN

+EFI_STATUS

 UnregisterMemoryProfileImage (

   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry

   );

@@ -2810,20 +2814,31 @@ UnregisterMemoryProfileImage (
   @param CallerAddress  Address of caller who call Allocate or Free.

   @param Action         This Allocate or Free action.

   @param MemoryType     Memory type.

+                        EfiMaxMemoryType means the MemoryType is unknown.

   @param Size           Buffer size.

   @param Buffer         Buffer address.

+  @param ActionString   String for memory profile action.

+                        Only needed for user defined allocate action.

 

-  @retval TRUE          Profile udpate success.

-  @retval FALSE         Profile update fail.

+  @return EFI_SUCCESS           Memory profile is updated.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported,

+                                or memory profile for the image is not required,

+                                or memory profile for the memory type is not required.

+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.

+  @return EFI_ABORTED           Memory profile recording is not enabled.

+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.

+  @return EFI_NOT_FOUND         No matched allocate info found for free action.

 

 **/

-BOOLEAN

+EFI_STATUS

+EFIAPI

 CoreUpdateProfile (

   IN EFI_PHYSICAL_ADDRESS   CallerAddress,

   IN MEMORY_PROFILE_ACTION  Action,

   IN EFI_MEMORY_TYPE        MemoryType,

   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool

-  IN VOID                   *Buffer

+  IN VOID                   *Buffer,

+  IN CHAR8                  *ActionString OPTIONAL

   );

 

 /**

diff --git a/MdeModulePkg/Core/Dxe/DxeMain.inf b/MdeModulePkg/Core/Dxe/DxeMain.inf
index e3e4d036db..450da579e0 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain.inf
+++ b/MdeModulePkg/Core/Dxe/DxeMain.inf
@@ -187,6 +187,7 @@
   gEfiMdeModulePkgTokenSpaceGuid.PcdMaxEfiSystemTablePointerAddress         ## CONSUMES

   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfileMemoryType                 ## CONSUMES

   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfilePropertyMask               ## CONSUMES

+  gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfileDriverPath                 ## CONSUMES

   gEfiMdeModulePkgTokenSpaceGuid.PcdPropertiesTableEnable                   ## CONSUMES

 

 # [Hob]

diff --git a/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c b/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
index 6626e10159..1ba8488a0e 100644
--- a/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
+++ b/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
@@ -17,6 +17,9 @@
 

 #define IS_UEFI_MEMORY_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0)

 

+#define GET_OCCUPIED_SIZE(ActualSize, Alignment) \

+  ((ActualSize) + (((Alignment) - ((ActualSize) & ((Alignment) - 1))) & ((Alignment) - 1)))

+

 typedef struct {

   UINT32                        Signature;

   MEMORY_PROFILE_CONTEXT        Context;

@@ -27,12 +30,14 @@ typedef struct {
   UINT32                        Signature;

   MEMORY_PROFILE_DRIVER_INFO    DriverInfo;

   LIST_ENTRY                    *AllocInfoList;

+  CHAR8                         *PdbString;

   LIST_ENTRY                    Link;

 } MEMORY_PROFILE_DRIVER_INFO_DATA;

 

 typedef struct {

   UINT32                        Signature;

   MEMORY_PROFILE_ALLOC_INFO     AllocInfo;

+  CHAR8                         *ActionString;

   LIST_ENTRY                    Link;

 } MEMORY_PROFILE_ALLOC_INFO_DATA;

 

@@ -58,7 +63,10 @@ GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA mMemoryProfileContext
 };

 GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mMemoryProfileContextPtr = NULL;

 

-BOOLEAN mMemoryProfileRecordingStatus = FALSE;

+BOOLEAN mMemoryProfileGettingStatus = FALSE;

+BOOLEAN mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;

+EFI_DEVICE_PATH_PROTOCOL *mMemoryProfileDriverPath;

+UINTN                    mMemoryProfileDriverPathSize;

 

 /**

   Get memory profile data.

@@ -69,6 +77,7 @@ BOOLEAN mMemoryProfileRecordingStatus = FALSE;
   @param[out]     ProfileBuffer     Profile buffer.

                       

   @return EFI_SUCCESS               Get the memory profile data successfully.

+  @return EFI_UNSUPPORTED           Memory profile is unsupported.

   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. 

                                     ProfileSize is updated with the size required.

 

@@ -90,7 +99,9 @@ ProfileProtocolGetData (
   @param[in] ImageSize          Image size.

   @param[in] FileType           File type of the image.

 

-  @return EFI_SUCCESS           Register success.

+  @return EFI_SUCCESS           Register successfully.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported,

+                                or memory profile for the image is not required.

   @return EFI_OUT_OF_RESOURCE   No enough resource for this register.

 

 **/

@@ -112,7 +123,9 @@ ProfileProtocolRegisterImage (
   @param[in] ImageBase          Image base address.

   @param[in] ImageSize          Image size.

 

-  @return EFI_SUCCESS           Unregister success.

+  @return EFI_SUCCESS           Unregister successfully.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported,

+                                or memory profile for the image is not required.

   @return EFI_NOT_FOUND         The image is not found.

 

 **/

@@ -125,10 +138,83 @@ ProfileProtocolUnregisterImage (
   IN UINT64                             ImageSize

   );

 

+/**

+  Get memory profile recording state.

+

+  @param[in]  This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.

+  @param[out] RecordingState    Recording state.

+

+  @return EFI_SUCCESS           Memory profile recording state is returned.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported.

+  @return EFI_INVALID_PARAMETER RecordingState is NULL.

+

+**/

+EFI_STATUS

+EFIAPI

+ProfileProtocolGetRecordingState (

+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,

+  OUT BOOLEAN                           *RecordingState

+  );

+

+/**

+  Set memory profile recording state.

+

+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.

+  @param[in] RecordingState     Recording state.

+

+  @return EFI_SUCCESS           Set memory profile recording state successfully.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported.

+

+**/

+EFI_STATUS

+EFIAPI

+ProfileProtocolSetRecordingState (

+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,

+  IN BOOLEAN                            RecordingState

+  );

+

+/**

+  Record memory profile of multilevel caller.

+

+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.

+  @param[in] CallerAddress      Address of caller.

+  @param[in] Action             Memory profile action.

+  @param[in] MemoryType         Memory type.

+                                EfiMaxMemoryType means the MemoryType is unknown.

+  @param[in] Buffer             Buffer address.

+  @param[in] Size               Buffer size.

+  @param[in] ActionString       String for memory profile action.

+                                Only needed for user defined allocate action.

+

+  @return EFI_SUCCESS           Memory profile is updated.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported,

+                                or memory profile for the image is not required,

+                                or memory profile for the memory type is not required.

+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.

+  @return EFI_ABORTED           Memory profile recording is not enabled.

+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.

+  @return EFI_NOT_FOUND         No matched allocate info found for free action.

+

+**/

+EFI_STATUS

+EFIAPI

+ProfileProtocolRecord (

+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,

+  IN PHYSICAL_ADDRESS                   CallerAddress,

+  IN MEMORY_PROFILE_ACTION              Action,

+  IN EFI_MEMORY_TYPE                    MemoryType,

+  IN VOID                               *Buffer,

+  IN UINTN                              Size,

+  IN CHAR8                              *ActionString OPTIONAL

+  );

+

 EDKII_MEMORY_PROFILE_PROTOCOL mProfileProtocol = {

   ProfileProtocolGetData,

   ProfileProtocolRegisterImage,

-  ProfileProtocolUnregisterImage

+  ProfileProtocolUnregisterImage,

+  ProfileProtocolGetRecordingState,

+  ProfileProtocolSetRecordingState,

+  ProfileProtocolRecord,

 };

 

 /**

@@ -307,13 +393,27 @@ BuildDriverInfo (
   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;

   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;

   VOID                              *EntryPointInImage;

+  CHAR8                             *PdbString;

+  UINTN                             PdbSize;

+  UINTN                             PdbOccupiedSize;

+

+  PdbSize = 0;

+  PdbOccupiedSize = 0;

+  PdbString = NULL;

+  if (ImageBase != 0) {

+    PdbString = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageBase);

+    if (PdbString != NULL) {

+      PdbSize = AsciiStrSize (PdbString);

+      PdbOccupiedSize = GET_OCCUPIED_SIZE (PdbSize, sizeof (UINT64));

+    }

+  }

 

   //

   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.

   //

   Status = CoreInternalAllocatePool (

              EfiBootServicesData,

-             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY),

+             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY) + PdbSize,

              (VOID **) &DriverInfoData

              );

   if (EFI_ERROR (Status)) {

@@ -325,7 +425,7 @@ BuildDriverInfo (
   DriverInfo = &DriverInfoData->DriverInfo;

   DriverInfoData->Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;

   DriverInfo->Header.Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;

-  DriverInfo->Header.Length = sizeof (MEMORY_PROFILE_DRIVER_INFO);

+  DriverInfo->Header.Length = (UINT16) (sizeof (MEMORY_PROFILE_DRIVER_INFO) + PdbOccupiedSize);

   DriverInfo->Header.Revision = MEMORY_PROFILE_DRIVER_INFO_REVISION;

   if (FileName != NULL) {

     CopyMem (&DriverInfo->FileName, FileName, sizeof (EFI_GUID));

@@ -349,6 +449,14 @@ BuildDriverInfo (
   DriverInfo->CurrentUsage = 0;

   DriverInfo->PeakUsage = 0;

   DriverInfo->AllocRecordCount = 0;

+  if (PdbSize != 0) {

+    DriverInfo->PdbStringOffset = (UINT16) sizeof (MEMORY_PROFILE_DRIVER_INFO);

+    DriverInfoData->PdbString = (CHAR8 *) (DriverInfoData->AllocInfoList + 1);

+    CopyMem (DriverInfoData->PdbString, PdbString, PdbSize);

+  } else {

+    DriverInfo->PdbStringOffset = 0;

+    DriverInfoData->PdbString = NULL;

+  }

 

   InsertTailList (ContextData->DriverInfoList, &DriverInfoData->Link);

   ContextData->Context.ImageCount ++;

@@ -358,6 +466,65 @@ BuildDriverInfo (
 }

 

 /**

+  Return if record for this driver is needed..

+

+  @param DriverFilePath     Driver file path.

+

+  @retval TRUE              Record for this driver is needed.

+  @retval FALSE             Record for this driver is not needed.

+

+**/

+BOOLEAN

+NeedRecordThisDriver (

+  IN EFI_DEVICE_PATH_PROTOCOL       *DriverFilePath

+  )

+{

+  EFI_DEVICE_PATH_PROTOCOL          *TmpDevicePath;

+  EFI_DEVICE_PATH_PROTOCOL          *DevicePathInstance;

+  UINTN                             DevicePathSize;

+  UINTN                             FilePathSize;

+

+  if (!IsDevicePathValid (mMemoryProfileDriverPath, mMemoryProfileDriverPathSize)) {

+    //

+    // Invalid Device Path means record all.

+    //

+    return TRUE;

+  }

+  

+  //

+  // Record FilePath without END node.

+  //

+  FilePathSize = GetDevicePathSize (DriverFilePath) - sizeof(EFI_DEVICE_PATH_PROTOCOL);

+

+  DevicePathInstance = mMemoryProfileDriverPath;

+  do {

+    //

+    // Find END node (it might be END_ENTIRE or END_INSTANCE).

+    //

+    TmpDevicePath = DevicePathInstance;

+    while (!IsDevicePathEndType (TmpDevicePath)) {

+      TmpDevicePath = NextDevicePathNode (TmpDevicePath);

+    }

+

+    //

+    // Do not compare END node.

+    //

+    DevicePathSize = (UINTN)TmpDevicePath - (UINTN)DevicePathInstance;

+    if ((FilePathSize == DevicePathSize) &&

+        (CompareMem (DriverFilePath, DevicePathInstance, DevicePathSize) == 0)) {

+      return TRUE;

+    }

+

+    //

+    // Get next instance.

+    //

+    DevicePathInstance = (EFI_DEVICE_PATH_PROTOCOL *)((UINTN)DevicePathInstance + DevicePathSize + DevicePathNodeLength(TmpDevicePath));

+  } while (DevicePathSubType (TmpDevicePath) != END_ENTIRE_DEVICE_PATH_SUBTYPE);

+

+  return FALSE;

+}

+

+/**

   Register DXE Core to memory profile.

 

   @param HobStart       The start address of the HOB.

@@ -376,6 +543,8 @@ RegisterDxeCore (
   EFI_PEI_HOB_POINTERS              DxeCoreHob;

   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;

   PHYSICAL_ADDRESS                  ImageBase;

+  UINT8                             TempBuffer[sizeof(MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof(EFI_DEVICE_PATH_PROTOCOL)];

+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;

 

   ASSERT (ContextData != NULL);

 

@@ -394,6 +563,14 @@ RegisterDxeCore (
   }

   ASSERT (DxeCoreHob.Raw != NULL);

 

+  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) TempBuffer;

+  EfiInitializeFwVolDevicepathNode (FilePath, &DxeCoreHob.MemoryAllocationModule->ModuleName);

+  SetDevicePathEndNode (FilePath + 1);

+

+  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *) FilePath)) {

+    return FALSE;

+  }

+

   ImageBase = DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryBaseAddress;

   DriverInfoData = BuildDriverInfo (

                      ContextData,

@@ -433,7 +610,14 @@ MemoryProfileInit (
     return;

   }

 

-  mMemoryProfileRecordingStatus = TRUE;

+  mMemoryProfileGettingStatus = FALSE;

+  if ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT7) != 0) {

+    mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;

+  } else {

+    mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_ENABLE;

+  }

+  mMemoryProfileDriverPathSize = PcdGetSize (PcdMemoryProfileDriverPath);

+  mMemoryProfileDriverPath = AllocateCopyPool (mMemoryProfileDriverPathSize, PcdGetPtr (PcdMemoryProfileDriverPath));

   mMemoryProfileContextPtr = &mMemoryProfileContext;

 

   RegisterDxeCore (HobStart, &mMemoryProfileContext);

@@ -504,11 +688,13 @@ GetFileNameFromFilePath (
   @param DriverEntry    Image info.

   @param FileType       Image file type.

 

-  @retval TRUE          Register success.

-  @retval FALSE         Register fail.

+  @return EFI_SUCCESS           Register successfully.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported,

+                                or memory profile for the image is not required.

+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.

 

 **/

-BOOLEAN

+EFI_STATUS

 RegisterMemoryProfileImage (

   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,

   IN EFI_FV_FILETYPE            FileType

@@ -518,12 +704,16 @@ RegisterMemoryProfileImage (
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;

 

   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {

-    return FALSE;

+    return EFI_UNSUPPORTED;

+  }

+

+  if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {

+    return EFI_UNSUPPORTED;

   }

 

   ContextData = GetMemoryProfileContext ();

   if (ContextData == NULL) {

-    return FALSE;

+    return EFI_UNSUPPORTED;

   }

 

   DriverInfoData = BuildDriverInfo (

@@ -536,10 +726,10 @@ RegisterMemoryProfileImage (
                      FileType

                      );

   if (DriverInfoData == NULL) {

-    return FALSE;

+    return EFI_OUT_OF_RESOURCES;

   }

 

-  return TRUE;

+  return EFI_SUCCESS;

 }

 

 /**

@@ -587,44 +777,8 @@ GetMemoryProfileDriverInfoByFileNameAndAddress (
 }

 

 /**

-  Search dummy image from memory profile.

-

-  @param ContextData    Memory profile context.

-

-  @return Pointer to memory profile driver info.

-

-**/

-MEMORY_PROFILE_DRIVER_INFO_DATA *

-FindDummyImage (

-  IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData

-  )

-{

-  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;

-  LIST_ENTRY                        *DriverLink;

-  LIST_ENTRY                        *DriverInfoList;

-

-  DriverInfoList = ContextData->DriverInfoList;

-

-  for (DriverLink = DriverInfoList->ForwardLink;

-       DriverLink != DriverInfoList;

-       DriverLink = DriverLink->ForwardLink) {

-    DriverInfoData = CR (

-                   DriverLink,

-                   MEMORY_PROFILE_DRIVER_INFO_DATA,

-                   Link,

-                   MEMORY_PROFILE_DRIVER_INFO_SIGNATURE

-                   );

-    if (CompareGuid (&gZeroGuid, &DriverInfoData->DriverInfo.FileName)) {

-      return DriverInfoData;

-    }

-  }

-

-  return BuildDriverInfo (ContextData, &gZeroGuid, 0, 0, 0, 0, 0);

-}

-

-/**

   Search image from memory profile.

-  It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize)

+  It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize).

 

   @param ContextData    Memory profile context.

   @param Address        Image or Function address.

@@ -661,10 +815,7 @@ GetMemoryProfileDriverInfoFromAddress (
     }

   }

 

-  //

-  // Should never come here.

-  //

-  return FindDummyImage (ContextData);

+  return NULL;

 }

 

 /**

@@ -672,11 +823,13 @@ GetMemoryProfileDriverInfoFromAddress (
 

   @param DriverEntry    Image info.

 

-  @retval TRUE          Unregister success.

-  @retval FALSE         Unregister fail.

+  @return EFI_SUCCESS           Unregister successfully.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported,

+                                or memory profile for the image is not required.

+  @return EFI_NOT_FOUND         The image is not found.

 

 **/

-BOOLEAN

+EFI_STATUS

 UnregisterMemoryProfileImage (

   IN LOADED_IMAGE_PRIVATE_DATA      *DriverEntry

   )

@@ -689,12 +842,16 @@ UnregisterMemoryProfileImage (
   VOID                              *EntryPointInImage;

 

   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {

-    return FALSE;

+    return EFI_UNSUPPORTED;

+  }

+

+  if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {

+    return EFI_UNSUPPORTED;

   }

 

   ContextData = GetMemoryProfileContext ();

   if (ContextData == NULL) {

-    return FALSE;

+    return EFI_UNSUPPORTED;

   }

 

   DriverInfoData = NULL;

@@ -716,12 +873,13 @@ UnregisterMemoryProfileImage (
     DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);

   }

   if (DriverInfoData == NULL) {

-    return FALSE;

+    return EFI_NOT_FOUND;

   }

 

   ContextData->Context.TotalImageSize -= DriverInfoData->DriverInfo.ImageSize;

 

-  DriverInfoData->DriverInfo.ImageBase = 0;

+  // Keep the ImageBase for RVA calculation in Application.

+  //DriverInfoData->DriverInfo.ImageBase = 0;

   DriverInfoData->DriverInfo.ImageSize = 0;

 

   if (DriverInfoData->DriverInfo.PeakUsage == 0) {

@@ -733,7 +891,7 @@ UnregisterMemoryProfileImage (
     CoreInternalFreePool (DriverInfoData, NULL);

   }

 

-  return TRUE;

+  return EFI_SUCCESS;

 }

 

 /**

@@ -803,55 +961,80 @@ GetProfileMemoryIndex (
   @param MemoryType     Memory type.

   @param Size           Buffer size.

   @param Buffer         Buffer address.

+  @param ActionString   String for memory profile action.

 

-  @retval TRUE          Profile udpate success.

-  @retval FALSE         Profile update fail.

+  @return EFI_SUCCESS           Memory profile is updated.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported,

+                                or memory profile for the image is not required.

+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.

 

 **/

-BOOLEAN

+EFI_STATUS

 CoreUpdateProfileAllocate (

   IN PHYSICAL_ADDRESS       CallerAddress,

   IN MEMORY_PROFILE_ACTION  Action,

   IN EFI_MEMORY_TYPE        MemoryType,

   IN UINTN                  Size,

-  IN VOID                   *Buffer

+  IN VOID                   *Buffer,

+  IN CHAR8                  *ActionString OPTIONAL

   )

 {

   EFI_STATUS                        Status;

-  MEMORY_PROFILE_CONTEXT           *Context;

-  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;

-  MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;

+  MEMORY_PROFILE_CONTEXT            *Context;

+  MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;

+  MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;

   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;

   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;

   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;

   UINTN                             ProfileMemoryIndex;

+  MEMORY_PROFILE_ACTION             BasicAction;

+  UINTN                             ActionStringSize;

+  UINTN                             ActionStringOccupiedSize;

 

-  AllocInfoData = NULL;

+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;

 

   ContextData = GetMemoryProfileContext ();

   if (ContextData == NULL) {

-    return FALSE;

+    return EFI_UNSUPPORTED;

   }

 

   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);

-  ASSERT (DriverInfoData != NULL);

+  if (DriverInfoData == NULL) {

+    return EFI_UNSUPPORTED;

+  }

+

+  ActionStringSize = 0;

+  ActionStringOccupiedSize = 0;

+  if (ActionString != NULL) {

+    ActionStringSize = AsciiStrSize (ActionString);

+    ActionStringOccupiedSize = GET_OCCUPIED_SIZE (ActionStringSize, sizeof (UINT64));

+  }

 

   //

   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.

   //

+  AllocInfoData = NULL;

   Status = CoreInternalAllocatePool (

              EfiBootServicesData,

-             sizeof (*AllocInfoData),

+             sizeof (*AllocInfoData) + ActionStringSize,

              (VOID **) &AllocInfoData

              );

   if (EFI_ERROR (Status)) {

-    return FALSE;

+    return EFI_OUT_OF_RESOURCES;

   }

   ASSERT (AllocInfoData != NULL);

+

+  //

+  // Only update SequenceCount if and only if it is basic action.

+  //

+  if (Action == BasicAction) {

+    ContextData->Context.SequenceCount ++;

+  }

+

   AllocInfo = &AllocInfoData->AllocInfo;

   AllocInfoData->Signature      = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;

   AllocInfo->Header.Signature   = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;

-  AllocInfo->Header.Length      = sizeof (MEMORY_PROFILE_ALLOC_INFO);

+  AllocInfo->Header.Length      = (UINT16) (sizeof (MEMORY_PROFILE_ALLOC_INFO) + ActionStringOccupiedSize);

   AllocInfo->Header.Revision    = MEMORY_PROFILE_ALLOC_INFO_REVISION;

   AllocInfo->CallerAddress      = CallerAddress;

   AllocInfo->SequenceId         = ContextData->Context.SequenceCount;

@@ -859,50 +1042,64 @@ CoreUpdateProfileAllocate (
   AllocInfo->MemoryType         = MemoryType;

   AllocInfo->Buffer             = (PHYSICAL_ADDRESS) (UINTN) Buffer;

   AllocInfo->Size               = Size;

+  if (ActionString != NULL) {

+    AllocInfo->ActionStringOffset = (UINT16) sizeof (MEMORY_PROFILE_ALLOC_INFO);

+    AllocInfoData->ActionString = (CHAR8 *) (AllocInfoData + 1);

+    CopyMem (AllocInfoData->ActionString, ActionString, ActionStringSize);

+  } else {

+    AllocInfo->ActionStringOffset = 0;

+    AllocInfoData->ActionString = NULL;

+  }

 

   InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);

 

-  ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);

-

+  Context = &ContextData->Context;

   DriverInfo = &DriverInfoData->DriverInfo;

-  DriverInfo->CurrentUsage += Size;

-  if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {

-    DriverInfo->PeakUsage = DriverInfo->CurrentUsage;

-  }

-  DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;

-  if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {

-    DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];

-  }

   DriverInfo->AllocRecordCount ++;

 

-  Context = &ContextData->Context;

-  Context->CurrentTotalUsage += Size;

-  if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {

-    Context->PeakTotalUsage = Context->CurrentTotalUsage;

-  }

-  Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;

-  if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {

-    Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];

+  //

+  // Update summary if and only if it is basic action.

+  //

+  if (Action == BasicAction) {

+    ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);

+

+    DriverInfo->CurrentUsage += Size;

+    if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {

+      DriverInfo->PeakUsage = DriverInfo->CurrentUsage;

+    }

+    DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;

+    if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {

+      DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];

+    }

+

+    Context->CurrentTotalUsage += Size;

+    if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {

+      Context->PeakTotalUsage = Context->CurrentTotalUsage;

+    }

+    Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;

+    if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {

+      Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];

+    }

   }

-  Context->SequenceCount ++;

 

-  return TRUE;

+  return EFI_SUCCESS;

 }

 

 /**

-  Get memory profile alloc info from memory profile

+  Get memory profile alloc info from memory profile.

 

-  @param DriverInfoData     Driver info

-  @param Action             This Free action

-  @param Size               Buffer size

-  @param Buffer             Buffer address

+  @param DriverInfoData     Driver info.

+  @param BasicAction        This Free basic action.

+  @param Size               Buffer size.

+  @param Buffer             Buffer address.

 

   @return Pointer to memory profile alloc info.

+

 **/

 MEMORY_PROFILE_ALLOC_INFO_DATA *

 GetMemoryProfileAllocInfoFromAddress (

   IN MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData,

-  IN MEMORY_PROFILE_ACTION              Action,

+  IN MEMORY_PROFILE_ACTION              BasicAction,

   IN UINTN                              Size,

   IN VOID                               *Buffer

   )

@@ -924,10 +1121,10 @@ GetMemoryProfileAllocInfoFromAddress (
                       MEMORY_PROFILE_ALLOC_INFO_SIGNATURE

                       );

     AllocInfo = &AllocInfoData->AllocInfo;

-    if (AllocInfo->Action != Action) {

+    if ((AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK) != BasicAction) {

       continue;

     }

-    switch (Action) {

+    switch (BasicAction) {

       case MemoryProfileActionAllocatePages:

         if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS) (UINTN) Buffer) &&

             ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size))) {

@@ -956,11 +1153,12 @@ GetMemoryProfileAllocInfoFromAddress (
   @param Size           Buffer size.

   @param Buffer         Buffer address.

 

-  @retval TRUE          Profile udpate success.

-  @retval FALSE         Profile update fail.

+  @return EFI_SUCCESS           Memory profile is updated.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported.

+  @return EFI_NOT_FOUND         No matched allocate info found for free action.

 

 **/

-BOOLEAN

+EFI_STATUS

 CoreUpdateProfileFree (

   IN PHYSICAL_ADDRESS       CallerAddress,

   IN MEMORY_PROFILE_ACTION  Action,

@@ -978,114 +1176,140 @@ CoreUpdateProfileFree (
   MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;

   MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;

   UINTN                            ProfileMemoryIndex;

+  MEMORY_PROFILE_ACTION            BasicAction;

+  BOOLEAN                          Found;

+

+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;

 

   ContextData = GetMemoryProfileContext ();

   if (ContextData == NULL) {

-    return FALSE;

+    return EFI_UNSUPPORTED;

   }

 

   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);

-  ASSERT (DriverInfoData != NULL);

 

-  switch (Action) {

-    case MemoryProfileActionFreePages:

-      AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);

-      break;

-    case MemoryProfileActionFreePool:

-      AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);

-      break;

-    default:

-      ASSERT (FALSE);

-      AllocInfoData = NULL;

-      break;

-  }

-  if (AllocInfoData == NULL) {

-    //

-    // Legal case, because driver A might free memory allocated by driver B, by some protocol.

-    //

-    DriverInfoList = ContextData->DriverInfoList;

-

-    for (DriverLink = DriverInfoList->ForwardLink;

-         DriverLink != DriverInfoList;

-         DriverLink = DriverLink->ForwardLink) {

-      ThisDriverInfoData = CR (

-                             DriverLink,

-                             MEMORY_PROFILE_DRIVER_INFO_DATA,

-                             Link,

-                             MEMORY_PROFILE_DRIVER_INFO_SIGNATURE

-                             );

-      switch (Action) {

+  //

+  // Do not return if DriverInfoData == NULL here,

+  // because driver A might free memory allocated by driver B.

+  //

+

+  //

+  // Need use do-while loop to find all possible records,

+  // because one address might be recorded multiple times.

+  //

+  Found = FALSE;

+  AllocInfoData = NULL;

+  do {

+    if (DriverInfoData != NULL) {

+      switch (BasicAction) {

         case MemoryProfileActionFreePages:

-          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);

+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);

           break;

         case MemoryProfileActionFreePool:

-          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);

+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);

           break;

         default:

           ASSERT (FALSE);

           AllocInfoData = NULL;

           break;

       }

-      if (AllocInfoData != NULL) {

-        DriverInfoData = ThisDriverInfoData;

-        break;

-      }

     }

-

     if (AllocInfoData == NULL) {

       //

-      // No matched allocate operation is found for this free operation.

-      // It is because the specified memory type allocate operation has been

-      // filtered by CoreNeedRecordProfile(), but free operations have no

-      // memory type information, they can not be filtered by CoreNeedRecordProfile().

-      // Then, they will be filtered here.

+      // Legal case, because driver A might free memory allocated by driver B, by some protocol.

       //

-      return FALSE;

-    }

-  }

+      DriverInfoList = ContextData->DriverInfoList;

+

+      for (DriverLink = DriverInfoList->ForwardLink;

+           DriverLink != DriverInfoList;

+           DriverLink = DriverLink->ForwardLink) {

+        ThisDriverInfoData = CR (

+                               DriverLink,

+                               MEMORY_PROFILE_DRIVER_INFO_DATA,

+                               Link,

+                               MEMORY_PROFILE_DRIVER_INFO_SIGNATURE

+                               );

+        switch (BasicAction) {

+          case MemoryProfileActionFreePages:

+            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);

+            break;

+          case MemoryProfileActionFreePool:

+            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);

+            break;

+          default:

+            ASSERT (FALSE);

+            AllocInfoData = NULL;

+            break;

+        }

+        if (AllocInfoData != NULL) {

+          DriverInfoData = ThisDriverInfoData;

+          break;

+        }

+      }

 

-  Context = &ContextData->Context;

-  DriverInfo = &DriverInfoData->DriverInfo;

-  AllocInfo = &AllocInfoData->AllocInfo;

+      if (AllocInfoData == NULL) {

+        //

+        // If (!Found), no matched allocate info is found for this free action.

+        // It is because the specified memory type allocate actions have been filtered by

+        // CoreNeedRecordProfile(), but free actions may have no memory type information,

+        // they can not be filtered by CoreNeedRecordProfile(). Then, they will be

+        // filtered here.

+        //

+        // If (Found), it is normal exit path.

+        return (Found ? EFI_SUCCESS : EFI_NOT_FOUND);

+      }

+    }

 

-  ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);

+    Found = TRUE;

 

-  Context->CurrentTotalUsage -= AllocInfo->Size;

-  Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;

+    Context = &ContextData->Context;

+    DriverInfo = &DriverInfoData->DriverInfo;

+    AllocInfo = &AllocInfoData->AllocInfo;

 

-  DriverInfo->CurrentUsage -= AllocInfo->Size;

-  DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;

-  DriverInfo->AllocRecordCount --;

+    DriverInfo->AllocRecordCount --;

+    //

+    // Update summary if and only if it is basic action.

+    //

+    if (AllocInfo->Action == (AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK)) {

+      ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);

 

-  RemoveEntryList (&AllocInfoData->Link);

+      Context->CurrentTotalUsage -= AllocInfo->Size;

+      Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;

 

-  if (Action == MemoryProfileActionFreePages) {

-    if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {

-      CoreUpdateProfileAllocate (

-        AllocInfo->CallerAddress,

-        MemoryProfileActionAllocatePages,

-        AllocInfo->MemoryType,

-        (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),

-        (VOID *) (UINTN) AllocInfo->Buffer

-        );

+      DriverInfo->CurrentUsage -= AllocInfo->Size;

+      DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;

     }

-    if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {

-      CoreUpdateProfileAllocate (

-        AllocInfo->CallerAddress,

-        MemoryProfileActionAllocatePages,

-        AllocInfo->MemoryType,

-        (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),

-        (VOID *) ((UINTN) Buffer + Size)

-        );

-    }

-  }

 

-  //

-  // Use CoreInternalFreePool() that will not update profile for this FreePool action.

-  //

-  CoreInternalFreePool (AllocInfoData, NULL);

+    RemoveEntryList (&AllocInfoData->Link);

+

+    if (BasicAction == MemoryProfileActionFreePages) {

+      if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {

+        CoreUpdateProfileAllocate (

+          AllocInfo->CallerAddress,

+          AllocInfo->Action,

+          AllocInfo->MemoryType,

+          (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),

+          (VOID *) (UINTN) AllocInfo->Buffer,

+          AllocInfoData->ActionString

+          );

+      }

+      if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {

+        CoreUpdateProfileAllocate (

+          AllocInfo->CallerAddress,

+          AllocInfo->Action,

+          AllocInfo->MemoryType,

+          (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),

+          (VOID *) ((UINTN) Buffer + Size),

+          AllocInfoData->ActionString

+          );

+      }

+    }

 

-  return TRUE;

+    //

+    // Use CoreInternalFreePool() that will not update profile for this FreePool action.

+    //

+    CoreInternalFreePool (AllocInfoData, NULL);

+  } while (TRUE);

 }

 

 /**

@@ -1094,62 +1318,90 @@ CoreUpdateProfileFree (
   @param CallerAddress  Address of caller who call Allocate or Free.

   @param Action         This Allocate or Free action.

   @param MemoryType     Memory type.

+                        EfiMaxMemoryType means the MemoryType is unknown.

   @param Size           Buffer size.

   @param Buffer         Buffer address.

-

-  @retval TRUE          Profile udpate success.

-  @retval FALSE         Profile update fail.

+  @param ActionString   String for memory profile action.

+                        Only needed for user defined allocate action.

+

+  @return EFI_SUCCESS           Memory profile is updated.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported,

+                                or memory profile for the image is not required,

+                                or memory profile for the memory type is not required.

+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.

+  @return EFI_ABORTED           Memory profile recording is not enabled.

+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.

+  @return EFI_NOT_FOUND         No matched allocate info found for free action.

 

 **/

-BOOLEAN

+EFI_STATUS

+EFIAPI

 CoreUpdateProfile (

   IN PHYSICAL_ADDRESS       CallerAddress,

   IN MEMORY_PROFILE_ACTION  Action,

   IN EFI_MEMORY_TYPE        MemoryType,

   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool

-  IN VOID                   *Buffer

+  IN VOID                   *Buffer,

+  IN CHAR8                  *ActionString OPTIONAL

   )

 {

+  EFI_STATUS                    Status;

   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;

+  MEMORY_PROFILE_ACTION         BasicAction;

 

   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {

-    return FALSE;

+    return EFI_UNSUPPORTED;

   }

 

-  if (!mMemoryProfileRecordingStatus) {

-    return FALSE;

+  if (mMemoryProfileGettingStatus) {

+    return EFI_ACCESS_DENIED;

+  }

+

+  if (!mMemoryProfileRecordingEnable) {

+    return EFI_ABORTED;

   }

 

   //

-  // Only record limited MemoryType.

+  // Get the basic action to know how to process the record

   //

-  if (!CoreNeedRecordProfile (MemoryType)) {

-    return FALSE;

+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;

+

+  //

+  // EfiMaxMemoryType means the MemoryType is unknown.

+  //

+  if (MemoryType != EfiMaxMemoryType) {

+    //

+    // Only record limited MemoryType.

+    //

+    if (!CoreNeedRecordProfile (MemoryType)) {

+      return EFI_UNSUPPORTED;

+    }

   }

 

   ContextData = GetMemoryProfileContext ();

   if (ContextData == NULL) {

-    return FALSE;

+    return EFI_UNSUPPORTED;

   }

 

-  switch (Action) {

+  switch (BasicAction) {

     case MemoryProfileActionAllocatePages:

-      CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);

+      Status = CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);

       break;

     case MemoryProfileActionFreePages:

-      CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);

+      Status = CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);

       break;

     case MemoryProfileActionAllocatePool:

-      CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);

+      Status = CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);

       break;

     case MemoryProfileActionFreePool:

-      CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);

+      Status = CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);

       break;

     default:

       ASSERT (FALSE);

+      Status = EFI_UNSUPPORTED;

       break;

   }

-  return TRUE;

+  return Status;

 }

 

 ////////////////////

@@ -1167,8 +1419,11 @@ MemoryProfileGetDataSize (
 {

   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;

   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;

+  MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;

   LIST_ENTRY                        *DriverInfoList;

   LIST_ENTRY                        *DriverLink;

+  LIST_ENTRY                        *AllocInfoList;

+  LIST_ENTRY                        *AllocLink;

   UINTN                             TotalSize;

 

 

@@ -1178,7 +1433,6 @@ MemoryProfileGetDataSize (
   }

 

   TotalSize = sizeof (MEMORY_PROFILE_CONTEXT);

-  TotalSize += sizeof (MEMORY_PROFILE_DRIVER_INFO) * (UINTN) ContextData->Context.ImageCount;

 

   DriverInfoList = ContextData->DriverInfoList;

   for (DriverLink = DriverInfoList->ForwardLink;

@@ -1190,7 +1444,20 @@ MemoryProfileGetDataSize (
                        Link,

                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE

                        );

-    TotalSize += sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfoData->DriverInfo.AllocRecordCount;

+    TotalSize += DriverInfoData->DriverInfo.Header.Length;

+

+    AllocInfoList = DriverInfoData->AllocInfoList;

+    for (AllocLink = AllocInfoList->ForwardLink;

+         AllocLink != AllocInfoList;

+         AllocLink = AllocLink->ForwardLink) {

+      AllocInfoData = CR (

+                        AllocLink,

+                        MEMORY_PROFILE_ALLOC_INFO_DATA,

+                        Link,

+                        MEMORY_PROFILE_ALLOC_INFO_SIGNATURE

+                        );

+      TotalSize += AllocInfoData->AllocInfo.Header.Length;

+    }

   }

 

   return TotalSize;

@@ -1217,6 +1484,8 @@ MemoryProfileCopyData (
   LIST_ENTRY                        *DriverLink;

   LIST_ENTRY                        *AllocInfoList;

   LIST_ENTRY                        *AllocLink;

+  UINTN                             PdbSize;

+  UINTN                             ActionStringSize;

 

   ContextData = GetMemoryProfileContext ();

   if (ContextData == NULL) {

@@ -1238,7 +1507,11 @@ MemoryProfileCopyData (
                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE

                        );

     CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));

-    AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) (DriverInfo + 1);

+    if (DriverInfo->PdbStringOffset != 0) {

+      PdbSize = AsciiStrSize (DriverInfoData->PdbString);

+      CopyMem ((VOID *) ((UINTN) DriverInfo + DriverInfo->PdbStringOffset), DriverInfoData->PdbString, PdbSize);

+    }

+    AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) ((UINTN) DriverInfo + DriverInfo->Header.Length);

 

     AllocInfoList = DriverInfoData->AllocInfoList;

     for (AllocLink = AllocInfoList->ForwardLink;

@@ -1251,10 +1524,14 @@ MemoryProfileCopyData (
                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE

                         );

       CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));

-      AllocInfo += 1;

+      if (AllocInfo->ActionStringOffset != 0) {

+        ActionStringSize = AsciiStrSize (AllocInfoData->ActionString);

+        CopyMem ((VOID *) ((UINTN) AllocInfo + AllocInfo->ActionStringOffset), AllocInfoData->ActionString, ActionStringSize);

+      }

+      AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) ((UINTN) AllocInfo + AllocInfo->Header.Length);

     }

 

-    DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) ((UINTN) (DriverInfo + 1) + sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfo->AllocRecordCount);

+    DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *)  AllocInfo;

   }

 }

 

@@ -1267,6 +1544,7 @@ MemoryProfileCopyData (
   @param[out]     ProfileBuffer     Profile buffer.

                       

   @return EFI_SUCCESS               Get the memory profile data successfully.

+  @return EFI_UNSUPPORTED           Memory profile is unsupported.

   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. 

                                     ProfileSize is updated with the size required.

 

@@ -1281,28 +1559,28 @@ ProfileProtocolGetData (
 {

   UINTN                                 Size;

   MEMORY_PROFILE_CONTEXT_DATA           *ContextData;

-  BOOLEAN                               MemoryProfileRecordingStatus;

+  BOOLEAN                               MemoryProfileGettingStatus;

 

   ContextData = GetMemoryProfileContext ();

   if (ContextData == NULL) {

     return EFI_UNSUPPORTED;

   }

 

-  MemoryProfileRecordingStatus = mMemoryProfileRecordingStatus;

-  mMemoryProfileRecordingStatus = FALSE;

+  MemoryProfileGettingStatus = mMemoryProfileGettingStatus;

+  mMemoryProfileGettingStatus = TRUE;

 

   Size = MemoryProfileGetDataSize ();

 

   if (*ProfileSize < Size) {

     *ProfileSize = Size;

-    mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;

+    mMemoryProfileGettingStatus = MemoryProfileGettingStatus;

     return EFI_BUFFER_TOO_SMALL;

   }

 

   *ProfileSize = Size;

   MemoryProfileCopyData (ProfileBuffer);

 

-  mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;

+  mMemoryProfileGettingStatus = MemoryProfileGettingStatus;

   return EFI_SUCCESS;

 }

 

@@ -1315,8 +1593,10 @@ ProfileProtocolGetData (
   @param[in] ImageSize          Image size.

   @param[in] FileType           File type of the image.

 

-  @return EFI_SUCCESS           Register success.

-  @return EFI_OUT_OF_RESOURCE   No enough resource for this register.

+  @return EFI_SUCCESS           Register successfully.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported,

+                                or memory profile for the image is not required.

+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.

 

 **/

 EFI_STATUS

@@ -1342,7 +1622,7 @@ ProfileProtocolRegisterImage (
   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;

   DriverEntry.ImageContext.ImageType = InternalPeCoffGetSubsystem ((VOID *) (UINTN) ImageBase);

 

-  return RegisterMemoryProfileImage (&DriverEntry, FileType) ? EFI_SUCCESS: EFI_OUT_OF_RESOURCES;

+  return RegisterMemoryProfileImage (&DriverEntry, FileType);

 }

 

 /**

@@ -1353,7 +1633,9 @@ ProfileProtocolRegisterImage (
   @param[in] ImageBase          Image base address.

   @param[in] ImageSize          Image size.

 

-  @return EFI_SUCCESS           Unregister success.

+  @return EFI_SUCCESS           Unregister successfully.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported,

+                                or memory profile for the image is not required.

   @return EFI_NOT_FOUND         The image is not found.

 

 **/

@@ -1378,7 +1660,105 @@ ProfileProtocolUnregisterImage (
   ASSERT_EFI_ERROR (Status);

   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;

 

-  return UnregisterMemoryProfileImage (&DriverEntry) ? EFI_SUCCESS: EFI_NOT_FOUND;

+  return UnregisterMemoryProfileImage (&DriverEntry);

+}

+

+/**

+  Get memory profile recording state.

+

+  @param[in]  This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.

+  @param[out] RecordingState    Recording state.

+

+  @return EFI_SUCCESS           Memory profile recording state is returned.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported.

+  @return EFI_INVALID_PARAMETER RecordingState is NULL.

+

+**/

+EFI_STATUS

+EFIAPI

+ProfileProtocolGetRecordingState (

+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,

+  OUT BOOLEAN                           *RecordingState

+  )

+{

+  MEMORY_PROFILE_CONTEXT_DATA       *ContextData;

+

+  ContextData = GetMemoryProfileContext ();

+  if (ContextData == NULL) {

+    return EFI_UNSUPPORTED;

+  }

+

+  if (RecordingState == NULL) {

+    return EFI_INVALID_PARAMETER;

+  }

+  *RecordingState = mMemoryProfileRecordingEnable;

+  return EFI_SUCCESS;

+}

+

+/**

+  Set memory profile recording state.

+

+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.

+  @param[in] RecordingState     Recording state.

+

+  @return EFI_SUCCESS           Set memory profile recording state successfully.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported.

+

+**/

+EFI_STATUS

+EFIAPI

+ProfileProtocolSetRecordingState (

+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,

+  IN BOOLEAN                            RecordingState

+  )

+{

+  MEMORY_PROFILE_CONTEXT_DATA       *ContextData;

+

+  ContextData = GetMemoryProfileContext ();

+  if (ContextData == NULL) {

+    return EFI_UNSUPPORTED;

+  }

+

+  mMemoryProfileRecordingEnable = RecordingState;

+  return EFI_SUCCESS;

+}

+

+/**

+  Record memory profile of multilevel caller.

+

+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.

+  @param[in] CallerAddress      Address of caller.

+  @param[in] Action             Memory profile action.

+  @param[in] MemoryType         Memory type.

+                                EfiMaxMemoryType means the MemoryType is unknown.

+  @param[in] Buffer             Buffer address.

+  @param[in] Size               Buffer size.

+  @param[in] ActionString       String for memory profile action.

+                                Only needed for user defined allocate action.

+

+  @return EFI_SUCCESS           Memory profile is updated.

+  @return EFI_UNSUPPORTED       Memory profile is unsupported,

+                                or memory profile for the image is not required,

+                                or memory profile for the memory type is not required.

+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.

+  @return EFI_ABORTED           Memory profile recording is not enabled.

+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.

+  @return EFI_NOT_FOUND         No matched allocate info found for free action.

+

+**/

+EFI_STATUS

+EFIAPI

+ProfileProtocolRecord (

+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,

+  IN PHYSICAL_ADDRESS                   CallerAddress,

+  IN MEMORY_PROFILE_ACTION              Action,

+  IN EFI_MEMORY_TYPE                    MemoryType,

+  IN VOID                               *Buffer,

+  IN UINTN                              Size,

+  IN CHAR8                              *ActionString OPTIONAL

+  )

+{

+  return CoreUpdateProfile (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);

 }

 

 ////////////////////

diff --git a/MdeModulePkg/Core/Dxe/Mem/Page.c b/MdeModulePkg/Core/Dxe/Mem/Page.c
index 898b722a43..b02bafb4be 100644
--- a/MdeModulePkg/Core/Dxe/Mem/Page.c
+++ b/MdeModulePkg/Core/Dxe/Mem/Page.c
@@ -1335,7 +1335,14 @@ CoreAllocatePages (
 

   Status = CoreInternalAllocatePages (Type, MemoryType, NumberOfPages, Memory);

   if (!EFI_ERROR (Status)) {

-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePages, MemoryType, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) *Memory);

+    CoreUpdateProfile (

+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),

+      MemoryProfileActionAllocatePages,

+      MemoryType,

+      EFI_PAGES_TO_SIZE (NumberOfPages),

+      (VOID *) (UINTN) *Memory,

+      NULL

+      );

     InstallMemoryAttributesTableOnMemoryAllocation (MemoryType);

   }

   return Status;

@@ -1444,7 +1451,14 @@ CoreFreePages (
 

   Status = CoreInternalFreePages (Memory, NumberOfPages, &MemoryType);

   if (!EFI_ERROR (Status)) {

-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePages, MemoryType, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) Memory);

+    CoreUpdateProfile (

+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),

+      MemoryProfileActionFreePages,

+      MemoryType,

+      EFI_PAGES_TO_SIZE (NumberOfPages),

+      (VOID *) (UINTN) Memory,

+      NULL

+      );

     InstallMemoryAttributesTableOnMemoryAllocation (MemoryType);

   }

   return Status;

diff --git a/MdeModulePkg/Core/Dxe/Mem/Pool.c b/MdeModulePkg/Core/Dxe/Mem/Pool.c
index 2980e2293e..6ef5fba395 100644
--- a/MdeModulePkg/Core/Dxe/Mem/Pool.c
+++ b/MdeModulePkg/Core/Dxe/Mem/Pool.c
@@ -276,7 +276,14 @@ CoreAllocatePool (
 

   Status = CoreInternalAllocatePool (PoolType, Size, Buffer);

   if (!EFI_ERROR (Status)) {

-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePool, PoolType, Size, *Buffer);

+    CoreUpdateProfile (

+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),

+      MemoryProfileActionAllocatePool,

+      PoolType,

+      Size,

+      *Buffer,

+      NULL

+      );

     InstallMemoryAttributesTableOnMemoryAllocation (PoolType);

   }

   return Status;

@@ -505,7 +512,14 @@ CoreFreePool (
 

   Status = CoreInternalFreePool (Buffer, &PoolType);

   if (!EFI_ERROR (Status)) {

-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePool, PoolType, 0, Buffer);

+    CoreUpdateProfile (

+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),

+      MemoryProfileActionFreePool,

+      PoolType,

+      0,

+      Buffer,

+      NULL

+      );

     InstallMemoryAttributesTableOnMemoryAllocation (PoolType);

   }

   return Status;