DataFlow: primary and secondary configurations

For now I've only implemented what XSS.qll needs

Author: jbj

java/ql/lib/semmle/code/java/security/XSS.qll

@@ -71,30 +71,25 @@ private module XssVulnerableWriterSourceToWritingMethodFlowConfig implements Dat
       sink.asExpr() = ma.getQualifier() and ma.getMethod() instanceof WritingMethod
     )
   }
+}
 
-  predicate observeDiffInformedIncrementalMode() {
-    // Since this configuration is for finding sinks to be used in a main
-    // data-flow configuration, this configuration should only restrict the
-    // sinks to be found if there are no main-configuration sources in the
-    // diff range. That's because if there is such a source, we need to
-    // report query results for it even with sinks outside the diff range.
-    exists(DataFlow::DiffInformedQuery q | not q.hasSourceInDiffRange())
-  }
-
-  // The sources are not exposed outside this file module, so we know the
-  // query will not select them.
-  Location getASelectedSourceLocation(DataFlow::Node source) { none() }
-
-  Location getASelectedSinkLocation(DataFlow::Node sink) {
+private module XssVulnerableWriterSourceToWritingMethodFlowSecondaryConfig implements
+  DataFlow::SecondaryConfig
+{
+  DataFlow::Node getPrimaryOfSecondaryNode(
+    DataFlow::IsSourceOrSink sourceOrSink, DataFlow::Node sink
+  ) {
+    sourceOrSink instanceof DataFlow::IsSink and
     // This code mirrors `DefaultXssSink()`.
-    exists(MethodCall ma | result = ma.getAnArgument().getLocation() |
+    exists(MethodCall ma | result.asExpr() = ma.getAnArgument() |
       sink.asExpr() = ma.getQualifier() and ma.getMethod() instanceof WritingMethod
     )
   }
 }
 
 private module XssVulnerableWriterSourceToWritingMethodFlow =
-  TaintTracking::Global<XssVulnerableWriterSourceToWritingMethodFlowConfig>;
+  TaintTracking::FindSinks<XssVulnerableWriterSourceToWritingMethodFlowConfig,
+    XssVulnerableWriterSourceToWritingMethodFlowSecondaryConfig>;
 
 /** A method that can be used to output data to an output stream or writer. */
 private class WritingMethod extends Method {

java/ql/lib/semmle/code/java/security/XssQuery.qll

@@ -20,9 +20,7 @@ module XssConfig implements DataFlow::ConfigSig {
   predicate isAdditionalFlowStep(DataFlow::Node node1, DataFlow::Node node2) {
     any(XssAdditionalTaintStep s).step(node1, node2)
   }
-
-  predicate observeDiffInformedIncrementalMode() { exists(DataFlow::DiffInformedQuery q) }
 }
 
 /** Tracks flow from remote sources to cross site scripting vulnerabilities. */
-module XssFlow = TaintTracking::Global<XssConfig>;
+module XssFlow = TaintTracking::Primary<XssConfig>;

shared/dataflow/codeql/dataflow/DataFlow.qll

@@ -691,6 +691,8 @@ module DataFlowMake<LocationSig Location, InputSig<Location> Lang> {
 
   class DiffInformedQuery = DiffInformedQueryImpl;
 
+  import SecondaryConfigHelpers
+
   /**
    * Constructs a global data flow computation.
    */
@@ -744,6 +746,33 @@ module DataFlowMake<LocationSig Location, InputSig<Location> Lang> {
     import Flow
   }
 
+  module Primary<ConfigSig Config> implements GlobalFlowSig {
+    private module Config0 implements FullStateConfigSig {
+      import DefaultState<Config>
+      import Config
+
+      predicate accessPathLimit = Config::accessPathLimit/0;
+
+      predicate isAdditionalFlowStep(Node node1, Node node2, string model) {
+        Config::isAdditionalFlowStep(node1, node2) and model = "Config"
+      }
+    }
+
+    private module C implements FullStateConfigSig {
+      import MakePrimaryDiffInformed<Config0>
+
+      predicate accessPathLimit = Config0::accessPathLimit/0;
+    }
+
+    private module Stage1 = ImplStage1<C>;
+
+    import Stage1::PartialFlow
+
+    private module Flow = Impl<C, Stage1::Stage1NoState>;
+
+    import Flow
+  }
+
   signature class PathNodeSig {
     /** Gets a textual representation of this element. */
     string toString();

shared/dataflow/codeql/dataflow/TaintTracking.qll

@@ -143,6 +143,61 @@ module TaintFlowMake<
     import Flow
   }
 
+  /**
+   * Constructs a global taint tracking computation.
+   */
+  module Primary<DataFlow::ConfigSig Config> implements DataFlow::GlobalFlowSig {
+    private module Config0 implements DataFlowInternal::FullStateConfigSig {
+      import DataFlowInternal::DefaultState<Config>
+      import Config
+
+      predicate isAdditionalFlowStep(
+        DataFlowLang::Node node1, DataFlowLang::Node node2, string model
+      ) {
+        Config::isAdditionalFlowStep(node1, node2) and model = "Config"
+      }
+    }
+
+    private module C implements DataFlowInternal::FullStateConfigSig {
+      import DataFlowInternal::MakePrimaryDiffInformed<AddTaintDefaults<Config0>>
+    }
+
+    private module Stage1 = DataFlowInternalStage1::ImplStage1<C>;
+
+    import Stage1::PartialFlow
+
+    private module Flow = DataFlowInternal::Impl<C, Stage1::Stage1NoState>;
+
+    import Flow
+  }
+
+  module FindSinks<DataFlow::ConfigSig Config, DataFlow::SecondaryConfig SC> implements
+    DataFlow::GlobalFlowSig
+  {
+    private module Config0 implements DataFlowInternal::FullStateConfigSig {
+      import DataFlowInternal::DefaultState<Config>
+      import Config
+
+      predicate isAdditionalFlowStep(
+        DataFlowLang::Node node1, DataFlowLang::Node node2, string model
+      ) {
+        Config::isAdditionalFlowStep(node1, node2) and model = "Config"
+      }
+    }
+
+    private module C implements DataFlowInternal::FullStateConfigSig {
+      import DataFlowInternal::MakeSinkFinder<AddTaintDefaults<Config0>, SC>
+    }
+
+    private module Stage1 = DataFlowInternalStage1::ImplStage1<C>;
+
+    import Stage1::PartialFlow
+
+    private module Flow = DataFlowInternal::Impl<C, Stage1::Stage1NoState>;
+
+    import Flow
+  }
+
   signature int speculationLimitSig();
 
   private module AddSpeculativeTaintSteps<

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

@@ -184,7 +184,6 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
 
     string toString() { result = "DataFlow::DiffInformedQuery" }
 
-    // TODO: how to wire this up? It overlaps with the data-flow configuration signature!
     Location getASelectedSourceLocation(Node source) { result = source.getLocation() }
 
     Location getASelectedSinkLocation(Node sink) { result = sink.getLocation() }
@@ -200,6 +199,90 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
     }
   }
 
+  module MakePrimaryDiffInformed<FullStateConfigSig Config> implements FullStateConfigSig {
+    // Workaround for name clash
+    predicate accessPathLimit = Config::accessPathLimit/0;
+
+    import Config
+
+    predicate observeDiffInformedIncrementalMode() {
+      // Add to existing configuration to support composition of config transformers
+      Config::observeDiffInformedIncrementalMode()
+      or
+      exists(DiffInformedQueryImpl q)
+    }
+
+    Location getASelectedSourceLocation(Node source) {
+      result = Config::getASelectedSourceLocation(source)
+      or
+      exists(DiffInformedQueryImpl q | result = q.getASelectedSourceLocation(source))
+    }
+
+    Location getASelectedSinkLocation(Node sink) {
+      result = Config::getASelectedSinkLocation(sink)
+      or
+      exists(DiffInformedQueryImpl q | result = q.getASelectedSinkLocation(sink))
+    }
+  }
+
+  module SecondaryConfigHelpers {
+    newtype IsSourceOrSink =
+      IsSource() or
+      IsSink()
+
+    signature module SecondaryConfig {
+      /**
+       * Gets the source/sink node from the primary configuration that is
+       * informed by a given source/sink node from the secondary configuration.
+       * Whether the secondary node is a source or a sink is determined by
+       * `sourceOrSink`.
+       */
+      bindingset[sourceOrSink, secondaryNode]
+      Node getPrimaryOfSecondaryNode(IsSourceOrSink sourceOrSink, Node secondaryNode);
+    }
+  }
+
+  module MakeSinkFinder<FullStateConfigSig Config, SecondaryConfig SC> implements FullStateConfigSig
+  {
+    // Workaround for name clash
+    predicate accessPathLimit = Config::accessPathLimit/0;
+
+    import Config
+
+    predicate observeDiffInformedIncrementalMode() {
+      // Add to existing configuration to support composition of config transformers
+      Config::observeDiffInformedIncrementalMode()
+      or
+      // Because this configuration is for finding sinks to be used in a main
+      // data-flow configuration, this configuration should only restrict the
+      // sinks to be found if there are no main-configuration sources in the
+      // diff range. That's because if there is such a source, we need to
+      // report query results for it even with sinks outside the diff range.
+      //
+      // The `MakeSinkFinder` and `MakeSourceFinder` modules are separate
+      // because each can only call one of `hasSourceInDiffRange` or
+      // `hasSinkInDiffRange`. Otherwise it would look like a non-monotonic
+      // recursion.
+      exists(DiffInformedQuery q | not q.hasSourceInDiffRange())
+    }
+
+    Location getASelectedSourceLocation(Node source) {
+      result = Config::getASelectedSourceLocation(source)
+      or
+      exists(DiffInformedQueryImpl q, IsSource s |
+        result = q.getASelectedSourceLocation(SC::getPrimaryOfSecondaryNode(s, source))
+      )
+    }
+
+    Location getASelectedSinkLocation(Node sink) {
+      result = Config::getASelectedSinkLocation(sink)
+      or
+      exists(DiffInformedQueryImpl q, IsSink s |
+        result = q.getASelectedSinkLocation(SC::getPrimaryOfSecondaryNode(s, sink))
+      )
+    }
+  }
+
   /**
    * Constructs a data flow computation given a full input configuration, and
    * an initial stage 1 pruning.