using isAdditionalTaintStep

haby0 · haby0 · commit bfe0d4098795 · 2021-06-18T21:36:44.000+08:00
diff --git a/java/ql/src/experimental/Security/CWE/CWE-094/JShellInjection.ql b/java/ql/src/experimental/Security/CWE/CWE-094/JShellInjection.ql
@@ -12,6 +12,7 @@
 
 import java
 import JShellInjection
+import semmle.code.java.dataflow.DataFlow2
 import semmle.code.java.dataflow.FlowSources
 import DataFlow::PathGraph
 
@@ -23,15 +24,12 @@ class JShellInjectionConfiguration extends TaintTracking::Configuration {
   override predicate isSink(DataFlow::Node sink) { sink instanceof JShellInjectionSink }
 
   override predicate isAdditionalTaintStep(DataFlow::Node pred, DataFlow::Node succ) {
-    exists(MethodAccess ma |
-      ma.getMethod().hasName("analyzeCompletion") and
-      ma.getMethod().getNumberOfParameters() = 1 and
-      ma.getMethod()
-          .getDeclaringType()
-          .getASupertype*()
-          .hasQualifiedName("jdk.jshell", "SourceCodeAnalysis") and
-      ma.getArgument(0) = pred.asExpr() and
-      ma = succ.asExpr()
+    exists(
+      SourceCodeAnalysisAnalyzeCompletionCall scaacc, CompletionInfoSourceOrRemainingCall cisorc
+    |
+      scaacc.getArgument(0) = pred.asExpr() and
+      cisorc = succ.asExpr() and
+      DataFlow2::localExprFlow(scaacc, cisorc.getQualifier())
     )
   }
 }
diff --git a/java/ql/src/experimental/Security/CWE/CWE-094/JShellInjection.qll b/java/ql/src/experimental/Security/CWE/CWE-094/JShellInjection.qll
@@ -7,26 +7,11 @@ class JShellInjectionSink extends DataFlow::Node {
     this.asExpr() = any(JShellEvalCall jsec).getArgument(0)
     or
     this.asExpr() = any(SourceCodeAnalysisWrappersCall scawc).getArgument(0)
-    or
-    exists(MethodAccess ma |
-      ma.getMethod().hasName("source") and
-      ma.getMethod().getNumberOfParameters() = 0 and
-      ma.getMethod()
-          .getDeclaringType()
-          .getASupertype*()
-          .hasQualifiedName("jdk.jshell", "SourceCodeAnalysis$CompletionInfo") and
-      ma.getQualifier() = this.asExpr() and
-      (
-        ma = any(JShellEvalCall jsec).getArgument(0)
-        or
-        ma = any(SourceCodeAnalysisWrappersCall scawc).getArgument(0)
-      )
-    )
   }
 }
 
 /** A call to `JShell.eval`. */
-class JShellEvalCall extends MethodAccess {
+private class JShellEvalCall extends MethodAccess {
   JShellEvalCall() {
     this.getMethod().hasName("eval") and
     this.getMethod().getDeclaringType().hasQualifiedName("jdk.jshell", "JShell") and
@@ -35,10 +20,34 @@ class JShellEvalCall extends MethodAccess {
 }
 
 /** A call to `SourceCodeAnalysis.wrappers`. */
-class SourceCodeAnalysisWrappersCall extends MethodAccess {
+private class SourceCodeAnalysisWrappersCall extends MethodAccess {
   SourceCodeAnalysisWrappersCall() {
     this.getMethod().hasName("wrappers") and
     this.getMethod().getDeclaringType().hasQualifiedName("jdk.jshell", "SourceCodeAnalysis") and
     this.getMethod().getNumberOfParameters() = 1
   }
 }
+
+/** A call to `SourceCodeAnalysis.analyzeCompletion`. */
+class SourceCodeAnalysisAnalyzeCompletionCall extends MethodAccess {
+  SourceCodeAnalysisAnalyzeCompletionCall() {
+    this.getMethod().hasName("analyzeCompletion") and
+    this.getMethod()
+        .getDeclaringType()
+        .getASupertype*()
+        .hasQualifiedName("jdk.jshell", "SourceCodeAnalysis") and
+    this.getMethod().getNumberOfParameters() = 1
+  }
+}
+
+/** A call to `CompletionInfo.source` or `CompletionInfo.remaining`. */
+class CompletionInfoSourceOrRemainingCall extends MethodAccess {
+  CompletionInfoSourceOrRemainingCall() {
+    this.getMethod().getName() in ["source", "remaining"] and
+    this.getMethod()
+        .getDeclaringType()
+        .getASupertype*()
+        .hasQualifiedName("jdk.jshell", "SourceCodeAnalysis$CompletionInfo") and
+    this.getMethod().getNumberOfParameters() = 0
+  }
+}
