Add support for multiple sink types per query

java/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/ATMConfig.qll

@@ -57,15 +57,23 @@ abstract class AtmConfig extends TaintTracking::Configuration {
   /**
    * EXPERIMENTAL. This API may change in the future.
    *
-   * Holds if `sink` is a known sink of flow.
+   * Holds if `sink` is a known sink of for this query
    */
   final predicate isKnownSink(DataFlow::Node sink) {
+    // If the list of characteristics includes positive indicators with maximal confidence for this class, then it's a
+    // known sink for the class.
+    isKnownSink(sink, this.getASinkEndpointType())
+  }
+
+  /**
+   * Holds if `sink` is a known sink for this query of type `sinkType`.
+   */
+  final predicate isKnownSink(DataFlow::Node sink, EndpointType sinkType) {
     // If the list of characteristics includes positive indicators with maximal confidence for this class, then it's a
     // known sink for the class.
     exists(EndpointCharacteristics::EndpointCharacteristic characteristic |
       characteristic.appliesToEndpoint(sink) and
-      characteristic
-          .hasImplications(this.getASinkEndpointType(), true, characteristic.maximalConfidence())
+      characteristic.hasImplications(sinkType, true, characteristic.maximalConfidence())
     )
   }
 
@@ -121,7 +129,7 @@ abstract class AtmConfig extends TaintTracking::Configuration {
   /**
    * EXPERIMENTAL. This API may change in the future.
    *
-   * Get an endpoint type for the sinks of this query. A query may have multiple applicable
+   * Get all sink types that can be sinks for this query. A query may have multiple applicable
    * endpoint types for its sinks.
    */
   abstract EndpointType getASinkEndpointType();

java/ql/experimental/adaptivethreatmodeling/src/ExtractPositiveExamples.ql

@@ -10,6 +10,7 @@
 private import java
 import semmle.code.java.dataflow.TaintTracking
 private import experimental.adaptivethreatmodeling.EndpointCharacteristics as EndpointCharacteristics
+private import experimental.adaptivethreatmodeling.EndpointTypes
 private import experimental.adaptivethreatmodeling.ATMConfig as AtmConfig
 private import experimental.adaptivethreatmodeling.SqlInjectionATM as SqlInjectionAtm
 private import experimental.adaptivethreatmodeling.TaintedPathATM as TaintedPathAtm
@@ -21,9 +22,9 @@ private import experimental.adaptivethreatmodeling.RequestForgeryATM as RequestF
  * the ML-gnerarated, noisy sinks will end up poluting the positive examples used in the prompt!
  */
 
-from DataFlow::Node sink, AtmConfig::AtmConfig config, string message
+from DataFlow::Node sink, AtmConfig::AtmConfig config, EndpointType sinkType, string message
 where
-  config.isKnownSink(sink) and
+  config.isKnownSink(sink, sinkType) and
   // If there are _any_ erroneous endpoints, return nothing. This will prevent us from accidentally running this query
   // when there's a codex-generated data extension file in `java/ql/lib/ext`.
   not EndpointCharacteristics::erroneousEndpoints(_, _, _, _, _) and
@@ -32,7 +33,7 @@ where
   // `isSink(n) and not isSanitizer(n)`. We don't want to include such nodes as positive examples in the prompt.
   not config.isSanitizer(sink) and
   message =
-    config.getASinkEndpointType().getDescription() + "\n" +
+    sinkType.getDescription() + "\n" +
       // Extract the needed metadata for this endpoint.
       any(string metadata | EndpointCharacteristics::hasMetadata(sink, metadata))
 select sink, message

java/ql/experimental/adaptivethreatmodeling/src/ExtractSinkCandidatesWithFlow.ql

@@ -28,7 +28,7 @@ where
         config.isSinkCandidateWithFlow(sinkPathNode) and
         sinkPathNode.getNode() = sink
       |
-        config.getASinkEndpointType().getDescription(), ", "
+        config, ", "
       ) + "\n" +
       // Extract the needed metadata for this endpoint.
       any(string metadata | EndpointCharacteristics::hasMetadata(sink, metadata))