stos: pool: Add pool allocation header + tag usage

Signed-off-by: Chloe M. <chloe@mensia.org>

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

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