.NET Internals and Code Injection

// rbcoree.cpp : Defines the entry point for the DLL application.
//

#include "stdafx.h"
#include <CommCtrl.h>
#include <CommDlg.h>
#include <tlhelp32.h> 
#include <tchar.h>
#include <CorHdr.h>
#include "corinfo.h"
#include "corjit.h"
#include "RebelDotNET.h"
#include "resource.h"

#ifndef PAGE_SIZE 
#define PAGE_SIZE 0x1000
#endif

#define IS_FLAG(Value, Flag) ((Value & Flag) == Flag)

HINSTANCE hInstance;

extern "C" __declspec(dllexport) void HookJIT();

VOID ListThread();


BOOL APIENTRY DllMain( HMODULE hModule,
                       DWORD  dwReason,
                       LPVOID lpReserved
					 )
{
	hInstance = (HINSTANCE) hModule;

	HookJIT();

	if (dwReason == DLL_PROCESS_ATTACH)
	{
		CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE) ListThread, 
			NULL, 0, NULL);
	}

	return TRUE;
}

// unimportant
extern "C" __declspec(dllexport) int __stdcall _CorExeMain(void)
{
	return 0;
}


//
// Hook JIT's compileMethod
//

BOOL bHooked = FALSE;

ULONG_PTR *(__stdcall *p_getJit)();
typedef int (__stdcall *compileMethod_def)(ULONG_PTR classthis, ICorJitInfo *comp, 
										   CORINFO_METHOD_INFO *info, unsigned flags,         
										   BYTE **nativeEntry, ULONG  *nativeSizeOfCode);
struct JIT
{
	compileMethod_def compileMethod;
};

compileMethod_def compileMethod;

int __stdcall my_compileMethod(ULONG_PTR classthis, ICorJitInfo *comp,
							   CORINFO_METHOD_INFO *info, unsigned flags,         
							   BYTE **nativeEntry, ULONG  *nativeSizeOfCode);

extern "C" __declspec(dllexport) void HookJIT()
{
	if (bHooked) return;

	LoadLibrary(_T("mscoree.dll"));

	HMODULE hJitMod = LoadLibrary(_T("mscorjit.dll"));

	if (!hJitMod)
		return;

	p_getJit = (ULONG_PTR *(__stdcall *)()) GetProcAddress(hJitMod, "getJit");

	if (p_getJit)
	{
		JIT *pJit = (JIT *) *((ULONG_PTR *) p_getJit());

		if (pJit)
		{
			DWORD OldProtect;
			VirtualProtect(pJit, sizeof (ULONG_PTR), PAGE_READWRITE, &OldProtect);
			compileMethod =  pJit->compileMethod;
			pJit->compileMethod = &my_compileMethod;
			VirtualProtect(pJit, sizeof (ULONG_PTR), OldProtect, &OldProtect);
			bHooked = TRUE;
		}
	}
}

//
// Logging
//

struct AssemblyInfo
{
	CORINFO_MODULE_HANDLE hCorModule;

	WCHAR AssemblyName[MAX_PATH];

	VOID *ImgBase;
	UINT ImgSize;

	BOOL bIdentified;

	HANDLE hRebReport;

	BOOL bDump;
	TCHAR DumpFileName[MAX_PATH];

} LoggedAssemblies[100];

UINT NumberOfLoggedAssemblies = 0;

VOID LogAssembly(ICorJitInfo *comp, CORINFO_METHOD_INFO *info);

DWORD GetTokenFromMethodHandle(ICorJitInfo *comp, CORINFO_METHOD_INFO *info);

VOID AddMethod(CORINFO_METHOD_INFO *mi);
BOOL CreateRebFile(AssemblyInfo *ai);
//
// hooked compileMethod
//
/*__declspec (naked) */
int __stdcall my_compileMethod(ULONG_PTR classthis, ICorJitInfo *comp,
							   CORINFO_METHOD_INFO *info, unsigned flags,         
							   BYTE **nativeEntry, ULONG  *nativeSizeOfCode)
{
	// in case somebody hooks us (x86 only)
#ifdef _M_IX86
	__asm 
	{
		nop
		nop
		nop
		nop
		nop
		nop
		nop
		nop
		nop
		nop
		nop
		nop
		nop
		nop
	}
#endif

	//
	// check if it's the dump process
	//

	if (comp == NULL)
	{
		AddMethod(info);
		return 0;
	}

	LogAssembly(comp, info);

	// call original method
	// I'm not using the naked + jmp approach to avoid x64 incompatibilities
	int nRet = compileMethod(classthis, comp, info, flags, nativeEntry, nativeSizeOfCode);
	
	return nRet;
}

//
// convert an address to its module ImgBase and Name (if possible)
//

VOID AddressToModuleInfo(VOID *pAddress, WCHAR *AssemblyName,
						 VOID **pImgBase, UINT *pImgSize,
						 BOOL *pbIdentified) 
{ 
	DWORD dwPID = GetCurrentProcessId();
	HANDLE hModuleSnap = INVALID_HANDLE_VALUE; 
	MODULEENTRY32 me32; 

	static BOOL bFirstUnkAsm = TRUE;

	hModuleSnap = CreateToolhelp32Snapshot( TH32CS_SNAPMODULE, dwPID ); 

	if (hModuleSnap == INVALID_HANDLE_VALUE) 
		return;

	me32.dwSize = sizeof (MODULEENTRY32); 

	if (!Module32First(hModuleSnap, &me32 )) 
	{ 
		CloseHandle(hModuleSnap);    
		return; 
	} 

	do 
	{
		if (((ULONG_PTR) pAddress) > ((ULONG_PTR) me32.modBaseAddr) &&
			((ULONG_PTR) pAddress) < (((ULONG_PTR) me32.modBaseAddr) + 
			me32.modBaseSize))
		{
			if (pImgBase) *pImgBase = (VOID *) me32.modBaseAddr;
			if (pImgSize) *pImgSize = me32.modBaseSize;
			wcscpy_s(AssemblyName, MAX_PATH, me32.szExePath);
			if (pbIdentified) *pbIdentified = TRUE;
			return;
		}

	} while (Module32Next(hModuleSnap, &me32)); 

	CloseHandle(hModuleSnap); 

	if (pbIdentified) *pbIdentified = FALSE;

	MEMORY_BASIC_INFORMATION mbi = { 0 };
	VirtualQuery(pAddress, &mbi, sizeof (MEMORY_BASIC_INFORMATION));

	if (pImgBase) *pImgBase = mbi.AllocationBase;

	DWORD ImgSize = 0;

	__try
	{
		IMAGE_DOS_HEADER *pDosHeader = (IMAGE_DOS_HEADER *)
			mbi.AllocationBase;

		if (pDosHeader->e_magic == IMAGE_DOS_SIGNATURE)
		{
			IMAGE_NT_HEADERS *pNtHeaders = (IMAGE_NT_HEADERS *)
				(pDosHeader->e_lfanew + (ULONG_PTR) pDosHeader);

			if (pNtHeaders->Signature == IMAGE_NT_SIGNATURE)
			{
				ImgSize = pNtHeaders->OptionalHeader.SizeOfImage;
			}
		}
	}
	__except (EXCEPTION_EXECUTE_HANDLER)
	{
		goto endinfo;
	}

endinfo:

	if (pImgSize) *pImgSize = ImgSize;

	if (bFirstUnkAsm)
	{
		wsprintfW(AssemblyName, L"Base: %p - Size: %08X - Primary Assembly", 
			mbi.AllocationBase, ImgSize);
		bFirstUnkAsm = FALSE;
	}
	else
	{
		wsprintfW(AssemblyName, L"Base: %p - Size: %08X - unidentfied", 
			mbi.AllocationBase, ImgSize);
	}
	
} 

VOID LogAssembly(ICorJitInfo *comp, CORINFO_METHOD_INFO *info)
{
	// already in the list?
	for (UINT x = 0; x < NumberOfLoggedAssemblies; x++)
	{
		if (LoggedAssemblies[x].hCorModule == info->scope)
			return;
	}

	//
	// Add assembly to the logged list
	//

	AddressToModuleInfo(info->ILCode, 
		LoggedAssemblies[NumberOfLoggedAssemblies].AssemblyName,
		&LoggedAssemblies[NumberOfLoggedAssemblies].ImgBase,
		&LoggedAssemblies[NumberOfLoggedAssemblies].ImgSize,
		&LoggedAssemblies[NumberOfLoggedAssemblies].bIdentified);

	LoggedAssemblies[NumberOfLoggedAssemblies].hCorModule = info->scope;
	LoggedAssemblies[NumberOfLoggedAssemblies].bDump = FALSE;

	NumberOfLoggedAssemblies++;
}

//
// Listing
//

LRESULT CALLBACK ListDlgProc(HWND hDlg, UINT uMsg, WPARAM wParam, LPARAM lParam);

VOID ListThread()
{
	Sleep(2000);

	InitCommonControls();

	DialogBox(hInstance, MAKEINTRESOURCE(IDD_ASMLIST), NULL, (DLGPROC) ListDlgProc);
}

LRESULT CALLBACK ListDlgProc(HWND hDlg, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
	switch (uMsg)
	{

	case WM_INITDIALOG:
		{
			HWND hList = GetDlgItem(hDlg, LST_ASMS);

			LV_COLUMN lvc;

			ZeroMemory(&lvc, sizeof (LV_COLUMN));

			lvc.mask = LVCF_FMT | LVCF_WIDTH | LVCF_TEXT | LVCF_SUBITEM;
			lvc.fmt = LVCFMT_LEFT;

			lvc.cx = 500;
			lvc.pszText = _T("Assembly Path");
			ListView_InsertColumn(hList, 0, &lvc);

			SendMessage(hList, LVM_SETEXTENDEDLISTVIEWSTYLE,
				LVS_EX_FULLROWSELECT | LVS_EX_INFOTIP,
				LVS_EX_FULLROWSELECT | LVS_EX_INFOTIP);


			SendMessage(hDlg, WM_COMMAND, IDC_REFRESH, 0);

			break;
		}

	case WM_CLOSE:
		{
			EndDialog(hDlg, 0);
			break;
		}

	case WM_COMMAND:
		{
			switch (LOWORD(wParam))
			{

			case IDC_REFRESH:
				{
					HWND hList = GetDlgItem(hDlg, LST_ASMS);

					ListView_DeleteAllItems(hList);

					LV_ITEM lvi;

					ZeroMemory(&lvi, sizeof (LV_ITEM));

					lvi.mask = LVIF_TEXT | LVIF_STATE | LVIF_PARAM;

					for (UINT x = 0; x < NumberOfLoggedAssemblies; x++)
					{
						lvi.lParam = (LPARAM) LoggedAssemblies[x].hCorModule;
						lvi.pszText = LoggedAssemblies[x].AssemblyName;
						ListView_InsertItem(hList, &lvi);
					}

					break;
				}

			case IDC_DUMPASM:
				{
					HWND hList = GetDlgItem(hDlg, LST_ASMS);

					int nSel = ListView_GetNextItem(hList, -1, LVNI_SELECTED);

					if (nSel == -1) break;

					OPENFILENAME SaveFileName;

					TCHAR DumpFileName[MAX_PATH];

					ZeroMemory(DumpFileName, MAX_PATH * sizeof (TCHAR));

					ZeroMemory(&SaveFileName, sizeof (OPENFILENAME));

					SaveFileName.lStructSize = sizeof (OPENFILENAME);
					SaveFileName.hwndOwner = hDlg;
					SaveFileName.lpstrFilter = _T("All Files (*.*)\0*.*\0");
					SaveFileName.lpstrFile = DumpFileName;
					SaveFileName.nMaxFile = MAX_PATH;
					SaveFileName.lpstrTitle = _T("Save Assembly As...");

					if (!GetSaveFileName(&SaveFileName)) 
						break;

					LV_ITEM lvi;
					ZeroMemory(&lvi, sizeof (LV_ITEM));

					lvi.mask = LVIF_PARAM;
					lvi.iItem = nSel;

					ListView_GetItem(hList, &lvi);

					for (UINT x = 0; x < NumberOfLoggedAssemblies; x++)
					{
						if (LoggedAssemblies[x].hCorModule == 
							(CORINFO_MODULE_HANDLE) lvi.lParam)
						{
							HANDLE hFile = CreateFile(DumpFileName, GENERIC_WRITE,
								FILE_SHARE_READ, NULL, CREATE_ALWAYS, 0, NULL);

							if (hFile == INVALID_HANDLE_VALUE)
								break;

							DWORD dwOldProtect;

							VirtualProtect(LoggedAssemblies[x].ImgBase, 
								LoggedAssemblies[x].ImgSize, PAGE_EXECUTE_READ,
								&dwOldProtect);

							for (UINT nPage = 0; 
								nPage < (LoggedAssemblies[x].ImgSize / PAGE_SIZE);
								nPage++)
							{
								DWORD BW;

								__try
								{
									VOID *pPage = (VOID *) ((nPage * PAGE_SIZE) + 
										(ULONG_PTR) LoggedAssemblies[x].ImgBase);

									WriteFile(hFile, pPage, PAGE_SIZE, &BW, NULL);
								}

								__except (EXCEPTION_EXECUTE_HANDLER)
								{
									SetFilePointer(hFile, PAGE_SIZE, NULL, FILE_CURRENT);
									SetEndOfFile(hFile);
								}
							}
							
							CloseHandle(hFile);
							
							MessageBox(hDlg, _T("Assembly successfully dumped."), 
								_T("Dumped"), MB_ICONINFORMATION);
						}
					}

					break;
				}

			case IDC_REBFILE:
				{
					HWND hList = GetDlgItem(hDlg, LST_ASMS);

					int nSel = ListView_GetNextItem(hList, -1, LVNI_SELECTED);

					if (nSel == -1) break;

					LV_ITEM lvi;
					ZeroMemory(&lvi, sizeof (LV_ITEM));

					lvi.mask = LVIF_PARAM;
					lvi.iItem = nSel;

					ListView_GetItem(hList, &lvi);

					for (UINT x = 0; x < NumberOfLoggedAssemblies; x++)
					{
						if (LoggedAssemblies[x].hCorModule == 
							(CORINFO_MODULE_HANDLE) lvi.lParam)
						{
							OPENFILENAME OpenFileName;

							TCHAR ReportFileName[MAX_PATH];

							ZeroMemory(ReportFileName, MAX_PATH * sizeof (TCHAR));

							ZeroMemory(&OpenFileName, sizeof (OPENFILENAME));

							OpenFileName.lStructSize = sizeof (OPENFILENAME);
							OpenFileName.hwndOwner = hDlg;
							OpenFileName.lpstrFilter = _T("Report Rebel File (*.rebel)\0*.rebel\0");
							OpenFileName.lpstrFile = ReportFileName;
							OpenFileName.nMaxFile = MAX_PATH;
							OpenFileName.lpstrTitle = _T("Select a Report Rebel File...");

							if (!GetOpenFileName(&OpenFileName)) 
								break;

							LoggedAssemblies[x].hRebReport = CreateFile(ReportFileName,
								GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);

							if (LoggedAssemblies[x].hRebReport == INVALID_HANDLE_VALUE)
								break;

							OPENFILENAME SaveFileName;

							ZeroMemory(LoggedAssemblies[x].DumpFileName, MAX_PATH * sizeof (TCHAR));

							ZeroMemory(&SaveFileName, sizeof (OPENFILENAME));

							SaveFileName.lStructSize = sizeof (OPENFILENAME);
							SaveFileName.hwndOwner = hDlg;
							SaveFileName.lpstrFilter = _T("Rebel File (*.rebel)\0*.rebel\0");
							SaveFileName.lpstrFile = LoggedAssemblies[x].DumpFileName;
							SaveFileName.nMaxFile = MAX_PATH;
							SaveFileName.lpstrTitle = _T("Save Rebel File As...");
							SaveFileName.lpstrDefExt = _T("rebel");

							if (!GetSaveFileName(&SaveFileName)) 
							{
								CloseHandle(LoggedAssemblies[x].hRebReport);
								break;
							}

							// dump
							CreateRebFile(&LoggedAssemblies[x]);

							break;
						}
					}

					break;
				}
			}

			break;
		}
	}

	return FALSE;
}

//
// Dumping
//

DWORD RvaToOffset(VOID *pBase, DWORD Rva)
{
	__try
	{
		DWORD Offset = Rva, Limit;

		IMAGE_DOS_HEADER *pDosHeader = (IMAGE_DOS_HEADER *) pBase;

		if (pDosHeader->e_magic != IMAGE_DOS_SIGNATURE)
			return 0;

		IMAGE_NT_HEADERS *pNtHeaders = (IMAGE_NT_HEADERS *) (
			pDosHeader->e_lfanew + (ULONG_PTR) pDosHeader);

		if (pNtHeaders->Signature != IMAGE_NT_SIGNATURE)
			return 0;

		IMAGE_SECTION_HEADER *Img = IMAGE_FIRST_SECTION(pNtHeaders);

		if (Rva < Img->PointerToRawData)
			return Rva;

		for (WORD i = 0; i < pNtHeaders->FileHeader.NumberOfSections; i++)
		{
			if (Img[i].SizeOfRawData)
				Limit = Img[i].SizeOfRawData;
			else
				Limit = Img[i].Misc.VirtualSize;

			if (Rva >= Img[i].VirtualAddress &&
				Rva < (Img[i].VirtualAddress + Limit))
			{
				if (Img[i].PointerToRawData != 0)
				{
					Offset -= Img[i].VirtualAddress;
					Offset += Img[i].PointerToRawData;
				}

				return Offset;
			}
		}
	}
	__except (EXCEPTION_EXECUTE_HANDLER)
	{
		return 0;
	}

	return 0;
}

UINT GetMethodSize(REBEL_METHOD *rbMethod)
{
	UINT nMethodSize = sizeof (REBEL_METHOD);

	if (!IS_FLAG(rbMethod->Mask, REBEL_METHOD_MASK_NAMEOFFSET))
		nMethodSize += rbMethod->NameOffsetOrSize;

	if (!IS_FLAG(rbMethod->Mask, REBEL_METHOD_MASK_SIGOFFSET))
		nMethodSize += rbMethod->SignatureOffsetOrSize;

	if (!IS_FLAG(rbMethod->Mask, REBEL_METHOD_MASK_LOCVARSIGOFFSET))
		nMethodSize += rbMethod->LocalVarSigOffsetOrSize;

	nMethodSize += rbMethod->CodeSize;
	nMethodSize += rbMethod->ExtraSectionsSize;

	return nMethodSize;
}

static HANDLE hRebuildDump = INVALID_HANDLE_VALUE;

VOID AddMethod(CORINFO_METHOD_INFO *mi)
{
	REBEL_METHOD rbMethod;

	ZeroMemory(&rbMethod, sizeof (REBEL_METHOD));

	rbMethod.Token = mi->locals.token;
	rbMethod.CodeSize = mi->ILCodeSize;
	
	DWORD BRW;

	SetFilePointer(hRebuildDump, 0, NULL, FILE_END);

	WriteFile(hRebuildDump, &rbMethod, sizeof (REBEL_METHOD), &BRW, NULL);
	WriteFile(hRebuildDump, mi->ILCode, mi->ILCodeSize, &BRW, NULL);

	//
	// Increase number of methods
	//

	SetFilePointer(hRebuildDump, 0, NULL, FILE_BEGIN);

	REBEL_NET_BASE rbBase;

	ReadFile(hRebuildDump, &rbBase, sizeof (REBEL_NET_BASE), &BRW, NULL);

	rbBase.NumberOfMethods++;

	SetFilePointer(hRebuildDump, 0, NULL, FILE_BEGIN);

	WriteFile(hRebuildDump, &rbBase, sizeof (REBEL_NET_BASE), &BRW, NULL);
}

BOOL CreateRebFile(AssemblyInfo *ai)
{
	DWORD BRW;

	hRebuildDump = CreateFile(ai->DumpFileName, GENERIC_READ | 
		GENERIC_WRITE, FILE_SHARE_READ, NULL,
		CREATE_ALWAYS, 0, NULL);

	if (hRebuildDump == INVALID_HANDLE_VALUE)
		return FALSE;

	REBEL_NET_BASE rbBase;

	ZeroMemory(&rbBase, sizeof (REBEL_NET_BASE));

	rbBase.Signature = REBEL_NET_SIGNATURE;
	rbBase.MethodsOffset = sizeof (REBEL_NET_BASE);

	WriteFile(hRebuildDump, &rbBase, sizeof (REBEL_NET_BASE), &BRW, NULL);

	//
	// Get the new JIT compileMethod which by now should be
	// hooked by the code protection
	//

	HMODULE hJitMod = LoadLibrary(_T("mscorjit.dll"));

	p_getJit = (ULONG_PTR *(__stdcall *)()) GetProcAddress(hJitMod, "getJit");

	if (!p_getJit)
	{
		CloseHandle(hRebuildDump);
		hRebuildDump = INVALID_HANDLE_VALUE;
		return FALSE;
	}

	JIT *pJit = (JIT *) *((ULONG_PTR *) p_getJit());

	//

	REBEL_NET_BASE repBase;

	if (!ReadFile(ai->hRebReport, &repBase, sizeof (REBEL_NET_BASE), &BRW, NULL))
	{
		CloseHandle(hRebuildDump);
		hRebuildDump = INVALID_HANDLE_VALUE;
		return FALSE;
	}

	UINT CurMethodOffset = repBase.MethodsOffset;

	for (UINT x = 0; x < repBase.NumberOfMethods; x++)
	{
		SetFilePointer(ai->hRebReport, CurMethodOffset, NULL, FILE_BEGIN);

		REBEL_METHOD rbRepMethod;

		ReadFile(ai->hRebReport, &rbRepMethod, sizeof (REBEL_METHOD), &BRW, NULL);

		//
		// Calculate current method's code location
		// Use RvaToOffset only when the module wasn't mapped
		//

		BYTE *pMethodCode;

		if (ai->bIdentified)
		{
			pMethodCode = (BYTE *) (rbRepMethod.RVA + (ULONG_PTR) ai->ImgBase);
		}
		else
		{
			DWORD Offset = RvaToOffset(ai->ImgBase, rbRepMethod.RVA);

			if (Offset == 0) continue;

			pMethodCode = (BYTE *) (Offset + (ULONG_PTR) ai->ImgBase);
		}

		//
		// we should check the validity of the memory
		// pointer by pMethodCode
		//

		// TODO

		//
		// skips the method header
		//

		BYTE HeaderFormat = *pMethodCode;

		HeaderFormat &= 3;

		if (HeaderFormat == 2)			// Tiny = 2 (CorILMethod_TinyFormat)
			pMethodCode++;
		else							// Fat = 3 (CorILMethod_FatFormat)
			pMethodCode += (sizeof (DWORD) * 3);

		//
		// Do the fake compileMethod request
		//

		CORINFO_METHOD_INFO mi = { 0 };

		mi.ILCode = pMethodCode;
		mi.ILCodeSize = rbRepMethod.CodeSize;
		mi.scope = ai->hCorModule;
		// use this to pass the token to our AddMethod
		mi.locals.token = rbRepMethod.Token;

		pJit->compileMethod((ULONG_PTR) pJit, NULL, &mi, 0, NULL, NULL);

		//
		// next method
		//

		CurMethodOffset += GetMethodSize(&rbRepMethod);
	}

	//
	// Close file and notify the user
	//

	CloseHandle(hRebuildDump);
	hRebuildDump = INVALID_HANDLE_VALUE;

	MessageBox(0, _T("Assembly code successfully dumped."), _T("JIT Dumper"),
		MB_ICONINFORMATION);

	return TRUE;
}