From 13ab94deff011c77fdb271e6f08390325d13ed42 Mon Sep 17 00:00:00 2001
From: Sergey <sergey.gorbaty@salesforce.com>
Date: Tue, 28 Feb 2017 12:48:56 -0800
Subject: [PATCH] Improving CRUD detection for Aggregate Results

---
 .../rule/security/ApexCRUDViolationRule.java  | 19 +++++++++--
 .../rule/security/xml/ApexCRUDViolation.xml   | 34 +++++++++++++++++--
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/pmd-apex/src/main/java/net/sourceforge/pmd/lang/apex/rule/security/ApexCRUDViolationRule.java b/pmd-apex/src/main/java/net/sourceforge/pmd/lang/apex/rule/security/ApexCRUDViolationRule.java
index 8ab246afbf..5c655fb3d6 100644
--- a/pmd-apex/src/main/java/net/sourceforge/pmd/lang/apex/rule/security/ApexCRUDViolationRule.java
+++ b/pmd-apex/src/main/java/net/sourceforge/pmd/lang/apex/rule/security/ApexCRUDViolationRule.java
@@ -480,6 +480,7 @@ public class ApexCRUDViolationRule extends AbstractApexRule {
 
     private void checkForAccessibility(final ASTSoqlExpression node, Object data) {
         final boolean isCount = node.getNode().getCanonicalQuery().startsWith("SELECT COUNT()");
+        final String typeFromSOQL = getTypeFromSOQLQuery(node);
 
         final HashSet<ASTMethodCallExpression> prevCalls = getPreviousMethodCalls(node);
         for (ASTMethodCallExpression prevCall : prevCalls) {
@@ -510,7 +511,7 @@ public class ApexCRUDViolationRule extends AbstractApexRule {
                     .append(":").append(type);
 
             if (!isGetter) {
-                validateCRUDCheckPresent(node, data, ANY, typeCheck.toString());
+                validateCRUDCheckPresent(node, data, ANY, typeFromSOQL == null ? typeCheck.toString() : typeFromSOQL);
             }
 
         }
@@ -523,7 +524,7 @@ public class ApexCRUDViolationRule extends AbstractApexRule {
                 if (varToTypeMapping.containsKey(variableWithClass)) {
                     String type = varToTypeMapping.get(variableWithClass);
                     if (!isGetter) {
-                        validateCRUDCheckPresent(node, data, ANY, type);
+                        validateCRUDCheckPresent(node, data, ANY, typeFromSOQL == null ? type : typeFromSOQL);
                     }
                 }
             }
@@ -533,11 +534,23 @@ public class ApexCRUDViolationRule extends AbstractApexRule {
         final ASTReturnStatement returnStatement = node.getFirstParentOfType(ASTReturnStatement.class);
         if (returnStatement != null) {
             if (!isGetter) {
-                validateCRUDCheckPresent(node, data, ANY, returnType == null ? "" : returnType);
+                String retType = typeFromSOQL == null ? returnType : typeFromSOQL;
+                validateCRUDCheckPresent(node, data, ANY, retType == null ? "" : retType);
             }
         }
     }
 
+    private String getTypeFromSOQLQuery(final ASTSoqlExpression node) {
+        final String canonQuery = node.getNode().getCanonicalQuery();
+
+        Matcher m = Pattern.compile("^[\\S|\\s]+?FROM[\\s]+?(\\S+)", Pattern.CASE_INSENSITIVE).matcher(canonQuery);
+        while (m.find()) {
+            return new StringBuffer().append(node.getNode().getDefiningType().getApexName()).append(":")
+                    .append(m.group(1)).toString();
+        }
+        return null;
+    }
+
     private String getReturnType(final ASTMethod method) {
         return new StringBuilder().append(method.getNode().getDefiningType().getApexName()).append(":")
                 .append(method.getNode().getMethodInfo().getEmitSignature().getReturnType().getApexName()).toString();
diff --git a/pmd-apex/src/test/resources/net/sourceforge/pmd/lang/apex/rule/security/xml/ApexCRUDViolation.xml b/pmd-apex/src/test/resources/net/sourceforge/pmd/lang/apex/rule/security/xml/ApexCRUDViolation.xml
index 67588ae088..14df490c81 100644
--- a/pmd-apex/src/test/resources/net/sourceforge/pmd/lang/apex/rule/security/xml/ApexCRUDViolation.xml
+++ b/pmd-apex/src/test/resources/net/sourceforge/pmd/lang/apex/rule/security/xml/ApexCRUDViolation.xml
@@ -2,6 +2,37 @@
 
 <test-data>
 
+	<test-code>
+		<description>Proper CRUD checks for Aggregate Result</description>
+		<expected-problems>0</expected-problems>
+		<code><![CDATA[ 
+public class Foo { 	
+	public void justGiveMeFoo() {
+		if (Opportunity.sObjectType.getDescribe().isAccessible()) {
+	 		return;
+	 	}
+	 	AggregateResult[] test = [SELECT Id FROM Opportunity];
+	} 
+} 
+]]></code>
+	</test-code>
+
+
+	<test-code>
+		<description>Proper CRUD checks for Aggregate Result return</description>
+		<expected-problems>0</expected-problems>
+		<code><![CDATA[ 
+public class Foo { 	
+	public AggregateResult[] justGiveMeFoo() {
+		if (Opportunity.sObjectType.getDescribe().isAccessible()) {
+	 		return null;
+	 	}
+	 	return [SELECT Id FROM Opportunity];
+	} 
+} 
+]]></code>
+	</test-code>
+
 	<test-code>
 		<description>Not a getter</description>
 		<expected-problems>1</expected-problems>
@@ -688,7 +719,7 @@ public class Foo {
 } 
 ]]></code>
 	</test-code>
-	
+
 	<test-code>
 		<description>Field detection</description>
 		<expected-problems>1</expected-problems>
@@ -703,5 +734,4 @@ public class MyProfilePageController {
 	
 ]]></code>
 	</test-code>
-	
 </test-data>