Using Javassist to insert try/finally logic that wraps the original method logic

1.9k views Asked by At

I have a Java 8 project that rewrites bytecode at runtime using Javassist 3.20.0-GA via a custom Java Agent. The goal is to instrument the method such that the original body is wrapped by a try/finally block with print statements. For example, given this trivial method:

public class TimeService {

    public long getCurrentTime(){

        long time = 0;
        try {
            time = System.currentTimeMillis();
        }
        catch(Throwable t){
            time = -1;
        }

        System.out.println("The current time is "+time);

        return time;
    }
}

I would like to output the following modified code (without the comments):

public class TimeService {

    public long getCurrentTime(){
        try {
            System.out.println("Start");

            // BEGIN original code
            long time = 0;
            try {
                time = System.currentTimeMillis();
            }
            catch(Throwable t){
                time = -1;
            }

            System.out.println("The current time is "+time);

            return time;
            // END original code
        }
        finally {
            System.out.println("Stop");
        }
    }
}

First I tried the following code:

CtClass ctClass =  ClassPool.getDefault().get("com.example.test.TimeService");

CtMethod ctMethod = ctClass.getDeclaredMethod("getCurrentTime");

ctMethod.insertBefore("System.out.println(\"Start\");");
ctMethod.insertAfter("System.out.println(\"Stop\");", true);

This appeared to work well for simple void methods without preexisting try/catch/finally blocks, but the output for the getCurrentTime() method was strange, and the System.out.println("Stop") statement was repeated twice:

public long getCurrentTime() {
    boolean var10 = false;

    long var10000;
    long var5;
    try {
        var10 = true;
        System.out.println("Start");
        long time = 0L;

        try {
            time = System.currentTimeMillis();
        } catch (Throwable var11) {
            time = -1L;
        }

        System.out.println("The current time is " + time);
        var10000 = time;
        var10 = false;
    } finally {
        if(var10) {
            var5 = 0L;
            System.out.println("Stop");
        }
    }

    var5 = var10000;
    System.out.println("Stop");
    return var5;
}

The above code "works" when there isn't an error, but pretend System.currentTimeMillis() throws an Exception, then var10 would never be set to false, and thusly the System.out.println("Stop") statement would be executed twice. Using a boolean flag as a safeguard doesn't work well here, hence I would prefer to simply insert a try/finally at the very beginning and end of the method.

Next I tried instrumenting and replacing the method directly like this:

CtClass ctClass = ClassPool.getDefault().get("com.example.test.TimeService");

CtMethod ctMethod = ctClass.getDeclaredMethod("getCurrentTime");

ctMethod.instrument(new ExprEditor(){
    public void edit(MethodCall m) throws CannotCompileException {
        String start = "{ System.out.println(\"Start\"); }";
        String stop = "{ System.out.println(\"Stop\"); }";

        m.replace("{ try {"+start+" $_ = $proceed($$); } finally { "+stop+" } }");
    }
});

But that turned into a garbled mess that repeated the "Start" and "Stop" printouts several times across many try/catch/finally blocks (too messy to paste here).

I'm not sure what to try next, or if I'm using the wrong part of the Javassist API. It seems like it should be very straightforward--and it is, with simple void methods--but when values are returned, especially in combination with preexisting try/catch/finally blocks the output becomes unpredictable.

Any thoughts or workarounds?

Thanks in advance.

1

There are 1 answers

2
light_keeper On

The generated getCurrentTime is correct. You enter if(var10) statement only if exception is thrown. But in such case you will never reach the code below finally statement. finally does not catch exceptions, it just execute some code and throws further

Generated code is pretty weird but it eventually does what it is supposed to do. Message "Stop" will be printed exactly once.