injector: allow method hooking to end of methods with multiple returns

runescape-client-injector-plugin/src/main/java/net/runelite/injector/InjectHookMethod.java

@@ -24,6 +24,8 @@
  */
 package net.runelite.injector;
 
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
 import net.runelite.asm.ClassFile;
@@ -46,6 +48,7 @@
 
 public class InjectHookMethod
 {
+
 	private static final Logger logger = LoggerFactory.getLogger(InjectHookMethod.class);
 
 	private final Inject inject;
@@ -108,47 +111,56 @@ private void injectHookMethod(Annotation hook, Method deobMethod, Method vanilla
 			builder.addArgument(inject.deobfuscatedTypeToApiType(type));
 		}
 
-		int insertPos = findHookLocation(hook, vanillaMethod);
-
 		assert deobMethod.isStatic() == vanillaMethod.isStatic();
 
 		if (!deobMethod.isStatic())
 		{
 			// Add variable to signature
 			builder.addArgument(0, inject.deobfuscatedTypeToApiType(new Type(deobMethod.getClassFile().getName())));
-			instructions.addInstruction(insertPos++, new ALoad(instructions, 0));
 		}
 
 		Signature signature = builder.build();
-		int index = deobMethod.isStatic() ? 0 : 1; // current variable index
 
-		for (int i = index; i < signature.size(); ++i)
+		List<Integer> insertIndexes = findHookLocations(hook, vanillaMethod);
+		insertIndexes.sort((a, b) -> Integer.compare(b, a));
+
+		for (int insertPos : insertIndexes)
 		{
-			Type type = signature.getTypeOfArg(i);
+			if (!deobMethod.isStatic())
+			{
+				instructions.addInstruction(insertPos++, new ALoad(instructions, 0));
+			}
 
-			Instruction load = inject.createLoadForTypeIndex(instructions, type, index);
-			instructions.addInstruction(insertPos++, load);
+			int index = deobMethod.isStatic() ? 0 : 1; // current variable index
 
-			index += type.getSize();
-		}
+			for (int i = index; i < signature.size(); ++i)
+			{
+				Type type = signature.getTypeOfArg(i);
+
+				Instruction load = inject.createLoadForTypeIndex(instructions, type, index);
+				instructions.addInstruction(insertPos++, load);
 
-		// Invoke callback
-		InvokeStatic invoke = new InvokeStatic(instructions,
-			new net.runelite.asm.pool.Method(
-				new net.runelite.asm.pool.Class(HOOKS),
-				hookName,
-				signature
-			)
-		);
+				index += type.getSize();
+			}
 
-		instructions.addInstruction(insertPos++, invoke);
+			// Invoke callback
+			InvokeStatic invoke = new InvokeStatic(instructions,
+				new net.runelite.asm.pool.Method(
+					new net.runelite.asm.pool.Class(HOOKS),
+					hookName,
+					signature
+				)
+			);
+
+			instructions.addInstruction(insertPos++, invoke);
+		}
 
 		logger.info("Injected method hook {} in {} with {} args: {}",
 			hookName, vanillaMethod, signature.size(),
 			signature.getArguments());
 	}
 
-	private int findHookLocation(Annotation hook, Method vanillaMethod) throws InjectionException
+	private List<Integer> findHookLocations(Annotation hook, Method vanillaMethod) throws InjectionException
 	{
 		Instructions instructions = vanillaMethod.getCode().getInstructions();
 
@@ -159,20 +171,21 @@ private int findHookLocation(Annotation hook, Method vanillaMethod) throws Injec
 			List<Instruction> returns = instructions.getInstructions().stream()
 				.filter(i -> i instanceof ReturnInstruction)
 				.collect(Collectors.toList());
+			List<Integer> indexes = new ArrayList<>();
 
-			if (returns.size() != 1)
+			for (Instruction ret : returns)
 			{
-				throw new InjectionException("returns != 1");
+				int idx = instructions.getInstructions().indexOf(ret);
+				assert idx != -1;
+				indexes.add(idx);
 			}
 
-			int idx = instructions.getInstructions().indexOf(returns.get(0));
-			assert idx != -1;
-			return idx;
+			return indexes;
 		}
 
 		if (!vanillaMethod.getName().equals("<init>"))
 		{
-			return 0;
+			return Arrays.asList(0);
 		}
 
 		// Find index after invokespecial
@@ -182,7 +195,7 @@ private int findHookLocation(Annotation hook, Method vanillaMethod) throws Injec
 
 			if (in.getType() == InstructionType.INVOKESPECIAL)
 			{
-				return i + 1; // one after
+				return Arrays.asList(i + 1); // one after
 			}
 		}
 

runescape-client-injector-plugin/src/test/java/net/runelite/injector/InjectHookMethodTest.java

@@ -26,7 +26,8 @@
 
 import java.io.IOException;
 import java.io.InputStream;
-import java.util.Optional;
+import java.util.List;
+import java.util.stream.Collectors;
 import net.runelite.asm.ClassFile;
 import net.runelite.asm.ClassGroup;
 import net.runelite.asm.ClassUtil;
@@ -46,8 +47,16 @@
 
 abstract class Obfuscated
 {
-	public void foo(Obfuscated o, int i)
+	public int foo(Obfuscated o, int i)
 	{
+		if (i > 0)
+		{
+			return i;
+		}
+		else
+		{
+			return -i;
+		}
 	}
 }
 
@@ -56,12 +65,12 @@ class Actor
 {
 	@ObfuscatedName("foo")
 	@ObfuscatedSignature(
-		signature = "(LObfuscated;I)V"
+		signature = "(LObfuscated;I)I"
 	)
-	@Hook("test")
-	public void bar(Actor actor, int i)
+	@Hook(value = "test", end = true)
+	public int bar(Actor actor, int i)
 	{
-
+		throw new IllegalStateException();
 	}
 }
 
@@ -74,15 +83,15 @@ public void testProcess() throws IOException, InjectionException
 		InputStream in = getClass().getResourceAsStream("Actor.class");
 		ClassFile cf = ClassUtil.loadClass(in);
 		cf.setName("Actor");
-		cf.findMethod("bar").setDescriptor(new Signature("(LActor;I)V"));
+		cf.findMethod("bar").setDescriptor(new Signature("(LActor;I)I"));
 
 		ClassGroup deobfuscated = new ClassGroup();
 		deobfuscated.addClass(cf);
 
 		in = getClass().getResourceAsStream("Obfuscated.class");
 		ClassFile obcf = ClassUtil.loadClass(in);
 		obcf.setName("Obfuscated");
-		obcf.findMethod("foo").setDescriptor(new Signature("(LObfuscated;I)V"));
+		obcf.findMethod("foo").setDescriptor(new Signature("(LObfuscated;I)I"));
 
 		ClassGroup obfuscated = new ClassGroup();
 		obfuscated.addClass(obcf);
@@ -97,14 +106,15 @@ public void testProcess() throws IOException, InjectionException
 		method = obcf.findMethod("foo");
 		assert method != null;
 		Code code = method.getCode();
-		Optional<InvokeStatic> invokeIns = code.getInstructions().getInstructions().stream()
+		List<InvokeStatic> invokeIns = code.getInstructions().getInstructions().stream()
 			.filter(i -> i instanceof InvokeStatic)
 			.map(i -> (InvokeStatic) i)
 			.filter(i -> i.getMethod().getClazz().getName().equals(HOOKS))
-			.findFirst();
-		assertTrue(invokeIns.isPresent());
+			.collect(Collectors.toList());
+		assertTrue(!invokeIns.isEmpty());
+		assertEquals(2, invokeIns.size());
 
-		InvokeStatic invokeStatic = invokeIns.get();
+		InvokeStatic invokeStatic = invokeIns.get(0);
 		Signature signature = invokeStatic.getMethod().getType();
 		assertEquals(3, signature.size()); // this + patamers