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

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

@@ -29,4 +29,6 @@ module CsharpDataFlow implements InputSig<Location> {
   predicate neverSkipInPathGraph(Node n) {
     exists(n.(AssignableDefinitionNode).getDefinition().getTargetAccess())
   }
+
+  DataFlowType getSourceContextParameterNodeType() { result.isSourceContextParameterType() }
 }

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>"
   }
 }
 
@@ -2469,6 +2475,11 @@ private predicate compatibleTypesDelegateLeft(DataFlowType dt1, DataFlowType dt2
   )
 }
 
+pragma[nomagic]
+private predicate compatibleTypesSourceContextParameterTypeLeft(DataFlowType dt1, DataFlowType dt2) {
+  dt1.isSourceContextParameterType() and not exists(dt2.asDelegate())
+}
+
 /**
  * Holds if `t1` and `t2` are compatible, that is, whether data can flow from
  * a node of type `t1` to a node of type `t2`.
@@ -2499,6 +2510,10 @@ predicate compatibleTypes(DataFlowType dt1, DataFlowType dt2) {
   compatibleTypesDelegateLeft(dt2, dt1)
   or
   dt1.asDelegate() = dt2.asDelegate()
+  or
+  compatibleTypesSourceContextParameterTypeLeft(dt1, dt2)
+  or
+  compatibleTypesSourceContextParameterTypeLeft(dt2, dt1)
 }
 
 pragma[nomagic]

csharp/ql/test/utils/modelgenerator/dataflow/Summaries.cs

@@ -536,6 +536,12 @@ public void Apply(Action<object> a, object o)
     {
         a(o);
     }
+
+    private void CallApply()
+    {
+        // Test that this call to `Apply` does not interfere with the flow summaries generated for `Apply`
+        Apply(x => x, null);
+    }
 }
 
 public static class HigherOrderExtensionMethods

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,16 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
         private module FwdTypeFlowInput implements TypeFlowInput {
           predicate enableTypeFlow = Param::enableTypeFlow/0;
 
+          pragma[nomagic]
+          predicate isParameterNodeInSourceCallContext(ParamNode p) {
+            hasSourceCallCtx() and
+            exists(Node source, DataFlowCallable c |
+              Config::isSource(pragma[only_bind_into](source), _) and
+              nodeEnclosingCallable(source, c) and
+              nodeEnclosingCallable(p, c)
+            )
+          }
+
           predicate relevantCallEdgeIn = PrevStage::relevantCallEdgeIn/2;
 
           predicate relevantCallEdgeOut = PrevStage::relevantCallEdgeOut/2;
@@ -1410,6 +1420,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 |