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

1
#include <windows.h>
2
#include <stdlib.h>
3

4
#include "byakugan.h"
5
#include "tenketsu.h"
6
#include "heapStructs.h"
7

8
#define BUFSIZE		4096
9

10
// UNDOCUMENTED FUNCTIONS TO ADD
11
//
12
//		RtlAllocateMemoryBlockLookaside
13
//		RtlFreeMemoryBlockLookaside
14
//		
15
//		RtlCreateMemoryBlockLookaside
16
//		RtlDestroyMemoryBlockLookaside
17
//		
18
//		RtlExtendMemoryBlockLookaside
19
//		RtlResetMemoryBlockLookaside
20

21
PCSTR	undocdFunc[] = { "ntdll!RtlpCoalesceFreeBlocks", NULL };
22
ULONG	undocdAddr[sizeof (undocdFunc)+1];
23

24
struct HeapState		heapModel;
25
BOOLEAN					running = FALSE;
26

27
// Two things that fucking rock? Bunnies and Jaguars. Werd.
28

29

30
int hookRtlHeap(BYTE type, char *fileName) {
31
    HRESULT     Status;
32
    HANDLE      process;
33
    DWORD       pid;
34
    HANDLE processHandle = 0;
35
    HANDLE threadHandle = 0;
36
    LPVOID stringAddress = NULL;
37
    LPCSTR dllName = "C:\\windbg\\injectsu.dll";
38
    ULONG64     funcAddr64;
39
    ULONG       *funcAddr, i;
40

41
	heapModel.state = heapModel.state ^ type;
42

43
	if (running) {
44
		dprintf("[Byakugan] Hooks are already injected.\n");
45
		return (0);
46
	}
47

48
	running = TRUE;
49

50
	dprintf("[Byakugan] Beginning data gathering thread... ");
51
	if(tenkListener()) {
52
		dprintf("\n[Byakugan] Failed to create heap info back channel :(\n");
53
		VirtualFreeEx(processHandle, stringAddress, strlen(dllName), MEM_DECOMMIT);
54
        CloseHandle(processHandle); 
55
        return (-1);	
56
	}
57
   	dprintf("Success!\n[Byakugan] Injecting Tenketsu Heap Monitoring DLL... ");
58
    
59
	Status = g_ExtSystem->GetCurrentProcessSystemId(&pid);
60
    if (Status != S_OK)
61
        return (-1);
62

63
    if(!(processHandle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid))){
64
        dprintf("\n[Byakugan] Unable to OpenProcess().\n");
65
        return (-1);
66
    }
67

68
    if(!(stringAddress = VirtualAllocEx(processHandle, NULL, strlen(dllName), MEM_COMMIT, PAGE_EXECUTE_READWRITE))){
69
        printf("\n[Byakugan] VirtualAllocEx() failed.\n");
70
        CloseHandle(processHandle);
71
        return (-1);
72
    }
73

74
    if(!WriteProcessMemory(processHandle, (LPVOID)stringAddress, dllName, strlen(dllName), NULL)){
75
        dprintf("\n[Byakugan] WriteProcessMemory() failed.\n");
76
        VirtualFreeEx(processHandle, stringAddress, strlen(dllName), MEM_DECOMMIT);
77
        CloseHandle(processHandle);
78
        return (-1);
79
    }
80

81
    dprintf("Success!\n");
82

83
	// Resolve undocumented functions!	
84
    i = 0;
85
	g_ExtSymbols->Reload("/f ntdll.dll");
86
	dprintf("[Byakugan] Resolving undocumented Heap functions...\n");
87
    while (undocdFunc[i] != NULL) {
88
		if (g_ExtSymbols->GetOffsetByName(undocdFunc[i], &funcAddr64) == E_FAIL)
89
			funcAddr64 = NULL;
90
		funcAddr = (ULONG *) &funcAddr64;
91
        if (*funcAddr != NULL)
92
            dprintf("[T] Resolved undocumented function '%s' @ 0x%08x.\n", undocdFunc[i], *funcAddr);
93
        else
94
            dprintf("[T] Unable to resolve undocumented function '%s' :(\n", undocdFunc[i]);
95
        undocdAddr[i] = *funcAddr;
96
        i++;
97
    }
98
    undocdAddr[i] = NULL;
99

100

101
    if(!(threadHandle = CreateRemoteThread(processHandle, NULL, 0, (LPTHREAD_START_ROUTINE)GetProcAddress(LoadLibrary("kernel32.dll"), "LoadLibraryA"), \
102
                (LPVOID)stringAddress, 0, NULL))){
103
        dprintf("\n[Byakugan] CreateRemoteThread() failed.\n");
104
        VirtualFreeEx(processHandle, stringAddress, strlen(dllName), MEM_DECOMMIT);
105
        CloseHandle(processHandle);
106
        return (-1);
107
    }
108

109
	CloseHandle(threadHandle);
110
    CloseHandle(processHandle);
111

112
	// Set up the log file
113
	if (type == 2) {
114
		heapModel.hLogFile = CreateFileA(fileName, GENERIC_WRITE, FILE_SHARE_READ, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
115
		if (!heapModel.hLogFile) {
116
			dprintf("[T] Unable to open file \"%s\" for writing. :(\n");
117
		}
118
	}
119

120
    return (0);
121
}
122

123
int tenkListener(void) {
124
   	BOOL fConnected; 
125
   	DWORD dwThreadId; 
126
   	HANDLE hPipe, hThread; 
127
   	LPTSTR lpszPipename = TEXT("\\\\.\\pipe\\tenketsu"); 
128

129
	
130
	hPipe = CreateNamedPipe(	lpszPipename,
131
								PIPE_ACCESS_DUPLEX,
132
								PIPE_TYPE_MESSAGE | PIPE_READMODE_MESSAGE | PIPE_WAIT,
133
								PIPE_UNLIMITED_INSTANCES,
134
								BUFSIZE,
135
								BUFSIZE,
136
								NMPWAIT_USE_DEFAULT_WAIT,
137
								NULL);
138
	
139
	if (hPipe == INVALID_HANDLE_VALUE) {
140
		dprintf("[Byakugan] CreateNamedPipe() failed.\n");
141
		return (-1);
142
	}
143
    
144

145
	hThread = CreateThread(		NULL,
146
								0,
147
								tenkBackChannel,
148
								(LPVOID) hPipe,
149
								0,
150
								&dwThreadId);
151

152
    if (hThread == NULL) {
153
        dprintf("[Byakugan] CreateThread() failed.\n");
154
        return (-1);
155
    } 
156

157

158
	CloseHandle(hThread);
159

160

161

162
	return (0);
163
}
164

165
DWORD WINAPI tenkBackChannel(LPVOID lpvParam) {
166
	HANDLE		hPipe;
167
	TCHAR		buf[BUFSIZE+1];
168
	ULONG		bytesRead, bytesWritten;
169
	BOOL		fSuccess;
170
	ULONG64		funcAddr64;
171
	ULONG		*funcAddr, i;
172
	
173

174
	struct AllocateStruct	*aStruct = (struct AllocateStruct *)buf;
175
	struct ReallocateStruct	*rStruct = (struct ReallocateStruct *)buf;
176
	struct FreeStruct		*fStruct = (struct FreeStruct *)buf;
177
	struct CreateStruct		*cStruct = (struct CreateStruct *)buf;
178
	struct DestroyStruct	*dStruct = (struct DestroyStruct *)buf;
179
	struct CoalesceStruct	*cfbStruct = (struct CoalesceStruct *)buf;
180

181
	
182
	hPipe = (HANDLE) lpvParam;
183
	
184
	dprintf("[Byakugan] Waiting for named pipe connection...\n");
185

186
	for(;!ConnectNamedPipe(hPipe, NULL) == TRUE;) { dprintf("[B] waiting for connection...\n"); }
187

188
	dprintf("[Byakugan] Connected to back channel. :)\n");
189

190
	// Load addresses from symbols if possible for undoumented interfaces
191
	i = 0;
192
	while (undocdFunc[i] != NULL) {
193
		dprintf("[T] Sending address of %s\n", undocdFunc[i]);
194
		fSuccess = WriteFile(hPipe, &(undocdAddr[i]), sizeof(ULONG), &bytesWritten, NULL);
195
		if (!fSuccess || bytesWritten != sizeof(ULONG))
196
			dprintf("[T] Failed to send address of %s\n", undocdFunc[i]);
197
		i++;
198
	}
199
	//FlushFileBuffers(hPipe);
200
	dprintf("[T] Sent addresses of %d undocumented functions.\n", i);
201

202
	initializeHeapModel(&heapModel);
203

204
#undef THREEDHEAPFU_ENABLED //Place this in setup.bat
205

206
#if 0 //#ifdef THREEDHEAPFU_ENABLED
207
//Create a heap event proxy, play these back out to 3dheapfu
208
LPTSTR	lpszProxyPipename = TEXT("\\\\.\\pipe\\tenketsuProxy");
209
BOOL	fProxySuccess;
210
DWORD	dwProxyMode	= PIPE_READMODE_MESSAGE;
211
ULONG   bytesProxyWritten;
212
static BOOL	fDontAttemptProxyWrite = false;
213

214
HANDLE hProxyPipe = CreateFile(lpszProxyPipename,
215
								GENERIC_READ | GENERIC_WRITE,
216
								0,
217
								NULL,
218
								OPEN_EXISTING,
219
								0,
220
								NULL);
221

222
			if (hProxyPipe == INVALID_HANDLE_VALUE)
223
					dprintf("hProxyPipe == invalid handle\n");
224
			else
225
					dprintf("hProxyPipe == good\n");
226
			SetNamedPipeHandleState(hProxyPipe, &dwProxyMode, NULL, NULL); //?
227

228
#endif
229

230
	while (1) {
231
		fSuccess = ReadFile(	hPipe,
232
								buf,
233
								BUFSIZE*sizeof(TCHAR),
234
								&bytesRead,
235
								NULL);
236
		if (!fSuccess || bytesRead == 0) {
237
			dprintf("[Byakugan] ReadFile failed, or read 0 bytes.\n");
238
			continue;
239
		}
240
#if 0 //#ifdef THREEDHEAPFU_ENABLED
241
		//dprintf("jc: receieved an event of size %d. Forwarding on to ProxyPipe\n", bytesRead);
242
		//WriteFile(hPipe, &freeinfo, sizeof(struct FreeStruct), &bytesWritten, NULL);
243

244
		if (!fDontAttemptProxyWrite)
245
		{
246
			fProxySuccess = WriteFile(hProxyPipe, buf, bytesRead, &bytesProxyWritten, NULL); 
247
			if (bytesRead != bytesProxyWritten)
248
			{
249
				dprintf("Partial write to proxy on last event! ;(\n");
250
				dprintf("event size was %d. wrote %d\n", bytesRead, bytesProxyWritten);
251
				dprintf("Disabling message proxying until explicitly enabled.\n");
252
				fDontAttemptProxyWrite = true;
253
			}
254
		}
255

256
#endif
257
		switch ( *((BYTE *) buf) ) {
258
			case ALLOCATESTRUCT:
259
				//dprintf("[T] New Chunk @ 0x%08x\n", aStruct->ret);
260
				//dprintf("Heap: 0x%08x\tFlags: 0x%08x\tSize: 0x%08x\n\n", 
261
				//		aStruct->heapHandle, aStruct->flags, aStruct->size);
262
				if (heapModel.state & MODEL) heapAllocate(&heapModel, aStruct);
263
				if (heapModel.state & LOG) logAllocate(&heapModel, aStruct);
264
				break;
265

266
			case REALLOCATESTRUCT:
267
				//dprintf("[T] Realloc'd Chunk @ 0x%08x\n", rStruct->ret);
268
				//dprintf("Heap: 0x%08x\tFlags: 0x%08x\tSize: 0x%08x\n", 
269
				//		rStruct->heapHandle, rStruct->flags, rStruct->size);
270
				//if (rStruct->ret !=  (PVOID) rStruct->memoryPointer)
271
				//	dprintf("Replaces chunk @ 0x%08x\n", rStruct->memoryPointer);
272
				//dprintf("\n");
273
				if (heapModel.state & MODEL) heapReallocate(&heapModel, rStruct);
274
				if (heapModel.state & LOG) logReallocate(&heapModel, rStruct);
275
				break;
276

277
			case FREESTRUCT:
278
				//dprintf("[T] Free'd Chunk @ 0x%08x\n", fStruct->memoryPointer);
279
				//dprintf("Heap: 0x%08x\tFlags: 0x%08x\n\n", fStruct->heapHandle, fStruct->flags);
280
				if (heapModel.state & MODEL) heapFree(&heapModel, fStruct);
281
				if (heapModel.state & LOG) logFree(&heapModel, fStruct);
282
				break;
283

284
			case CREATESTRUCT:
285
				dprintf("[T] New Heap: 0x%08x\n", cStruct->ret);
286
				dprintf("Base: 0x%08x\tReserve: 0x%08x\tFlags: 0x%08x\n",
287
						cStruct->base, cStruct->reserve, cStruct->flags);
288
				//dprintf("Commit: 0x%08x\tLock: 0x%08x\n\n", cStruct->commit, cStruct->lock);
289
				if (heapModel.state & MODEL) heapCreate(&heapModel, cStruct);
290
				break;
291

292
			case DESTROYSTRUCT:
293
				dprintf("[T] Heap Destroyed: 0x%08x\n\n", dStruct->heapHandle);
294
				if (heapModel.state & MODEL) heapDestroy(&heapModel, dStruct);
295
				break;
296
			
297
			case COALESCESTRUCT:
298
				//dprintf("[T] Free Block Consolidation (returned 0x%08x)\n", cfbStruct->ret);
299
				//dprintf("Heap: 0x%08x\tArg2: 0x%08x\tArg3: 0x%08x\tArg4: 0x%08x\n\n",
300
				//		cfbStruct->heapHandle, cfbStruct->arg2, cfbStruct->arg3, cfbStruct->arg4);
301
				if (heapModel.state & MODEL) heapCoalesce(&heapModel, cfbStruct);
302
				break;
303

304
			default:
305
				dprintf("[Byakugan] Tenketsu: Unrecognized data was returned.\n");
306
		}
307

308
	}
309

310

311
	return (0);
312
}
313

314
void tenkHelp() {
315
	dprintf(HELPSTRING);
316
	dprintf("Tenketsu Commands:\n");
317
	dprintf("\tmodel\t- Load tenketsu heap visualization libraries and begin modeling\n");
318
	dprintf("\tlog\t- Load tenketsu heap visualization libraries and begin logging\n");
319
	dprintf("\tlistHeaps\t- List all currently tracked heaps and their information\n");
320
	dprintf("\tlistChunks <heap base>\t- List all chunks associated with a givend heap\n");
321
	dprintf("\tvalidate <heap base> - check the chunk chain and find corrupted chunk headers\n");
322
}
323

324
void tenkListHeaps() {
325
	struct HPool	*curHeap;
326
	ULONG			i;
327

328
	dprintf("[T] Currently tracking %d heaps:\n", heapModel.numHeaps);
329
	for (i = 0; i < heapModel.numHeaps; i++) {
330
		curHeap = &(heapModel.heaps[i]);
331
		if (curHeap->inUse == FALSE)
332
			continue;
333
		dprintf("\tBase: 0x%08x\tNumber of Chunks: %d\n", curHeap->base, curHeap->numChunks);
334
		dprintf("\tFlags: 0x%08x\tReserve: 0x%08x\n", 
335
				curHeap->flags, curHeap->reserve, curHeap->commit);
336
		dprintf("\tLock: 0x%08x\n\n", curHeap->lock);
337
	}
338
}
339

340
void tenkValidate(PVOID heapHandle) {
341
	struct HPool            *heap;
342
    struct DestroyStruct    dStruct;
343
    struct HeapChunk        *curChunk;
344
	ULONG					chunkPtr;
345
    ULONG                   i, nextIndex;
346
	BOOL					screwed = FALSE;
347

348
	heap = getHeap(&heapModel, heapHandle);
349

350
	i = heap->inUseHead;
351
	while (i != NULLNODE) {
352
		if (CHUNK(i).free) {
353
			// CHUNK(i).nextInUse must be equal to the next ptr
354
			if(!ReadMemory((ULONG64)(CHUNK(i).addr)+4, (PVOID) &chunkPtr, 4, NULL)) {
355
				dprintf("[T] Unable to read memory at address 0x%08x\n!");
356
				return;
357
			}
358

359
			// Find next free chunk - continue if there are no more
360
			nextIndex = CHUNK(i).nextInUse;
361
			while (nextIndex != NULLNODE && !(CHUNK(nextIndex).free))
362
				nextIndex = CHUNK(nextIndex).nextInUse;
363
			if (nextIndex == NULLNODE) {
364
				i = CHUNK(i).nextInUse;
365
				continue;
366
			}
367

368
			// Validate next free chunk
369
			if (CHUNK(nextIndex).addr != (PVOID) chunkPtr) {
370
				dprintf("[T] Corruped next pointer for chunk at 0x%08x\n", CHUNK(i).addr);
371
				dprintf(">\tGot: 0x%08x\tExpected: 0x%08x\n", chunkPtr, CHUNK(nextIndex).addr);
372
				screwed = TRUE;
373
			}
374
			
375
			// next free chunk prev, must equal CHUNK(i).addr
376
			if(!ReadMemory((ULONG64)CHUNK(nextIndex).addr, (PVOID) &chunkPtr, 4, NULL)) {
377
                dprintf("[T] Unable to read memory at address 0x%08x\n!");
378
                return; 
379
            }
380
			if ((PVOID) chunkPtr != CHUNK(i).addr) {
381
                dprintf("[T] Corruped prev pointer for chunk at 0x%08x\n", CHUNK(nextIndex).addr);
382
                dprintf(">\tGot: 0x%08x\tExpected: 0x%08x\n", chunkPtr, CHUNK(i).addr);
383
				screwed = TRUE;
384
			}
385
		
386
		
387
		} else {
388
		}
389
		i = CHUNK(i).nextInUse;
390
	}
391
	
392
	dprintf("[T] Validation complete: ");
393
	if (!screwed)
394
		dprintf("all known free chunks are correct\n");
395
	else
396
		dprintf("errors found\n");
397
}
398

399
void tenkListChunks(PVOID heapHandle) {
400
	struct HPool			*heap;
401
	struct DestroyStruct	dStruct;
402
	struct HeapChunk		*curChunk;
403
	ULONG					i;
404

405
	heap = getHeap(&heapModel, heapHandle);
406
	dprintf("[T] Currently tracking %d chunks for heap 0x%08x\n", 
407
			heap->numChunks, heap->base);
408
	
409
	i = heap->inUseHead;
410
	while (i != NULLNODE) {
411
		if (CHUNK(i).inUse) {
412
			dprintf("\tAddress: 0x%08x\tSize: 0x%08x", CHUNK(i).addr, CHUNK(i).size);
413
			dprintf("\tFlags: 0x%08x\t%s\n\n", CHUNK(i).flags, 
414
					(CHUNK(i).free)?"FREE'D":"IN USE");
415
		}
416
		i = CHUNK(i).nextInUse;
417
	}
418

419
	if (heap->numChunks == 0) {
420
		dStruct.heapHandle	= heap->base;
421
		heapDestroy(&heapModel, &dStruct);
422
	}
423
}
424

425