From dac2510de75839c7617c2fb648905541496494e0 Mon Sep 17 00:00:00 2001
From: zaaarf <zaaarf@proton.me>
Date: Wed, 29 Mar 2023 11:59:06 +0200
Subject: [PATCH] fix: map descriptors too

---
 .../tools/containers/FieldContainer.java      |  5 +-
 .../tools/containers/MethodContainer.java     |  5 +-
 .../tools/obfuscation/ObfuscationMapper.java  | 51 +++++++++++++++++++
 3 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/src/main/java/ftbsc/lll/processor/tools/containers/FieldContainer.java b/src/main/java/ftbsc/lll/processor/tools/containers/FieldContainer.java
index 23cd0c6..0b39a1c 100644
--- a/src/main/java/ftbsc/lll/processor/tools/containers/FieldContainer.java
+++ b/src/main/java/ftbsc/lll/processor/tools/containers/FieldContainer.java
@@ -63,11 +63,12 @@ public class FieldContainer {
 				throw new AmbiguousDefinitionException("Cannot use name-based lookups for fields of unverifiable classes!");
 			this.elem = null;
 			this.name = name;
-			this.descriptor = descriptor;
+			this.descriptor = mapper == null ? descriptor : mapper.obfuscateMethodDescriptor(descriptor);
 		} else {
 			this.elem = (VariableElement) findMember(parent, name, descriptor, descriptor != null, true);
 			this.name = this.elem.getSimpleName().toString();
-			this.descriptor = descriptorFromType(this.elem.asType());
+			String validatedDescriptor = descriptorFromType(this.elem.asType());
+			this.descriptor = mapper == null ? descriptor : mapper.obfuscateMethodDescriptor(validatedDescriptor);
 		}
 		this.nameObf = findMemberName(parent.fqnObf, name, descriptor, mapper);
 	}
diff --git a/src/main/java/ftbsc/lll/processor/tools/containers/MethodContainer.java b/src/main/java/ftbsc/lll/processor/tools/containers/MethodContainer.java
index 8b54a93..d8ab04f 100644
--- a/src/main/java/ftbsc/lll/processor/tools/containers/MethodContainer.java
+++ b/src/main/java/ftbsc/lll/processor/tools/containers/MethodContainer.java
@@ -67,11 +67,12 @@ public class MethodContainer {
 				throw new AmbiguousDefinitionException("Cannot use name-based lookups for methods of unverifiable classes!");
 			this.elem = null;
 			this.name = name;
-			this.descriptor = descriptor;
+			this.descriptor = mapper == null ? descriptor : mapper.obfuscateMethodDescriptor(descriptor);
 		} else {
 			this.elem = (ExecutableElement) findMember(parent, name, descriptor, descriptor != null && strict, false);
 			this.name = this.elem.getSimpleName().toString();
-			this.descriptor = descriptorFromExecutableElement(this.elem);
+			String validatedDescriptor = descriptorFromExecutableElement(this.elem);
+			this.descriptor = mapper == null ? descriptor : mapper.obfuscateMethodDescriptor(validatedDescriptor);
 		}
 		this.nameObf = findMemberName(parent.fqnObf, name, descriptor, mapper);
 	}
diff --git a/src/main/java/ftbsc/lll/processor/tools/obfuscation/ObfuscationMapper.java b/src/main/java/ftbsc/lll/processor/tools/obfuscation/ObfuscationMapper.java
index 0c42983..e4e11c4 100644
--- a/src/main/java/ftbsc/lll/processor/tools/obfuscation/ObfuscationMapper.java
+++ b/src/main/java/ftbsc/lll/processor/tools/obfuscation/ObfuscationMapper.java
@@ -2,6 +2,8 @@ package ftbsc.lll.processor.tools.obfuscation;
 
 import ftbsc.lll.exceptions.AmbiguousDefinitionException;
 import ftbsc.lll.exceptions.MappingNotFoundException;
+import ftbsc.lll.tools.DescriptorBuilder;
+import org.objectweb.asm.Type;
 
 import java.util.HashMap;
 import java.util.List;
@@ -93,6 +95,55 @@ public class ObfuscationMapper {
 		return member;
 	}
 
+	/**
+	 * Obfuscates a method descriptor, replacing its class references
+	 * with their obfuscated counterparts.
+	 * @param descriptor a {@link String} containing the descriptor
+	 * @return the obfuscated descriptor
+	 * @since 0.5.1
+	 */
+	public String obfuscateMethodDescriptor(String descriptor) {
+		Type method = Type.getMethodType(descriptor);
+		Type[] arguments = method.getArgumentTypes();
+		Type returnType = method.getReturnType();
+
+		Type[] obfArguments = new Type[arguments.length];
+		for(int i = 0; i < obfArguments.length; i++)
+			obfArguments[i] = this.obfuscateType(arguments[i]);
+
+		return Type.getMethodDescriptor(this.obfuscateType(returnType), obfArguments);
+	}
+
+	/**
+	 * Given a {@link Type} it returns its obfuscated counterpart.
+	 * @param type the type in question
+	 * @return the obfuscated type
+	 * @since 0.5.1
+	 */
+	public Type obfuscateType(Type type) {
+		//unwrap arrays
+		Type unwrapped = type;
+		int arrayLevel = 0;
+		while(unwrapped.getSort() == org.objectweb.asm.Type.ARRAY) {
+			unwrapped = unwrapped.getElementType();
+			arrayLevel++;
+		}
+
+		//if it's a primitive no operation is needed
+		if(type.getSort() < org.objectweb.asm.Type.ARRAY)
+			return type;
+
+		String internalName = type.getInternalName();
+
+		String internalNameObf;
+		try {
+			internalNameObf = this.obfuscateClass(internalName);
+			return Type.getType(DescriptorBuilder.nameToDescriptor(internalNameObf, arrayLevel));
+		} catch(MappingNotFoundException e) {
+			return type;
+		}
+	}
+
 	/**
 	 * Gets the unobfuscated name of the given member.
 	 * Due to how it's implemented, it's considerably less efficient than its