Merge pull request #21592 from hvitved/dataflow/source-call-context-type-flow

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

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 node `p` 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(Node p) { 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
@@ -2104,8 +2115,12 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
      * context.
      */
     private predicate typeFlowParamType(ParamNode p, Type t, boolean cc) {
-      exists(Callable c |
-        Input::dataFlowNonCallEntry(c, cc) and
+      exists(Callable c | Input::dataFlowNonCallEntry(c, cc) |
+        cc = true and
+        nodeEnclosingCallable(p, c) and
+        t = getSourceContextParameterNodeType(p)
+        or
+        (cc = false or not exists(getSourceContextParameterNodeType(p))) and
         trackedParamWithType(p, t, c)
       )
       or