Data flow: Add hook for preventing lambda dispatch in source call contexts

Author: hvitved

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

@@ -1179,7 +1179,8 @@ private module Cached {
   cached
   newtype TDataFlowType =
     TGvnDataFlowType(Gvn::GvnType t) or
-    TDelegateDataFlowType(Callable lambda) { lambdaCreationExpr(_, lambda) }
+    TDelegateDataFlowType(Callable lambda) { lambdaCreationExpr(_, lambda) } or
+    TSourceContextParameterType()
 }
 
 import Cached
@@ -2394,6 +2395,8 @@ class DataFlowType extends TDataFlowType {
 
   Callable asDelegate() { this = TDelegateDataFlowType(result) }
 
+  predicate isSourceContextParameterType() { this = TSourceContextParameterType() }
+
   /**
    * Gets an expression that creates a delegate of this type.
    *
@@ -2412,6 +2415,9 @@ class DataFlowType extends TDataFlowType {
     result = this.asGvnType().toString()
     or
     result = this.asDelegate().toString()
+    or
+    this.isSourceContextParameterType() and
+    result = "<source context parameter type>"
   }
 }
 
@@ -2427,6 +2433,8 @@ DataFlowType getNodeType(Node n) {
   ] = result.getADelegateCreation()
 }
 
+DataFlowType getSourceContextParameterNodeType() { result.isSourceContextParameterType() }
+
 private class DataFlowNullType extends Gvn::GvnType {
   DataFlowNullType() { this = Gvn::getGlobalValueNumber(any(NullType nt)) }
 
@@ -2499,6 +2507,10 @@ predicate compatibleTypes(DataFlowType dt1, DataFlowType dt2) {
   compatibleTypesDelegateLeft(dt2, dt1)
   or
   dt1.asDelegate() = dt2.asDelegate()
+  or
+  dt1.isSourceContextParameterType() and not exists(dt2.asDelegate())
+  or
+  dt2.isSourceContextParameterType() and not exists(dt1.asDelegate())
 }
 
 pragma[nomagic]

csharp/ql/test/utils/modelgenerator/dataflow/CaptureContentSummaryModels.expected

@@ -1,3 +1,2 @@
 unexpectedModel
-| Unexpected contentbased-summary found: Models;HigherOrderParameters;false;Apply;(System.Func<System.Object,System.Object>,System.Object);;Argument[1];ReturnValue;value;dfc-generated |
 expectedModel

shared/dataflow/codeql/dataflow/DataFlow.qll

@@ -63,6 +63,35 @@ signature module InputSig<LocationSig Location> {
 
   DataFlowType getNodeType(Node node);
 
+  /**
+   * Gets a special type to use for parameter nodes belonging to callables with a
+   * source node where a source call context `FlowFeature` is used, if any.
+   *
+   * This can be used to prevent lambdas from being resolved, when a concrete call
+   * context is needed. Example:
+   *
+   * ```csharp
+   * void Foo(Action<string> a)
+   * {
+   *     var x = Source();
+   *     a(x);              // (1)
+   *     a = s => Sink(s);  // (2)
+   *     a(x);              // (3)
+   * }
+   *
+   * void Bar()
+   * {
+   *     Foo(s => Sink(s)); // (4)
+   * }
+   * ```
+   *
+   * If a source call context flow feature is used, `a` can be assigned a special
+   * type that is incompatible with the type of _any_ lambda expression, which will
+   * prevent the call edge from (1) to (4). Note that the call edge from (3) to (2)
+   * will still be valid.
+   */
+  default DataFlowType getSourceContextParameterNodeType() { none() }
+
   predicate nodeIsHidden(Node node);
 
   class DataFlowExpr;

shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll

@@ -1103,6 +1103,13 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
         private module FwdTypeFlowInput implements TypeFlowInput {
           predicate enableTypeFlow = Param::enableTypeFlow/0;
 
+          predicate isParameterNodeInSourceCallContext(ParamNode p) {
+            hasSourceCallCtx() and
+            Config::isSource(any(Node source |
+                getNodeEnclosingCallable(source) = getNodeEnclosingCallable(p)
+              ), _)
+          }
+
           predicate relevantCallEdgeIn = PrevStage::relevantCallEdgeIn/2;
 
           predicate relevantCallEdgeOut = PrevStage::relevantCallEdgeOut/2;
@@ -1410,6 +1417,8 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
         private module RevTypeFlowInput implements TypeFlowInput {
           predicate enableTypeFlow = Param::enableTypeFlow/0;
 
+          predicate isParameterNodeInSourceCallContext(ParamNode p) { none() }
+
           predicate relevantCallEdgeIn(Call call, Callable c) {
             flowOutOfCallAp(call, c, _, _, _, _, _)
           }

shared/dataflow/codeql/dataflow/internal/DataFlowImplCommon.qll

@@ -1893,6 +1893,9 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
   signature module TypeFlowInput {
     predicate enableTypeFlow();
 
+    /** Holds if `p` is a parameter of a callable with a source node that has a call context. */
+    predicate isParameterNodeInSourceCallContext(ParamNode p);
+
     /** Holds if the edge is possibly needed in the direction `call` to `c`. */
     predicate relevantCallEdgeIn(Call call, Callable c);
 
@@ -1953,6 +1956,9 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
     /**
      * Holds if a sequence of calls may propagate the value of `arg` to some
      * argument-to-parameter call edge that strengthens the static type.
+     *
+     * This predicate is a reverse flow computation, starting at calls that
+     * strengthen the type and then following relevant call edges backwards.
      */
     pragma[nomagic]
     private predicate trackedArgTypeCand(ArgNode arg) {
@@ -1987,6 +1993,9 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
      * Holds if `p` is part of a value-propagating call path where the
      * end-points have stronger types than the intermediate parameter and
      * argument nodes.
+     *
+     * This predicate is a forward flow computation, intersecting with the
+     * reverse flow computation done in `trackedArgTypeCand`.
      */
     private predicate trackedParamType(ParamNode p) {
       exists(Call call1, Callable c1, ArgNode argOut, Call call2, Callable c2, ArgNode argIn |
@@ -2013,6 +2022,8 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
         typeStrongerThanFilter(at, pt)
       )
       or
+      Input::isParameterNodeInSourceCallContext(p)
+      or
       exists(ArgNode arg |
         trackedArgType(arg) and
         relevantCallEdge(_, _, arg, p) and
@@ -2106,7 +2117,9 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
     private predicate typeFlowParamType(ParamNode p, Type t, boolean cc) {
       exists(Callable c |
         Input::dataFlowNonCallEntry(c, cc) and
-        trackedParamWithType(p, t, c)
+        if cc = true and exists(getSourceContextParameterNodeType())
+        then t = getSourceContextParameterNodeType()
+        else trackedParamWithType(p, t, c)
       )
       or
       exists(Type t1, Type t2 |