From dda7e0a0f5de7fdd676a83d482dbdd5ae4f82f7f Mon Sep 17 00:00:00 2001
From: Cedric Champeau <cedric.champeau@gmail.com>
Date: Mon, 6 Feb 2012 11:57:13 +0100
Subject: [PATCH] GROOVY-5277: SecureASTCustomizer fails to visit classes

---
 .../customizers/SecureASTCustomizer.java      | 44 +++++++++++++++----
 .../SecureASTCustomizerTest.groovy            | 31 +++++++++++++
 2 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/src/main/org/codehaus/groovy/control/customizers/SecureASTCustomizer.java b/src/main/org/codehaus/groovy/control/customizers/SecureASTCustomizer.java
index ea7d963d77..b2641c4ee6 100644
--- a/src/main/org/codehaus/groovy/control/customizers/SecureASTCustomizer.java
+++ b/src/main/org/codehaus/groovy/control/customizers/SecureASTCustomizer.java
@@ -516,10 +516,7 @@ public void call(final SourceUnit source, final GeneratorContext context, final
         if (!isPackageAllowed && ast.getPackage() != null) {
             throw new SecurityException("Package definitions are not allowed");
         }
-        final List<MethodNode> methods = ast.getMethods();
-        if (!isMethodDefinitionAllowed && methods !=null && methods.size()>0) {
-            throw new SecurityException("Method definitions are not allowed");
-        }
+        checkMethodDefinitionAllowed(classNode);
 
         // verify imports
         if (importsBlacklist != null || importsWhitelist != null || starImportsBlacklist != null || starImportsWhitelist != null) {
@@ -547,13 +544,42 @@ public void call(final SourceUnit source, final GeneratorContext context, final
 
         final SecuringCodeVisitor visitor = new SecuringCodeVisitor();
         ast.getStatementBlock().visit(visitor);
+        for (ClassNode clNode : ast.getClasses()) {
+            if (clNode!=classNode) {
+                checkMethodDefinitionAllowed(clNode);
+                for (MethodNode methodNode : clNode.getMethods()) {
+                    if (!methodNode.isSynthetic()) {
+                        methodNode.getCode().visit(visitor);
+                    }
+                }
+            }
+        }
 
-        if (isMethodDefinitionAllowed && methods !=null) {
+        List<MethodNode> methods = filterMethods(classNode);
+        if (isMethodDefinitionAllowed) {
             for (MethodNode method : methods) {
-                method.getCode().visit(visitor);
+                if (method.getDeclaringClass()==classNode) method.getCode().visit(visitor);
             }
         }
     }
+    
+    private void checkMethodDefinitionAllowed(ClassNode owner) {
+        if (isMethodDefinitionAllowed) return;
+        List<MethodNode> methods = filterMethods(owner);
+        if (!methods.isEmpty()) throw new SecurityException("Method definitions are not allowed");
+    }
+    
+    private List<MethodNode> filterMethods(ClassNode owner) {
+        List<MethodNode> result = new LinkedList<MethodNode>();
+        List<MethodNode> methods = owner.getMethods();
+        for (MethodNode method : methods) {
+            if (method.getDeclaringClass() == owner && !method.isSynthetic()) {
+                if ("main".equals(method.getName()) || "run".equals(method.getName()) && owner.isScriptBody()) continue;
+                result.add(method);
+            }
+        }
+        return result;
+    }
 
     private void assertStarImportIsAllowed(final String packageName) {
         if (starImportsWhitelist != null && !starImportsWhitelist.contains(packageName)) {
@@ -649,13 +675,13 @@ private void assertStatementAuthorized(final Statement statement) throws Securit
         private void assertExpressionAuthorized(final Expression expression) throws SecurityException {
             final Class<? extends Expression> clazz = expression.getClass();
             if (expressionsBlacklist != null && expressionsBlacklist.contains(clazz)) {
-                throw new SecurityException(clazz.getSimpleName() + "s are not allowed");
+                throw new SecurityException(clazz.getSimpleName() + "s are not allowed: " + expression.getText());
             } else if (expressionsWhitelist != null && !expressionsWhitelist.contains(clazz)) {
-                throw new SecurityException(clazz.getSimpleName() + "s are not allowed");
+                throw new SecurityException(clazz.getSimpleName() + "s are not allowed: " + expression.getText());
             }
             for (ExpressionChecker expressionChecker : expressionCheckers) {
                 if (!expressionChecker.isAuthorized(expression)) {
-                    throw new SecurityException("Expression [" + clazz.getSimpleName() + "] is not allowed");
+                    throw new SecurityException("Expression [" + clazz.getSimpleName() + "] is not allowed: " + expression.getText());
                 }
             }
             if (isIndirectImportCheckEnabled) {
diff --git a/src/test/org/codehaus/groovy/control/customizers/SecureASTCustomizerTest.groovy b/src/test/org/codehaus/groovy/control/customizers/SecureASTCustomizerTest.groovy
index 71d0e21d99..b2134f24d6 100644
--- a/src/test/org/codehaus/groovy/control/customizers/SecureASTCustomizerTest.groovy
+++ b/src/test/org/codehaus/groovy/control/customizers/SecureASTCustomizerTest.groovy
@@ -81,6 +81,37 @@ class SecureASTCustomizerTest extends GroovyTestCase {
         }
     }
 
+    void testMethodDefinitionInClass() {
+        def shell = new GroovyShell(configuration)
+        String script = """
+            class A {
+                def method() {
+                    true
+                }
+            }
+            new A()
+        """
+        shell.evaluate(script)
+        // no error means success
+        customizer.methodDefinitionAllowed = false
+        assert hasSecurityException {
+            shell.evaluate(script)
+        }
+    }
+
+    void testClassExtendingClassWithMethods() {
+        def shell = new GroovyShell(configuration)
+        String script = """
+            class A extends LinkedList {
+            }
+            new A()
+        """
+        shell.evaluate(script)
+        // no error means success
+        customizer.methodDefinitionAllowed = false
+        shell.evaluate(script)
+    }
+
     void testExpressionWhiteList() {
         def shell = new GroovyShell(configuration)
         customizer.expressionsWhitelist = [BinaryExpression, ConstantExpression]