src/win/patch.c

#include <stdint.h>
#include <string.h>
#include <stdbool.h>
#include <windows.h>
#include <psapi.h>
#include "dbg.h"
#include "patch.h"

static CRITICAL_SECTION s_cs;

static IMAGE_IMPORT_DESCRIPTOR *
imports(IMAGE_DOS_HEADER *dh)
{
	char *		      base = (char *)dh;
	IMAGE_NT_HEADERS *    nth = (IMAGE_NT_HEADERS *)(base + dh->e_lfanew);
	IMAGE_DATA_DIRECTORY *impdd =
	    nth->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_IMPORT;
	ASSERT(dh->e_lfanew);
	ASSERT(nth);
	ASSERT(impdd);
	return impdd && impdd->VirtualAddress ?
		  (IMAGE_IMPORT_DESCRIPTOR *)(base + impdd->VirtualAddress) :
		  0;
}

static IMAGE_THUNK_DATA *
lookup2(char *base, IMAGE_THUNK_DATA *td, IMAGE_THUNK_DATA *otd, const char *nm)
{
	while (otd->u1.AddressOfData) {
		IMAGE_IMPORT_BY_NAME *name =
		    (IMAGE_IMPORT_BY_NAME *)(otd->u1.AddressOfData + base);
		if (otd->u1.Ordinal & IMAGE_ORDINAL_FLAG)
			dbg("   ordinal1\n");
		else {
			dbg("  name: %s %p\n", name->Name, td->u1.Function);
			if (0 == strcmp((char *)name->Name, nm))
				return td;
		}
		otd++;
		td++;
	}
	return 0;
}

static IMAGE_THUNK_DATA *
lookup(IMAGE_DOS_HEADER *dh, const char *nm)
{
	char *			 base = (char *)dh;
	IMAGE_IMPORT_DESCRIPTOR *id = imports(dh);
	if (!id)
		return 0;
	while (id->Name) {
		if (id->FirstThunk && id->OriginalFirstThunk) {
			IMAGE_THUNK_DATA *d = lookup2(base,
			    (IMAGE_THUNK_DATA *)(id->FirstThunk + base),
			    (IMAGE_THUNK_DATA *)(id->OriginalFirstThunk + base),
			    nm);
			dbg(" import %s\n", id->Name + base);
			if (d)
				return d;
		}
		id++;
	}
	return 0;
}

static void
modpatch(
    IMAGE_DOS_HEADER *dh, void *orig, void *hook, void **preal, const char *nm)
{
	IMAGE_THUNK_DATA *td = lookup(dh, nm);
	if (td && orig == (void *)td->u1.Function) {
		DWORD o;
		dbg("   patching %s %p %p -> %p\n", nm, td->u1.Function, orig,
		    hook);
		*preal = (void *)td->u1.Function;
		CHK(VirtualProtect(
		    td, sizeof(*td), PAGE_EXECUTE_READWRITE, &o));
		td->u1.Function = (uintptr_t)hook;
		CHK(VirtualProtect(td, sizeof(*td), o, &o));
	}
}

void
patchInstall(void *orig, void *hook, void **preal, const char *nm)
{
	HMODULE			mod[4096];
	DWORD			n;
	DWORD			i;
	extern IMAGE_DOS_HEADER __ImageBase;
	CHK(EnumProcessModules(GetCurrentProcess(), mod, sizeof(mod), &n));
	n /= sizeof(HMODULE);
	dbg("orig %s %p\n", nm, orig);
	for (i = 0; i < n; i++) {
		HMODULE m = mod[i];
		char	mname[4096];
		CHK(GetModuleFileNameA(m, mname, sizeof(mname)));
		dbg("module %s\n", mname);
		if (m != (HMODULE)&__ImageBase)
			modpatch((IMAGE_DOS_HEADER *)m, orig, hook, preal, nm);
	}
	dbg("modules patched\n");
}

bool
patchInstalled(DWORD pid)
{
	static uint8_t pids[0x1000] = { 0 };
	unsigned       index = pid >> 2;
	unsigned       byte = (index >> 3) & 0xfff;
	unsigned       mask = 1 << (index & 7);
	bool	       ret;
	EnterCriticalSection(&s_cs);
	ret = (pids[byte] & mask) != 0;
	pids[byte] |= mask;
	LeaveCriticalSection(&s_cs);
	return ret;
}

void
patchInit()
{
	InitializeCriticalSection(&s_cs);
}

void
patchTerm()
{
	DeleteCriticalSection(&s_cs);
}