Implement FNativeHookManagerInternal::UnregisterHookFunction

satisfactorymodding · Nov 2, 2023 · 80d1ff4 · 80d1ff4
1 parent e9de79c
commit 80d1ff4
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 11 deletions.
diff --git a/Mods/SML/Source/SML/Private/Patching/NativeHookManager.cpp b/Mods/SML/Source/SML/Private/Patching/NativeHookManager.cpp
@@ -13,6 +13,9 @@ static TMap<void*, void*> RegisteredListenerMap;
 //Map of the function implementation pointer to the trampoline function pointer. Used to ensure one hook per function installed
 static TMap<void*, void*> InstalledHookMap;
 
+//Store the funchook instance used to hook each function
+static TMap<void*, funchook_t*> FunchookMap;
+
 void* FNativeHookManagerInternal::GetHandlerListInternal( const void* RealFunctionAddress ) {
 	void** ExistingMapEntry = RegisteredListenerMap.Find(RealFunctionAddress);
 	return ExistingMapEntry ? *ExistingMapEntry : nullptr;
@@ -52,6 +55,7 @@ bool HookStandardFunction(const FString& DebugSymbolName, void* OriginalFunction
 	CHECK_FUNCHOOK_ERR(funchook_prepare(funchook, OutTrampolineFunction, HookFunctionPointer));
 	CHECK_FUNCHOOK_ERR(funchook_install(funchook, 0));
 	InstalledHookMap.Add(OriginalFunctionPointer, *OutTrampolineFunction);
+	FunchookMap.Add(OriginalFunctionPointer, funchook);
 	return true;
 }
 
@@ -88,3 +92,16 @@ SML_API void* FNativeHookManagerInternal::RegisterHookFunction(const FString& De
 	return ResolvedHookingFunctionPointer;
 }
 
+void FNativeHookManagerInternal::UnregisterHookFunction(const FString& DebugSymbolName, const void* RealFunctionAddress) {
+	funchook_t** funchookPtr = FunchookMap.Find(RealFunctionAddress);
+	if (funchookPtr == nullptr) {
+		UE_LOG(LogNativeHookManager, Warning, TEXT("Attempt to unregister hook for function %s at %p which was not registered"), *DebugSymbolName, RealFunctionAddress);
+		return;
+	}
+	funchook_t* funchook = *funchookPtr;
+	CHECK_FUNCHOOK_ERR(funchook_uninstall(funchook, 0));
+	CHECK_FUNCHOOK_ERR(funchook_destroy(funchook));
+	FunchookMap.Remove(RealFunctionAddress);
+	InstalledHookMap.Remove(RealFunctionAddress);
+	UE_LOG(LogNativeHookManager, Display, TEXT("Successfully unregistered hook for function %s at %p"), *DebugSymbolName, RealFunctionAddress);
+}
diff --git a/Mods/SML/Source/SML/Public/Patching/NativeHookManager.h b/Mods/SML/Source/SML/Public/Patching/NativeHookManager.h
@@ -51,7 +51,7 @@ class SML_API FNativeHookManagerInternal {
 	static void* GetHandlerListInternal( const void* RealFunctionAddress);
 	static void SetHandlerListInstanceInternal(void* RealFunctionAddress, void* HandlerList);
 	static void* RegisterHookFunction(const FString& DebugSymbolName, void* OriginalFunctionPointer, const void* SampleObjectInstance, int ThisAdjustment, void* HookFunctionPointer, void** OutTrampolineFunction);
-	static void UnregisterHookFunction( const void* RealFunctionAddress );
+	static void UnregisterHookFunction( const FString& DebugSymbolName, const void* RealFunctionAddress );
 };
 
 template <typename T, typename E>
@@ -272,11 +272,11 @@ struct HookInvokerExecutorGlobalFunction {
 	}
 
 	// Uninstalls the hook. Also frees the handler lists object.
-	static void UninstallHook()
+	static void UninstallHook(const FString& DebugSymbolName)
 	{
 		if (bHookInitialized)
 		{
-			FNativeHookManagerInternal::UnregisterHookFunction( RealFunctionAddress );
+			FNativeHookManagerInternal::UnregisterHookFunction( DebugSymbolName, RealFunctionAddress );
 			DestroyHandlerLists<Handler, HandlerAfter>( RealFunctionAddress );
 			bHookInitialized = false;
 			RealFunctionAddress = nullptr;
@@ -308,7 +308,7 @@ struct HookInvokerExecutorGlobalFunction {
 		return NewDelegateHandle;
 	}
 
-	static void RemoveHandler( FDelegateHandle InHandlerHandle )
+	static void RemoveHandler(const FString& DebugSymbolName, FDelegateHandle InHandlerHandle )
 	{
 		if ( HandlerBeforeReferences->Contains( InHandlerHandle ) )
 		{
@@ -323,7 +323,7 @@ struct HookInvokerExecutorGlobalFunction {
 
 		if ( HandlersAfter->IsEmpty() && HandlersBefore->IsEmpty() )
 		{
-			UninstallHook();
+			UninstallHook(DebugSymbolName);
 		}
 	}
 };
@@ -449,11 +449,11 @@ struct HookInvokerExecutorMemberFunction {
 	}
 
 	// Uninstalls the hook. Also frees the handler lists object.
-	static void UninstallHook()
+	static void UninstallHook(const FString& DebugSymbolName)
 	{
 		if (bHookInitialized)
 		{
-			FNativeHookManagerInternal::UnregisterHookFunction( RealFunctionAddress );
+			FNativeHookManagerInternal::UnregisterHookFunction( DebugSymbolName, RealFunctionAddress );
 			DestroyHandlerLists<Handler, HandlerAfter>( RealFunctionAddress );
 			bHookInitialized = false;
 			RealFunctionAddress = nullptr;
@@ -485,7 +485,7 @@ struct HookInvokerExecutorMemberFunction {
 		return NewDelegateHandle;
 	}
 
-	static void RemoveHandler( FDelegateHandle InHandlerHandle )
+	static void RemoveHandler(const FString& DebugSymbolName, FDelegateHandle InHandlerHandle )
 	{
 		if ( HandlerBeforeReferences->Contains( InHandlerHandle ) )
 		{
@@ -500,7 +500,7 @@ struct HookInvokerExecutorMemberFunction {
 
 		if ( HandlersAfter->IsEmpty() && HandlersBefore->IsEmpty() )
 		{
-			UninstallHook();
+			UninstallHook(DebugSymbolName);
 		}
 	}
 };
@@ -568,7 +568,7 @@ using CallScope = TCallScope<T>;
 	} )
 
 #define UNSUBSCRIBE_METHOD(MethodReference, HandlerHandle) \
-	HookInvoker<decltype(&MethodReference), &MethodReference>::RemoveHandler( HandlerHandle )
+	HookInvoker<decltype(&MethodReference), &MethodReference>::RemoveHandler(TEXT(#MethodReference), HandlerHandle )
 
 #define UNSUBSCRIBE_UOBJECT_METHOD(ObjectClass, MethodName, HandlerHandle) \
-	HookInvoker<decltype(&ObjectClass::MethodName), &ObjectClass::MethodName>::RemoveHandler( HandlerHandle )
+	HookInvoker<decltype(&ObjectClass::MethodName), &ObjectClass::MethodName>::RemoveHandler(TEXT(#ObjectClass "::" #MethodName), HandlerHandle )