First (rough) implementation, not tested yet

bbepis · Apr 18, 2019 · 4c7f7c7 · 4c7f7c7
1 parent 60cb59e
commit 4c7f7c7
Show file tree

Hide file tree

Showing 6 changed files with 164 additions and 18 deletions.
diff --git a/Harmony/Internal/MethodCopier.cs b/Harmony/Internal/MethodCopier.cs
@@ -291,8 +291,7 @@ internal void FinalizeILCodes(List<MethodInfo> transpilers, List<Label> endLabel
 				// TODO: we ignore the resulting label because we have no way to use it
 				//
 				codeInstruction.blocks.Do(block => {
-					Label? label;
-					Emitter.MarkBlockBefore(generator, block, out label);
+					Emitter.MarkBlockBefore(generator, block, out var label);
 				});
 
 				var code = codeInstruction.opcode;

diff --git a/Harmony/Internal/MethodPatcher.cs b/Harmony/Internal/MethodPatcher.cs
@@ -14,6 +14,7 @@ internal static class MethodPatcher
 		public static string ORIGINAL_METHOD_PARAM = "__originalMethod";
 		public static string RESULT_VAR = "__result";
 		public static string STATE_VAR = "__state";
+		public static string EXCEPTION_VAR = "__exception";
 		public static string PARAM_INDEX_PREFIX = "__";
 		public static string INSTANCE_FIELD_PREFIX = "___";
 
@@ -26,22 +27,23 @@ public static DynamicMethod CreatePatchedMethod(MethodBase original, List<Method
 		[UpgradeToLatestVersion(1)]
 		public static DynamicMethod CreatePatchedMethod(MethodBase original, string harmonyInstanceID, List<MethodInfo> prefixes, List<MethodInfo> postfixes, List<MethodInfo> transpilers, List<MethodInfo> finalizers)
 		{
-			Memory.MarkForNoInlining(original);
-
-			if (original == null)
-				throw new ArgumentNullException(nameof(original), "Original method is null. Did you specify it correctly?");
-
 			try
 			{
+				if (original == null)
+					throw new ArgumentNullException(nameof(original));
+
+				Memory.MarkForNoInlining(original);
+
 				if (HarmonyInstance.DEBUG)
 				{
 					FileLog.LogBuffered("### Patch " + original.DeclaringType + ", " + original);
 					FileLog.FlushBuffer();
 				}
 
-				var idx = prefixes.Count() + postfixes.Count();
+				var idx = prefixes.Count() + postfixes.Count() + finalizers.Count();
 				var firstArgIsReturnBuffer = NativeThisPointer.NeedsNativeThisPointerFix(original);
 				var returnType = AccessTools.GetReturnedType(original);
+				var hasFinalizers = finalizers.Any();
 				var patch = DynamicTools.CreateDynamicMethod(original, "_Patch" + idx);
 				if (patch == null)
 					return null;
@@ -72,8 +74,22 @@ public static DynamicMethod CreatePatchedMethod(MethodBase original, string harm
 					}
 				});
 
-				if (firstArgIsReturnBuffer)
-					Emitter.Emit(il, original.IsStatic ? OpCodes.Ldarg_0 : OpCodes.Ldarg_1);
+				LocalBuilder finalizedVariable = null;
+				LocalBuilder exceptionVariable = null;
+				if (hasFinalizers)
+				{
+					finalizedVariable = DynamicTools.DeclareLocalVariable(il, typeof(bool));
+
+					exceptionVariable = DynamicTools.DeclareLocalVariable(il, typeof(Exception));
+					privateVars[EXCEPTION_VAR] = exceptionVariable;
+
+					// begin try
+					Emitter.MarkBlockBefore(il, new ExceptionBlock(ExceptionBlockType.BeginExceptionBlock), out var label);
+				}
+
+				// TODO: verify correct usage
+				// if (firstArgIsReturnBuffer)
+				// 	Emitter.Emit(il, original.IsStatic ? OpCodes.Ldarg_0 : OpCodes.Ldarg_1);
 
 				var skipOriginalLabel = il.DefineLabel();
 				var canHaveJump = AddPrefixes(il, original, prefixes, privateVars, skipOriginalLabel);
@@ -99,10 +115,61 @@ public static DynamicMethod CreatePatchedMethod(MethodBase original, string harm
 
 				AddPostfixes(il, original, postfixes, privateVars, true);
 
-				if (firstArgIsReturnBuffer)
-					Emitter.Emit(il, OpCodes.Stobj, returnType);
+				if (hasFinalizers)
+				{
+					AddFinalizers(il, original, finalizers, privateVars, false);
+					Emitter.Emit(il, OpCodes.Ldc_I4_1);
+					Emitter.Emit(il, OpCodes.Stloc, finalizedVariable);
+					Emitter.Emit(il, OpCodes.Ldloc, exceptionVariable);
+					var noExceptionLabel = il.DefineLabel();
+					Emitter.Emit(il, OpCodes.Brfalse, noExceptionLabel);
+					Emitter.Emit(il, OpCodes.Ldloc, exceptionVariable);
+					Emitter.Emit(il, OpCodes.Throw);
+					il.MarkLabel(noExceptionLabel);
+
+					// end try (includes leave)
+					Emitter.MarkBlockAfter(il, new ExceptionBlock(ExceptionBlockType.EndExceptionBlock));
+				}
+				else
+				{
+					// TODO: verify correct usage
+					// if (firstArgIsReturnBuffer)
+					// 	Emitter.Emit(il, OpCodes.Stobj, returnType);
+
+					Emitter.Emit(il, OpCodes.Ret);
+				}
+
+				if (hasFinalizers)
+				{
+					// begin catch
+					Emitter.MarkBlockBefore(il, new ExceptionBlock(ExceptionBlockType.BeginCatchBlock), out var label);
+					Emitter.Emit(il, OpCodes.Stloc, exceptionVariable);
+
+					Emitter.Emit(il, OpCodes.Ldloc, finalizedVariable);
+					var endFinalizerLabel = il.DefineLabel();
+					Emitter.Emit(il, OpCodes.Brtrue, endFinalizerLabel);
+
+					var rethrowPossible = AddFinalizers(il, original, finalizers, privateVars, true);
 
-				Emitter.Emit(il, OpCodes.Ret);
+					il.MarkLabel(endFinalizerLabel);
+
+					if (rethrowPossible)
+						Emitter.Emit(il, OpCodes.Rethrow);
+					else
+					{
+						Emitter.Emit(il, OpCodes.Ldloc, exceptionVariable);
+						Emitter.Emit(il, OpCodes.Throw);
+					}
+
+					// end catch
+					Emitter.MarkBlockAfter(il, new ExceptionBlock(ExceptionBlockType.EndExceptionBlock));
+
+					// TODO: verify correct usage
+					// if (firstArgIsReturnBuffer)
+					// 	Emitter.Emit(il, OpCodes.Stobj, returnType);
+
+					Emitter.Emit(il, OpCodes.Ret);
+				}
 
 				if (HarmonyInstance.DEBUG)
 				{
@@ -445,6 +512,34 @@ static void AddPostfixes(ILGenerator il, MethodBase original, List<MethodInfo> p
 					}
 				});
 		}
+
+		static bool AddFinalizers(ILGenerator il, MethodBase original, List<MethodInfo> finalizers, Dictionary<string, LocalBuilder> variables, bool catchExceptions)
+		{
+			var rethrowPossible = true;
+			finalizers
+				.Do(fix =>
+				{
+					if (catchExceptions)
+						Emitter.MarkBlockBefore(il, new ExceptionBlock(ExceptionBlockType.BeginExceptionBlock), out var label);
+
+					EmitCallParameter(il, original, fix, variables, true);
+					Emitter.Emit(il, OpCodes.Call, fix);
+					if (fix.ReturnType != typeof(void))
+					{
+						Emitter.Emit(il, OpCodes.Stloc, variables[EXCEPTION_VAR]);
+						rethrowPossible = false;
+					}
+
+					if (catchExceptions)
+					{
+						Emitter.MarkBlockAfter(il, new ExceptionBlock(ExceptionBlockType.EndExceptionBlock));
+						Emitter.MarkBlockBefore(il, new ExceptionBlock(ExceptionBlockType.BeginCatchBlock), out var label);
+						Emitter.Emit(il, OpCodes.Pop);
+						Emitter.MarkBlockAfter(il, new ExceptionBlock(ExceptionBlockType.EndExceptionBlock));
+					}
+				});
+			return rethrowPossible;
+		}
 	}
 }
 
diff --git a/Harmony/Public/ExceptionBlock.cs b/Harmony/Public/ExceptionBlock.cs
@@ -32,10 +32,10 @@ public class ExceptionBlock
 		/// <param name="blockType">Block type</param>
 		/// <param name="catchType">Catch type</param>
 		///
-		public ExceptionBlock(ExceptionBlockType blockType, Type catchType)
+		public ExceptionBlock(ExceptionBlockType blockType, Type catchType = null)
 		{
 			this.blockType = blockType;
-			this.catchType = catchType;
+			this.catchType = catchType ?? typeof(object);
 		}
 	}
 }
diff --git a/HarmonyTests/Patching/Assets/PatchClasses.cs b/HarmonyTests/Patching/Assets/PatchClasses.cs
@@ -308,7 +308,7 @@ public static void Prefix(out object __state)
 			__state = null;
 		}
 
-		public static void Postfix(int __state)
+		public static void Postfix(object __state)
 		{
 
 		}
@@ -440,4 +440,23 @@ public static void Postfix(bool __result)
 			postfixed = true;
 		}
 	}
+
+	public class SimpleFinalizerClass
+	{
+		[MethodImpl(MethodImplOptions.NoInlining)]
+		public void Method()
+		{
+			Console.WriteLine("Method");
+		}
+	}
+
+	public class SimpleFinalizerPatch
+	{
+		public static bool finalized = false;
+
+		public static void Finalizer()
+		{
+			finalized = true;
+		}
+	}
 }
diff --git a/HarmonyTests/Patching/FinalizerPatches.cs b/HarmonyTests/Patching/FinalizerPatches.cs
@@ -0,0 +1,35 @@
+using Harmony;
+using HarmonyTests.Assets;
+using NUnit.Framework;
+using System;
+using System.Collections.Generic;
+using System.Reflection;
+
+namespace HarmonyTests
+{
+	[TestFixture]
+	public class FinalizerPatches
+	{
+		[Test]
+		public void TestSimpleFinalizer()
+		{
+			var originalClass = typeof(SimpleFinalizerClass);
+			Assert.IsNotNull(originalClass);
+			var originalMethod = originalClass.GetMethod("Method");
+			Assert.IsNotNull(originalMethod);
+
+			var patchClass = typeof(SimpleFinalizerPatch);
+			var finalizer = patchClass.GetMethod("Finalizer");
+			Assert.IsNotNull(finalizer);
+
+			var instance = HarmonyInstance.Create("test");
+			Assert.IsNotNull(instance);
+			var patcher = new PatchProcessor(instance, new List<MethodBase> { originalMethod }, null, null, null, new HarmonyMethod(finalizer));
+			Assert.IsNotNull(patcher);
+			patcher.Patch();
+
+			new SimpleFinalizerClass().Method();
+			Assert.IsTrue(SimpleFinalizerPatch.finalized);
+		}
+	}
+}
diff --git a/HarmonyTests/Patching/StaticPatches.cs b/HarmonyTests/Patching/StaticPatches.cs
@@ -186,7 +186,6 @@ public void TestMethod5()
 			Assert.IsTrue(Class5Patch.postfixed, "Prefix was not executed");
 		}
 
-		/* fails during travis-ci testing, so disabled for now 
 		[Test]
 		public void TestPatchUnpatch()
 		{
@@ -213,7 +212,6 @@ public void TestPatchUnpatch()
 
 			instanceB.UnpatchAll("test");
 		}
-		*/
 
 		[Test]
 		public void TestAttributes()