diff --git a/paw/stos/ex/pool.c b/paw/stos/ex/pool.c index 59b17c6..d9c8d3a 100644 --- a/paw/stos/ex/pool.c +++ b/paw/stos/ex/pool.c @@ -171,14 +171,12 @@ AllocatePoolWithTag(MEMORY_POOL *Pool, POOL_TYPE Type, USIZE ByteCount, ULONG Ta return AllocateFromBlock(BlockIndex, Block, ByteCount); } -/* - * TODO: Handle tag assignment and prepend a pool header - */ VOID * ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag) { KPCR *ThisCore; MEMORY_POOL *AllocPool; + MEMORY_HEADER *Header; VOID *Base; if (ByteCount == 0) { @@ -203,7 +201,7 @@ ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag) * that we can allocate. Furthermore, ensure that it does not exceed * the size of a single page. */ - ByteCount = ALIGN_UP(ByteCount, LEVEL_GRAN(0)); + ByteCount = ALIGN_UP(ByteCount + sizeof(*Header), LEVEL_GRAN(0)); if (ByteCount >= PAGESIZE) { DTRACE( "allocation in pool #%d exceeds page size\n", @@ -216,7 +214,16 @@ ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag) Base = AllocatePoolWithTag(AllocPool, Type, ByteCount, Tag); KeSpinLockRelease(&AllocPool->Lock); - return Base; + + if (Base == NULL) { + return NULL; + } + + Header = (MEMORY_HEADER *)Base; + Header->Magic = POOL_MAGIC; + Header->MemoryTag = Tag; + Header->Length = ByteCount; + return PTR_OFFSET(Base, sizeof(MEMORY_HEADER)); } ST_STATUS diff --git a/paw/stos/head/ex/pool.h b/paw/stos/head/ex/pool.h index dd19f68..ef8edc3 100644 --- a/paw/stos/head/ex/pool.h +++ b/paw/stos/head/ex/pool.h @@ -32,6 +32,9 @@ #define POOL_REAL_PAGESZ \ (PAGESIZE - sizeof(MEMORY_PAGE)) +/* Used to prevent corruption */ +#define POOL_MAGIC 0xCA7F00D + /* * Valid pool types * @@ -41,6 +44,20 @@ typedef enum { POOL_NON_PAGED } POOL_TYPE; +/* + * Represents the header that goes before each returned + * and allocated block. + * + * @Magic: Magic numbers + * @MemoryTag: Tag assigned at allocation + * @Length: Length of memory allocation + */ +typedef struct PACKED { + ULONG Magic; + ULONG MemoryTag; + USIZE Length; +} MEMORY_HEADER; + /* * Page data structure used internally to link together page * lists