From ef875ced1a6534d0213bfa417b207fa89985f7e9 Mon Sep 17 00:00:00 2001 From: Lenni0451 <20379977+Lenni0451@users.noreply.github.com> Date: Sat, 22 Jun 2024 19:12:14 +0200 Subject: [PATCH] Added method mapper --- .../lenni0451/reflect/proxy/ProxyBuilder.java | 75 +++++++++++++------ .../reflect/proxy/ProxyMethodBuilder.java | 18 +++-- .../lenni0451/reflect/proxy/ProxyUtils.java | 17 +++++ .../reflect/proxy/impl/ProxyRuntime.java | 4 +- .../lenni0451/reflect/proxy/ProxyTest.java | 41 ++++++++++ 5 files changed, 123 insertions(+), 32 deletions(-) diff --git a/src/main/java/net/lenni0451/reflect/proxy/ProxyBuilder.java b/src/main/java/net/lenni0451/reflect/proxy/ProxyBuilder.java index 1ecf42f..e171996 100644 --- a/src/main/java/net/lenni0451/reflect/proxy/ProxyBuilder.java +++ b/src/main/java/net/lenni0451/reflect/proxy/ProxyBuilder.java @@ -16,6 +16,7 @@ import java.lang.reflect.Field; import java.lang.reflect.Method; import java.util.Arrays; +import java.util.function.Function; import java.util.function.Predicate; import static net.lenni0451.reflect.bytecode.BytecodeUtils.*; @@ -35,11 +36,20 @@ public class ProxyBuilder { @Nullable private Class[] interfaces; private Predicate methodFilter = m -> true; + private Function methodMapper = Function.identity(); private InvocationHandler invocationHandler = InvocationHandler.forwarding(); private ProxyClassDefiner classDefiner = ProxyClassDefiner.loader(ProxyBuilder.class.getClassLoader()); private Class proxyClass; + /** + * @return The super class of the proxy class + */ + @Nullable + public Class getSuperClass() { + return this.superClass; + } + /** * Set the super class of the proxy class.
* The super class must be public and not final. @@ -56,11 +66,11 @@ public ProxyBuilder setSuperClass(@Nullable final Class superClass) { } /** - * @return The super class of the proxy class + * @return The interfaces of the proxy class */ @Nullable - public Class getSuperClass() { - return this.superClass; + public Class[] getInterfaces() { + return this.interfaces; } /** @@ -101,15 +111,16 @@ public ProxyBuilder setInterfaces(@Nullable final Class... interfaces) { } /** - * @return The interfaces of the proxy class + * @return The current method filter */ - @Nullable - public Class[] getInterfaces() { - return this.interfaces; + public Predicate getMethodFilter() { + return this.methodFilter; } /** - * Set the filter for the methods that should be overridden by the proxy class. + * Set the filter for the methods that should be overridden by the proxy class.
+ * Filtered methods will not be overridden by the proxy class.
+ * Abstract methods will always be overridden and cannot be filtered. * * @param methodFilter The method filter * @return This builder @@ -121,21 +132,23 @@ public ProxyBuilder setMethodFilter(@Nonnull final Predicate methodFilte } /** - * @return The current method filter + * @return The current method mapper */ - public Predicate getMethodFilter() { - return this.methodFilter; + public Function getMethodMapper() { + return this.methodMapper; } /** - * Set the invocation handler for the proxy class. + * Set the mapper for the methods that should be overridden by the proxy class.
+ * The mapper can be used to change the owner of an overridden method.
+ * The owner must be a super class of the original owner. * - * @param invocationHandler The invocation handler + * @param methodMapper The method mapper * @return This builder */ - public ProxyBuilder setInvocationHandler(@Nonnull final InvocationHandler invocationHandler) { + public ProxyBuilder setMethodMapper(@Nonnull final Function methodMapper) { this.reset(); - this.invocationHandler = invocationHandler; + this.methodMapper = methodMapper; return this; } @@ -147,14 +160,14 @@ public InvocationHandler getInvocationHandler() { } /** - * Set the class definer for the proxy class. + * Set the invocation handler for the proxy class. * - * @param classDefiner The class definer + * @param invocationHandler The invocation handler * @return This builder */ - public ProxyBuilder setClassDefiner(@Nonnull final ProxyClassDefiner classDefiner) { + public ProxyBuilder setInvocationHandler(@Nonnull final InvocationHandler invocationHandler) { this.reset(); - this.classDefiner = classDefiner; + this.invocationHandler = invocationHandler; return this; } @@ -165,6 +178,18 @@ public ProxyClassDefiner getClassDefiner() { return this.classDefiner; } + /** + * Set the class definer for the proxy class. + * + * @param classDefiner The class definer + * @return This builder + */ + public ProxyBuilder setClassDefiner(@Nonnull final ProxyClassDefiner classDefiner) { + this.reset(); + this.classDefiner = classDefiner; + return this; + } + /** * Build the proxy class. * @@ -173,21 +198,22 @@ public ProxyClassDefiner getClassDefiner() { public ProxyClass build() { if (this.proxyClass == null) { Reference methodsReference = new Reference<>(); - BuiltClass builtClass = this.buildClass(methodsReference); + Reference originalMethodsReference = new Reference<>(); + BuiltClass builtClass = this.buildClass(methodsReference, originalMethodsReference); this.proxyClass = this.classDefiner.defineProxyClass(builtClass, this.superClass, this.interfaces); //Set the static fields Field methods = Fields.getDeclaredField(this.proxyClass, METHODS_FIELD); Fields.setObject(null, methods, methodsReference.value); - Class[] proxyMethodClasses = this.buildProxyMethodClasses(methodsReference.value); + Class[] proxyMethodClasses = this.buildProxyMethodClasses(methodsReference.value, originalMethodsReference.value); Field proxyMethodClassesField = Fields.getDeclaredField(this.proxyClass, PROXY_METHOD_CLASSES_FIELD); Fields.setObject(null, proxyMethodClassesField, proxyMethodClasses); } return new ProxyClass(this.proxyClass, this.invocationHandler); } - private BuiltClass buildClass(final Reference methodsReference) { + private BuiltClass buildClass(final Reference methodsReference, final Reference originalMethodsReference) { String className = "net/lenni0451/reflect/proxy/ProxyImpl$" + System.nanoTime(); Class[] interfaces = this.interfaces; if (interfaces == null) { @@ -208,6 +234,7 @@ private BuiltClass buildClass(final Reference methodsReference) { Method[] methods = ProxyUtils.getOverridableMethod(this.superClass, this.interfaces, this.methodFilter); methodsReference.value = methods; + originalMethodsReference.value = ProxyUtils.mapMethods(methods, this.methodMapper); this.addFields(cb, methods); this.addMethods(cb, methods); this.addDefaultMethods(cb); @@ -329,10 +356,10 @@ private void addDefaultMethods(final ClassBuilder cb) { }); } - private Class[] buildProxyMethodClasses(final Method[] methods) { + private Class[] buildProxyMethodClasses(final Method[] methods, final Method[] originalMethods) { Class[] proxyMethodClasses = new Class[methods.length]; for (int i = 0; i < methods.length; i++) { - proxyMethodClasses[i] = ProxyMethodBuilder.buildProxyMethodClass(this.proxyClass, methods[i]); + proxyMethodClasses[i] = ProxyMethodBuilder.buildProxyMethodClass(this.proxyClass, methods[i], originalMethods[i]); } return proxyMethodClasses; } diff --git a/src/main/java/net/lenni0451/reflect/proxy/ProxyMethodBuilder.java b/src/main/java/net/lenni0451/reflect/proxy/ProxyMethodBuilder.java index c8b39f6..c7e2d53 100644 --- a/src/main/java/net/lenni0451/reflect/proxy/ProxyMethodBuilder.java +++ b/src/main/java/net/lenni0451/reflect/proxy/ProxyMethodBuilder.java @@ -7,6 +7,7 @@ import net.lenni0451.reflect.proxy.impl.ProxyMethod; import net.lenni0451.reflect.proxy.impl.ProxyRuntime; import org.jetbrains.annotations.ApiStatus; +import org.jetbrains.annotations.Nullable; import java.lang.invoke.MethodHandle; import java.lang.reflect.Method; @@ -22,14 +23,14 @@ class ProxyMethodBuilder { private static final BytecodeBuilder BUILDER = BytecodeBuilder.get(); - public static Class buildProxyMethodClass(final Class proxyClass, final Method method) { + public static Class buildProxyMethodClass(final Class proxyClass, final Method method, @Nullable final Method originalMethod) { BuiltClass builtClass = BUILDER.class_(BUILDER.opcode("ACC_PUBLIC"), slash(proxyClass) + "$ProxyMethodImpl", null, slash(Object.class), new String[]{slash(ProxyMethod.class)}, cb -> { addFields(proxyClass, cb); - addStaticBlock(method, cb); + addStaticBlock(method, originalMethod, cb); addConstructor(proxyClass, cb); addMethodGetter(cb); addInvokeWith(method, cb); - addInvokeSuper(proxyClass, method, cb); + addInvokeSuper(proxyClass, originalMethod == null ? method : originalMethod, cb); addCancel(method, cb); }); @@ -43,10 +44,15 @@ private static void addFields(final Class proxyClass, final ClassBuilder cb) cb.field(BUILDER.opcode("ACC_PRIVATE", "ACC_FINAL"), "method", desc(Method.class), null, null); } - private static void addStaticBlock(final Method method, final ClassBuilder cb) { + private static void addStaticBlock(final Method method, @Nullable final Method originalMethod, final ClassBuilder cb) { cb.method(BUILDER.opcode("ACC_PUBLIC", "ACC_STATIC"), "", mdesc(void.class), null, null, mb -> { + mb.typeLdc(BUILDER, method.getDeclaringClass()); //Instance class + if (originalMethod != null) { + mb.typeLdc(BUILDER, originalMethod.getDeclaringClass()); //Super class + } else { + mb.insn(BUILDER.opcode("DUP")); + } mb - .typeLdc(BUILDER, method.getDeclaringClass()) //Instance class .ldc(method.getName()) //Method name .intPush(BUILDER, method.getParameterCount()) .type(BUILDER.opcode("ANEWARRAY"), slash(Class.class)); //Parameter array @@ -60,7 +66,7 @@ private static void addStaticBlock(final Method method, final ClassBuilder cb) { } mb .typeLdc(BUILDER, method.getReturnType()) //Return type - .method(BUILDER.opcode("INVOKESTATIC"), slash(ProxyRuntime.class), "getMethodHandles", mdesc(MethodHandle[].class, Class.class, String.class, Class[].class, Class.class), false); + .method(BUILDER.opcode("INVOKESTATIC"), slash(ProxyRuntime.class), "getMethodHandles", mdesc(MethodHandle[].class, Class.class, Class.class, String.class, Class[].class, Class.class), false); mb //Store the method handles .insn(BUILDER.opcode("DUP")) .intPush(BUILDER, 0) diff --git a/src/main/java/net/lenni0451/reflect/proxy/ProxyUtils.java b/src/main/java/net/lenni0451/reflect/proxy/ProxyUtils.java index 221ed15..33077fa 100644 --- a/src/main/java/net/lenni0451/reflect/proxy/ProxyUtils.java +++ b/src/main/java/net/lenni0451/reflect/proxy/ProxyUtils.java @@ -12,6 +12,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.Function; import java.util.function.Predicate; /** @@ -64,4 +65,20 @@ public static void getOverridableMethod(final Class clazz, final Map inter : clazz.getInterfaces()) getOverridableMethod(inter, methods); } + public static Method[] mapMethods(final Method[] methods, final Function methodMapper) { + Method[] originalMethods = new Method[methods.length]; + for (int i = 0; i < methods.length; i++) { + Method method = methods[i]; + Method mappedMethod = methodMapper.apply(method); + if (method.equals(mappedMethod)) continue; + + if (!mappedMethod.getDeclaringClass().isAssignableFrom(method.getDeclaringClass())) { + throw new IllegalArgumentException("The method '" + method.getName() + "' in class '" + method.getDeclaringClass().getName() + "' is not assignable to the method in the proxy class"); + } + methods[i] = mappedMethod; + originalMethods[i] = method; + } + return originalMethods; + } + } diff --git a/src/main/java/net/lenni0451/reflect/proxy/impl/ProxyRuntime.java b/src/main/java/net/lenni0451/reflect/proxy/impl/ProxyRuntime.java index 688e3ca..ff4eb92 100644 --- a/src/main/java/net/lenni0451/reflect/proxy/impl/ProxyRuntime.java +++ b/src/main/java/net/lenni0451/reflect/proxy/impl/ProxyRuntime.java @@ -28,11 +28,11 @@ public class ProxyRuntime { * @throws NoSuchMethodException If the method does not exist * @throws IllegalAccessException If the method is not accessible */ - public static MethodHandle[] getMethodHandles(final Class owner, final String name, final Class[] parameters, final Class returnType) throws NoSuchMethodException, IllegalAccessException { + public static MethodHandle[] getMethodHandles(final Class owner, final Class superOwner, final String name, final Class[] parameters, final Class returnType) throws NoSuchMethodException, IllegalAccessException { MethodHandle[] methodHandles = new MethodHandle[2]; methodHandles[0] = TRUSTED_LOOKUP.findVirtual(owner, name, MethodType.methodType(returnType, parameters)); try { - methodHandles[1] = TRUSTED_LOOKUP.findSpecial(owner, name, MethodType.methodType(returnType, parameters), owner); + methodHandles[1] = TRUSTED_LOOKUP.findSpecial(superOwner, name, MethodType.methodType(returnType, parameters), superOwner); } catch (Throwable ignored) { } return methodHandles; diff --git a/src/test/java/net/lenni0451/reflect/proxy/ProxyTest.java b/src/test/java/net/lenni0451/reflect/proxy/ProxyTest.java index 84852ea..491ae29 100644 --- a/src/test/java/net/lenni0451/reflect/proxy/ProxyTest.java +++ b/src/test/java/net/lenni0451/reflect/proxy/ProxyTest.java @@ -1,10 +1,13 @@ package net.lenni0451.reflect.proxy; +import net.lenni0451.reflect.Methods; import net.lenni0451.reflect.proxy.impl.Proxy; import net.lenni0451.reflect.proxy.test.*; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import java.lang.reflect.Method; + import static org.junit.jupiter.api.Assertions.*; class ProxyTest { @@ -99,4 +102,42 @@ void testSuperImplementsProxy() { assertDoesNotThrow(instance::get); } + @Test + void testInvokeWrongObject() { + ProxyClass proxyClass = new ProxyBuilder() + .setSuperClass(Class2.class) + .setInvocationHandler((thiz, proxyMethod, args) -> proxyMethod.invokeWith(new Class1(), args)) + .build(); + Class2 proxy = proxyClass.allocateInstance(); + //The method found by the proxy builder belongs to Class2 (the proxy super class) + //Invoking it with an instance of Class1 should throw an exception because Class1 does not extend Class2, even tho it is the other way round + //The cast is done by the proxy to allow calling invokeExact of the method handle + assertThrows(ClassCastException.class, proxy::test); + } + + @Test + void testMethodMapper() { + ProxyClass proxyClass = new ProxyBuilder() + .setSuperClass(Class2.class) + .setMethodMapper(method -> { + Method class1Method = Methods.getDeclaredMethod(Class1.class, method.getName(), method.getParameterTypes()); + if (class1Method != null) return class1Method; + return method; + }) + .setInvocationHandler((thiz, proxyMethod, args) -> proxyMethod.invokeWith(new Class1(), args)) + .build(); + Class2 proxy = proxyClass.allocateInstance(); + //The invocation is the same as in testInvokeWrongObject but this time the method mapper maps the method to one owned by Class1 + //This should work because Class2 extends Class1 and the correct method should be invoked + assertEquals(1, proxy.test()); + + ((Proxy) proxy).setInvocationHandler(InvocationHandler.forwarding()); + //Here it should return 2 because the super method is still the one from Class2 + assertEquals(2, proxy.test()); + + ((Proxy) proxy).setInvocationHandler((thiz, proxyMethod, args) -> proxyMethod.invokeWith(new Class2(), args)); + //This should also work because it's forwarded to Class2 + assertEquals(2, proxy.test()); + } + }