C#: Restrict multi-body dataflow dispatch based on file-system distance

csharp/ql/lib/semmle/code/csharp/dataflow/internal/DataFlowDispatch.qll

@@ -8,6 +8,7 @@ private import semmle.code.csharp.dispatch.Dispatch
 private import semmle.code.csharp.dispatch.RuntimeCallable
 private import semmle.code.csharp.frameworks.system.Collections
 private import semmle.code.csharp.frameworks.system.collections.Generic
+private import semmle.code.csharp.internal.Location
 
 /**
  * Gets a source declaration of callable `c` that has a body and is
@@ -24,6 +25,21 @@ newtype TReturnKind =
   TOutReturnKind(int i) { i = any(Parameter p | p.isOut()).getPosition() } or
   TRefReturnKind(int i) { i = any(Parameter p | p.isRef()).getPosition() }
 
+private predicate hasMultipleSourceLocations(Callable c) { strictcount(getASourceLocation(c)) > 1 }
+
+private module NearestBodyLocationInput implements NearestLocationInputSig {
+  class C = ControlFlowElement;
+
+  predicate relevantLocations(ControlFlowElement body, Location l1, Location l2) {
+    exists(Callable c |
+      hasMultipleSourceLocations(c) and
+      l1 = getASourceLocation(c) and
+      body = c.getBody() and
+      l2 = body.getLocation()
+    )
+  }
+}
+
 cached
 private module Cached {
   /**
@@ -33,7 +49,18 @@ private module Cached {
    */
   cached
   newtype TDataFlowCallable =
-    TCallable(Callable c) { c.isUnboundDeclaration() } or
+    TCallable(Callable c, Location l) {
+      c.isUnboundDeclaration() and
+      l = [c.getLocation(), getASourceLocation(c)] and
+      (
+        not hasMultipleSourceLocations(c)
+        or
+        // when `c` has multiple source locations, only use those with a body;
+        // for example, `partial` methods may have multiple source locations but
+        // we are only interested in the one with a body
+        NearestLocation<NearestBodyLocationInput>::nearestLocation(_, l, _)
+      )
+    } or
     TSummarizedCallable(FlowSummary::SummarizedCallable sc) or
     TFieldOrPropertyCallable(FieldOrProperty f) or
     TCapturedVariableCallable(LocalScopeVariable v) { v.isCaptured() }
@@ -82,17 +109,23 @@ private module DispatchImpl {
    */
   predicate mayBenefitFromCallContext(DataFlowCall call) { mayBenefitFromCallContext(call, _) }
 
+  bindingset[dc, result]
+  pragma[inline_late]
+  private Callable viableImplInCallContext0(DispatchCall dc, NonDelegateDataFlowCall ctx) {
+    result = dc.getADynamicTargetInCallContext(ctx.getDispatchCall()).getUnboundDeclaration()
+  }
+
   /**
    * Gets a viable dispatch target of `call` in the context `ctx`. This is
    * restricted to those `call`s for which a context might make a difference.
    */
   DataFlowCallable viableImplInCallContext(DataFlowCall call, DataFlowCall ctx) {
-    exists(DispatchCall dc | dc = call.(NonDelegateDataFlowCall).getDispatchCall() |
-      result.getUnderlyingCallable() =
-        getCallableForDataFlow(dc.getADynamicTargetInCallContext(ctx.(NonDelegateDataFlowCall)
-                .getDispatchCall()).getUnboundDeclaration())
+    exists(DispatchCall dc, Callable c | dc = call.(NonDelegateDataFlowCall).getDispatchCall() |
+      result = call.getARuntimeTarget() and
+      getCallableForDataFlow(c) = result.asCallable(_) and
+      c = viableImplInCallContext0(dc, ctx)
       or
-      exists(Callable c, DataFlowCallable encl |
+      exists(DataFlowCallable encl |
         result.asSummarizedCallable() = c and
         mayBenefitFromCallContext(call, encl) and
         encl = ctx.getARuntimeTarget() and
@@ -159,7 +192,51 @@ class RefReturnKind extends OutRefReturnKind, TRefReturnKind {
 /** A callable used for data flow. */
 class DataFlowCallable extends TDataFlowCallable {
   /** Gets the underlying source code callable, if any. */
-  Callable asCallable() { this = TCallable(result) }
+  Callable asCallable(Location l) { this = TCallable(result, l) }
+
+  /** Holds if this callable is multi-bodied. */
+  pragma[nomagic]
+  predicate isMultiBodied() { hasMultipleSourceLocations(this.asCallable(_)) }
+
+  private ControlFlow::Nodes::ElementNode getAMultiBodyEntryNode() {
+    exists(ControlFlowElement body, Location l |
+      body = this.asCallable(l).getBody() and
+      NearestLocation<NearestBodyLocationInput>::nearestLocation(body, l, _) and
+      result = body.getAControlFlowEntryNode()
+    )
+  }
+
+  pragma[nomagic]
+  private ControlFlow::Nodes::ElementNode getAMultiBodyControlFlowNodePred() {
+    result = this.getAMultiBodyEntryNode().getAPredecessor()
+    or
+    result = this.getAMultiBodyControlFlowNodePred().getAPredecessor()
+  }
+
+  pragma[nomagic]
+  private ControlFlow::Nodes::ElementNode getAMultiBodyControlFlowNodeSucc() {
+    result = this.getAMultiBodyEntryNode().getASuccessor()
+    or
+    result = this.getAMultiBodyControlFlowNodeSucc().getASuccessor()
+  }
+
+  pragma[nomagic]
+  private ControlFlow::Nodes::ElementNode getAMultiBodyControlFlowNode() {
+    result =
+      [
+        this.getAMultiBodyEntryNode(), this.getAMultiBodyControlFlowNodePred(),
+        this.getAMultiBodyControlFlowNodeSucc()
+      ]
+  }
+
+  /** Gets a control flow node belonging to this callable. */
+  pragma[inline]
+  ControlFlow::Node getAControlFlowNode() {
+    result = this.getAMultiBodyControlFlowNode()
+    or
+    not this.isMultiBodied() and
+    result.getEnclosingCallable() = this.asCallable(_)
+  }
 
   /** Gets the underlying summarized callable, if any. */
   FlowSummary::SummarizedCallable asSummarizedCallable() { this = TSummarizedCallable(result) }
@@ -171,7 +248,7 @@ class DataFlowCallable extends TDataFlowCallable {
 
   /** Gets the underlying callable. */
   Callable getUnderlyingCallable() {
-    result = this.asCallable() or result = this.asSummarizedCallable()
+    result = this.asCallable(_) or result = this.asSummarizedCallable()
   }
 
   /** Gets a textual representation of this dataflow callable. */
@@ -185,7 +262,9 @@ class DataFlowCallable extends TDataFlowCallable {
 
   /** Get the location of this dataflow callable. */
   Location getLocation() {
-    result = this.getUnderlyingCallable().getLocation()
+    this = TCallable(_, result)
+    or
+    result = this.asSummarizedCallable().getLocation()
     or
     result = this.asFieldOrProperty().getLocation()
     or
@@ -256,6 +335,26 @@ abstract class DataFlowCall extends TDataFlowCall {
   }
 }
 
+private predicate relevantFolder(Folder f) {
+  exists(NonDelegateDataFlowCall call, Location l | f = l.getFile().getParentContainer() |
+    l = call.getLocation() and
+    call.getARuntimeTargetCandidate(_, _).isMultiBodied()
+    or
+    call.getARuntimeTargetCandidate(l, _).isMultiBodied()
+  )
+}
+
+private predicate adjacentFolders(Folder f1, Folder f2) {
+  f1 = f2.getParentContainer()
+  or
+  f2 = f1.getParentContainer()
+}
+
+bindingset[f1, f2]
+pragma[inline_late]
+private predicate folderDist(Folder f1, Folder f2, int i) =
+  shortestDistances(relevantFolder/1, adjacentFolders/2)(f1, f2, i)
+
 /** A non-delegate C# call relevant for data flow. */
 class NonDelegateDataFlowCall extends DataFlowCall, TNonDelegateCall {
   private ControlFlow::Nodes::ElementNode cfn;
@@ -266,25 +365,63 @@ class NonDelegateDataFlowCall extends DataFlowCall, TNonDelegateCall {
   /** Gets the underlying call. */
   DispatchCall getDispatchCall() { result = dc }
 
-  override DataFlowCallable getARuntimeTarget() {
-    result.asCallable() = getCallableForDataFlow(dc.getADynamicTarget())
-    or
+  pragma[nomagic]
+  private predicate hasSourceTarget() { dc.getADynamicTarget().fromSource() }
+
+  pragma[nomagic]
+  private FlowSummary::SummarizedCallable getASummarizedCallableTarget() {
     // Only use summarized callables with generated summaries in case
     // we are not able to dispatch to a source declaration.
-    exists(FlowSummary::SummarizedCallable sc, boolean static |
-      result.asSummarizedCallable() = sc and
-      sc = this.getATarget(static) and
+    exists(boolean static |
+      result = this.getATarget(static) and
       not (
-        sc.applyGeneratedModel() and
-        dc.getADynamicTarget().getUnboundDeclaration().getFile().fromSource()
+        result.applyGeneratedModel() and
+        this.hasSourceTarget()
       )
     |
       static = false
       or
-      static = true and not sc instanceof RuntimeCallable
+      static = true and not result instanceof RuntimeCallable
     )
   }
 
+  pragma[nomagic]
+  DataFlowCallable getARuntimeTargetCandidate(Location l, Callable c) {
+    c = result.asCallable(l) and
+    c = getCallableForDataFlow(dc.getADynamicTarget())
+  }
+
+  pragma[nomagic]
+  private DataFlowCallable getAMultiBodiedRuntimeTargetCandidate(Callable c, int distance) {
+    result.isMultiBodied() and
+    exists(Location l | result = this.getARuntimeTargetCandidate(l, c) |
+      inSameFile(l, this.getLocation()) and
+      distance = -1
+      or
+      folderDist(l.getFile().getParentContainer(),
+        this.getLocation().getFile().getParentContainer(), distance)
+    )
+  }
+
+  pragma[nomagic]
+  override DataFlowCallable getARuntimeTarget() {
+    // For calls to multi-bodied methods, we restrict the viable targets to those
+    // that are closest to the callsite, measured by file-system distance.
+    exists(Callable c |
+      result =
+        min(DataFlowCallable cand, int distance |
+          cand = this.getAMultiBodiedRuntimeTargetCandidate(c, distance)
+        |
+          cand order by distance
+        )
+    )
+    or
+    result = this.getARuntimeTargetCandidate(_, _) and
+    not result.isMultiBodied()
+    or
+    result.asSummarizedCallable() = this.getASummarizedCallableTarget()
+  }
+
   /** Gets a static or dynamic target of this call. */
   Callable getATarget(boolean static) {
     result = dc.getADynamicTarget().getUnboundDeclaration() and static = false
@@ -296,9 +433,7 @@ class NonDelegateDataFlowCall extends DataFlowCall, TNonDelegateCall {
 
   override DataFlow::ExprNode getNode() { result.getControlFlowNode() = cfn }
 
-  override DataFlowCallable getEnclosingCallable() {
-    result.asCallable() = cfn.getEnclosingCallable()
-  }
+  override DataFlowCallable getEnclosingCallable() { result.getAControlFlowNode() = cfn }
 
   override string toString() { result = cfn.toString() }
 
@@ -326,9 +461,7 @@ class ExplicitDelegateLikeDataFlowCall extends DelegateDataFlowCall, TExplicitDe
 
   override DataFlow::ExprNode getNode() { result.getControlFlowNode() = cfn }
 
-  override DataFlowCallable getEnclosingCallable() {
-    result.asCallable() = cfn.getEnclosingCallable()
-  }
+  override DataFlowCallable getEnclosingCallable() { result.getAControlFlowNode() = cfn }
 
   override string toString() { result = cfn.toString() }