ptos: pool: Add pool allocator groundwork

Signed-off-by: Chloe M. <chloe@mensia.org>
2026-07-13 23:40:24 +00:00
parent 895521c510
commit 6472e67fe2
2 changed files with 236 additions and 13 deletions
@@ -8,32 +8,231 @@

 #include <ex/pool.h>
 #include <ke/bugcheck.h>
+#include <hal/kpcr.h>
 #include <mm/vm.h>
 #include <ptapi/status.h>
+#include <string.h>
+
+/*
+ * Compute the segment level index from the number of
+ * bytes requested to be allocated. We use the following
+ * formula to do so:
+ *
+ * ((log2(count) - POOL_MIN_LOG2) - 1) % POOL_LEVEL_MAX
+ */
+static UCHAR
+PoolComputeLevelIndex(USIZE ByteCount)
+{
+    USIZE TmpCount;
+    USIZE Log2Sum = 0;
+
+    /* Compute the base 2 logarithm */
+    TmpCount = ByteCount;
+    while (TmpCount != 0) {
+       TmpCount >>= 1;
+       ++Log2Sum;
+    }
+
+    /* Subtract the right side */
+    TmpCount -= POOL_MIN_LOG2 + 1;
+    TmpCount %= POOL_LEVEL_MAX;
+    return TmpCount;
+}
+
+/*
+ * Obtain a reference to the current pool region
+ */
+static POOL_REGION *
+PoolGetCurrentRegion(VOID)
+{
+    KPCR *Kpcr;
+
+    Kpcr = HalKpcrCurrent();
+    return &Kpcr->PoolRegion;
+}
+
+/*
+ * Allocate a page descriptor
+ *
+ * @GranBytes:  Granularity in bytes of level
+ */
+static POOL_PAGE *
+PoolAllocatePage(USIZE GranBytes)
+{
+    PT_STATUS Status;
+    POOL_PAGE *PageDesc;
+    MM_RANGE Range;
+    ULONG PageSizeBitCount;
+    ULONG BitmapReserveMask;
+
+    Status = MmAllocatePages(1, &Range);
+    if (Status != STATUS_SUCCESS) {
+        return NULL;
+    }
+
+    PageDesc = RANGE_PTR(&Range);
+    PageSizeBitCount = sizeof(*PageDesc) * 8;
+    PageSizeBitCount /= GranBytes;
+    BitmapReserveMask = (1 << PageSizeBitCount) - 1;
+    PageDesc->Bitmap |= BitmapReserveMask;
+    return PageDesc;
+}

 /*
 * Initialize a pool segment
 *
 * @Segment: Pool segment to initialize
+ * @Level:   Segment level
 */
 static PT_STATUS
-PoolSegmentInit(POOL_SEGMENT *Segment)
+PoolSegmentInit(POOL_SEGMENT *Segment, UCHAR Level)
 {
-    PT_STATUS Status;
-    MM_RANGE Range;
+    USIZE Gran;

    if (Segment == NULL) {
        return STATUS_INVALID_PARAM;
    }

-    Status = MmAllocatePages(1, &Range);
-    if (Status != STATUS_SUCCESS) {
-        return Status;
+    Gran = POOL_LEVEL_GRAN(Level);
+    Segment->PageCount = 1;
+    Segment->FirstPage = PoolAllocatePage(Gran);
+    return STATUS_SUCCESS;
 }

-    Segment->PageCount = 1;
-    Segment->LastPage = RANGE_PTR(&Range);
-    return STATUS_SUCCESS;
+/*
+ * Allocate memory from a page descriptor
+ *
+ * @Gran:       Granularity in bytes of level
+ * @Page:       Page to allocate from
+ * @ByteCount:  Number of bytes to allocate
+ */
+static VOID *
+PoolAllocateFromPage(USIZE Gran, POOL_PAGE *Page, USIZE ByteCount)
+{
+    USIZE ChunksNeeded, ChunksFound;
+    USIZE SetIdx, Idx;
+    SSIZE IdxBase = -1;
+
+    if (Page == NULL) {
+        return NULL;
+    }
+
+    ChunksNeeded = ByteCount / Gran;
+
+    /* Scan the page for free chunks */
+    for (Idx = 0; Idx < sizeof(Page->Bitmap) * 8; ++Idx) {
+       if (TESTBIT(&Page->Bitmap, Idx)) {
+           ChunksFound = 0;
+           IdxBase = -1;
+           continue;
+       }
+
+       if (IdxBase < 0) {
+           IdxBase = Idx;
+       }
+
+       ++ChunksFound;
+       if (ChunksFound >= ChunksNeeded) {
+           break;
+       }
+    }
+
+    if (ChunksFound < ChunksNeeded) {
+        return NULL;
+    }
+
+    /* Mark all as allocated */
+    for (SetIdx = IdxBase; SetIdx < Idx + 1; ++SetIdx) {
+        SETBIT(&Page->Bitmap, SetIdx);
+    }
+
+    return PTR_OFFSET(Page, IdxBase * Gran);
+}
+
+/*
+ * Allocate memory from a pool segment
+ *
+ * @ByteCount:  Number of bytes to allocate
+ * @LevelIndex: Segment level index
+ * @Segment:    Segment to allocate from
+ */
+static VOID *
+PoolAllocateFromSegment(USIZE ByteCount, UCHAR LevelIndex, POOL_SEGMENT *Segment)
+{
+    USIZE ChunkCount, Idx;
+    ULONG Gran;
+    POOL_PAGE *Page;
+    VOID *Ptr;
+
+    if (ByteCount == 0 || Segment == NULL) {
+        return NULL;
+    }
+
+    Gran = POOL_LEVEL_GRAN(LevelIndex);
+    Page = Segment->FirstPage;
+
+    if (Page == NULL) {
+        return NULL;
+    }
+
+    while (Page != NULL) {
+        Ptr = PoolAllocateFromPage(Gran, Page, ByteCount);
+        if (Ptr != NULL) {
+            return Ptr;
+        }
+
+        if (Page->Next == NULL) {
+            Page->Next = PoolAllocatePage(Gran);
+            if (Page->Next == NULL)
+                return NULL;
+        }
+
+        Page = Page->Next;
+    }
+
+    return NULL;
+}
+
+/*
+ * TODO: We need per-region locking
+ */
+VOID *
+ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag)
+{
+    UCHAR SegmentIndex;
+    POOL_REGION *Region;
+    POOL_SEGMENT *Segment;
+    VOID *Base;
+
+    if (ByteCount == 0) {
+        return NULL;
+    }
+
+    ByteCount = ALIGN_UP(ByteCount, POOL_LEVEL_GRAN(0));
+
+    /* Verify that the pool type is valid */
+    switch (Type) {
+    case POOL_NON_PAGED:
+        break;
+    default:
+        return NULL;
+    }
+
+    Region = PoolGetCurrentRegion();
+    SegmentIndex = PoolComputeLevelIndex(ByteCount);
+    Segment = &Region->Level[SegmentIndex];
+    Base = PoolAllocateFromSegment(
+        ByteCount,
+        SegmentIndex,
+        Segment
+    );
+
+    if (Base == NULL) {
+        return NULL;
+    }
+
+    RtlMemSet(Base, 0, ByteCount);
+    return Base;
 }

 VOID
@@ -53,7 +252,7 @@ ExPoolRegionInit(POOL_REGION *Region)
    /* Initialize each segment level */
    for (Level = 0; Level < POOL_LEVEL_MAX; ++Level) {
        Segment = &Region->Level[Level];
-        Status = PoolSegmentInit(Segment);
+        Status = PoolSegmentInit(Segment, Level);

        if (PT_ERROR(Status)) {
            KeBugCheck(
@@ -21,11 +21,20 @@
 #define POOL_LEVEL_GRAN(X) \
    (1 << ((X) + POOL_MIN_LOG2))

+/*
+ * Valid types that can be assigned to an allocated region
+ *
+ * @POOL_NON_PAGED:  This type is not paged to disk
+ */
+typedef enum {
+    POOL_NON_PAGED
+} POOL_TYPE;
+
 /*
 * A pool page represents a single allocatable
 * unit within a pool segment.
 *
- * @Bitmap: Bitmap within pool page
+ * @Bitmap: Bitmap within pool page [0=free, 1=allocated]
 * @Next:   Next page within list
 */
 typedef struct _POOL_PAGE {
@@ -37,11 +46,11 @@ typedef struct _POOL_PAGE {
 * A pool segment represents a list of pages at a specific
 * granularity.
 *
- * @LastPage:  Last page in segment
+ * @FirstPage:  Last page in segment
 * @PageCount: Number of pages in this segment
 */
 typedef struct {
-    POOL_PAGE *LastPage;
+    POOL_PAGE *FirstPage;
    USIZE PageCount;
 } POOL_SEGMENT;

@@ -62,4 +71,19 @@ typedef struct {
 */
 VOID ExPoolRegionInit(POOL_REGION *Region);

+/*
+ * Allocate a memory pool of a specific type and assign a tag
+ * to it for verifying frees among other debugging needs.
+ *
+ * @Type:       Pool type to assign
+ * @ByteCount:  Number of bytes to allocate
+ * @Tag:        Tag to assign to allocated memory
+ *
+ * Returns NULL upon failure.
+ */
+VOID *ExAllocatePoolWithTag(
+    POOL_TYPE Type, USIZE ByteCount,
+    ULONG Tag
+);
+
 #endif  /* !_EX_POOL_H_ */