From 65923ed2c15b69a643d8be1eadd97a1089c53f4a Mon Sep 17 00:00:00 2001
From: tiferet <tiferet@github.com>
Date: Thu, 2 Feb 2023 12:48:54 -0800
Subject: [PATCH] Add support for multiple sink types per query

---
 .../adaptivethreatmodeling/ATMConfig.qll         | 16 ++++++++++++----
 .../src/ExtractPositiveExamples.ql               |  7 ++++---
 .../src/ExtractSinkCandidatesWithFlow.ql         |  2 +-
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/java/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/ATMConfig.qll b/java/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/ATMConfig.qll
index 4614a5fbae3..c6ff581439d 100644
--- a/java/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/ATMConfig.qll
+++ b/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();
diff --git a/java/ql/experimental/adaptivethreatmodeling/src/ExtractPositiveExamples.ql b/java/ql/experimental/adaptivethreatmodeling/src/ExtractPositiveExamples.ql
index dc58251b10e..494b674b348 100644
--- a/java/ql/experimental/adaptivethreatmodeling/src/ExtractPositiveExamples.ql
+++ b/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
diff --git a/java/ql/experimental/adaptivethreatmodeling/src/ExtractSinkCandidatesWithFlow.ql b/java/ql/experimental/adaptivethreatmodeling/src/ExtractSinkCandidatesWithFlow.ql
index a7d296e243a..b204a96262d 100644
--- a/java/ql/experimental/adaptivethreatmodeling/src/ExtractSinkCandidatesWithFlow.ql
+++ b/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))