From 40529d1e51b57f079a16ac6a4265ca71b0303222 Mon Sep 17 00:00:00 2001
From: Jonas Jensen <jbj@github.com>
Date: Mon, 7 Oct 2024 13:16:10 +0200
Subject: [PATCH] DataFlow: Add the concept of selected locations

This extension allows queries to be diff-informed even when the elements
they select are different from the sources and sinks found by data flow.
---
 shared/dataflow/codeql/dataflow/DataFlow.qll  | 44 +++++++++++++++++++
 .../codeql/dataflow/internal/DataFlowImpl.qll |  8 +++-
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/shared/dataflow/codeql/dataflow/DataFlow.qll b/shared/dataflow/codeql/dataflow/DataFlow.qll
index 400ed5746eb..7c437adabb8 100644
--- a/shared/dataflow/codeql/dataflow/DataFlow.qll
+++ b/shared/dataflow/codeql/dataflow/DataFlow.qll
@@ -442,6 +442,28 @@ module Configs<LocationSig Location, InputSig<Location> Lang> {
      * are used directly in a query result.
      */
     default predicate observeDiffInformedIncrementalMode() { none() }
+
+    /**
+     * Gets a location that will be associated with the given `source` in a
+     * diff-informed query that uses this configuration (see
+     * `observeDiffInformedIncrementalMode`). By default, this is the location
+     * of the source itself, but this predicate should include any locations
+     * that are reported as the primary-location of the query or as an
+     * additional location ("$@" interpolation). For a query that doesn't
+     * report the source at all, this predicate can be `none()`.
+     */
+    default Location getASelectedSourceLocation(Node source) { result = source.getLocation() }
+
+    /**
+     * Gets a location that will be associated with the given `sink` in a
+     * diff-informed query that uses this configuration (see
+     * `observeDiffInformedIncrementalMode`). By default, this is the location
+     * of the sink itself, but this predicate should include any locations
+     * that are reported as the primary-location of the query or as an
+     * additional location ("$@" interpolation). For a query that doesn't
+     * report the sink at all, this predicate can be `none()`.
+     */
+    default Location getASelectedSinkLocation(Node sink) { result = sink.getLocation() }
   }
 
   /** An input configuration for data flow using flow state. */
@@ -569,6 +591,28 @@ module Configs<LocationSig Location, InputSig<Location> Lang> {
      * are used directly in a query result.
      */
     default predicate observeDiffInformedIncrementalMode() { none() }
+
+    /**
+     * Gets a location that will be associated with the given `source` in a
+     * diff-informed query that uses this configuration (see
+     * `observeDiffInformedIncrementalMode`). By default, this is the location
+     * of the source itself, but this predicate should include any locations
+     * that are reported as the primary-location of the query or as an
+     * additional location ("$@" interpolation). For a query that doesn't
+     * report the source at all, this predicate can be `none()`.
+     */
+    default Location getASelectedSourceLocation(Node source) { result = source.getLocation() }
+
+    /**
+     * Gets a location that will be associated with the given `sink` in a
+     * diff-informed query that uses this configuration (see
+     * `observeDiffInformedIncrementalMode`). By default, this is the location
+     * of the sink itself, but this predicate should include any locations
+     * that are reported as the primary-location of the query or as an
+     * additional location ("$@" interpolation). For a query that doesn't
+     * report the sink at all, this predicate can be `none()`.
+     */
+    default Location getASelectedSinkLocation(Node sink) { result = sink.getLocation() }
   }
 }
 
diff --git a/shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll b/shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll
index 1373345423f..6534aaa640c 100644
--- a/shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll
+++ b/shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll
@@ -137,6 +137,10 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
      * are used directly in a query result.
      */
     predicate observeDiffInformedIncrementalMode();
+
+    Location getASelectedSourceLocation(Node source);
+
+    Location getASelectedSinkLocation(Node sink);
   }
 
   /**
@@ -177,7 +181,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
       private predicate isFilteredSource(Node source) {
         Config::isSource(source, _) and
         if Config::observeDiffInformedIncrementalMode()
-        then AlertFiltering::filterByLocation(source.getLocation())
+        then AlertFiltering::filterByLocation(Config::getASelectedSourceLocation(source))
         else any()
       }
 
@@ -188,7 +192,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
           Config::isSink(sink)
         ) and
         if Config::observeDiffInformedIncrementalMode()
-        then AlertFiltering::filterByLocation(sink.getLocation())
+        then AlertFiltering::filterByLocation(Config::getASelectedSinkLocation(sink))
         else any()
       }