thunks.cpp

#include "stdafx.h"

__pragma(optimize("", off))
__pragma(runtime_checks("", off))
__pragma(check_stack(off))
__pragma(strict_gs_check(push, off))

// anonymous namespace to ensure that the function names are not exported,
// and hence that I can take their address without suffering indirections
namespace {

__pragma(code_seg(push, "thunks"))

	struct functions {
		template<typename R, typename... A>
		__pragma(runtime_checks("", off))
		static R __declspec(code_seg("thunks")) __stdcall fn_1_stdcall(A... args) {
			std::function<R(A...)>* f(reinterpret_cast<std::function<R(A...)>*>(static_cast<size_t>(0xdeadbeefbaadf00d)));
			union converter_t {
				size_t raw_pointer;
				decltype(&std::function<R(A...)>::operator()) cooked_pointer;
			};
			converter_t converter = { 0xcafebabed15ea5e5 };
			return (f->*(converter.cooked_pointer))(args...);
		}
		__pragma(runtime_checks("", restore))

		template<typename R, typename... A>
		__pragma(runtime_checks("", off))
		static R __declspec(code_seg("thunks")) __fastcall fn_2_fastcall(A... args) {
			std::function<R(A...)>* f(reinterpret_cast<std::function<R(A...)>*>(static_cast<size_t>(0xdeadbeefbaadf00d)));
			union converter_t {
				size_t raw_pointer;
				decltype(&std::function<R(A...)>::operator()) cooked_pointer;
			};
			converter_t converter = { 0xcafebabed15ea5e5 };
			return (f->*(converter.cooked_pointer))(args...);
		}
		__pragma(runtime_checks("", restore))

		template<typename R, typename... A>
		__pragma(runtime_checks("", off))
		static R __declspec(code_seg("thunks")) __cdecl fn_3_cdecl(A... args) {
			std::function<R(A...)>* f(reinterpret_cast<std::function<R(A...)>*>(static_cast<size_t>(0xdeadbeefbaadf00d)));
			union converter_t {
				size_t raw_pointer;
				decltype(&std::function<R(A...)>::operator()) cooked_pointer;
			};
			converter_t converter = { 0xcafebabed15ea5e5 };
			return (f->*(converter.cooked_pointer))(args...);
		}
		__pragma(runtime_checks("", restore))

		template<typename R, typename... A>
		__pragma(runtime_checks("", off))
		static R __declspec(code_seg("thunks")) __vectorcall fn_4_vectorcall(A... args) {
			std::function<R(A...)>* f(reinterpret_cast<std::function<R(A...)>*>(static_cast<size_t>(0xdeadbeefbaadf00d)));
			union converter_t {
				size_t raw_pointer;
				decltype(&std::function<R(A...)>::operator()) cooked_pointer;
			};
			converter_t converter = { 0xcafebabed15ea5e5 };
			return (f->*(converter.cooked_pointer))(args...);
		}
		__pragma(runtime_checks("", restore))
	};
__pragma(code_seg(pop))

}

__pragma(strict_gs_check(pop))
__pragma(check_stack)
__pragma(runtime_checks("", restore))
__pragma(optimize("", on))

template<typename R, typename... A>
struct thunk {
	thunk(std::function<R(A...)> fn) : the_thunk(new thunk_holder(fn)) {
	}

	typedef R   (__stdcall    *stdcall_type   )(A...);
	typedef R   (__fastcall   *fastcall_type  )(A...);
	typedef R   (__cdecl      *cdecl_type     )(A...);
	typedef R   (__vectorcall *vectorcall_type)(A...);
	typedef void(             *void_type      )(    );

	stdcall_type as_stdcall() const {
		return the_thunk->as_stdcall();
	}

	fastcall_type as_fastcall() const {
		return the_thunk->as_fastcall();
	}

	cdecl_type as_cdecl() const {
		return the_thunk->as_cdecl();
	}

	vectorcall_type as_vectorcall() const {
		return the_thunk->as_vectorcall();
	}

private:
	struct thunk_holder {
		thunk_holder(std::function<R(A...)> f_) : f(f_), page(nullptr) {
			generate_thunk();
		}

		~thunk_holder() {
			decltype(&::SetProcessValidCallTargets) spvct = reinterpret_cast<decltype(&::SetProcessValidCallTargets)>(::GetProcAddress(::GetModuleHandleW(L"kernelbase.dll"), "SetProcessValidCallTargets"));
			if (spvct) {
				call_targets[std_call   ].Flags = 0;
				call_targets[fast_call  ].Flags = 0;
				call_targets[c_decl     ].Flags = 0;
				call_targets[vector_call].Flags = 0;

				spvct(::GetCurrentProcess(), page, function_size, calling_conventions_max, call_targets);
			}
			::VirtualFree(page, function_size, MEM_RELEASE);
		}

		stdcall_type as_stdcall() const {
			return reinterpret_cast<stdcall_type>(static_cast<char*>(page) + call_targets[std_call].Offset);
		}

		fastcall_type as_fastcall() const {
			return reinterpret_cast<fastcall_type>(static_cast<char*>(page) + call_targets[fast_call].Offset);
		}

		cdecl_type as_cdecl() const {
			return reinterpret_cast<cdecl_type>(static_cast<char*>(page) + call_targets[c_decl].Offset);
		}

		vectorcall_type as_vectorcall() const {
			return reinterpret_cast<cdecl_type>(static_cast<char*>(page) + call_targets[vector_call].Offset);
		}

	private:
		enum calling_conventions {
			std_call,
			fast_call,
			c_decl,
			vector_call,
			calling_conventions_max
		};
	
		::CFG_CALL_TARGET_INFO call_targets[calling_conventions_max];
	
		void generate_thunk() {
			::MEMORY_BASIC_INFORMATION mbi = { 0 };
			::VirtualQuery(static_cast<stdcall_type>(&functions::fn_1_stdcall<R, A...>), &mbi, sizeof(mbi));
			function_size = mbi.RegionSize;
		
			page = ::VirtualAlloc(nullptr, function_size, MEM_COMMIT, PAGE_READWRITE);
			std::memcpy(page, mbi.BaseAddress, function_size);

			call_targets[std_call   ].Offset = reinterpret_cast<const char*>(static_cast<stdcall_type>(&functions::fn_1_stdcall<R, A...>)) - reinterpret_cast<const char*>(mbi.BaseAddress);
			call_targets[std_call   ].Flags = CFG_CALL_TARGET_VALID;
			call_targets[fast_call  ].Offset = reinterpret_cast<const char*>(static_cast<fastcall_type>(&functions::fn_2_fastcall<R, A...>)) - reinterpret_cast<const char*>(mbi.BaseAddress);
			call_targets[fast_call  ].Flags = CFG_CALL_TARGET_VALID;
			call_targets[c_decl     ].Offset = reinterpret_cast<const char*>(static_cast<cdecl_type>(&functions::fn_3_cdecl<R, A...>)) - reinterpret_cast<const char*>(mbi.BaseAddress);
			call_targets[c_decl     ].Flags = CFG_CALL_TARGET_VALID;
			call_targets[vector_call].Offset = reinterpret_cast<const char*>(static_cast<vectorcall_type>(&functions::fn_4_vectorcall<R, A...>)) - reinterpret_cast<const char*>(mbi.BaseAddress);
			call_targets[vector_call].Flags = CFG_CALL_TARGET_VALID;
			
			char* start(static_cast<char*>(page));
			for(size_t i(0); i < function_size - sizeof(size_t); ++i) {
				if(*reinterpret_cast<size_t*>(start + i) == static_cast<size_t>(0xdeadbeefbaadf00d)) {
					*reinterpret_cast<size_t*>(start + i) = reinterpret_cast<size_t>(&f);
					i += sizeof(size_t) - 1;
				}
				else if(*reinterpret_cast<size_t*>(start + i) == static_cast<size_t>(0xcafebabed15ea5e5)) {
					union converter_t {
						decltype(&std::function<R(A...)>::operator()) cooked_pointer;
						size_t raw_pointer;
					};
					converter_t converter = { &std::function<R(A...)>::operator() };
					*reinterpret_cast<size_t*>(start + i) = converter.raw_pointer;
					i += sizeof(size_t) - 1;
				}
			}

			DWORD old_protection(0);
			::VirtualProtect(page, function_size, PAGE_EXECUTE_READ, &old_protection);
			
			decltype(&::SetProcessValidCallTargets) spvct = reinterpret_cast<decltype(&::SetProcessValidCallTargets)>(::GetProcAddress(::GetModuleHandleW(L"kernelbase.dll"), "SetProcessValidCallTargets"));
			if (spvct) {
				spvct(::GetCurrentProcess(), page, function_size, calling_conventions_max, call_targets);
			}
		
			::FlushInstructionCache(::GetCurrentProcess(), page, function_size);
		}

		std::function<R(A...)> f;
		void* page;
		size_t function_size;
	};

	std::shared_ptr<thunk_holder> the_thunk;
};

template<typename R, typename... A>
thunk<R, A...> make_thunk(std::function<R(A...)> f) {
	return thunk<R, A...>(f);
}

// motivating example

struct window {
	window() : wndproc(thunk<LRESULT, HWND, UINT, WPARAM, LPARAM>([this](HWND window, UINT msg, WPARAM wp, LPARAM lp) -> LRESULT {
			return this->window_procedure(window, msg, wp, lp);
		})) {
		WNDCLASSEX clazz = { 0 };
		clazz.cbSize = sizeof(WNDCLASSEX);
		clazz.lpszClassName = L"window-class";
		clazz.lpfnWndProc = wndproc.as_stdcall();
		clazz.style = CS_DBLCLKS;
		clazz.hInstance = ::GetModuleHandle(NULL);
		clazz.hCursor = ::LoadCursor(NULL, IDC_ARROW);
		clazz.hbrBackground = (HBRUSH)(COLOR_WINDOW + 1);

		cls = ::RegisterClassEx(&clazz);
		wnd = ::CreateWindow(clazz.lpszClassName, L"Window", WS_OVERLAPPEDWINDOW, 100, 100, 640, 480, 0, 0, ::GetModuleHandle(nullptr), nullptr);
		::ShowWindow(wnd, SW_RESTORE);
	}

	~window() {
		::UnregisterClass(reinterpret_cast<const wchar_t*>(cls), ::GetModuleHandle(nullptr));
	}

	int pump_messages() {
		MSG msg = { 0 };
		while(::GetMessageW(&msg, NULL, 0, 0)) {
			::TranslateMessage(&msg);
			::DispatchMessageW(&msg);
		}
		return static_cast<int>(msg.wParam);
	}

private:
	LRESULT window_procedure(HWND w, UINT message, WPARAM wp, LPARAM lp) {
		// NB regular method, not static.
		// NB no need to pass 'this' via WM_NCCREATE/CREATESTRUCTW
		switch (message) {
		case WM_NCDESTROY:
			::PostQuitMessage(0);
			return 0;
		default:
			return ::DefWindowProc(w, message, wp, lp);
		}
	}

	thunk<LRESULT, HWND, UINT, WPARAM, LPARAM> wndproc;
	ATOM cls;
	HWND wnd;
};

//int __stdcall WinMain(HINSTANCE, HINSTANCE, LPSTR, int)
int main()
{
	window w;

	return w.pump_messages();
}