gistfile1.java

import java.lang.invoke.CallSite;
import java.lang.invoke.ConstantCallSite;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodHandles.Lookup;
import java.lang.invoke.MethodType;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

import jdk.internal.org.objectweb.asm.ClassReader;
import jdk.internal.org.objectweb.asm.ClassVisitor;
import jdk.internal.org.objectweb.asm.ClassWriter;
import jdk.internal.org.objectweb.asm.Handle;
import jdk.internal.org.objectweb.asm.MethodVisitor;
import jdk.internal.org.objectweb.asm.Opcodes;
import jdk.internal.org.objectweb.asm.Type;

@SuppressWarnings("restriction")
public class AnonymousLoader extends ClassLoader {
  private final Lookup hostLookup;
  final Map<String, Object> cpPatches;
  private final String anonClassName;
  private final Class<?> anonClass;

  private AnonymousLoader(Lookup hostLookup, Map<String, Object> cpPatches, String anonClassName, byte[] data) {
    this.hostLookup = hostLookup;
    this.cpPatches = cpPatches;
    this.anonClassName = anonClassName;
    this.anonClass = createClass(anonClassName, data);
  }

  public static Class<?> defineAnonymousClass(Lookup hostLookup, byte[] data, Map<String, Object> cpPatches) {
    String className = hostLookup.lookupClass().getName() + /*'\\' +*/ COUNTER.getAndIncrement();
    return new AnonymousLoader(hostLookup, cpPatches, className, data).anonClass;
  }
  
  /*@Override
  protected java.lang.Class<?> findClass(String name) throws ClassNotFoundException {
    if (!name.equals(anonClassName)) {
      throw new ClassNotFoundException(name);
    }
    return anonClass;
  }*/
  
  private Class<?> createClass(String anonClassName, byte[] data) {
    ClassReader reader = new ClassReader(data);
    ClassWriter writer = new ClassWriter(reader, 0);
    
    String hostOwner = hostLookup.lookupClass().getName().replace('.', '/');
    reader.accept(new ClassVisitor(Opcodes.ASM5, writer) {
      @Override
      public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
        super.visit(version, access, anonClassName.replace('.',  '/'), signature, superName, interfaces);
      }
      @Override
      public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
        MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
        return new MethodVisitor(Opcodes.ASM5, mv) {
          @Override
          public void visitLdcInsn(Object cst) {
            if (cpPatches.containsKey(cst)) {
              super.visitInvokeDynamicInsn("const", "()Ljava/lang/Object;", BSM_CONST, cst);
              return;
            }
            super.visitLdcInsn(cst);
          }
          @Override
          public void visitFieldInsn(int opcode, String owner, String name, String desc) {
            if (owner.equals(hostOwner)) {
              String fieldDesc = fieldDesc(opcode, owner, desc);
              super.visitInvokeDynamicInsn(name, fieldDesc, BSM, getTag(opcode), Type.getObjectType(owner));
              return;
            }
            super.visitFieldInsn(opcode, owner, name, desc);
          }
          @Override
          public void visitMethodInsn(int opcode, String owner, String name, String desc, boolean itf) {
            if (!name.equals("<init>") && !name.equals("<clinit>") && owner.equals(hostOwner)) {
              super.visitInvokeDynamicInsn(name, methodDesc(opcode, owner, desc), BSM, getTag(opcode), Type.getObjectType(owner));
              return;
            }
            super.visitMethodInsn(opcode, owner, name, desc, itf);
          }
        };
      }
    }, 0);
    byte[] bytecode = writer.toByteArray();
    
    //CheckClassAdapter.verify(new ClassReader(bytecode), true, new PrintWriter(System.out));
    
    return defineClass(anonClassName, bytecode, 0, bytecode.length);
  }
  
  static String fieldDesc(int opcode, String owner, String desc) {
    switch(opcode) {
    case Opcodes.GETFIELD:
      return "(L" + owner + ";)" + desc; 
    case Opcodes.PUTFIELD:
      return "(L" + owner + ';' + desc + ")V"; 
    case Opcodes.GETSTATIC:
      return "()" + desc; 
    default: //case Opcodes.PUTSTATIC:
      return '(' + desc + ")V"; 
    }
  }
  
  static String methodDesc(int opcode, String owner, String desc) {
    switch(opcode) {
    case Opcodes.INVOKEINTERFACE:
    case Opcodes.INVOKESPECIAL:
    case Opcodes.INVOKEVIRTUAL:
      return "(L" + owner + ';' + desc.substring(1); 
    default: //case Opcodes.INVOKESTATIC:
      return desc;
    }
  }
  
  static int getTag(int opcode) {
    switch(opcode) {
    case Opcodes.INVOKEINTERFACE:
      return Opcodes.H_INVOKEINTERFACE;
    case Opcodes.INVOKESPECIAL:
      return Opcodes.H_INVOKESPECIAL;
    case Opcodes.INVOKESTATIC:
      return Opcodes.H_INVOKESTATIC;
    case Opcodes.INVOKEVIRTUAL:
      return Opcodes.H_INVOKEVIRTUAL;
    case Opcodes.GETFIELD:
      return Opcodes.H_GETFIELD;
    case Opcodes.PUTFIELD:
      return Opcodes.H_PUTFIELD;
    case Opcodes.GETSTATIC:
      return Opcodes.H_GETSTATIC;
    case Opcodes.PUTSTATIC:
      return Opcodes.H_PUTSTATIC;
    default:
      throw new AssertionError();
    }
  }
  
  private static MethodHandle retarget(Lookup hostLookup, int tag, Class<?> declaringClass, String name, MethodType type) throws ReflectiveOperationException {
    switch(tag) {
    case Opcodes.H_INVOKEINTERFACE:
    case Opcodes.H_INVOKEVIRTUAL:
      return hostLookup.findVirtual(declaringClass, name, type.dropParameterTypes(0, 1));
    case Opcodes.H_INVOKESPECIAL:
      return hostLookup.findSpecial(declaringClass, name, type.dropParameterTypes(0, 1), hostLookup.lookupClass());
    case Opcodes.H_INVOKESTATIC:
      return hostLookup.findStatic(declaringClass, name, type);
    case Opcodes.H_GETFIELD:
      return hostLookup.findGetter(declaringClass, name, type.returnType());
    case Opcodes.H_PUTFIELD:
      return hostLookup.findSetter(declaringClass, name, type.parameterType(1));
    case Opcodes.H_GETSTATIC:
      return hostLookup.findStaticGetter(declaringClass, name, type.returnType());
    case Opcodes.H_PUTSTATIC:
      return hostLookup.findStaticSetter(declaringClass, name, type.parameterType(0));
    default:
      throw new AssertionError();
    }
  }
  
  static final Handle BSM, BSM_CONST;
  static {
    String loaderInternalName = AnonymousLoader.class.getSimpleName().replace('.', '/');
    BSM = new Handle(Opcodes.H_INVOKESTATIC,
        loaderInternalName,
        "bsm_call",
        MethodType.methodType(CallSite.class, Lookup.class, String.class, MethodType.class,
            int.class, Class.class).toMethodDescriptorString());
    BSM_CONST = new Handle(Opcodes.H_INVOKESTATIC,
        loaderInternalName,
        "bsm_const",
        MethodType.methodType(CallSite.class, Lookup.class, String.class, MethodType.class,
            String.class).toMethodDescriptorString());
  }
  static final AtomicInteger COUNTER = new AtomicInteger();
  
  public static CallSite bsm_call(Lookup lookup, String name, MethodType methodType, int tag, Class<?> owner) throws Throwable {
    AnonymousLoader loader = (AnonymousLoader)lookup.lookupClass().getClassLoader();
    MethodHandle target = retarget(loader.hostLookup, tag, owner, name, methodType);
    return new ConstantCallSite(target);
  }
  
  public static CallSite bsm_const(Lookup lookup, String name, MethodType methodType, String constant) {
    AnonymousLoader loader = (AnonymousLoader)lookup.lookupClass().getClassLoader();
    return new ConstantCallSite(MethodHandles.constant(Object.class, loader.cpPatches.get(constant)));
  }

  // ------------------ test
  public static class A {
    private static int foo = 666;
    
    static Lookup lookup() {
      return MethodHandles.lookup();
    }
  }
  
  public static class B {
    public static int foo() {
      System.out.println((Object)"<PLACEHOLDER>");
      return 0;
    }
  }
  
  public static void main(String[] args) throws Throwable {
    ClassReader reader = new ClassReader(B.class.getResourceAsStream(B.class.getName().replace('.', '/') + ".class"));
    ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS|ClassWriter.COMPUTE_FRAMES);
    reader.accept(new ClassVisitor(Opcodes.ASM5, writer) {
      @Override
      public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
        MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
        if (!"foo".equals(name)) {
          return mv;
        }
        return new MethodVisitor(Opcodes.ASM5, mv) {
          @Override
          public void visitInsn(int opcode) {
            if (opcode == Opcodes.ICONST_0) {
              super.visitFieldInsn(Opcodes.GETSTATIC, "AnonymousLoader$A", "foo", "I");
              return;
            }
            super.visitInsn(opcode);
          }
        };
      }
    }, 0);
    
    //CheckClassAdapter.verify(new ClassReader(writer.toByteArray()), true, new PrintWriter(System.out));
    
    HashMap<String, Object> cpPatches = new HashMap<>();
    cpPatches.put("<PLACEHOLDER>", new Integer(111));
    Class<?> clazz = AnonymousLoader.defineAnonymousClass(A.lookup(), writer.toByteArray(), cpPatches);
    
    MethodHandle mh = MethodHandles.publicLookup().findStatic(clazz, "foo", MethodType.methodType(int.class));
    System.out.println((int)mh.invokeExact());
  }
}