Groovy AST: Writing an annotation

Sandbox

In my last post I explained how to run a script in a sandbox, and how the DSL implementation can regain full access by wrapping calls inside an AccessController.doPrivileged call. For example:

String javaClassVersion() {
    AccessController.doPrivileged({
        System.getProperty("java.class.version")
    } as PrivilegedAction)
}

This doesn't look very nice though. Even though we have already simplified it using Groovy's asType operator to convert a Closure into a PrivilegedAction, it is still too much boilerplate code.

Helper method

We can create our own doPriviliged helper method like this:

String javaClassVersion() {
    doPrivileged {
        System.getProperty("java.class.version")
    }
}

static doPrivileged(Closure cl) {
    AccessController.doPrivileged(cl as PrivilegedAction)
}

But the new helper method should be in some helper class and we need to remember that class whenever we need to use the method.

Also, AccessController.doPrivileged is overloaded. It has 4 variants that accept either a PrivilegedAction or a PrivilegedExceptionAction, and optionally an AccessControlContext. We don't need to pass an AccessControlContext in our case, because we want to have full access, so we can skip supporting those overloads.

However, we do need to know if we should cast the closure to a PrivilegedAction or a PrivilegedExceptionAction. When our code can throw a checked exception, we should use a PrivilegedExceptionAction, catch PrivilegedActionException, and rethrow the wrapped checked exception. For example:

String readFile(String path) throws IOException {
    try {
        AccessController.doPrivileged({
            new File(path).text
        } as PrivilegedExceptionAction)
    } catch (PrivilegedActionException e) {
        throw e.cause
    }
}

We would need another helper method for this, with a different name, and we would need to know which one to call. Or we should always use PrivilegedExceptionAction.

Annotation

It would be better if we could simply tag some code to be privileged and do all of this automatically. For this, we can use an annotation:

@Privileged
String javaClassVersion() {
    System.getProperty("java.class.version")
}

This tells the compiler to wrap the code inside the annotated method to be run inside an AccessController.doPrivileged call. It should also handle checked exceptions, so we can write the other example like this:

@Privileged
String readFile(String path) throws IOException {
    new File(path).text
}

Unfortunately, there is no such annotation in the GDK yet. Which gives us the opportunity to write our own and learn about Groovy AST transformations!

Groovy local AST transformation

Groovy has a nice feature that allows us to hook into the Groovy compiler and transform Groovy code before it is compiled to bytecode. It is called Groovy AST transformations, and has support for global and local transformations. For our annotation, we need to use a local AST transformation. 

The first step is to write the annotation interface. Here is the full source:

package blogsample

import org.codehaus.groovy.transform.GroovyASTTransformationClass
import java.lang.annotation.*

@Retention(RetentionPolicy.SOURCE)
@Target([ElementType.METHOD])
@GroovyASTTransformationClass(classes = [PrivilegedASTTransformation])
@interface Privileged {
}

As you can see, the annotation is itself annotated by multiple annotations. The most important one is @GroovyASTTransformationClass,  which points to our to be made PrivilegedASTTransformation class. 

It's important to put the annotation interface and AST transformation class into a separate module that is compiled before code that uses it, as pointed out by Christian Oestreich in his blog (issue 1). This is because the transformation is applied at compile time during the CANONICALIZATION phase. Also make sure to explicitly recompile the code that uses the @Privileged annotation when testing changes in the transformation class, else the transformation will not be triggered.

Before we're going to implement anything, we first need to write unit tests. I will leave this as an exercise for the reader, to keep this post length limited.

Now we're ready to implement the transformation class. An initial skeleton will look like this:

@GroovyASTTransformation(phase = CompilePhase.CANONICALIZATION)
class PrivilegedASTTransformation implements ASTTransformation {
    private static final ClassNode MY_TYPE = ClassHelper.make(Privileged)

    @Override
    void visit(ASTNode[] nodes, SourceUnit sourceUnit) {
        if (nodes.length != 2 ||
                !(nodes[0] instanceof AnnotationNode) ||
                !(nodes[1] instanceof AnnotatedNode)) {
            throw new RuntimeException(
                    "Internal error: " +
                    "expecting [AnnotationNode, AnnotatedNode] " +
                    "but got: $nodes")
        }

        AnnotatedNode parent = (AnnotatedNode) nodes[1]
        AnnotationNode node = (AnnotationNode) nodes[0]
        if (MY_TYPE == node.classNode &&
                parent instanceof MethodNode &&
                parent.code instanceof BlockStatement) {
            transformMethod(parent)
        }
    }

    private void transformMethod(MethodNode methodNode) {
        // todo
    }
}

The visit method is called by the Groovy compiler for each method annotated with @Privileged. Inside it we verify if we are called in the right way. We assume the code inside the method is a BlockStatement, and we'll make sure we're gonna replace it with a new BlockStatement. 

Now we can actually start to do some transformation work, inside transformMethod.

Transformation implementation

Recall that we are going to wrap the method code inside a call to one of the AccessController.doPrivileged overloads.

The first step is to check which overload to use. If the method has a throws clause, we shall assume it is throwing checked exceptions, in which case we have to use a PrivilegedExceptionAction. Otherwise, we can use PrivilegedAction:

private void transformMethod(MethodNode methodNode) {
    Statement s
    if (methodNode.exceptions) {
        s = createPrivilegedExceptionActionStatement(methodNode.code)

    } else {
        s = createPrivilegedActionStatement(methodNode.code)

    }
    methodNode.code = new BlockStatement([s], null)
}

Let's start with the simplest case, when there are no checked exceptions. It should be transformed into something like this:

AccessController.doPrivileged({-> <code> } as PrivilegedAction)


Here's the code to do it:

private Statement createPrivilegedActionStatement(Statement code) {
    createDoPrivilegedStatement(
            wrapInClosureAsExpression(
                    code, PrivilegedAction))
}

private Statement createDoPrivilegedStatement(Expression actionExpression) {
    new ExpressionStatement(
            new StaticMethodCallExpression(
                    ClassHelper.make(AccessController),
                    'doPrivileged',
                    new ArgumentListExpression(actionExpression)))
}

private Expression wrapInClosureAsExpression(Statement code, Class asClass) {
    // Construct following code: {-> <code> } as <asClass>
    CastExpression.asExpression(
            ClassHelper.make(asClass),
            wrapInClosureExpression(code))
}

private ClosureExpression wrapInClosureExpression(Statement code) {
    // Construct following code: {-> <code> }
    new ClosureExpression(Parameter.EMPTY_ARRAY, code)
}

To make the code readable and reusable, each transformation part is inside a separate method. Some of these will be quite generic and can be moved to a generic helper class. Note that CastExpression has a factory method to create an as expression.

For the exceptions case, we need to do a bit more. Normally, in Java, we'd need to catch a PrivilegedActionException, cast the wrapped exception to the checked exception, and rethrow it. If there are multiple checked exceptions in the throws clause, we'd need to explicitly check and cast each exception type.

Fortunately, this is Groovy, and we can just rethrow the exception cause without checking or casting any of the exception types. The transformed code should look like this:

try {
     AccessController.doPrivileged({-> <code> } as PrivilegedExceptionAction)
} catch (PrivilegedActionException e) {
     throw e.cause
}

Here is the transformation code:

private Statement createPrivilegedExceptionActionStatement(Statement code) {
    def actionExpression = createPrivilegedActionExpression(
            code, PrivilegedExceptionAction)
    def doPrivilegedStatement = createDoPrivilegedStatement(actionExpression)
    def exceptionType = ClassHelper.make(PrivilegedActionException)
    def exceptionParameter = new Parameter(exceptionType, 'e')
    def throwStatement = new ThrowStatement(
            new MethodCallExpression(
                    new VariableExpression(exceptionParameter),
                    'getCause',
                    MethodCallExpression.NO_ARGUMENTS))
    return createTryCatchStatement(
            doPrivilegedStatement,
            new CatchStatement(exceptionParameter, throwStatement))
}

private Statement createTryCatchStatement(Statement tryStatement,
                                          CatchStatement... catchStatements) {
    def finallyStatement = new EmptyStatement()
    def tryCatchStatement = new TryCatchStatement(tryStatement, finallyStatement)
    catchStatements.each { tryCatchStatement.addCatch(it) }
    return tryCatchStatement
}

VariableScope

We are almost finished with our AST transformation class. The last part is the most tricky one, and it is about variable scopes. Classes like MethodNode, BlockStatement and ClosureExpression have a VariableScope field, which is primarily used to handle access of variables from closures that are declared outside the closure.

Since we have wrapped the original method code inside a closure, any parameters of the method should become accessible from the closure. This means we somehow have to change the existing and new variable scopes to make this work. This proved to be quite tricky, and luckily I found a simple solution on the aforementioned blog of Christian Oestreich (issue 2).

To fix the variable scopes we will have to apply a VariableScopeVisitor on the source unit. We can simply add this in the visit method, after the transformMethod call:

transformMethod(parent)
VariableScopeVisitor scopeVisitor = new VariableScopeVisitor(sourceUnit)
sourceUnit.AST.classes.each {
    scopeVisitor.visitClass(it)
}

This will magically make any parameters of the method accessible from the new closure!

Summary

We have created a method annotation that will wrap the method's code inside a closure and pass it to some other code. This is almost the same as using the & operator on the method to convert it to a MethodClosure, except that we're keeping everything local to the method.

With the @Privileged annotation, all the method's code will become privileged. For security reasons, care should be taken to make the method body as small as possible.