Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place.

1
#include "byakugan.h"
2
#include "heapStructs.h"
3

4
// Get offset into the hash map
5
ULONG hash(PVOID addr, struct HPool *heap) {
6
    ULONG   offset;
7

8
    // assume 8 bit alignment, and use the golden ratio of 2^32
9
    offset  = (((ULONG) addr >> 3) * (2654435761)) % (0x007fffff/8);
10
    //dprintf("[XXX] Hashed to bucket %d\n", offset);
11

12
    return (offset);
13
}
14

15
/*
16
 * added by jc. I need to be able to find the offset of a chunk inide the chunks array quickly.
17
 * returns < 0 on error.
18
 */
19
int FindOffsetForChunk(struct HPool *heap, PVOID memoryPointer)
20
{
21
    ULONG                   i;
22

23
    i = hash(memoryPointer, heap);
24
    i = heap->map[i];
25
    while (i != NULLNODE) {
26
        if (CHUNK(i).addr == memoryPointer)
27
            return (i);
28
        i = CHUNK(i).nextBucket;
29
    }
30
    return (NULLNODE); // error
31
}
32

33
ULONG insertBucket(ULONG newC, struct HPool *heap, ULONG inAfter) {
34
    ULONG   offset, lastNum, checkNum;
35

36
    heap->numChunks++;
37
    CHUNK(newC).inUse       = TRUE;
38

39
    // Set up nextInUse list
40
	if (inAfter == NULLNODE) {
41
		if (heap->inUseHead == NULLNODE) heap->inUseHead = newC;
42
		else CHUNK(heap->lastInUse).nextInUse = newC;
43
		heap->lastInUse = newC;
44
		CHUNK(newC).nextInUse = NULLNODE;
45
		//dprintf("[T] Inserting new chunk %d at the end\n", newC);
46
	} else {
47
		//dprintf("[T] Inserting new chunk %d after chunk %d\n", newC, inAfter);
48
		if (heap->lastInUse == inAfter) heap->lastInUse = newC;
49
		CHUNK(newC).nextInUse = CHUNK(inAfter).nextInUse;
50
		CHUNK(inAfter).nextInUse = newC;
51
	}
52

53
    offset = hash(CHUNK(newC).addr, heap);
54

55
    if (heap->map[offset] == NULLNODE) {
56
        heap->map[offset] = newC;
57
        CHUNK(newC).nextBucket = NULLNODE;
58
    } else {
59
        offset = heap->map[offset];
60
        while (CHUNK(offset).nextBucket != NULLNODE)
61
            offset = CHUNK(offset).nextBucket;
62
        CHUNK(newC).nextBucket  = NULLNODE;
63
        CHUNK(offset).nextBucket = newC;
64
    }
65
    return (newC);
66
}
67

68
void removeBucket(PVOID addr, struct HPool *heap) {
69
    ULONG   offset, chunkNum, last = NULLNODE;
70

71
    if (heap == NULL)
72
        return;
73
	
74

75
    offset 		= hash(addr, heap);
76
	chunkNum 	= heap->map[offset];
77

78
    // Find our bucket and the bucket before it
79
    while (chunkNum != NULLNODE && CHUNK(chunkNum).addr != addr) {
80
        last = chunkNum;
81
        chunkNum = CHUNK(chunkNum).nextBucket;
82
    }
83

84
    // remove our bucket from the list if we found it
85
    if (chunkNum != NULLNODE) {
86
		//dprintf("[T] Deleting node %d\n", chunkNum);
87
    	CHUNK(chunkNum).inUse       = FALSE;
88
        
89
		if (last == NULLNODE)
90
            heap->map[offset] = NULLNODE;
91
        else
92
            CHUNK(last).nextBucket = CHUNK(chunkNum).nextBucket;
93
		
94
		heap->numChunks--;
95
    } else 
96
		dprintf("[T] Couldnt find chunk 0x%08x to delete :(\n", addr);
97

98

99
	// Head of the list case
100
	chunkNum = heap->inUseHead;
101
	if (CHUNK(chunkNum).addr == addr) {
102
		heap->inUseHead = CHUNK(chunkNum).nextInUse;
103
		return;
104
	}
105

106
	// search the list for the previous chunk
107
	while (chunkNum != NULLNODE && CHUNK(chunkNum).addr != addr) {
108
			last = chunkNum;
109
			chunkNum = CHUNK(chunkNum).nextInUse;
110
	}
111
	if (chunkNum != NULLNODE) {
112
		CHUNK(last).nextInUse = CHUNK(chunkNum).nextInUse;
113
		if (heap->lastInUse == chunkNum) heap->lastInUse = last;
114
	}
115
}
116

117
void initializeHeapModel(struct HeapState *heapModel) {
118
	
119
	memset(heapModel+1, 0, sizeof (struct HeapState));
120
	heapModel->hPoolListLen = 16;
121
	heapModel->heaps = (struct HPool *) calloc(heapModel->hPoolListLen, sizeof (struct HPool));
122
}
123

124
#define ALLOCHEAD	"<Allocation>\n"
125
#define ALLOCFOOT	"</Allocation>\n"
126
#define REALLOCHEAD "<Reallocation>\n"
127
#define REALLOCFOOT "</Reallocation>\n"
128
#define FREEHEAD	"<Free>\n"
129
#define FREEFOOT	"</Free>\n"
130
#define HEAP		"\t<Heap>0x%08x</Heap>\n"
131
#define ADDRESS		"\t<Address>0x%08x</Address>\n" 
132
#define RETURNPTR	"\t<NewAddress>0x%08x</NewAddress>\n"
133
#define SIZE		"\t<Size>0x%08x</Size>\n"
134
#define	TRACEHEAD	"\t<StackTrace>\n"
135
#define	TRACEFOOT	"\t</StackTrace>\n"
136
#define CALLADDR	"\t\t<Caller>0x%08x</Caller>\n"
137

138
void hprintf(HANDLE fHandle, char *format, ...) {
139
	char	ptr[1024];
140
	DWORD	out;
141
	va_list	args;
142

143
	va_start(args, format);
144

145

146
	StringCbVPrintf(ptr, 1024, format, args);
147
	WriteFile(fHandle, ptr, strlen(ptr), &out, NULL);
148
}
149

150
void logAllocate(struct HeapState *heapModel, struct AllocateStruct *aStruct) {
151
	ULONG	infoUsed;
152
	
153
	WriteFile(heapModel->hLogFile, ALLOCHEAD, strlen(ALLOCHEAD), &infoUsed, NULL);
154
	hprintf(heapModel->hLogFile, HEAP, aStruct->heapHandle);
155
	hprintf(heapModel->hLogFile, ADDRESS, aStruct->ret);
156
	hprintf(heapModel->hLogFile, SIZE, aStruct->size);
157
	WriteFile(heapModel->hLogFile, TRACEHEAD, strlen(TRACEHEAD), &infoUsed, NULL);
158

159
	hprintf(heapModel->hLogFile, CALLADDR, aStruct->caller);
160

161
	WriteFile(heapModel->hLogFile, TRACEFOOT, strlen(TRACEFOOT), &infoUsed, NULL);
162
	WriteFile(heapModel->hLogFile, ALLOCFOOT, strlen(ALLOCFOOT), &infoUsed, NULL);
163
}
164

165
void logReallocate(struct HeapState *heapModel, struct ReallocateStruct *rStruct) {
166
	ULONG	infoUsed;
167

168
	WriteFile(heapModel->hLogFile, REALLOCHEAD, strlen(REALLOCHEAD), &infoUsed, NULL);
169
	hprintf(heapModel->hLogFile, HEAP, rStruct->heapHandle);
170
	hprintf(heapModel->hLogFile, ADDRESS, rStruct->memoryPointer);
171
	hprintf(heapModel->hLogFile, RETURNPTR, rStruct->ret);
172
	hprintf(heapModel->hLogFile, SIZE, rStruct->size);
173
	WriteFile(heapModel->hLogFile, TRACEHEAD, strlen(TRACEHEAD), &infoUsed, NULL);
174

175
	hprintf(heapModel->hLogFile, CALLADDR, rStruct->caller);
176

177
	WriteFile(heapModel->hLogFile, TRACEFOOT, strlen(TRACEFOOT), &infoUsed, NULL);
178
	WriteFile(heapModel->hLogFile, ALLOCFOOT, strlen(REALLOCFOOT), &infoUsed, NULL);
179
}
180

181
void logFree(struct HeapState *heapModel, struct FreeStruct *fStruct) {
182
	ULONG	infoUsed;
183

184
	WriteFile(heapModel->hLogFile, FREEHEAD, strlen(FREEHEAD), &infoUsed, NULL);
185
	hprintf(heapModel->hLogFile, HEAP, fStruct->heapHandle);
186
	hprintf(heapModel->hLogFile, ADDRESS, fStruct->memoryPointer);
187
	WriteFile(heapModel->hLogFile, TRACEHEAD, strlen(TRACEHEAD), &infoUsed, NULL);
188
	
189
	hprintf(heapModel->hLogFile, CALLADDR, fStruct->caller);
190

191
	WriteFile(heapModel->hLogFile, TRACEFOOT, strlen(TRACEFOOT), &infoUsed, NULL);
192
	WriteFile(heapModel->hLogFile, FREEFOOT, strlen(FREEFOOT), &infoUsed, NULL);
193
}
194

195
void heapAllocate(struct HeapState *heapModel, struct AllocateStruct *aStruct) {
196
	struct	HPool		*myHeap;
197
	struct	HeapChunk	*myChunk, *remainder;
198

199
	if (aStruct->heapHandle == 0 || aStruct->ret == 0)
200
		return;
201

202
	myHeap	= getHeap(heapModel, aStruct->heapHandle);
203
	myChunk = getChunk(myHeap, aStruct->ret, NULLNODE);
204
	
205
	if (myChunk == NULL)
206
		return;
207

208
	if (myChunk->size > aStruct->size + SPACEBETWEEN) {
209
		//dprintf("[T] Splitting chunk at 0x%08x:\n", aStruct->ret);
210
		//dprintf("\tNew Chunk 0x%08x, size %d bytes, after chunk %u\n",
211
		//		((ULONG) aStruct->ret + aStruct->size + SPACEBETWEEN),
212
		//		(myChunk->size - aStruct->size - SPACEBETWEEN),
213
		//		FindOffsetForChunk(myHeap, myChunk->addr));
214
		remainder		=	getChunk(myHeap, 
215
							(PVOID) ((ULONG) aStruct->ret + aStruct->size + SPACEBETWEEN), 
216
							FindOffsetForChunk(myHeap, myChunk->addr));
217
		remainder->size = (myChunk->size - aStruct->size) - SPACEBETWEEN;
218
		remainder->free = TRUE;
219
	}
220
	
221
	myChunk->size		= aStruct->size;
222
	myChunk->flags		= aStruct->flags;
223
	myChunk->free		= FALSE;
224
}
225

226
void heapReallocate(struct HeapState *heapModel, struct ReallocateStruct *rStruct) {
227
    struct  HPool   *myHeap; 
228
	struct	HeapChunk	*oChunk, *nChunk, *remainder;
229
	    
230
    if (rStruct->heapHandle == 0 || rStruct->ret == 0)
231
        return;
232

233
	myHeap = getHeap(heapModel, rStruct->heapHandle);
234
	nChunk = getChunk(myHeap, rStruct->ret, NULLNODE);
235
	if (nChunk == NULL)
236
		return;
237

238
	if (rStruct->ret != rStruct->memoryPointer && rStruct->memoryPointer != 0) {
239
		oChunk = getChunk(myHeap, rStruct->memoryPointer, NULLNODE);
240
		if (oChunk != NULL)
241
			oChunk->free = TRUE;
242
	}
243

244
	// Split new chunk on realloc move if necessary
245
    if (nChunk->size > rStruct->size + SPACEBETWEEN) {
246
        remainder       =   getChunk(myHeap,
247
                            (PVOID) ((ULONG) rStruct->ret + rStruct->size + SPACEBETWEEN),
248
							FindOffsetForChunk(myHeap, nChunk->addr));
249
        remainder->size = (nChunk->size - rStruct->size) - SPACEBETWEEN;
250
        remainder->free = TRUE;
251
    }
252

253
	nChunk->size	= rStruct->size;
254
	nChunk->flags	= rStruct->flags;
255
	nChunk->free	= FALSE;
256
}
257

258
void heapFree(struct HeapState *heapModel, struct FreeStruct *fStruct) {
259
    struct  HPool   *myHeap; 
260
	struct	HeapChunk	*myChunk;
261
	
262
	//dprintf("[XXX] Freeing 0x%08x\n", fStruct->memoryPointer);
263

264
	if (fStruct->heapHandle == 0 || fStruct->memoryPointer == 0x00000000) { 
265
		// So many of these that it slows us down :(
266
		//dprintf("[T] Program attempted to free a NULL pointer.\n\n");
267
		return;
268
	}
269
	myHeap	= getHeap(heapModel, fStruct->heapHandle);
270
	myChunk	= getChunk(myHeap, fStruct->memoryPointer, NULLNODE);
271
	if (myChunk == NULL)
272
		return; // dupe free :(
273

274
	#if 0
275
	if (myChunk->free == 2)
276
		dprintf("[T] Possible 'double free' of chunk @ 0x%08x:0x%x\n\n", 
277
				myChunk->addr, myChunk->free);
278
	#endif
279

280
	myChunk->flags	= fStruct->flags;
281
	myChunk->free	+= 1;
282
}
283

284
void heapCreate(struct HeapState *heapModel, struct CreateStruct *cStruct) {
285
	struct HPool	*newHeap;
286
	ULONG			i;
287

288
    if (cStruct->base == 0)
289
		return;
290

291
	newHeap = getHeap(heapModel, cStruct->base);
292
    newHeap->commit		= cStruct->commit;
293
    newHeap->reserve	= cStruct->reserve;
294
    newHeap->lock		= cStruct->lock;
295
    newHeap->flags		= cStruct->flags;
296
}
297

298
void heapDestroy(struct HeapState *heapModel, struct DestroyStruct *dStruct) {
299
	// Only place we mark heaps as not in use
300
	struct HPool	*deadHeap;
301

302
	if (dStruct->heapHandle == 0)
303
		return;
304

305
	deadHeap = getHeap(heapModel, dStruct->heapHandle);
306
	free(deadHeap->chunks);
307
	memset(deadHeap, 0, sizeof(struct HPool));
308
	heapModel->numHeaps--;
309
}
310

311
void heapCoalesce(struct HeapState *heapModel, struct CoalesceStruct *cfbStruct) {
312
	// Only place we mark chunks as not in use
313
	ULONG                   i, j, latest, last, offset;
314
	struct HPool			*heap;
315

316
	if (heapModel->numHeaps == 0)
317
		return;
318

319
	// Get the heap handle from the CoalesceStruct
320
	heap = getHeap(heapModel, cfbStruct->heapHandle);
321

322
	if (heap->numChunks < 2)
323
		return;
324
 
325
	//dprintf("[T] Attempting to coalese heap 0x%08x...\n", cfbStruct->heapHandle);
326
	
327
	i = heap->inUseHead;
328

329
	//dprintf("[T] Starting with chunk %d at 0x%08x\n", i, CHUNK(i).addr);
330

331
	// Walk the list Coalescing consecutive free chunks
332
	// This assumes that two chunks in a row are actually consecutive...
333
	// This is wrong. If I alloc two large blocks, free the first, then alloc
334
	// a small one, then free the small one and the second one, then coalesce
335
	// the result will ignore the space in between, so we need to take addresses
336
	// into account, and/or we need to create free spaceholder chunks when creating
337
	// chunks in a smaller space than is currently existing.
338
	while (i != NULLNODE) {
339
		if (CHUNK(i).free) {
340
			//dprintf("[T] Found free chunk %d at 0x%08x\n", i, CHUNK(i).addr);
341
			j = FindOffsetForChunk(heap, NEXTADDR(i));
342
			while (j != NULLNODE && CHUNK(j).free) {
343

344
				//dprintf("[T] Found free chunk %d at 0x%08x\n", j, CHUNK(j).addr);
345
				
346
				// XXX Handle this by address for unknown chunks
347
				CHUNK(i).size += CHUNK(j).size + SPACEBETWEEN;	// This is only right on Vista... and only
348
															// most of the time - win2k is insane and 
349
															// goes backwards :(
350
				
351
				// fix up hash bucket list
352
				removeBucket(CHUNK(j).addr, heap);
353

354
				//dprintf("[T] Coalescing 0x%08x and 0x%08x. New size: %d bytes\n",
355
				//CHUNK(i).addr, CHUNK(j).addr, CHUNK(i).size);
356
				
357
				j = FindOffsetForChunk(heap, NEXTADDR(i));
358
			}
359
		} 
360
       	i = CHUNK(i).nextInUse;
361
		// Check Address ordering
362
		// If out of order, quicksort then start this loop again
363
		// from first in use 
364
    }
365

366
}
367

368
struct HPool *getHeap(struct HeapState *heapModel, PVOID heapHandle) {
369
	ULONG					i;
370
	struct HPool			*newHeap = NULL;
371

372
	for (i = 0; i < heapModel->hPoolListLen; i++) {
373
		if ((heapModel->heaps[i]).base == heapHandle)
374
			return (&(heapModel->heaps[i]));
375
		if (heapModel->numHeaps == heapModel->hPoolListLen)	// Add some short circuits here
376
			continue;
377
		if (newHeap != NULL)
378
			continue;
379
		if ((heapModel->heaps[i]).inUse == FALSE)
380
			newHeap = &(heapModel->heaps[i]);
381
	}
382

383
	// If we get here we haven't found a heap, so lets make one
384
	//dprintf("[T] Creating entry for heap: 0x%08x\n", heapHandle);
385

386
	if (newHeap == NULL) {
387
		heapModel->hPoolListLen = heapModel->hPoolListLen * 2;
388
		newHeap    = heapModel->heaps;
389
		heapModel->heaps = (struct HPool *) realloc(heapModel->heaps,
390
				            heapModel->hPoolListLen * sizeof(struct HPool));
391
		if (heapModel->heaps == NULL) {
392
            heapModel->heaps = newHeap;
393
            //crushModel(heapModel);    ADD ME XXX
394
			#ifndef THREEDHEAPFU_ENABLED
395
            dprintf("[T] OOM getting new heaps!\n\n");
396
			#endif
397
            return (NULL);
398
        }
399
		newHeap = &(heapModel->heaps[heapModel->numHeaps]);
400
		// Clean the newly allocated memory
401
		memset(newHeap, 0, (sizeof (struct HPool) * (heapModel->hPoolListLen - heapModel->numHeaps)));
402
	}
403

404
	heapModel->numHeaps++;
405
	newHeap->base			= heapHandle;
406
    newHeap->inUse			= TRUE;
407
	newHeap->chunkListLen	= 32;
408
	newHeap->chunks			= (struct HeapChunk *) calloc(newHeap->chunkListLen, 
409
								sizeof (struct HeapChunk));
410
	newHeap->map			= (ULONG *) malloc((0x007fffff/8) * sizeof (ULONG));
411
	memset(newHeap->map, 0xff, (0x007fffff/8) * sizeof (ULONG));
412
	newHeap->inUseHead = newHeap->lastInUse = NULLNODE;
413

414
	return (newHeap);
415
}
416

417
struct HeapChunk *getChunk(struct HPool *heap, PVOID memoryPointer, ULONG inAfter) {
418
	ULONG					i, *oldMap;
419
	struct HeapChunk		*newChunk, *niuChunk = NULL;
420

421
	i = hash(memoryPointer, heap);
422
	i = heap->map[i];
423
	while (i != NULLNODE) {
424
		if (CHUNK(i).addr == memoryPointer)
425
			return (&CHUNK(i));
426
		i = CHUNK(i).nextBucket;
427
	}
428

429
	if (niuChunk == NULL) {
430
		if (heap->numChunks < heap->chunkListLen) {
431
			if ((heap->chunks[heap->numChunks]).inUse == FALSE) {
432
				//dprintf("[XXX] FOUND CHUNK WHERE EXPECTED! :)\n\n");
433
				niuChunk = &(heap->chunks[heap->numChunks]);
434
    			niuChunk->addr			= memoryPointer;
435
				niuChunk->heapHandle	= heap->base;
436
				insertBucket(heap->numChunks, heap, inAfter);
437
				return (niuChunk);
438
			}
439
			// Heap has been coalesced and there are gaps
440
			// Work backwards in this case - should be MUCH faster
441
			// If this is noticably slow, we'll switch to a stack of unused chunks
442
			// (realloc stack with heap)
443
			for (i = 0; i < heap->chunkListLen;  i++) {
444
				if ((heap->chunks[i]).inUse == FALSE) {
445
					niuChunk = &(heap->chunks[i]);
446
    				niuChunk->addr  		= memoryPointer;
447
					niuChunk->heapHandle	= heap->base;
448
					insertBucket(i, heap, inAfter);
449
					return (niuChunk);
450
				}
451
			}
452
			dprintf("[XXX] Totally fucked up - chunk info doesnt jive with heap\n\n");
453
		}
454

455
		heap->chunkListLen = heap->chunkListLen * 2;
456
		niuChunk = heap->chunks;
457
		heap->chunks = (struct HeapChunk *) realloc(heap->chunks,
458
						heap->chunkListLen * sizeof (struct HeapChunk));
459
		
460
		if (heap->chunks == NULL) {
461
			heap->chunks = niuChunk;
462
			//crushHeap(heap);
463
			heap->chunkListLen = heap->chunkListLen/2;
464
			#ifndef THREEDHEAPFU_ENABLED
465
			dprintf("[T] OOM getting new chunks! %d -> %d\n\n",
466
					heap->chunkListLen, heap->chunkListLen*2);
467
			#endif
468
			return (NULL);
469
		}
470
        
471
		niuChunk = &(heap->chunks[heap->numChunks]);
472
		memset(niuChunk, 0, (sizeof (struct HeapChunk) * (heap->chunkListLen - heap->numChunks)));
473
	} 
474

475
	niuChunk->addr			= memoryPointer;
476
	niuChunk->heapHandle	= heap->base;
477
	insertBucket(heap->numChunks, heap, inAfter);
478
	return (niuChunk);
479
}
480

481