Fix return addresses after callbacks rebuilding

metamod/src/callback_jit.cpp

@@ -114,6 +114,7 @@ void CForwardCallbackJIT::naked_main()
 
 	// setup meta globals
 	mov(dword_ptr[globals + mg_mres], MRES_UNSET);
+	mov(dword_ptr[globals + mg_esp_save], esp);
 
 	// setup retval pointers
 	if (m_jitdata->has_ret) {
@@ -372,6 +373,23 @@ void CJit::clear_tramps()
 	m_tramp_allocator.deallocate_all();
 }
 
+size_t CJit::is_callback_retaddr(uint32 addr)
+{
+	if (m_callback_allocator.contain(addr)) {
+		// FF D1		call    ecx
+		// 83 C4 20		add     esp, 20h ; optional
+		// 8B 13		mov     edx, [ebx]
+		char *ptr = (char *)addr - 2;
+		return mem_compare(ptr, "\xFF\xD1\x83\xC4", 4) || mem_compare(ptr, "\xFF\xD1\x8B\x13", 4);
+	}
+	return false;
+}
+
+char* CJit::find_callback_pattern(char* pattern, size_t len)
+{
+	return m_callback_allocator.find_pattern(pattern, len);
+}
+
 bool CJit::is_hook_needed(jitdata_t* jitdata)
 {
 	if (jitdata->mm_hook)

metamod/src/callback_jit.h

@@ -88,7 +88,8 @@ public:
 	size_t compile_tramp(size_t ptr_to_func/*, size_t hook, size_t hook_time*/);
 	void clear_callbacks();
 	void clear_tramps();
-	size_t find_new_retaddr(size_t pfn);
+	size_t is_callback_retaddr(uint32 addr);
+	char* find_callback_pattern(char* pattern, size_t len);
 
 private:
 	static bool is_hook_needed(jitdata_t* jitdata);

metamod/src/dllapi.cpp

@@ -202,7 +202,6 @@ C_DLLEXPORT int GetNewDLLFunctions(NEW_DLL_FUNCTIONS *pNewFunctionTable, int *in
 		*interfaceVersion = NEW_DLL_FUNCTIONS_VERSION;
 		return FALSE;
 	}
-
 	memcpy(pNewFunctionTable, &sNewFunctionTable, sizeof(NEW_DLL_FUNCTIONS));
 	return TRUE;
 }

metamod/src/meta_api.h

@@ -49,7 +49,7 @@ struct meta_globals_t
 	void *override_ret;		// readable; return value from overriding/superceding plugin
 
 #ifdef METAMOD_CORE
-	int esp_save;
+	uint32* esp_save;
 #endif
 };
 

metamod/src/metamod.cpp

@@ -422,8 +422,90 @@ bool meta_load_gamedll()
 	return true;
 }
 
+struct fixdata_t
+{
+	uint32* pret;
+	uint32 callback;
+};
+
+static void meta_collect_fix_data(uint32* const esp, std::vector<fixdata_t>& data)
+{
+	uint32* pret = esp;
+
+	// scan for retaddr
+	do pret--;
+	while (!g_jit.is_callback_retaddr(*pret));
+
+	char* raddr = (char *)*pret;
+	size_t args_count = 0;
+
+	if ((raddr[0] == 0x83 && raddr[1] == 0xC4))
+		args_count = raddr[2] / 4;
+
+	// 8B 0D 4E 61 BC 00		mov     ecx, ds : 0BC614Eh
+	// 80 3D 4E 61 BC 00 05		cmp     byte ptr ds : 0BC614Eh, 5
+	char pattern[] = "\x8B\x0D\x2A\x2A\x2A\x2A\x80\x3D";
+
+	// scan for callback addr
+	do raddr--;
+	while (!mem_compare(raddr, pattern, sizeof pattern - 1));
+
+	uint32 callback = *(uint32 *)(raddr + 2);
+	data.push_back({pret, callback});
+
+	// 0F 29 44 24 1C				movaps  xmmword ptr[esp + 1Ch], xmm0
+	// 66 0F D6 4C 24 14			movq    mmword ptr[esp + 14h], xmm1
+
+	// 0F 29 84 24 36 05 00 00		movaps  xmmword ptr[esp + 536h], xmm0
+	// 66 0F D6 8C 24 36 05 00 00	movq    qword ptr[esp + 536h], xmm1
+	char pattern2[] = "\x66\x0F\xD6\x2A\x24";
+
+	// scan for mg_backup
+	do raddr--;
+	while (!mem_compare(raddr, pattern2, sizeof pattern2 - 1));
+
+	uint32 mg_backup;
+	if (raddr[3] == 0x4C)
+		mg_backup = *(uint8 *)&raddr[sizeof pattern2 - 1];
+	else
+		mg_backup = *(uint32 *)&raddr[sizeof pattern2 - 1];
+
+	uint32* esp_save = *(uint32 **)(uint32(esp) + mg_backup + sizeof(int));
+	if (esp_save)
+		meta_collect_fix_data(esp_save, data);
+}
+
+static void meta_apply_fix_data(std::vector<fixdata_t>& data)
+{
+	for (auto& d : data) {
+		// 8B 0D 4E 61 BC 00	mov     ecx, ds:0BC614Eh
+		char pattern[6] = "\x8B\x0D";
+		*(uint32 *)(pattern + 2) = d.callback;
+
+		char* ptr = g_jit.find_callback_pattern(pattern, sizeof pattern);
+		if (!ptr) {
+			META_ERROR("Failed to fix callback retaddr.\n Bye bye...\n");
+			do_exit(666);
+		}
+
+		// FF D1		call    ecx
+		// 83 C4 20		add     esp, 20h ; optional
+		// 8B 13		mov     edx, [ebx]
+		do ptr++;
+		while (!mem_compare(ptr, "\xFF\xD1\x83\xC4", 4) && !mem_compare(ptr, "\xFF\xD1\x8B\x13", 4));
+
+		*d.pret = uint32(ptr + 2);
+	}
+}
+
 void meta_rebuild_callbacks()
 {
+	std::vector<fixdata_t> fixdata;
+	if (g_metaGlobals.esp_save) {
+		META_LOG("dll: Begin scan to collect callback fix data...");
+		meta_collect_fix_data(g_metaGlobals.esp_save, fixdata);
+	}
+
 	META_LOG("dll: Rebuilding callbacks...");
 
 #ifdef JIT_DEBUG_FILE
@@ -434,4 +516,11 @@ void meta_rebuild_callbacks()
 
 	compile_engine_callbacks();
 	compile_gamedll_callbacks();
+
+	if (fixdata.size()) {
+		META_LOG("dll: Begin callbacks fixing...");
+		meta_apply_fix_data(fixdata);
+	}
+
+	META_LOG("dll: Callbacks rebuilded.");
 }

metamod/src/utils.cpp

@@ -65,6 +65,26 @@ size_t static_allocator::memory_used() const
 	return (m_pages.size() - 1) * Pagesize + m_used;
 }
 
+bool static_allocator::contain(uint32 addr)
+{
+	for (auto p : m_pages) {
+		if (uint32(p) <= addr && addr < uint32(p) + Pagesize)
+			return true;
+	}
+	return false;
+}
+
+char* static_allocator::find_pattern(char* pattern, size_t len)
+{
+	for (auto p : m_pages) {
+		for (char* c = (char *)p, *e = c + Pagesize - len; c < e; c++) {
+			if (mem_compare(c, pattern, len))
+				return c;
+		}
+	}
+	return nullptr;
+}
+
 void static_allocator::allocate_page()
 {
 #ifdef WIN32
@@ -275,3 +295,13 @@ char* full_gamedir_path(const char* path, char* fullpath)
 	normalize_path(fullpath);
 	return fullpath;
 }
+
+bool mem_compare(const char* addr, const char* pattern, size_t len)
+{
+	for (auto c = pattern, pattern_end = pattern + len; c < pattern_end; ++c, ++addr) {
+		if (*c == *addr || *c == '\x2A')
+			continue;
+		return false;
+	}
+	return true;
+}

metamod/src/utils.h

@@ -35,6 +35,8 @@ public:
 	char* strdup(const char* string);
 	void deallocate_all();
 	size_t memory_used() const;
+	bool contain(uint32 addr);
+	char* find_pattern(char* pattern, size_t len);
 
 	template<typename T>
 	T* allocate()
@@ -77,5 +79,6 @@ void __declspec(noreturn) do_exit(int exitval);
 
 bool is_file_exists_in_gamedir(const char *path);
 char *full_gamedir_path(const char *path, char *fullpath);
+bool mem_compare(const char* addr, const char* pattern, size_t len);
 
 extern const char* g_platform_postfixes[4];