Fix data flow through ExtensionMethodAccess

tamasvajk · igfoo · commit 95cb0149a3c4 · 2022-05-10T19:51:27.000+01:00
diff --git a/java/ql/lib/semmle/code/java/ControlFlowGraph.qll b/java/ql/lib/semmle/code/java/ControlFlowGraph.qll
@@ -558,7 +558,7 @@ private module ControlFlowGraphImpl {
       or
       exists(ExtensionMethodAccess e | e = this |
         // the actual qualifier of the expression method access
-        index = -1 and result.(Expr).isNthChildOf(this, index) and not result instanceof TypeAccess
+        index = -1 and result = e.getImplicitQualifier() and not result instanceof TypeAccess
         or
         // the extension receiver
         index = 0 and result = e.getQualifier()
diff --git a/java/ql/lib/semmle/code/java/Expr.qll b/java/ql/lib/semmle/code/java/Expr.qll
@@ -1990,6 +1990,12 @@ class MethodAccess extends Expr, Call, @methodaccess {
 class ExtensionMethodAccess extends MethodAccess {
   ExtensionMethodAccess() { this.getMethod() instanceof ExtensionMethod }
 
+  /**
+   * Gets the implicit qualifier of this extension method access. This expression
+   * is the qualifier of this method access in bytecode.
+   */
+  Expr getImplicitQualifier() { result.isNthChildOf(this, -1) }
+
   // The syntactic qualifier of an extension method is its receiver (arg 0),
   // whereas the actual arguments begin at index 1.
   override Expr getQualifier() { result.isNthChildOf(this, 0) }
@@ -2240,8 +2246,11 @@ private module Qualifier {
 
     /** Gets the qualifier of this member access, if any. */
     Expr getQualifier() {
-      result = this.(FieldAccess).getQualifier() or
-      result = this.(MethodAccess).getQualifier()
+      result = this.(FieldAccess).getQualifier()
+      or
+      result = this.(MethodAccess).getQualifier() and not this instanceof ExtensionMethodAccess
+      or
+      result = this.(ExtensionMethodAccess).getImplicitQualifier()
     }
   }
 
diff --git a/java/ql/lib/semmle/code/java/dataflow/internal/DataFlowNodes.qll b/java/ql/lib/semmle/code/java/dataflow/internal/DataFlowNodes.qll
@@ -49,10 +49,16 @@ newtype TNode =
 
 private predicate explicitInstanceArgument(Call call, Expr instarg) {
   call instanceof MethodAccess and
-  instarg = call.getQualifier() and
+  instarg = getQualifier(call) and
   not call.getCallee().isStatic()
 }
 
+private Expr getQualifier(MethodAccess ma) {
+  not ma instanceof ExtensionMethodAccess and result = ma.getQualifier()
+  or
+  result = ma.(ExtensionMethodAccess).getImplicitQualifier()
+}
+
 private predicate implicitInstanceArgument(Call call, InstanceAccessExt ia) {
   ia.isImplicitMethodQualifier(call) or
   ia.isImplicitThisConstructorArgument(call)
@@ -360,6 +366,9 @@ module Private {
     ArgumentNode() {
       exists(Argument arg | this.asExpr() = arg | not arg.isVararg())
       or
+      // In JVM bytecode the qualifier of an extension method call is the 0th argument
+      exists(ExtensionMethodAccess ema | this.asExpr() = ema.getQualifier())
+      or
       this instanceof ImplicitVarargsArray
       or
       this = getInstanceArgument(_)
@@ -373,7 +382,20 @@ module Private {
      */
     predicate argumentOf(DataFlowCall call, int pos) {
       exists(Argument arg | this.asExpr() = arg |
-        call.asCall() = arg.getCall() and pos = arg.getPosition()
+        call.asCall() = arg.getCall() and
+        (
+          pos = arg.getPosition() and not call.asCall() instanceof ExtensionMethodAccess
+          or
+          // Shift the positions of an extension method call by 1 to make space for the qualifier at index 0.
+          pos = arg.getPosition() + 1 and call.asCall() instanceof ExtensionMethodAccess
+        )
+      )
+      or
+      // In JVM bytecode the qualifier of an extension method call is the 0th argument
+      exists(ExtensionMethodAccess ema |
+        call.asCall() = ema and
+        pos = 0 and
+        this.asExpr() = ema.getQualifier()
       )
       or
       call.asCall() = this.(ImplicitVarargsArray).getCall() and
diff --git a/java/ql/lib/semmle/code/java/dispatch/VirtualDispatch.qll b/java/ql/lib/semmle/code/java/dispatch/VirtualDispatch.qll
@@ -273,8 +273,14 @@ private module Dispatch {
     not e instanceof FunctionalExpr and result = e.getType()
   }
 
+  private Expr getQualifier(MethodAccess ma) {
+    not ma instanceof ExtensionMethodAccess and result = ma.getQualifier()
+    or
+    result = ma.(ExtensionMethodAccess).getImplicitQualifier()
+  }
+
   private predicate hasQualifierType(VirtualMethodAccess ma, RefType t, boolean exact) {
-    exists(Expr src | src = variableTrack(ma.getQualifier()) |
+    exists(Expr src | src = variableTrack(getQualifier(ma)) |
       // If we have a qualifier, then we track it through variable assignments
       // and take the type of the assigned value.
       exists(RefType srctype | srctype = getPreciseType(src) |
diff --git a/java/ql/test/kotlin/library-tests/dataflow/extensionMethod/test.expected b/java/ql/test/kotlin/library-tests/dataflow/extensionMethod/test.expected
@@ -1,2 +1,6 @@
 | test.kt:20:29:20:31 | new C(...) | test.kt:23:22:23:28 | self1(...) |
+| test.kt:20:29:20:31 | new C(...) | test.kt:26:22:26:28 | self2(...) |
 | test.kt:20:29:20:31 | new C(...) | test.kt:29:18:29:29 | fn1(...) |
+| test.kt:20:29:20:31 | new C(...) | test.kt:32:18:32:29 | fn2(...) |
+| test.kt:20:29:20:31 | new C(...) | test.kt:35:18:35:31 | call1(...) |
+| test.kt:20:29:20:31 | new C(...) | test.kt:38:22:38:28 | call2(...) |
