Search code examples
javabytecodejavassist

Modify a method declared in the superclass with Javassist


I'm trying to modify a method with CtMethod#insertBefore that is declared in the superclass. However, it seems to not be possible with Javassist.

private class AbstractTestDataSource {
    public Connection getConnection() throws SQLException {
        return connection;
    }
}

private class TestDataSource extends AbstractTestDataSource implements DataSource { 
    public Connection getConnection(String username, String password) throws SQLException {
        return connection;
    }
    // other methods omitted
}

This is my ClassFileTransformer

public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                                  ProtectionDomain protectionDomain, byte[] classfileBuffer)
        throws Exception {
    if (!className.equals("org/example/TestDataSource")) {
        return classfileBuffer;
    }
    final CtClass ctClass = createCtClass(loader, classfileBuffer);
    for (CtMethod method : ctClass.getMethods()) {
        if (method.getName().equals("getConnection")) {
            System.out.print(method.getName());
            System.out.println(method.getSignature());
            method.insertBefore("System.out.println(\"foo\");");
        }
    }
    return ctClass.toBytecode();
}

When I'm calling the getConnection(String, String) method, foo is printed to the console, but if I call the getConnection() method that is declared in AbstractTestDataSource nothing happens. What am I doing wrong?

Edit

I can confirm that both methods are instrumented because this is what is printed to the console:

getConnection(Ljava/lang/String;Ljava/lang/String;)Ljava/sql/Connection;
getConnection()Ljava/sql/Connection;

Solution

  • My solution is to check whether the getConnection method is declared in a class other than the current one.

    if (!ctClass.equals(method.getDeclaringClass())) {
        method = overrideMethod(ctClass, method);
    }
    

    If so, I create (and thus override) the getConnection method and delegate to the superclass.

    private CtMethod overrideMethod(CtClass ctClass, CtMethod getConnectionMethodOfSuperclass)
            throws NotFoundException, CannotCompileException {
        final CtMethod m = CtNewMethod.delegator(getConnectionMethodOfSuperclass, ctClass);
        ctClass.addMethod(m);
        return m;
    }
    

    It doesn't feel like the ideal solution, but it is working fine.