fix(java): References within InvocationHandler, Issue #1364

java/fury-core/src/main/java/org/apache/fury/serializer/JdkProxySerializer.java

@@ -23,13 +23,39 @@
 import java.lang.reflect.Proxy;
 import org.apache.fury.Fury;
 import org.apache.fury.memory.MemoryBuffer;
+import org.apache.fury.resolver.RefResolver;
+import org.apache.fury.util.GraalvmSupport;
+import org.apache.fury.util.Platform;
 import org.apache.fury.util.Preconditions;
 import org.apache.fury.util.ReflectionUtils;
 
 /** Serializer for jdk {@link Proxy}. */
 @SuppressWarnings({"rawtypes", "unchecked"})
 public class JdkProxySerializer extends Serializer {
 
+  // Make offset compatible with graalvm native image.
+  private static final long PROXY_HANDLER_FIELD_OFFSET;
+
+  static {
+    if (GraalvmSupport.isGraalBuildtime()) {
+      try {
+        // Make offset compatible with graalvm native image.
+        PROXY_HANDLER_FIELD_OFFSET = Platform.objectFieldOffset(Proxy.class.getDeclaredField("h"));
+      } catch (NoSuchFieldException e) {
+        throw new RuntimeException(e);
+      }
+    } else {
+      // not all JVM implementations use 'h' as internal InvocationHandler name
+      PROXY_HANDLER_FIELD_OFFSET =
+          ReflectionUtils.getFieldOffset(Proxy.class, InvocationHandler.class);
+    }
+  }
+
+  private static final InvocationHandler STUB_HANDLER =
+      (proxy, method, args) -> {
+        throw new IllegalStateException("Deserialization stub handler still active");
+      };
+
   public JdkProxySerializer(Fury fury, Class cls) {
     super(fury, cls);
     if (cls != ReplaceStub.class) {
@@ -39,17 +65,22 @@ public JdkProxySerializer(Fury fury, Class cls) {
 
   @Override
   public void write(MemoryBuffer buffer, Object value) {
-    fury.writeRef(buffer, Proxy.getInvocationHandler(value));
     fury.writeRef(buffer, value.getClass().getInterfaces());
+    fury.writeRef(buffer, Proxy.getInvocationHandler(value));
   }
 
   @Override
   public Object read(MemoryBuffer buffer) {
-    InvocationHandler invocationHandler = (InvocationHandler) fury.readRef(buffer);
-    Preconditions.checkNotNull(invocationHandler);
+    final RefResolver resolver = fury.getRefResolver();
+    final int refId = resolver.lastPreservedRefId();
     final Class<?>[] interfaces = (Class<?>[]) fury.readRef(buffer);
     Preconditions.checkNotNull(interfaces);
-    return Proxy.newProxyInstance(fury.getClassLoader(), interfaces, invocationHandler);
+    Object proxy = Proxy.newProxyInstance(fury.getClassLoader(), interfaces, STUB_HANDLER);
+    resolver.setReadObject(refId, proxy);
+    InvocationHandler invocationHandler = (InvocationHandler) fury.readRef(buffer);
+    Preconditions.checkNotNull(invocationHandler);
+    Platform.putObject(proxy, PROXY_HANDLER_FIELD_OFFSET, invocationHandler);
+    return proxy;
   }
 
   public static class ReplaceStub {}

java/fury-core/src/main/java/org/apache/fury/util/ReflectionUtils.java

@@ -348,6 +348,40 @@ public static long getFieldOffset(Class<?> cls, String fieldName) {
     return getFieldOffset(field);
   }
 
+  /**
+   * Gets the offset of the only field in the class matching the required type.
+   *
+   * @param clazz the class in which the field should be declared
+   * @param fieldType the required type of the field
+   * @return offset of the field
+   * @throws IllegalStateException if there are multiple fields of the required type
+   * @throws IllegalArgumentException if there are no fields of the required type
+   */
+  public static long getFieldOffset(Class<?> clazz, Class<?> fieldType) {
+    Field f = null;
+    for (Field fi : clazz.getDeclaredFields()) {
+      if (fieldType.equals(fi.getType())) {
+        if (f != null) {
+          throw new IllegalStateException(
+              "Found multiple field s matching "
+                  + fieldType
+                  + " in "
+                  + clazz
+                  + ": "
+                  + f
+                  + " and "
+                  + fi);
+        }
+        f = fi;
+      }
+    }
+    if (f == null) {
+      throw new IllegalArgumentException(
+          "Found no field matching " + fieldType.getName() + " in " + clazz + "!");
+    }
+    return Platform.objectFieldOffset(f);
+  }
+
   public static long getFieldOffsetChecked(Class<?> cls, String fieldName) {
     long offset = getFieldOffset(cls, fieldName);
     Preconditions.checkArgument(offset != -1);

java/fury-core/src/main/resources/META-INF/native-image/org.apache.fury/fury-core/native-image.properties

@@ -23,6 +23,7 @@ Args=--initialize-at-build-time=org.apache.fury.memory.MemoryBuffer,\
     org.apache.fury.serializer.collection.SynchronizedSerializers$Offset,\
     org.apache.fury.serializer.collection.CollectionSerializers$ArraysAsListSerializer,\
     org.apache.fury.serializer.collection.MapSerializers$EnumMapSerializer,\
+    org.apache.fury.serializer.JdkProxySerializer,\
     org.apache.fury.serializer.StringSerializer$Offset,\
     org.apache.fury.serializer.StringSerializer,\
     org.apache.fury.serializer.Serializers,\

java/fury-core/src/test/java/org/apache/fury/serializer/JdkProxySerializerTest.java

@@ -57,4 +57,49 @@ public void testJdkProxy(boolean referenceTracking) {
     Function deserializedFunction = (Function) fury.deserialize(fury.serialize(function));
     assertEquals(deserializedFunction.apply(null), 1);
   }
+
+  private static class RefTestInvocationHandler implements InvocationHandler, Serializable {
+
+    private Function proxy;
+
+    @Override
+    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+      if (method.getName().equals("equals")) {
+        return args[0] == this.proxy;
+      }
+      return "Hello world from "
+          + (proxy == null
+              ? "null"
+              : proxy.getClass().getName() + "@" + System.identityHashCode(proxy));
+    }
+
+    private void setProxy(Function myProxy) {
+      this.proxy = myProxy;
+    }
+
+    private Function getProxy() {
+      return proxy;
+    }
+  }
+
+  @Test
+  public void testJdkProxyRef() {
+    Fury fury =
+        Fury.builder()
+            .withLanguage(Language.JAVA)
+            .withRefTracking(true)
+            .requireClassRegistration(false)
+            .build();
+    RefTestInvocationHandler hdlr = new RefTestInvocationHandler();
+    Function function =
+        (Function)
+            Proxy.newProxyInstance(fury.getClassLoader(), new Class[] {Function.class}, hdlr);
+    hdlr.setProxy(function);
+    assertEquals(hdlr.getProxy(), function);
+
+    Function deserializedFunction = (Function) fury.deserialize(fury.serialize(function));
+    RefTestInvocationHandler deserializedHandler =
+        (RefTestInvocationHandler) Proxy.getInvocationHandler(deserializedFunction);
+    assertEquals(deserializedHandler.getProxy(), deserializedFunction);
+  }
 }