From 6472e67fe221d5baace9c570c786f8d04010fcff Mon Sep 17 00:00:00 2001 From: "Chloe M." Date: Mon, 13 Jul 2026 23:40:24 +0000 Subject: [PATCH] ptos: pool: Add pool allocator groundwork Signed-off-by: Chloe M. --- service/ptos/ex/pool.c | 219 ++++++++++++++++++++++++++++++++++-- service/ptos/head/ex/pool.h | 30 ++++- 2 files changed, 236 insertions(+), 13 deletions(-) diff --git a/service/ptos/ex/pool.c b/service/ptos/ex/pool.c index 598b75e..6db18ee 100644 --- a/service/ptos/ex/pool.c +++ b/service/ptos/ex/pool.c @@ -8,32 +8,231 @@ #include #include +#include #include #include +#include + +/* + * 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; +} + +/* + * 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; } - Segment->PageCount = 1; - Segment->LastPage = RANGE_PTR(&Range); - return STATUS_SUCCESS; + 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( diff --git a/service/ptos/head/ex/pool.h b/service/ptos/head/ex/pool.h index fb2ca2c..fafab35 100644 --- a/service/ptos/head/ex/pool.h +++ b/service/ptos/head/ex/pool.h @@ -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_ */