AltSnap/thunks.c at main · RamonUnch/AltSnap · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

#include <windows.h>

enum {
	TOT_THUNK_SIZE=4096,
#if defined(_M_AMD64)
	THUNK_SIZE=22,
#elif defined(_M_IX86)
	THUNK_SIZE=14, // 13 aligned to USHORT
#else
	#error Unsupported processor type, determine THUNK_SIZE here...
#endif
	MAX_THUNKS=TOT_THUNK_SIZE/THUNK_SIZE,
};

static BYTE *thunks_array = NULL;
static WORD numthunks = 0;
static WORD free_idx = 0; // Free index mem;

static void *AllocThunk(LONG_PTR tProc, void *vParam)
{
	if (!thunks_array) {
		thunks_array = (BYTE*)VirtualAlloc( NULL, TOT_THUNK_SIZE, MEM_COMMIT, PAGE_EXECUTE_READWRITE );
		if (!thunks_array)
			return NULL;

		free_idx = 0;
		for (WORD i=0; i < MAX_THUNKS; i++)
			*(WORD*)&thunks_array[i*THUNK_SIZE] = i+1;
	}

	if (numthunks >= MAX_THUNKS) {
		return NULL;
	}
	++numthunks;

	// get the free thunk
	BYTE *thunk = &thunks_array[free_idx * THUNK_SIZE];
	free_idx = *(WORD*)thunk; // Pick up the free index

	DWORD oldprotect;
	VirtualProtect(thunk, THUNK_SIZE, PAGE_EXECUTE_READWRITE, &oldprotect);
	// x86
	//   | mov dword ptr [esp+4] this
	//   | jmp MainProc
	// AMD64
	//   | mov rcx this
	//   | mov rax MainProc
	//   | jmp rax
	//
	#if defined(_M_AMD64)
		*(WORD*)   (thunk + 0) = 0xb948;
		*(void**)  (thunk + 2) = vParam;
		*(WORD*)   (thunk +10) = 0xb848;
		*(void**)  (thunk +12) = tProc;
		*(WORD*)   (thunk +20) = 0xe0ff;
	#elif defined(_M_IX86)
		*(DWORD*)  (thunk +0) = 0x042444C7;
		*(void**)  (thunk +4) = vParam;
		*(BYTE*)   (thunk +8) = 0xE9;
		*(DWORD *) (thunk+9)  = (BYTE*)((BYTE*)tProc)-(thunk+13);
	#else
		#error Unsupported processor type, please implement assembly code...
	#endif

	// Make thuk execute only for safety.
	VirtualProtect(thunk, THUNK_SIZE, PAGE_EXECUTE, &oldprotect);
	FlushInstructionCache( GetCurrentProcess(), thunk, THUNK_SIZE );

	return thunk;
}
static void ReleaseThunk(BYTE *thunk)
{
	if (!thunk || thunks_array + TOT_THUNK_SIZE < thunk  || thunk < thunks_array )
		return;

	DWORD oldprotect;
	VirtualProtect(thunks_array, TOT_THUNK_SIZE, PAGE_EXECUTE_READWRITE, &oldprotect);

	*(WORD*)thunk = free_idx; // Set free index in the current thunk location
	free_idx = (WORD)( (thunk - thunks_array) / THUNK_SIZE );
	--numthunks;

	VirtualProtect(thunks_array, TOT_THUNK_SIZE, PAGE_EXECUTE_READ, &oldprotect);

	// Everything is free!
	if (numthunks == 0) {
		VirtualFree( thunks_array, 0, MEM_RELEASE );
		thunks_array = NULL;
	}
}
#ifdef TEST_MAIN
#include <stdio.h>
int main(void)
{
	enum { LEN = 1024 };
	BYTE *t[LEN];
	for (int i =0; i < LEN; i++)
	{
		t[i] = (BYTE*)AllocThunk((LONG_PTR)333, NULL);
		if (t[i])
			printf("t[%d] = %lx\n", i, (DWORD)t[i]);
	}

	for (int i = LEN-1; i >=0; i--)
	{
		size_t prevfreeidx = free_idx;
		ReleaseThunk(t[i]);
		if (t[i])
			printf("released t[%d] = %lx  free_idx=%d / prev=%d\n", i, (DWORD)t[i], (int)free_idx, prevfreeidx);
		if(i%3==0)
			t[i] = (BYTE*)AllocThunk((LONG_PTR)444, NULL);
	}

	return 0;
}
#endif