Demo for Groovy CompilerConfiguration and SecureASTCustomiser

demo_DemoDSLCompilerConfigProvider.java

package fr.zami.demo;

import groovy.util.DelegatingScript;
import org.codehaus.groovy.ast.ASTNode;
import org.codehaus.groovy.ast.expr.ArgumentListExpression;
import org.codehaus.groovy.ast.expr.BinaryExpression;
import org.codehaus.groovy.ast.expr.BooleanExpression;
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.StaticMethodCallExpression;
import org.codehaus.groovy.ast.expr.TupleExpression;
import org.codehaus.groovy.ast.expr.VariableExpression;
import org.codehaus.groovy.ast.stmt.BlockStatement;
import org.codehaus.groovy.ast.stmt.ExpressionStatement;
import org.codehaus.groovy.ast.stmt.IfStatement;
import org.codehaus.groovy.ast.stmt.ReturnStatement;
import org.codehaus.groovy.ast.stmt.Statement;
import org.codehaus.groovy.control.CompilerConfiguration;
import org.codehaus.groovy.control.customizers.CompilationCustomizer;
import org.codehaus.groovy.control.customizers.ImportCustomizer;
import org.codehaus.groovy.control.customizers.SecureASTCustomizer;

import java.lang.reflect.Method;
import java.math.BigDecimal;
import java.util.*;
import java.util.stream.Collectors;

import static java.util.Arrays.asList;
import static java.util.Collections.singletonList;
import static org.codehaus.groovy.syntax.Types.*;

class DemoDSLCompilerConfigProvider {

    private static final List<String> STATIC_STAR_IMPORTS = singletonList("java.lang.Math");
    private static final Set<String> MY_DSL_FUNCTIONS = getMyDSLFunctionNames();
    private static final Set<Class<?>> ALLOWED_EXPRESSION_TYPES = new HashSet<>(asList(
            MethodCallExpression.class,
            VariableExpression.class,
            ArgumentListExpression.class,
            BinaryExpression.class,
            ConstantExpression.class,
            BooleanExpression.class,
            StaticMethodCallExpression.class,
            TupleExpression.class
    ));
    private static final Set<Class<?>> ALLOWED_STATEMENT_TYPES = new HashSet<>(asList(
            BlockStatement.class,
            ExpressionStatement.class,
            IfStatement.class
    ));

    private DemoDSLCompilerConfigProvider() {}

    static CompilerConfiguration createCompilerConfiguration() {
        ImportCustomizer imports = new ImportCustomizer();
        imports.addStaticStars(STATIC_STAR_IMPORTS.toArray(new String[0]));
        CompilerConfiguration compilerConfiguration = new CompilerConfiguration();
        compilerConfiguration.addCompilationCustomizers(imports);
        compilerConfiguration.addCompilationCustomizers(createSecureASTCustomizer());
        compilerConfiguration.setScriptBaseClass(DelegatingScript.class.getName());
        return compilerConfiguration;
    }

    private static CompilationCustomizer createSecureASTCustomizer() {
        SecureASTCustomizer secureASTCustomizer = new SecureASTCustomizer();
        secureASTCustomizer.setClosuresAllowed(false);
        secureASTCustomizer.setMethodDefinitionAllowed(false);
        secureASTCustomizer.setImportsWhitelist(Collections.emptyList());
        secureASTCustomizer.setStaticImportsWhitelist(STATIC_STAR_IMPORTS);
        secureASTCustomizer.setTokensWhitelist(asList(
                EQUAL,
                PLUS,
                MINUS,
                MULTIPLY,
                DIVIDE,
                MOD,
                POWER,
                PLUS_PLUS,
                MINUS_MINUS,
                COMPARE_EQUAL,
                COMPARE_NOT_EQUAL,
                COMPARE_LESS_THAN,
                COMPARE_LESS_THAN_EQUAL,
                COMPARE_GREATER_THAN,
                COMPARE_GREATER_THAN_EQUAL,
                REGEX_PATTERN,
                FIND_REGEX,
                MATCH_REGEX,
                LOGICAL_AND,
                LOGICAL_OR,
                NOT
        ));
        secureASTCustomizer.setConstantTypesClassesWhiteList(asList(
                Object.class, String.class,
                Boolean.class, Integer.class, Float.class, Long.class, Double.class, BigDecimal.class,
                Boolean.TYPE, Integer.TYPE, Long.TYPE, Float.TYPE, Double.TYPE
        ));
        secureASTCustomizer.setReceiversClassesWhiteList(asList(
                Object.class, String.class, DemoDSLHelper.class,
                Math.class, Integer.class, Float.class, Double.class, Long.class, BigDecimal.class
        ));
        DSLASTChecker checker = new DSLASTChecker();
        secureASTCustomizer.addExpressionCheckers(checker);
        secureASTCustomizer.addStatementCheckers(checker);
        return secureASTCustomizer;
    }

    private static Set<String> getMyDSLFunctionNames() {
        Method[] methods = DemoDSLHelper.class.getDeclaredMethods();
        // Here you could also filter for methods with a particular signature, or annotation, etc.
        return Arrays.stream(methods).map(Method::getName).collect(Collectors.toSet());
    }

    private static class DSLASTChecker
            implements SecureASTCustomizer.ExpressionChecker, SecureASTCustomizer.StatementChecker {

        @Override
        public boolean isAuthorized(Expression expression) {
            Class<? extends Expression> expressionClass = expression.getClass();
            if (!ALLOWED_EXPRESSION_TYPES.contains(expressionClass)) {
                throw new DSLSyntaxError("Expression of type '"+expressionClass.getSimpleName()+"' is not allowed", expression);
            }
            if (expression instanceof MethodCallExpression) {
                MethodCallExpression methodCallExpression = (MethodCallExpression) expression;
                Expression method = methodCallExpression.getMethod();
                if (method instanceof ConstantExpression) {
                    ConstantExpression constantExpressionMethod = (ConstantExpression) method;
                    String methodName = constantExpressionMethod.getValue().toString();
                    if (!MY_DSL_FUNCTIONS.contains(methodName)) {
                        throw new DSLSyntaxError("Unknown function '" + methodName + "'", expression);
                    }
                }
            }
            return true;
        }

        @Override
        public boolean isAuthorized(Statement statement) {
            if (ALLOWED_STATEMENT_TYPES.contains(statement.getClass())) {
                return true;
            }
            if (statement instanceof ReturnStatement) {
                ReturnStatement returnStatement = (ReturnStatement) statement;
                boolean isImplicitReturn = returnStatement.getLineNumber() == -1;
                if (isImplicitReturn && returnStatement.isReturningNullOrVoid()) {
                    return true;
                }
                else {
                    throw new DSLSyntaxError("Cannot return from DemoDSL script", statement);
                }
            }
            throw new DSLSyntaxError("Invalid statement", statement);
        }

    }

    static class DSLSyntaxError extends RuntimeException {
        public DSLSyntaxError(String message, ASTNode node) {
            super(message + " (line " + node.getLineNumber() + ", column " + node.getColumnNumber() + ")");
        }
    }
}

demo_DemoDSLHelper.java

package fr.zami.demo;

class DemoDSLHelper {

    String hello(String name) {
        return "Hi " + name + "!";
    }

    int giveMeFive() {
        return 5;
    }

    void exit() {
        System.out.println("*** Exit called ***");
    }
}

demo_DemoDSLRunner.java

package fr.zami.demo;

import groovy.lang.GroovyShell;
import groovy.lang.Script;
import groovy.util.DelegatingScript;

public class DemoDSLRunner {

    public static void main(String[] args) {
        DemoDSLHelper delegate = new DemoDSLHelper();
        GroovyShell shell = new GroovyShell(DemoDSLCompilerConfigProvider.createCompilerConfiguration());
        Script script = shell.parse("hello(name) + ' ' + sin(giveMeFive())");
        ((DelegatingScript)script).setDelegate(delegate);
        script.setProperty("name", "Bob");
        Object result = script.run();
        System.out.println(result);

        script = shell.parse("System.exit(0)");
        ((DelegatingScript)script).setDelegate(delegate);
        script.run();
    }
}

First script demonstrates that one can execute DSL functions, use expressions with +, use methods from Math class, etc.

Second script shows an attempt at executing System.exit(0) from DSL is denied.

Output:

Hi Bob! -0.9589242746631385
Exception in thread "main" org.codehaus.groovy.control.MultipleCompilationErrorsException: startup failed:
General error during canonicalization: Method calls not allowed on [java.lang.System]

java.lang.SecurityException: Method calls not allowed on [java.lang.System]
	at org.codehaus.groovy.control.customizers.SecureASTCustomizer$SecuringCodeVisitor.visitMethodCallExpression(SecureASTCustomizer.java:857)
	at org.codehaus.groovy.ast.expr.MethodCallExpression.visit(MethodCallExpression.java:66)
	at org.codehaus.groovy.control.customizers.SecureASTCustomizer$SecuringCodeVisitor.visitExpressionStatement(SecureASTCustomizer.java:779)
	at org.codehaus.groovy.ast.stmt.ExpressionStatement.visit(ExpressionStatement.java:42)
	at org.codehaus.groovy.control.customizers.SecureASTCustomizer$SecuringCodeVisitor.visitBlockStatement(SecureASTCustomizer.java:739)
	at org.codehaus.groovy.ast.stmt.BlockStatement.visit(BlockStatement.java:71)
	at org.codehaus.groovy.control.customizers.SecureASTCustomizer.call(SecureASTCustomizer.java:554)
	at org.codehaus.groovy.control.CompilationUnit.applyToPrimaryClassNodes(CompilationUnit.java:1065)
	at org.codehaus.groovy.control.CompilationUnit.doPhaseOperation(CompilationUnit.java:603)
	at org.codehaus.groovy.control.CompilationUnit.processPhaseOperations(CompilationUnit.java:581)
	at org.codehaus.groovy.control.CompilationUnit.compile(CompilationUnit.java:558)
	at groovy.lang.GroovyClassLoader.doParseClass(GroovyClassLoader.java:298)
	at groovy.lang.GroovyClassLoader.parseClass(GroovyClassLoader.java:268)
	at groovy.lang.GroovyShell.parseClass(GroovyShell.java:688)
	at groovy.lang.GroovyShell.parse(GroovyShell.java:700)
	at groovy.lang.GroovyShell.parse(GroovyShell.java:736)
	at groovy.lang.GroovyShell.parse(GroovyShell.java:727)
	at fr.zami.demo.DemoDSLRunner.main(DemoDSLRunner.java:18)```