How do I replace a class with new one with Java instrumentation?

Question

I need to create a java agent that when is enabled it gets the path to a jar file as argument and then it replaces any loaded class the the one inside the jar file if their names are matched.

For example, we have an application with a class called com.something.ClassTest. Now if the mentioned jar (is not in the class path) has a class exactly the same name as com.something.ClassTest, I want to replace it with the one in the jar.

I have this class transformer but not sure if that's correct or not. I get IOException with message Class not found.

    @Override
public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

    if(classNames.contains(className.replace("/", "."))) {
        System.out.format("\n==> Found %s \n", className);
        try {
            Class c = urlClassLoader.loadClass(className.replace("/", "."));
            InputStream is = urlClassLoader.getResourceAsStream(className.replace("/", "."));
            System.out.println("Loaded class " + c);

            ClassReader reader = new ClassReader(is);
            ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
            byte[] content = writer.toByteArray();

            System.out.println("Redifned " + new String(content));
            System.out.println("Orig " + new String(classfileBuffer));
            ClassDefinition cd = new ClassDefinition(c, content);
            instrumentation.redefineClasses(cd);

            return content;
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (UnmodifiableClassException e) {
            e.printStackTrace();
        }

    }
    return classfileBuffer;
}

The error I get is at the line where the ClassReader is instantiated. I guess the error is because the urlClassloader is somehow under the hierarchy of the current class loader... but I don't know how else I can do it.

Here is the code where URL class loaded is initialized

    public SimpleClassTransformer(Instrumentation instrumentation, String jarFileName) {

    this.jarFileName = jarFileName;

    if(jarFileName != null) {
        JarFile jarFile = null;
        try {
            jarFile = new JarFile(this.jarFileName);
            Enumeration e = jarFile.entries();

            System.out.println("Jar file: " + this.jarFileName);
            URL[] urls = { new URL("jar:file:" + this.jarFileName+"!/") };
            urlClassLoader = URLClassLoader.newInstance(urls);

            while (e.hasMoreElements()) {
                JarEntry je = (JarEntry) e.nextElement();
                if(je.isDirectory() || !je.getName().endsWith(".class")){
                    continue;
                }
                // -6 because of .class
                String jarClassName = je.getName().substring(0,je.getName().length()-6);
                jarClassName = jarClassName.replace('/', '.');
                System.out.println("Adding class " + jarClassName);
                this.classNames.add(jarClassName);

            }
        } catch (IOException e) {
            e.printStackTrace();
        }

    }

}

Instrumentation is set while instantiating the transformer in the premain() method. I am trying to avoid using Javassist. Can you please help me on this.

This is the exception I get:

java.io.IOException: Class not found
at jdk.internal.org.objectweb.asm.ClassReader.readClass(ClassReader.java:484)
at jdk.internal.org.objectweb.asm.ClassReader.<init>(ClassReader.java:453)
at com.agent.SimpleClassTransformer.transform(SimpleClassTransformer.java:79)
at sun.instrument.TransformerManager.transform(TransformerManager.java:188)
at sun.instrument.InstrumentationImpl.transform(InstrumentationImpl.java:428)
at java.lang.ClassLoader.defineClass1(Native Method)
at java.lang.ClassLoader.defineClass(ClassLoader.java:760)
at java.security.SecureClassLoader.defineClass(SecureClassLoader.java:142)
at java.net.URLClassLoader.defineClass(URLClassLoader.java:467)
at java.net.URLClassLoader.access$100(URLClassLoader.java:73)
at java.net.URLClassLoader$1.run(URLClassLoader.java:368)
at java.net.URLClassLoader$1.run(URLClassLoader.java:362)
at java.security.AccessController.doPrivileged(Native Method)
at java.net.URLClassLoader.findClass(URLClassLoader.java:361)
at java.lang.ClassLoader.loadClass(ClassLoader.java:424)
at sun.misc.Launcher$AppClassLoader.loadClass(Launcher.java:331)
at java.lang.ClassLoader.loadClass(ClassLoader.java:357)
at org.springframework.util.ClassUtils.forName(ClassUtils.java:250)
at org.springframework.beans.factory.support.AbstractBeanDefinition.resolveBeanClass(AbstractBeanDefinition.java:394)
at org.springframework.beans.factory.support.AbstractBeanFactory.doResolveBeanClass(AbstractBeanFactory.java:1397)
at org.springframework.beans.factory.support.AbstractBeanFactory.resolveBeanClass(AbstractBeanFactory.java:1344)
at org.springframework.beans.factory.support.AbstractAutowireCapableBeanFactory.determineTargetType(AbstractAutowireCapableBeanFactory.java:628)
at org.springframework.beans.factory.support.AbstractAutowireCapableBeanFactory.predictBeanType(AbstractAutowireCapableBeanFactory.java:597)
at org.springframework.beans.factory.support.AbstractBeanFactory.isFactoryBean(AbstractBeanFactory.java:1445)
at org.springframework.beans.factory.support.DefaultListableBeanFactory.doGetBeanNamesForType(DefaultListableBeanFactory.java:445)
at org.springframework.beans.factory.support.DefaultListableBeanFactory.getBeanNamesForType(DefaultListableBeanFactory.java:415)
at org.springframework.boot.autoconfigure.web.DispatcherServletAutoConfiguration$DefaultDispatcherServletCondition.checkServlets(DispatcherServletAutoConfiguration.java:141)
at org.springframework.boot.autoconfigure.web.DispatcherServletAutoConfiguration$DefaultDispatcherServletCondition.getMatchOutcome(DispatcherServletAutoConfiguration.java:131)
at org.springframework.boot.autoconfigure.condition.SpringBootCondition.matches(SpringBootCondition.java:47)
at org.springframework.context.annotation.ConditionEvaluator.shouldSkip(ConditionEvaluator.java:102)
at org.springframework.context.annotation.ConfigurationClassParser.processConfigurationClass(ConfigurationClassParser.java:203)
at org.springframework.context.annotation.ConfigurationClassParser.processMemberClasses(ConfigurationClassParser.java:336)
at org.springframework.context.annotation.ConfigurationClassParser.doProcessConfigurationClass(ConfigurationClassParser.java:248)
at org.springframework.context.annotation.ConfigurationClassParser.processConfigurationClass(ConfigurationClassParser.java:231)
at org.springframework.context.annotation.ConfigurationClassParser.processImports(ConfigurationClassParser.java:509)
at org.springframework.context.annotation.ConfigurationClassParser.processDeferredImportSelectors(ConfigurationClassParser.java:454)
at org.springframework.context.annotation.ConfigurationClassParser.parse(ConfigurationClassParser.java:185)
at org.springframework.context.annotation.ConfigurationClassPostProcessor.processConfigBeanDefinitions(ConfigurationClassPostProcessor.java:321)
at org.springframework.context.annotation.ConfigurationClassPostProcessor.postProcessBeanDefinitionRegistry(ConfigurationClassPostProcessor.java:243)
at org.springframework.context.support.PostProcessorRegistrationDelegate.invokeBeanDefinitionRegistryPostProcessors(PostProcessorRegistrationDelegate.java:273)
at org.springframework.context.support.PostProcessorRegistrationDelegate.invokeBeanFactoryPostProcessors(PostProcessorRegistrationDelegate.java:98)
at org.springframework.context.support.AbstractApplicationContext.invokeBeanFactoryPostProcessors(AbstractApplicationContext.java:677)
at org.springframework.context.support.AbstractApplicationContext.refresh(AbstractApplicationContext.java:519)
at org.springframework.boot.context.embedded.EmbeddedWebApplicationContext.refresh(EmbeddedWebApplicationContext.java:118)
at org.springframework.boot.SpringApplication.refresh(SpringApplication.java:752)
at org.springframework.boot.SpringApplication.doRun(SpringApplication.java:347)
at org.springframework.boot.SpringApplication.run(SpringApplication.java:295)
at org.springframework.boot.SpringApplication.run(SpringApplication.java:1112)
at org.springframework.boot.SpringApplication.run(SpringApplication.java:1101)

==== EDIT ====

After fixing the issue with url class loader I now get this error that happens when Spring tries to refresh its context:

TestClass has been compiled by a more recent version of the Java Runtime (class file version 0.0), this version of the Java Runtime only recognizes class file versions up to 52.0

Answer 1

Another Approach to your problem is using endorsed dirs. Keep classes that you want to be loaded in a library and provide -Djava.endorsed.dirs=<directory_path> as JVM argument to the program.

While loading classes, JVM first checks the class availability in this directory and if not found then it will check the application classes. This works perfectly fine without any issue and without any coding.

Answer 2

I managed to fix the issue. In case someone has the same issue here was the problem:

I was using ClassReader and ClassWriter. For some reason, ClassWriter were stuffing the byte code, perhaps it was my mistake to pass already compiled class to class writer but anyway the following code:

    ClassReader reader = new ClassReader(is);
    ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
    byte[] content = writer.toByteArray();

Replaced with:

        InputStream is = urlClassLoader.getResourceAsStream(className + ".class");

        byte[] content = new byte[is.available()];
        is.read(content);

        System.out.println ("Original class version: " + ((classfileBuffer[6]&0xff)<<8 | (classfileBuffer[7]&0xff)));
        System.out.println ("Redefined class version: " + ((content[6]&0xff)<<8 | (content[7]&0xff)));

As you can see I am using InputStream to retrieve the byte code directly. That fixed the issue and in case you are interested, Spring detected the difference nicely and refreshed the context.

==== EDIT ====

I noticed using URLClassLoader here is not reliable as for some reason, it may return the already loaded class in the application itself and not the class inside the JAR file. It was random, sometimes returns the class inside the jar and sometimes the original class so I have decided to remove the URLClassLoader and instead get the class file as InputStream while traversing the jar file. This is the final code of my transformer for anyone who needs it:

public class JarFileClassTransformer implements ClassFileTransformer {

private String jarFileName = null;
protected Map<String, InputStream> classNames = new HashMap<>();
static Instrumentation instrumentation = null;

/**
 * Constructor.
 * @param jarFileName
 */
public JarFileClassTransformer(String jarFileName) {
    this.jarFileName = jarFileName;

    File file  = new File(jarFileName);
    System.out.println("Jar file '" + this.jarFileName + "' " + (file.exists() ? "exists" : "doesn't exists!"));

    if(file.exists()) {
        try {
            JarFile jarFile = new JarFile(file);
            Enumeration e = jarFile.entries();

            while (e.hasMoreElements()) {
                JarEntry je = (JarEntry) e.nextElement();
                if(je.isDirectory() || !je.getName().endsWith(".class")){
                    continue;
                }
                // -6 because of .class
                String jarClassName = je.getName().substring(0,je.getName().length()-6);
                jarClassName = jarClassName.replace('/', '.');
                System.out.println("Adding class " + jarClassName);
                this.classNames.put(jarClassName, jarFile.getInputStream(je));

            }
        } catch (IOException e) {
            e.printStackTrace();
        }

    }

}

@Override
public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

    if(classNames.containsKey(className.replace("/", "."))) {
        System.out.format("\n==> Found %s to replace with the existing version\n", className);
        try {

            Class c = loader.loadClass(className.replace("/", "."));
            System.out.println("Existing class: " + c);
            InputStream is = classNames.get(className.replace("/", "."));

            byte[] content = new byte[is.available()];
            is.read(content);

            System.out.println("Original class version: " + ((classfileBuffer[6]&0xff)<<8 | (classfileBuffer[7]&0xff)));
            System.out.println("Redefined class version: " + ((content[6]&0xff)<<8 | (content[7]&0xff)));

            System.out.println("Original bytecode: " + new String(classfileBuffer));
            System.out.println("Redefined byte code: " + new String(content));
            ClassDefinition cd = new ClassDefinition(c, content);
            instrumentation.redefineClasses(cd);

            return content;
        } catch (Throwable e) {
            e.printStackTrace();

        }

    }
    return classfileBuffer;
}

}