Skip to content

Commit ed709d1

Browse files
author
UserR00T
committed
Working on HasReturnValue, to return a value to injection Target.
Currently seems stable, still needs some refactoring and improvements. Preview: Injected [HookAttribute.Hook("Injected.Method3")] [HasReturnValue] public static bool InjectedMethod_M(Injected instance) { return true; } InjectionTarget public bool Method3() { Console.WriteLine("Magic v3"); return true; } After injected https://i.imgur.com/sZ7mpAT.png
1 parent 1ee54cf commit ed709d1

12 files changed

Lines changed: 205 additions & 28 deletions

‎HookAttribute/HasReturnValueAttribute.cs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
using System;
2+
using System.Collections.Generic;
3+
using System.Linq;
4+
using System.Text;
5+
using System.Threading.Tasks;
6+
7+
namespace HookAttribute
8+
{
9+
[AttributeUsage(AttributeTargets.Method, Inherited = false, AllowMultiple = false)]
10+
public sealed class HasReturnValueAttribute : Attribute
11+
{
12+
public HasReturnValueAttribute()
13+
{
14+
}
15+
}
16+
}

‎HookAttribute/HookAttribute.csproj‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
<Reference Include="System.Xml" />
4242
</ItemGroup>
4343
<ItemGroup>
44+
<Compile Include="HasReturnValueAttribute.cs" />
4445
<Compile Include="AddMethodAttribute.cs" />
4546
<Compile Include="HookAttribute.cs" />
4647
<Compile Include="Properties\AssemblyInfo.cs" />

‎TestProject/Core.cs‎

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -10,32 +10,28 @@ namespace TestProject
1010
{
1111
public class Core
1212
{
13-
public void InjectedMethod()
13+
[HookAttribute.Hook("Injected.Method3")]
14+
[HasReturnValue]
15+
public static bool InjectedMethod_M(Injected instance)
1416
{
15-
Console.WriteLine("Magic v1");
17+
return true;
1618
}
17-
public void InjectedMethod2()
19+
[HookAttribute.Hook("Injected.Method2")]
20+
public static void InjectedMethod2_M(Injected instance)
1821
{
19-
Console.WriteLine("Magic v2");
22+
Console.WriteLine("InjectedMethod2_M called");
2023
}
2124

2225

23-
24-
25-
[HookAttribute.Hook("Core.InjectedMethod")]
26-
public static bool InjectedMethod_M(Core instance)
27-
{
28-
return true;
29-
}
30-
[HookAttribute.Hook("Core.InjectedMethod2")]
31-
public static void InjectedMethod2_M(Core instance)
26+
[AddMethod(nameof(Injected), "TestMethod")]
27+
public static void OnTestMethod(Injected instance)
3228
{
33-
Console.WriteLine("InjectedMethod2_M called");
29+
Console.WriteLine("OnTestMethod Called");
3430
}
35-
[AddMethod(nameof(Core), "TestMethod")]
36-
public static void OnTestMethod(Core instance)
31+
[AddMethod(nameof(StaticInjected), "StaticTestMethod")]
32+
public static void OnStaticTestMethod()
3733
{
38-
Console.WriteLine("OnTestMethod Called");
34+
Console.WriteLine("OnStaticTestMethod Called");
3935
}
4036
}
4137
}

‎TestProject/Injected.cs‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
using System;
2+
using System.Collections.Generic;
3+
using System.Linq;
4+
using System.Text;
5+
using System.Threading.Tasks;
6+
7+
namespace TestProject
8+
{
9+
public class Injected
10+
{
11+
public void Method1()
12+
{
13+
Console.WriteLine("Magic v1");
14+
}
15+
public void Method2()
16+
{
17+
Console.WriteLine("Magic v2");
18+
}
19+
public bool Method3()
20+
{
21+
Console.WriteLine("Magic v3");
22+
return true;
23+
}
24+
public bool Method4()
25+
{
26+
Console.WriteLine("Magic v4");
27+
return false;
28+
}
29+
30+
public static void StaticMethod1()
31+
{
32+
Console.WriteLine("Magic v1");
33+
}
34+
public static void StaticMethod2()
35+
{
36+
Console.WriteLine("Magic v2");
37+
}
38+
public static bool StaticMethod3()
39+
{
40+
Console.WriteLine("Magic v3");
41+
return true;
42+
}
43+
public static bool StaticMethod4()
44+
{
45+
Console.WriteLine("Magic v4");
46+
return false;
47+
}
48+
}
49+
public static class StaticInjected
50+
{
51+
52+
}
53+
}

‎TestProject/TestProject.csproj‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
</ItemGroup>
4343
<ItemGroup>
4444
<Compile Include="Core.cs" />
45+
<Compile Include="Injected.cs" />
4546
<Compile Include="Properties\AssemblyInfo.cs" />
4647
</ItemGroup>
4748
<ItemGroup>

‎UniversalUnityHooks/Attributes.cs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ public static Dictionary<string, List<AttributeData>> FindAndInvokeAllAttributes
2121
{
2222
var typelist = GetTypesInNamespace(Assembly.GetExecutingAssembly(), "UniversalUnityHooks.Attributes");
2323
var returnData = new Dictionary<string, List<AttributeData>>();
24+
long lastElapsedMs = 0;
2425
foreach (var type in typelist)
2526
{
2627
if (type.BaseType != typeof(Attributes.Attribute))
@@ -32,7 +33,8 @@ public static Dictionary<string, List<AttributeData>> FindAndInvokeAllAttributes
3233
if (attributes == null)
3334
continue;
3435
returnData.Add(type.Name, attributes);
35-
ConsoleHelper.WriteMessage($"Attributes found of type {type.Name}: {attributes.Count}");
36+
ConsoleHelper.WriteMessage($"Attributes found of type {type.Name}: {attributes.Count} ({timer.GetElapsedMs - lastElapsedMs}ms)");
37+
lastElapsedMs = timer.GetElapsedMs;
3638
}
3739
return returnData;
3840
}
@@ -84,6 +86,7 @@ public AttributeData(TypeDefinition type, MethodDefinition method, CustomAttribu
8486
Attribute = attribute;
8587
Assembly = assembly;
8688
}
89+
public Cecil.ReturnData TargetData { get; set; }
8790
public AssemblyDefinition Assembly { get; set; }
8891
public TypeDefinition Type { get; set; }
8992
public MethodDefinition Method { get; set; }

‎UniversalUnityHooks/Attributes/AddMethodAttribute.cs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,18 +17,23 @@ public override AddAttributesResponse AddAllFound()
1717
TempData = AllAttributes.Where(x => x.Attribute.AttributeType.Name == nameof(HookAttribute.AddMethodAttribute)).ToList();
1818
if (TempData == null || TempData.Count == 0)
1919
{
20-
ConsoleHelper.WriteMessage(ConsoleHelper.MessageType.Warning, $"No add method attributes were found. ({Timer.GetElapsedMs}ms)\r\n");
20+
ConsoleHelper.WriteMessage(ConsoleHelper.MessageType.Warning, $"No add method attributes were found. ({Timer.GetElapsedMs}ms)");
2121
return AddAttributesResponse.Info;
2222
}
2323
return base.AddAllFound();
2424
}
2525
public static int InjectHooks(Dictionary<string, List<AttributeData>> attributes, AssemblyDefinition assemblyDefinition)
2626
{
27+
if (!attributes.ContainsKey(nameof(AddMethodAttribute)))
28+
return 0;
2729
var injectedCorrectly = 0;
30+
ConsoleHelper.WriteNewline();
2831
foreach (var hook in attributes[nameof(AddMethodAttribute)])
2932
{
3033
// Improve: Add check for multiple args here
3134
var typeDefinition = Cecil.ConvertStringToClass(hook.Attribute.ConstructorArguments[0].Value.ToString(), assemblyDefinition);
35+
if (typeDefinition == null)
36+
continue;
3237
var methodName = hook.Attribute.ConstructorArguments[1].Value.ToString();
3338
if (Cecil.MethodExists(typeDefinition, methodName))
3439
{

‎UniversalUnityHooks/Attributes/HasReturnValueAttribute.cs‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
using Mono.Cecil;
2+
using Mono.Cecil.Cil;
3+
using System;
4+
using System.Collections.Generic;
5+
using System.Linq;
6+
using System.Text;
7+
using System.Threading.Tasks;
8+
using static UniversalUnityHooks.AttributesHelper;
9+
using static UniversalUnityHooks.Util;
10+
11+
namespace UniversalUnityHooks.Attributes
12+
{
13+
public sealed class HasReturnValueAttribute : Attribute
14+
{
15+
public HasReturnValueAttribute(List<AttributeData> attributes, OperationTimer timer) : base(attributes, timer) { }
16+
public override AddAttributesResponse AddAllFound()
17+
{
18+
TempData = AllAttributes.Where(x => x.Attribute.AttributeType.Name == nameof(HookAttribute.HasReturnValueAttribute)).ToList();
19+
if (TempData == null || TempData.Count == 0)
20+
{
21+
ConsoleHelper.WriteMessage(ConsoleHelper.MessageType.Warning, $"No has return value attributes were found. ({Timer.GetElapsedMs}ms)");
22+
return AddAttributesResponse.Info;
23+
}
24+
if (TempData.Any(x=>x.Method.ReturnType == x.Assembly.MainModule.TypeSystem.Void))
25+
{
26+
ConsoleHelper.WriteMessage(ConsoleHelper.MessageType.Error, $"You cannot use void as return value for this attribute. ({Timer.GetElapsedMs}ms)");
27+
return AddAttributesResponse.Info;
28+
}
29+
return base.AddAllFound();
30+
}
31+
public static int InjectHooks(Dictionary<string, List<AttributeData>> attributes, AssemblyDefinition assemblyDefinition)
32+
{
33+
if (!attributes.ContainsKey(nameof(HookAttributes)) || !attributes.ContainsKey(nameof(HasReturnValueAttribute)))
34+
return 0;
35+
var injectedCorrectly = 0;
36+
var hookAttributes = new Dictionary<MethodDefinition, AttributeData>();
37+
hookAttributes = attributes[nameof(HookAttributes)].ToDictionary(x => x.Method, y => y);
38+
ConsoleHelper.WriteNewline();
39+
foreach (var hook in attributes[nameof(HasReturnValueAttribute)])
40+
{
41+
if (!hookAttributes.TryGetValue(hook.Method, out var hookAttribute))
42+
{
43+
ConsoleHelper.WriteError($"Cannot get hook attribute for method {hook.Method.Name}. Did you decorate that method with the [HookAttribute] method?");
44+
continue;
45+
}
46+
if (hook.Method.ReturnType != hookAttribute.TargetData.methodDefinition.ReturnType)
47+
{
48+
ConsoleHelper.WriteError($"Return types of {hook.Method.Name} and {hookAttribute.TargetData.methodDefinition.Name} do not match. They need to match to add the [HasReturnValue] attribute.");
49+
continue;
50+
}
51+
var il = hookAttribute.TargetData.methodDefinition.Body.GetILProcessor();
52+
var firstInstruction = il.Body.Instructions.First();
53+
54+
ConsoleHelper.WriteMessage($"Mdef: {hookAttribute.TargetData.methodDefinition} | Tdef: {hookAttribute.TargetData.typeDefinition}");
55+
}
56+
return injectedCorrectly;
57+
}
58+
}
59+
}

‎UniversalUnityHooks/Attributes/HookAttributes.cs‎

Lines changed: 40 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using Mono.Cecil;
2+
using Mono.Cecil.Cil;
23
using System;
34
using System.Collections.Generic;
45
using System.Linq;
@@ -16,22 +17,55 @@ public override AddAttributesResponse AddAllFound()
1617
TempData = AllAttributes.Where(x => x.Attribute.AttributeType.Name == nameof(HookAttribute.HookAttribute)).ToList();
1718
if (TempData == null || TempData.Count == 0)
1819
{
19-
ConsoleHelper.WriteMessage(ConsoleHelper.MessageType.Warning, $"No hook attributes were found. Without hook attributes it cannot inject anything. ({Timer.GetElapsedMs}ms)\r\n");
20+
ConsoleHelper.WriteMessage(ConsoleHelper.MessageType.Warning, $"No hook attributes were found. Without hook attributes it cannot inject anything. ({Timer.GetElapsedMs}ms)");
2021
return AddAttributesResponse.Info;
2122
}
2223
return base.AddAllFound();
2324
}
2425
public static int InjectHooks(Dictionary<string, List<AttributeData>> attributes, AssemblyDefinition assemblyDefinition)
2526
{
27+
if (!attributes.ContainsKey(nameof(HookAttributes)))
28+
return 0;
2629
var injectedCorrectly = 0;
30+
var returnValueAttributes = new Dictionary<MethodDefinition, AttributeData>();
31+
returnValueAttributes = attributes[nameof(HasReturnValueAttribute)].ToDictionary(x => x.Method, y => y);
2732
foreach (var hook in attributes[nameof(HookAttributes)])
2833
{
29-
// Improve: Add check for multiple args here
30-
var returnData = Cecil.ConvertStringToClassAndMethod(hook.Attribute.ConstructorArguments[0].Value.ToString(), assemblyDefinition);
31-
if (returnData == null)
34+
hook.TargetData = Cecil.ConvertStringToClassAndMethod(hook.Attribute.ConstructorArguments[0].Value.ToString(), assemblyDefinition);
35+
if (hook.TargetData == null)
3236
continue;
33-
if (Cecil.Inject(returnData, hook))
34-
++injectedCorrectly;
37+
if (!returnValueAttributes.ContainsKey(hook.Method))
38+
{
39+
if (Cecil.Inject(hook.TargetData, hook))
40+
++injectedCorrectly;
41+
continue;
42+
}
43+
ConsoleHelper.WriteMessage($"F1: {hook.Method.FullName}: {hook.Method.ReturnType.Name} | F2: {hook.TargetData.methodDefinition.FullName}: {hook.TargetData.methodDefinition.ReturnType.Name}");
44+
if (hook.Method.ReturnType.Name != hook.TargetData.methodDefinition.ReturnType.Name)
45+
{
46+
ConsoleHelper.WriteError($"Return types of {hook.Method.Name} and {hook.TargetData.methodDefinition.Name} do not match. They need to match to add the [HasReturnValue] attribute.");
47+
continue;
48+
}
49+
ConsoleHelper.WriteMessage($"Mdef: {hook.TargetData.methodDefinition} | Tdef: {hook.TargetData.typeDefinition}");
50+
var il = hook.TargetData.methodDefinition.Body.GetILProcessor();
51+
var firstInstruction = il.Body.Instructions.First();
52+
var method = hook.TargetData.methodDefinition;
53+
Instruction lastInstruction = null;
54+
Instruction currInstruction = null;
55+
if (!method.IsStatic && !method.DeclaringType.IsSealed)
56+
{
57+
currInstruction = il.Create(OpCodes.Ldarg_0);
58+
if (firstInstruction != null)
59+
il.InsertBefore(firstInstruction, currInstruction);
60+
else
61+
il.Append(currInstruction);
62+
}
63+
lastInstruction = currInstruction;
64+
currInstruction = il.Create(OpCodes.Call, assemblyDefinition.MainModule.Import(hook.Method));
65+
il.InsertAfter(lastInstruction ?? firstInstruction, currInstruction);
66+
lastInstruction = currInstruction;
67+
il.InsertAfter(lastInstruction, Instruction.Create(OpCodes.Ret));
68+
injectedCorrectly += 2; // HookAttribute and HasReturnValue both succeeded, add two
3569
}
3670
return injectedCorrectly;
3771
}

‎UniversalUnityHooks/Cecil.cs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,15 @@ public static bool InjectNewMethod(TypeDefinition type, AssemblyDefinition assem
4545
try
4646
{
4747
ConsoleHelper.WriteMessage(ConsoleHelper.MessageType.Wait, $"Injecting new method \"{type.Name}.{methodName}\"..");
48-
var methodDefinition = new MethodDefinition(methodName, MethodAttributes.Public, assembly.MainModule.TypeSystem.Void);
48+
var methodDefinition = new MethodDefinition(methodName, MethodAttributes.Public, assembly.MainModule.TypeSystem.Void)
49+
{
50+
IsStatic = type.IsSealed
51+
};
4952
type.Methods.Add(methodDefinition);
5053
var il = methodDefinition.Body.GetILProcessor();
51-
methodDefinition.Body.Instructions.Insert(0, il.Create(OpCodes.Ldarg_0));
52-
methodDefinition.Body.Instructions.Insert(1, il.Create(OpCodes.Call, assembly.MainModule.Import(hook.Method)));
54+
if (!type.IsSealed)
55+
methodDefinition.Body.Instructions.Add(il.Create(OpCodes.Ldarg_0));
56+
methodDefinition.Body.Instructions.Add(il.Create(OpCodes.Call, assembly.MainModule.Import(hook.Method)));
5357
methodDefinition.Body.Instructions.Add(Instruction.Create(OpCodes.Ret));
5458
ConsoleHelper.WriteMessage(ConsoleHelper.MessageType.Success, $"Injected new method \"{type.Name}.{methodName}\".");
5559
return true;

0 commit comments

Comments
 (0)