From abe3a2dae1b8f40f8dee164b98148760d1af2107 Mon Sep 17 00:00:00 2001
From: tiferet <tiferet@github.com>
Date: Thu, 16 Feb 2023 16:42:57 -0800
Subject: [PATCH] Improve positive prompt examples:

Include only sinks that are arguments to an external API call, because these are the sinks we are most interested in.
---
 .../adaptivethreatmodeling/src/ExtractPositiveExamples.ql     | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/java/ql/experimental/adaptivethreatmodeling/src/ExtractPositiveExamples.ql b/java/ql/experimental/adaptivethreatmodeling/src/ExtractPositiveExamples.ql
index c35d997f920..4df4af62f33 100644
--- a/java/ql/experimental/adaptivethreatmodeling/src/ExtractPositiveExamples.ql
+++ b/java/ql/experimental/adaptivethreatmodeling/src/ExtractPositiveExamples.ql
@@ -9,6 +9,7 @@
 
 private import java
 import semmle.code.java.dataflow.TaintTracking
+private import semmle.code.java.security.ExternalAPIs as ExternalAPIs
 private import experimental.adaptivethreatmodeling.EndpointCharacteristics as EndpointCharacteristics
 private import experimental.adaptivethreatmodeling.EndpointTypes
 private import experimental.adaptivethreatmodeling.ATMConfig as AtmConfig
@@ -34,6 +35,9 @@ where
     // treated by the actual query as a sanitizer, since the final logic is something like
     // `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
+    // Include only sinks that are arguments to an external API call, because these are the sinks we are most interested
+    // in.
+    sink instanceof ExternalAPIs::ExternalApiDataNode and
     message =
       sinkType.getDescription() + "\n" +
         // Extract the needed metadata for this endpoint.