Secure Groovy Expression / Method Invocations

SecureExpressionASTTransformation.groovy

import org.codehaus.groovy.ast.ASTNode
import org.codehaus.groovy.ast.ClassCodeVisitorSupport
import org.codehaus.groovy.ast.ClassHelper
import org.codehaus.groovy.ast.expr.ArgumentListExpression
import org.codehaus.groovy.ast.expr.BinaryExpression;
import org.codehaus.groovy.ast.expr.ConstantExpression
import org.codehaus.groovy.ast.expr.Expression
import org.codehaus.groovy.ast.expr.MethodCallExpression
import org.codehaus.groovy.ast.expr.MethodPointerExpression
import org.codehaus.groovy.ast.expr.PropertyExpression
import org.codehaus.groovy.ast.expr.StaticMethodCallExpression
import org.codehaus.groovy.control.CompilePhase
import org.codehaus.groovy.control.SourceUnit
import org.codehaus.groovy.syntax.Types;
import org.codehaus.groovy.transform.ASTTransformation
import org.codehaus.groovy.transform.GroovyASTTransformationq

import com.google.common.base.Predicate
import com.google.inject.Inject
import com.google.inject.Injector
import com.google.inject.Key
import com.google.inject.name.Names


/**
 * Attempts to wrap expressions that return objects in a predicate which allows a security check
 * against unsafe object access.
 * 
 * @author sskrla
 */
@GroovyASTTransformation(phase = CompilePhase.SEMANTIC_ANALYSIS)
class SecureExpressionASTTransformation implements ASTTransformation {		
	@Override
	void visit(ASTNode[] nodes, SourceUnit source) {
		source.getAST()?.classes?.each {
			new SecurePropertyVisitor(sourceUnit: source).visitClass(it)
	    }
	}
		
	private class SecurePropertyVisitor extends ClassCodeVisitorSupport {
		def SourceUnit sourceUnit;
		
		@Override
		public void visitBinaryExpression(BinaryExpression expression) {
			super.visitBinaryExpression(expression);
			switch(expression.operation.type) {
				// Both Sides
				case Types.COMPARE_NOT_EQUAL..Types.COMPARE_TO:
				case Types.NOT..Types.LOGICAL_AND_EQUAL:
				case Types.PLUS..Types.STAR_STAR:
				case Types.BITWISE_OR..Types.BITWISE_XOR:
					expression.leftExpression = wrap(expression.leftExpression);
				
				// Just the right
				case Types.PLUS_EQUAL..Types.POWER_EQUAL:
				case Types.LEFT_SQUARE_BRACKET:
				case Types.EQUAL:
				case Types.LEFT_SHIFT:
				case Types.LEFT_SHIFT_EQUAL:
				case Types.BITWISE_OR_EQUAL..Types.BITWISE_XOR_EQUAL:
					expression.rightExpression = wrap(expression.rightExpression);
					break;
				
				// Just the left
				case Types.KEYWORD_INSTANCEOF:
				case Types.KEYWORD_AS:
					expression.leftExpression = wrap(expression.leftExpression);
			}
		}
		
		@Override
		public void visitPropertyExpression(PropertyExpression expression) {
			super.visitPropertyExpression(expression)
			expression.objectExpression = wrap(expression.objectExpression)
		}
		
		@Override
		public void visitMethodCallExpression(MethodCallExpression expression) {
			super.visitMethodCallExpression(expression);
			expression.objectExpression = wrap(expression.objectExpression)
		}
		
		@Override
		public void visitMethodPointerExpression(MethodPointerExpression expression) {
			super.visitMethodPointerExpression(expression);
			expression.expression = wrap(expression.expression)
		}	
		
		def wrap(Expression expr) {
			if(expr instanceof ConstantExpression) 
				return expr;
			
			new StaticMethodCallExpression(
				ClassHelper.make(SecureExpressionASTTransformation.class), 
				'assertValidReturn', 
				new ArgumentListExpression(expr));
		}
	}
	
	/**
	 * This portion is specific to your platform, you will need to replace it with your own security checks.
	 **/
	def static assertValidReturn(Object result) {
		Predicate<Object> predicate = ... // Get ahold of some predicate
		if(result != null && !predicate.apply(result))
			throw new IllegalAccessException("You are not allowed to access objects of type ${result.class}");
		result
	}
}