From aae870e8744d62ff782f43ea07d85f086bcad6f6 Mon Sep 17 00:00:00 2001
From: Chloe M <chloe@mensia.org>
Date: Mon, 13 Jul 2026 20:09:53 -0400
Subject: [PATCH] ptos: pool: Propagate page information and prepend header

Signed-off-by: Chloe M <chloe@mensia.org>
---
 service/ptos/ex/pool.c      | 47 +++++++++++++++++++++++++++++++------
 service/ptos/head/ex/pool.h | 18 ++++++++++++++
 2 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/service/ptos/ex/pool.c b/service/ptos/ex/pool.c
index 6db18ee..c83fb2a 100644
--- a/service/ptos/ex/pool.c
+++ b/service/ptos/ex/pool.c
@@ -13,6 +13,20 @@
 #include <ptapi/status.h>
 #include <string.h>
 
+/*
+ * Used for storing in the header so that memory
+ * can be freed.
+ *
+ * @BitBase: Allocation bit base
+ * @Gran:    Granularity of page bitmap
+ * @Page:    Page we reside in
+ */
+typedef struct {
+    USHORT BitBase : 6;
+    USHORT Gran;
+    POOL_PAGE *Page;
+} PAGE_INFO;
+
 /*
  * Compute the segment level index from the number of
  * bytes requested to be allocated. We use the following
@@ -105,15 +119,16 @@ PoolSegmentInit(POOL_SEGMENT *Segment, UCHAR Level)
  * @Gran:       Granularity in bytes of level
  * @Page:       Page to allocate from
  * @ByteCount:  Number of bytes to allocate
+ * @Info:       Page information result
  */
 static VOID *
-PoolAllocateFromPage(USIZE Gran, POOL_PAGE *Page, USIZE ByteCount)
+PoolAllocateFromPage(USIZE Gran, POOL_PAGE *Page, USIZE ByteCount, PAGE_INFO *Info)
 {
     USIZE ChunksNeeded, ChunksFound;
     USIZE SetIdx, Idx;
     SSIZE IdxBase = -1;
 
-    if (Page == NULL) {
+    if (Info == NULL || Page == NULL) {
         return NULL;
     }
 
@@ -146,6 +161,9 @@ PoolAllocateFromPage(USIZE Gran, POOL_PAGE *Page, USIZE ByteCount)
         SETBIT(&Page->Bitmap, SetIdx);
     }
 
+    Info->BitBase = IdxBase;
+    Info->Page = Page;
+    Info->Gran = Gran;
     return PTR_OFFSET(Page, IdxBase * Gran);
 }
 
@@ -155,9 +173,10 @@ PoolAllocateFromPage(USIZE Gran, POOL_PAGE *Page, USIZE ByteCount)
  * @ByteCount:  Number of bytes to allocate
  * @LevelIndex: Segment level index
  * @Segment:    Segment to allocate from
+ * @Info:       Page information
  */
 static VOID *
-PoolAllocateFromSegment(USIZE ByteCount, UCHAR LevelIndex, POOL_SEGMENT *Segment)
+PoolAllocateFromSegment(USIZE ByteCount, UCHAR LevelIndex, POOL_SEGMENT *Segment, PAGE_INFO *Info)
 {
     USIZE ChunkCount, Idx;
     ULONG Gran;
@@ -176,7 +195,7 @@ PoolAllocateFromSegment(USIZE ByteCount, UCHAR LevelIndex, POOL_SEGMENT *Segment
     }
 
     while (Page != NULL) {
-        Ptr = PoolAllocateFromPage(Gran, Page, ByteCount);
+        Ptr = PoolAllocateFromPage(Gran, Page, ByteCount, Info);
         if (Ptr != NULL) {
             return Ptr;
         }
@@ -202,13 +221,18 @@ ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag)
     UCHAR SegmentIndex;
     POOL_REGION *Region;
     POOL_SEGMENT *Segment;
+    POOL_HEADER *Header;
+    PAGE_INFO PageInfo;
     VOID *Base;
 
     if (ByteCount == 0) {
         return NULL;
     }
 
-    ByteCount = ALIGN_UP(ByteCount, POOL_LEVEL_GRAN(0));
+    ByteCount = ALIGN_UP(
+        ByteCount,
+        POOL_LEVEL_GRAN(0)
+    );
 
     /* Verify that the pool type is valid */
     switch (Type) {
@@ -222,15 +246,24 @@ ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag)
     SegmentIndex = PoolComputeLevelIndex(ByteCount);
     Segment = &Region->Level[SegmentIndex];
     Base = PoolAllocateFromSegment(
-        ByteCount,
+        ByteCount + ALIGN_UP(sizeof(POOL_HEADER), POOL_LEVEL_GRAN(0)),
         SegmentIndex,
-        Segment
+        Segment,
+        &PageInfo
     );
 
     if (Base == NULL) {
         return NULL;
     }
 
+    Header = (POOL_HEADER *)Base;
+    Header->Tag = Tag;
+    Header->SizeBytes = ByteCount;
+    Header->Gran = PageInfo.Gran;
+    Header->BitBase = PageInfo.BitBase;
+    Header->Page = PageInfo.Page;
+
+    Base = PTR_OFFSET(Base, sizeof(POOL_HEADER));
     RtlMemSet(Base, 0, ByteCount);
     return Base;
 }
diff --git a/service/ptos/head/ex/pool.h b/service/ptos/head/ex/pool.h
index fafab35..88a350b 100644
--- a/service/ptos/head/ex/pool.h
+++ b/service/ptos/head/ex/pool.h
@@ -64,6 +64,24 @@ typedef struct {
     POOL_SEGMENT Level[POOL_LEVEL_MAX];
 } POOL_REGION;
 
+/*
+ * A pool header is prepended to the start of allocated
+ * memory.
+ *
+ * @Gran:       Granularity of page bitmap
+ * @BitBase:    Allocation bitbase
+ * @SizeBytes:  Allocation size in bytes
+ * @Tag:        Tag assigned to allocation
+ * @Page:       Page we reside in
+ */
+typedef struct {
+    USHORT Gran;
+    UCHAR BitBase : 6;
+    USIZE SizeBytes;
+    ULONG Tag;
+    POOL_PAGE *Page;
+} POOL_HEADER;
+
 /*
  * Initialize a pool region
  *