#include <CrtLibSupport.h>\r
#include <Library/MemoryAllocationLib.h>\r
\r
+//\r
+// Extra header to record the memory buffer size from malloc routine.\r
+//\r
+#define CRYPTMEM_HEAD_SIGNATURE SIGNATURE_32('c','m','h','d')\r
+typedef struct {\r
+ UINT32 Signature;\r
+ UINT32 Reserved;\r
+ UINTN Size;\r
+} CRYPTMEM_HEAD;\r
+\r
+#define CRYPTMEM_OVERHEAD sizeof(CRYPTMEM_HEAD)\r
+\r
//\r
// -- Memory-Allocation Routines --\r
//\r
/* Allocates memory blocks */\r
void *malloc (size_t size)\r
{\r
- return AllocatePool ((UINTN) size);\r
+ CRYPTMEM_HEAD *PoolHdr;\r
+ UINTN NewSize;\r
+ VOID *Data;\r
+\r
+ //\r
+ // Adjust the size by the buffer header overhead\r
+ //\r
+ NewSize = (UINTN)(size) + CRYPTMEM_OVERHEAD;\r
+\r
+ Data = AllocatePool (NewSize);\r
+ if (Data != NULL) {\r
+ PoolHdr = (CRYPTMEM_HEAD *)Data;\r
+ //\r
+ // Record the memory brief information\r
+ //\r
+ PoolHdr->Signature = CRYPTMEM_HEAD_SIGNATURE;\r
+ PoolHdr->Size = size;\r
+\r
+ return (VOID *)(PoolHdr + 1);\r
+ } else {\r
+ //\r
+ // The buffer allocation failed.\r
+ //\r
+ return NULL;\r
+ }\r
}\r
\r
/* Reallocate memory blocks */\r
void *realloc (void *ptr, size_t size)\r
{\r
- //\r
- // BUG: hardcode OldSize == size! We have no any knowledge about\r
- // memory size of original pointer ptr.\r
- //\r
- return ReallocatePool ((UINTN) size, (UINTN) size, ptr);\r
+ CRYPTMEM_HEAD *OldPoolHdr;\r
+ CRYPTMEM_HEAD *NewPoolHdr;\r
+ UINTN OldSize;\r
+ UINTN NewSize;\r
+ VOID *Data;\r
+\r
+ NewSize = (UINTN)size + CRYPTMEM_OVERHEAD;\r
+ Data = AllocatePool (NewSize);\r
+ if (Data != NULL) {\r
+ NewPoolHdr = (CRYPTMEM_HEAD *)Data;\r
+ NewPoolHdr->Signature = CRYPTMEM_HEAD_SIGNATURE;\r
+ NewPoolHdr->Size = size;\r
+ if (ptr != NULL) {\r
+ //\r
+ // Retrieve the original size from the buffer header.\r
+ //\r
+ OldPoolHdr = (CRYPTMEM_HEAD *)ptr - 1;\r
+ ASSERT (OldPoolHdr->Signature == CRYPTMEM_HEAD_SIGNATURE);\r
+ OldSize = OldPoolHdr->Size;\r
+\r
+ //\r
+ // Duplicate the buffer content.\r
+ //\r
+ CopyMem ((VOID *)(NewPoolHdr + 1), ptr, MIN (OldSize, size));\r
+ FreePool ((VOID *)OldPoolHdr);\r
+ }\r
+\r
+ return (VOID *)(NewPoolHdr + 1);\r
+ } else {\r
+ //\r
+ // The buffer allocation failed.\r
+ //\r
+ return NULL;\r
+ }\r
}\r
\r
/* De-allocates or frees a memory block */\r
void free (void *ptr)\r
{\r
+ CRYPTMEM_HEAD *PoolHdr;\r
+\r
//\r
// In Standard C, free() handles a null pointer argument transparently. This\r
// is not true of FreePool() below, so protect it.\r
//\r
if (ptr != NULL) {\r
- FreePool (ptr);\r
+ PoolHdr = (CRYPTMEM_HEAD *)ptr - 1;\r
+ ASSERT (PoolHdr->Signature == CRYPTMEM_HEAD_SIGNATURE);\r
+ FreePool (PoolHdr);\r
}\r
}\r