metasploit-framework/external/source/byakugan/heapSplay.cpp

324 lines
8.6 KiB
C++
Raw Normal View History

#include "byakugan.h"
#include "heapStructs.h"
void initializeHeapModel(struct HeapState *heapModel) {
memset(heapModel, 0, sizeof (struct HeapState));
heapModel->hPoolListLen = 16;
heapModel->heaps = (struct HPool *) calloc(heapModel->hPoolListLen, sizeof (struct HPool));
}
void heapAllocate(struct HeapState *heapModel, struct AllocateStruct *aStruct) {
struct HPool *myHeap;
struct HeapChunk *myChunk;
if (aStruct->heapHandle == 0 || aStruct->ret == 0)
return;
myHeap = getHeap(heapModel, aStruct->heapHandle);
myChunk = getChunk(myHeap, aStruct->ret);
if (myChunk == NULL)
return;
myChunk->size = aStruct->size;
myChunk->flags = aStruct->flags;
myChunk->free = FALSE;
}
void heapReallocate(struct HeapState *heapModel, struct ReallocateStruct *rStruct) {
struct HPool *myHeap;
struct HeapChunk *oChunk, *nChunk;
if (rStruct->heapHandle == 0 || rStruct->ret == 0)
return;
myHeap = getHeap(heapModel, rStruct->heapHandle);
nChunk = getChunk(myHeap, rStruct->ret);
if (nChunk == NULL)
return;
if (rStruct->ret != rStruct->memoryPointer && rStruct->memoryPointer != 0) {
oChunk = getChunk(myHeap, rStruct->memoryPointer);
if (oChunk != NULL)
oChunk->free = TRUE;
}
nChunk->size = rStruct->size;
nChunk->flags = rStruct->flags;
nChunk->free = FALSE;
}
void heapFree(struct HeapState *heapModel, struct FreeStruct *fStruct) {
struct HPool *myHeap;
struct HeapChunk *myChunk;
if (fStruct->heapHandle == 0 || fStruct->memoryPointer == 0x00000000) {
// So many of these that it slows us down :(
//dprintf("[T] Program attempted to free a NULL pointer.\n\n");
return;
}
myHeap = getHeap(heapModel, fStruct->heapHandle);
myChunk = getChunk(myHeap, fStruct->memoryPointer);
if (myChunk == NULL)
return;
if (myChunk->free == TRUE)
dprintf("[T] Possible 'double free' of chunk @ 0x%08x:0x%08x\n\n",
myChunk->addr, fStruct->memoryPointer);
myChunk->flags = fStruct->flags;
myChunk->free = TRUE;
}
void heapCreate(struct HeapState *heapModel, struct CreateStruct *cStruct) {
struct HPool *newHeap;
ULONG i;
if (cStruct->base == 0)
return;
newHeap = getHeap(heapModel, cStruct->base);
newHeap->commit = cStruct->commit;
newHeap->reserve = cStruct->reserve;
newHeap->lock = cStruct->lock;
newHeap->flags = cStruct->flags;
}
void heapDestroy(struct HeapState *heapModel, struct DestroyStruct *dStruct) {
// Only place we mark heaps as not in use
struct HPool *deadHeap;
if (dStruct->heapHandle == 0)
return;
deadHeap = getHeap(heapModel, dStruct->heapHandle);
free(deadHeap->chunks);
memset(deadHeap, 0, sizeof(struct HPool));
heapModel->numHeaps--;
}
void heapCoalesce(struct HeapState *heapModel, struct CoalesceStruct *cfbStruct) {
// Only place we mark chunks as not in use
}
struct HPool *getHeap(struct HeapState *heapModel, PVOID heapHandle) {
ULONG i;
struct HPool *newHeap = NULL;
for (i = 0; i < heapModel->hPoolListLen; i++) {
if ((heapModel->heaps[i]).base == heapHandle)
return (&(heapModel->heaps[i]));
if (heapModel->numHeaps == heapModel->hPoolListLen) // Add some short circuits here
continue;
if (newHeap != NULL)
continue;
if ((heapModel->heaps[i]).inUse == FALSE)
newHeap = &(heapModel->heaps[i]);
}
// If we get here we haven't found a heap, so lets make one
//dprintf("[T] Creating entry for heap: 0x%08x\n", heapHandle);
if (newHeap == NULL) {
heapModel->hPoolListLen = heapModel->hPoolListLen * 2;
newHeap = heapModel->heaps;
heapModel->heaps = (struct HPool *) realloc(heapModel->heaps,
heapModel->hPoolListLen * sizeof(struct HPool));
if (heapModel->heaps == NULL) {
heapModel->heaps = newHeap;
//crushModel(heapModel); ADD ME XXX
dprintf("[T] OOM getting new heaps!\n\n");
return (NULL);
}
newHeap = &(heapModel->heaps[heapModel->numHeaps]);
// Clean the newly allocated memory
memset(newHeap, 0, (sizeof (struct HPool) * (heapModel->hPoolListLen - heapModel->numHeaps)));
}
heapModel->numHeaps++;
newHeap->base = heapHandle;
newHeap->inUse = TRUE;
newHeap->chunkListLen = 0x007fffff/8;
newHeap->chunks = (struct HeapChunk *) calloc(newHeap->chunkListLen,
sizeof (struct HeapChunk));
return (newHeap);
}
// CRRRRAAAAZZYYY splay tree action on top of our chunk code...
// Many thanks to Sedgewick and Sleator :)
struct HeapChunk *splay(PVOID addr, struct HeapChunk *t) {
struct HeapChunk N, *l, *r, *y;
if (t == NULL) {
return (t);
}
memset(&N, 0, sizeof (HeapChunk));
l = r = &N;
for (;;) {
if (addr < t->addr) {
if (t->left == NULL) break;
if (addr < (t->left)->addr) {
y = t->left;
t->left = y->right;
y->right = t;
t = y;
if (t->left == NULL) break;
}
r->left = t;
r = t;
t = t->left;
} else if (addr > t->addr) {
if (t->right == NULL) break;
if (addr > (t->right)->addr) {
y = t->right;
t->right = y->left;
y->left = t;
t = y;
if (t->right == NULL) break;
}
l->right = t;
l = t;
t = t->right;
} else {
break;
}
}
l->right = t->left;
r->left = t->right;
t->left = N.right;
t->right = N.left;
return (t);
}
struct HeapChunk *insert(struct HeapChunk *newC, struct HPool *heap) {
/* Insert i into the tree t, unless it's already there. */
/* Return a pointer to the resulting tree. */
struct HeapChunk *t;
heap->numChunks++;
newC->inUse = TRUE;
t = heap->t;
if (t == NULL) {
heap->t = newC;
return (newC);
}
t = splay(newC->addr,t);
if (newC->addr < t->addr) {
newC->left = t->left;
newC->right = t;
t->left = NULL;
heap->t = newC;
return (newC);
} else if (newC->addr > t->addr) {
newC->right = t->right;
newC->left = t;
t->right = NULL;
heap->t = newC;
return (newC);
} else { /* We get here if it's already in the tree */
/* Don't add it again */
dprintf("[XXX] WTF?? Inserting the same address twice!\nBase: 0x%08x\tAddr: 0x%08x:0x%08x\n\n",
heap->base, newC->addr, t->addr);
t->inUse = TRUE;
heap->t = t;
return (t);
}
}
struct HeapChunk *getChunk(struct HPool *heap, PVOID memoryPointer) {
ULONG i;
struct HeapChunk *newChunk, *niuChunk = NULL;
// Old and slow - O(n)
#if 0
for (i = 0; i < heap->chunkListLen; i++) {
newChunk = &(heap->chunks[i]);
if (newChunk->addr == memoryPointer)
return (newChunk);
if (heap->numChunks == heap->chunkListLen)
continue;
if (niuChunk != NULL)
continue;
if (newChunk->inUse == FALSE)
niuChunk = newChunk;
}
// New splayness - O(log n)
newChunk = heap->t;
i = 0;
while (newChunk != NULL) {
if (newChunk->addr == memoryPointer) {
return (newChunk);
}
if (newChunk->inUse == FALSE) {
dprintf("[XXX] Unused chunk in the tree!! L%d 0x%08x\n\n",
i, newChunk->addr);
niuChunk = newChunk;
}
if (memoryPointer < newChunk->addr) {
newChunk = newChunk->left;
i++;
} else {
newChunk = newChunk->right;
i++;
}
}
#endif
heap->t = splay(memoryPointer, heap->t);
if (heap->t != NULL && (heap->t)->addr == memoryPointer)
return (heap->t);
if (niuChunk == NULL) {
if (heap->numChunks < heap->chunkListLen) {
if ((heap->chunks[heap->numChunks]).inUse == FALSE) {
//dprintf("[XXX] FOUND CHUNK WHERE EXPECTED! :)\n\n");
niuChunk = &(heap->chunks[heap->numChunks]);
niuChunk->addr = memoryPointer;
niuChunk->heapHandle = heap->base;
insert(niuChunk, heap);
return (niuChunk);
}
// Worst muddahfuggin case...
dprintf("[XXX] Unused chunks should be at the END of the list!\n\n");
for (i = 0; i < heap->chunkListLen; i++) {
if ((heap->chunks[i]).inUse == FALSE) {
niuChunk = &(heap->chunks[i]);
niuChunk->addr = memoryPointer;
niuChunk->heapHandle = heap->base;
insert(niuChunk, heap);
return (niuChunk);
}
}
}
heap->chunkListLen = heap->chunkListLen * 2;
niuChunk = heap->chunks;
heap->chunks = (struct HeapChunk *) realloc(heap->chunks,
heap->chunkListLen * sizeof (struct HeapChunk));
if (heap->chunks == NULL) {
heap->chunks = niuChunk;
//crushHeap(heap);
heap->chunkListLen = heap->chunkListLen/2;
dprintf("[T] OOM getting new chunks! %d -> %d\n\n",
heap->chunkListLen, heap->chunkListLen*2);
return (NULL);
}
niuChunk = &(heap->chunks[heap->numChunks]);
memset(niuChunk, 0, (sizeof (struct HeapChunk) * (heap->chunkListLen - heap->numChunks)));
}
niuChunk->addr = memoryPointer;
niuChunk->heapHandle = heap->base;
insert(niuChunk, heap);
return (niuChunk);
}