void testMemoryAllocatorFuzzy() {
const int memSize = 64 * 1024 * 1024;
u8 *memory = (u8*)malloc(memSize);
MemorySegment segment;
segment.base = memory;
segment.size = memSize;
segment.used = 0;
MemoryPool pool = InitializeAllocator(&segment);
void *resrv[128];
memset(resrv, 0, 128 * sizeof(void*));
for (int i = 0; i < 128; ++i) {
resrv[i] = Allocate(&pool, rand() % (memSize / 128));
}
for (int i = 0; i < 64; ++i) {
Free(&pool, resrv[rand() % 128]);
}
for (int i = 0; i < 128; ++i) {
Free(&pool, resrv[i]);
}
Assert(segment.used == memSize);
auto derp = (MemoryPoolRegion*)pool.base;
Assert(derp->used == 0);
free(memory);
}
void testMemoryAllocator() {
// Test memory
u8 *memory = (u8*)malloc(64 * 1024);
MemorySegment segment;
segment.base = memory;
segment.size = 64 * 1024;
MemoryPool pool = InitializeAllocator(&segment);
// Check that the memory is reused
void *test = Allocate(&pool, 1024);
Free(&pool, test);
void *test2 = Allocate(&pool, 1024);
Free(&pool, test);
Assert(test == test2);
// Check that two allocations does not overlap
void *overlap1 = Allocate(&pool, 512);
void *overlap2 = Allocate(&pool, 512);
Assert(overlap1 != overlap2);
Assert((u8*)overlap2 > (u8*)overlap1 + 512);
Free(&pool, overlap1);
Free(&pool, overlap2);
void *tooBig = Allocate(&pool, 128 * 1024);
Assert(tooBig == NULL);
// Scattered allocations
void *resrv[3];
resrv[0] = Allocate(&pool, 128);
resrv[1] = Allocate(&pool, 64);
resrv[2] = Allocate(&pool, 128);
Free(&pool, resrv[1]);
void *nxt = Allocate(&pool, 64);
Assert(nxt == resrv[1]);
Free(&pool, resrv[0]);
Free(&pool, resrv[1]);
Free(&pool, resrv[2]);
free(memory);
}
/*
@implemented
*/
KSDDKAPI NTSTATUS NTAPI
KsCreateDefaultAllocatorEx(
IN PIRP Irp,
IN PVOID InitializeContext OPTIONAL,
IN PFNKSDEFAULTALLOCATE DefaultAllocate OPTIONAL,
IN PFNKSDEFAULTFREE DefaultFree OPTIONAL,
IN PFNKSINITIALIZEALLOCATOR InitializeAllocator OPTIONAL,
IN PFNKSDELETEALLOCATOR DeleteAllocator OPTIONAL)
{
NTSTATUS Status;
PKSALLOCATOR_FRAMING AllocatorFraming;
PALLOCATOR Allocator;
PVOID Ctx;
/* first validate connect request */
Status = KsValidateAllocatorCreateRequest(Irp, &AllocatorFraming);
if (!NT_SUCCESS(Status))
return STATUS_INVALID_PARAMETER;
/* check the valid file alignment */
if (AllocatorFraming->FileAlignment > (PAGE_SIZE-1))
{
FreeItem(AllocatorFraming);
return STATUS_INVALID_PARAMETER;
}
/* allocate allocator struct */
Allocator = AllocateItem(NonPagedPool, sizeof(ALLOCATOR));
if (!Allocator)
{
FreeItem(AllocatorFraming);
return STATUS_INSUFFICIENT_RESOURCES;
}
/* allocate object header */
Status = KsAllocateObjectHeader((KSOBJECT_HEADER*)&Allocator->Header, 0, NULL, Irp, &DispatchTable);
if (!NT_SUCCESS(Status))
{
FreeItem(AllocatorFraming);
FreeItem(Allocator);
return Status;
}
/* set allocator type in object header */
Allocator->lpVtbl = &vt_IKsAllocator;
Allocator->Header->Unknown = (PUNKNOWN)&Allocator->lpVtbl;
Allocator->ref = 1;
if (DefaultAllocate)
{
/* use external allocator */
Allocator->Type = ALLOCATOR_CUSTOM;
Allocator->Allocate.DefaultAllocate = DefaultAllocate;
Allocator->Free.DefaultFree = DefaultFree;
Allocator->Delete.DefaultDelete = DeleteAllocator;
Ctx = InitializeAllocator(InitializeContext, AllocatorFraming, &Allocator->u.CustomList);
/* check for success */
if (!Ctx)
{
KsFreeObjectHeader(Allocator->Header);
FreeItem(Allocator);
return Status;
}
}
else if (AllocatorFraming->PoolType == NonPagedPool)
{
/* use non-paged pool allocator */
Allocator->Type = ALLOCATOR_NPAGED_LOOKASIDE;
Allocator->Allocate.NPagedPool = ExAllocateFromNPagedLookasideList;
Allocator->Free.NPagedPool = ExFreeToNPagedLookasideList;
Allocator->Delete.NPagedPool = ExDeleteNPagedLookasideList;
ExInitializeNPagedLookasideList(&Allocator->u.NPagedList, NULL, NULL, 0, AllocatorFraming->FrameSize, 0, 0);
}
else if (AllocatorFraming->PoolType == PagedPool)
{
/* use paged pool allocator */
Allocator->Allocate.PagedPool = ExAllocateFromPagedLookasideList;
Allocator->Free.PagedPool = ExFreeToPagedLookasideList;
Allocator->Delete.PagedPool = ExDeletePagedLookasideList;
Allocator->Type = ALLOCATOR_PAGED_LOOKASIDE;
ExInitializePagedLookasideList(&Allocator->u.PagedList, NULL, NULL, 0, AllocatorFraming->FrameSize, 0, 0);
}
/* backup allocator framing */
RtlMoveMemory(&Allocator->Status.Framing, AllocatorFraming, sizeof(KSALLOCATOR_FRAMING));
FreeItem(AllocatorFraming);
return Status;
}
