Java: automodel fr mode: add mad output to extraction metadata

java/ql/automodel/src/AutomodelFrameworkModeCharacteristics.qll

@@ -37,6 +37,20 @@ abstract class FrameworkModeEndpoint extends TFrameworkModeEndpoint {
    */
   abstract int getIndex();
 
+  /**
+   * Gets the input (if any) for this endpoint, eg.: `Argument[0]`.
+   *
+   * For endpoints that are source candidates, this will be `none()`.
+   */
+  abstract string getMaDInput();
+
+  /**
+   * Gets the output (if any) for this endpoint, eg.: `ReturnValue`.
+   *
+   * For endpoints that are sink candidates, this will be `none()`.
+   */
+  abstract string getMaDOutput();
+
   /**
    * Returns the name of the parameter of the endpoint.
    */
@@ -63,6 +77,10 @@ class ExplicitParameterEndpoint extends FrameworkModeEndpoint, TExplicitParamete
 
   override int getIndex() { result = param.getPosition() }
 
+  override string getMaDInput() { result = "Argument[" + param.getPosition() + "]" }
+
+  override string getMaDOutput() { none() }
+
   override string getParamName() { result = param.getName() }
 
   override Callable getEnclosingCallable() { result = param.getCallable() }
@@ -81,6 +99,10 @@ class QualifierEndpoint extends FrameworkModeEndpoint, TQualifier {
 
   override int getIndex() { result = -1 }
 
+  override string getMaDInput() { result = "Argument[this]" }
+
+  override string getMaDOutput() { none() }
+
   override string getParamName() { result = "this" }
 
   override Callable getEnclosingCallable() { result = callable }
@@ -100,10 +122,11 @@ class ReturnValue extends FrameworkModeEndpoint, TReturnValue {
     result = -1
   }
 
-  override string getParamName() {
-    // FIXME bogus value
-    result = "return value"
-  }
+  override string getMaDInput() { none() }
+
+  override string getMaDOutput() { result = "ReturnValue" }
+
+  override string getParamName() { none() }
 
   override Callable getEnclosingCallable() { result = callable }
 
@@ -163,7 +186,7 @@ module FrameworkCandidatesImpl implements SharedCharacteristics::CandidateSig {
     FrameworkModeGetCallable::getCallable(e).hasQualifiedName(package, type, name) and
     signature = ExternalFlow::paramsString(FrameworkModeGetCallable::getCallable(e)) and
     ext = "" and
-    input = AutomodelJavaUtil::getArgumentForIndex(e.getIndex())
+    input = e.getMaDInput()
   }
 
   /**
@@ -213,11 +236,12 @@ class FrameworkModeMetadataExtractor extends string {
 
   predicate hasMetadata(
     Endpoint e, string package, string type, string subtypes, string name, string signature,
-    string input, string parameterName
+    string input, string output, string parameterName
   ) {
-    parameterName = e.getParamName() and
+    (if exists(e.getParamName()) then parameterName = e.getParamName() else parameterName = "") and
     name = e.getEnclosingCallable().getName() and
-    input = AutomodelJavaUtil::getArgumentForIndex(e.getIndex()) and
+    (if exists(e.getMaDInput()) then input = e.getMaDInput() else input = "") and
+    (if exists(e.getMaDOutput()) then output = e.getMaDOutput() else output = "") and
     package = e.getEnclosingCallable().getDeclaringType().getPackage().getName() and
     type = e.getEnclosingCallable().getDeclaringType().getErasure().(RefType).nestedName() and
     subtypes = AutomodelJavaUtil::considerSubtypes(e.getEnclosingCallable()).toString() and
@@ -285,8 +309,8 @@ private class ExceptionCharacteristic extends CharacteristicsImpl::NotASinkChara
  * A characteristic that limits candidates to parameters of methods that are recognized as `ModelApi`, iow., APIs that
  * are considered worth modeling.
  */
-private class NotAModelApiParameter extends CharacteristicsImpl::UninterestingToModelCharacteristic {
-  NotAModelApiParameter() { this = "not a model API parameter" }
+private class NotAModelApi extends CharacteristicsImpl::UninterestingToModelCharacteristic {
+  NotAModelApi() { this = "not a model API" }
 
   override predicate appliesToEndpoint(Endpoint e) {
     not e.getEnclosingCallable() instanceof ModelExclusions::ModelApi

java/ql/automodel/src/AutomodelFrameworkModeExtractCandidates.ql

@@ -18,8 +18,8 @@ private import AutomodelJavaUtil
 from
   Endpoint endpoint, string message, FrameworkModeMetadataExtractor meta, DollarAtString package,
   DollarAtString type, DollarAtString subtypes, DollarAtString name, DollarAtString signature,
-  DollarAtString input, DollarAtString parameterName, DollarAtString alreadyAiModeled,
-  DollarAtString extensibleType
+  DollarAtString input, DollarAtString output, DollarAtString parameterName,
+  DollarAtString alreadyAiModeled, DollarAtString extensibleType
 where
   endpoint.getExtensibleType() = extensibleType and
   not exists(CharacteristicsImpl::UninterestingToModelCharacteristic u |
@@ -36,7 +36,7 @@ where
     alreadyAiModeled.matches("%ai-%") and
     CharacteristicsImpl::isSink(endpoint, _, alreadyAiModeled)
   ) and
-  meta.hasMetadata(endpoint, package, type, subtypes, name, signature, input, parameterName) and
+  meta.hasMetadata(endpoint, package, type, subtypes, name, signature, input, output, parameterName) and
   includeAutomodelCandidate(package, type, name, signature) and
   // The message is the concatenation of all sink types for which this endpoint is known neither to be a sink nor to be
   // a non-sink, and we surface only endpoints that have at least one such sink type.
@@ -48,7 +48,7 @@ where
       sinkType, ", "
     )
 select endpoint,
-  message + "\nrelated locations: $@, $@." + "\nmetadata: $@, $@, $@, $@, $@, $@, $@, $@.", //
+  message + "\nrelated locations: $@, $@." + "\nmetadata: $@, $@, $@, $@, $@, $@, $@, $@, $@.", //
   CharacteristicsImpl::getRelatedLocationOrCandidate(endpoint, MethodDoc()), "MethodDoc", //
   CharacteristicsImpl::getRelatedLocationOrCandidate(endpoint, ClassDoc()), "ClassDoc", //
   package, "package", //
@@ -57,6 +57,7 @@ select endpoint,
   name, "name", //
   signature, "signature", //
   input, "input", //
+  output, "output", //
   parameterName, "parameterName", //
   alreadyAiModeled, "alreadyAiModeled", //
   extensibleType, "extensibleType"

java/ql/automodel/src/AutomodelFrameworkModeExtractNegativeExamples.ql

@@ -16,7 +16,8 @@ from
   Endpoint endpoint, EndpointCharacteristic characteristic, float confidence,
   DollarAtString message, FrameworkModeMetadataExtractor meta, DollarAtString package,
   DollarAtString type, DollarAtString subtypes, DollarAtString name, DollarAtString signature,
-  DollarAtString input, DollarAtString parameterName, DollarAtString extensibleType
+  DollarAtString input, DollarAtString output, DollarAtString parameterName,
+  DollarAtString extensibleType
 where
   endpoint.getExtensibleType() = extensibleType and
   characteristic.appliesToEndpoint(endpoint) and
@@ -25,7 +26,7 @@ where
   // Exclude endpoints that have contradictory endpoint characteristics, because we only want examples we're highly
   // certain about in the prompt.
   not erroneousEndpoints(endpoint, _, _, _, _, false) and
-  meta.hasMetadata(endpoint, package, type, subtypes, name, signature, input, parameterName) and
+  meta.hasMetadata(endpoint, package, type, subtypes, name, signature, input, output, parameterName) and
   // It's valid for a node to satisfy the logic for both `isSink` and `isSanitizer`, but in that case it will be
   // 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 negative examples in the prompt, because
@@ -47,5 +48,6 @@ select endpoint,
   name, "name", //
   signature, "signature", //
   input, "input", //
+  output, "output", //
   parameterName, "parameterName", //
   extensibleType, "extensibleType"

java/ql/automodel/src/AutomodelFrameworkModeExtractPositiveExamples.ql

@@ -15,13 +15,13 @@ private import AutomodelJavaUtil
 from
   Endpoint endpoint, SinkType sinkType, FrameworkModeMetadataExtractor meta, DollarAtString package,
   DollarAtString type, DollarAtString subtypes, DollarAtString name, DollarAtString signature,
-  DollarAtString input, DollarAtString parameterName, DollarAtString extensibleType
+  DollarAtString input, DollarAtString output, DollarAtString parameterName, DollarAtString extensibleType
 where
   endpoint.getExtensibleType() = extensibleType and
   // Exclude endpoints that have contradictory endpoint characteristics, because we only want examples we're highly
   // certain about in the prompt.
   not erroneousEndpoints(endpoint, _, _, _, _, false) and
-  meta.hasMetadata(endpoint, package, type, subtypes, name, signature, input, parameterName) and
+  meta.hasMetadata(endpoint, package, type, subtypes, name, signature, input, output, parameterName) and
   // Extract positive examples of sinks belonging to the existing ATM query configurations.
   CharacteristicsImpl::isKnownAs(endpoint, sinkType, _)
 select endpoint,
@@ -34,5 +34,6 @@ select endpoint,
   name, "name", //
   signature, "signature", //
   input, "input", //
+  output, "output", //
   parameterName, "parameterName", //
   extensibleType, "extensibleType"

java/ql/automodel/src/AutomodelJavaUtil.qll

@@ -44,14 +44,6 @@ predicate isKnownKind(string kind, AutomodelEndpointTypes::EndpointType type) {
   type instanceof AutomodelEndpointTypes::RemoteSourceType
 }
 
-/** Gets the models-as-data description for the method argument with the index `index`. */
-bindingset[index]
-string getArgumentForIndex(int index) {
-  index = -1 and result = "Argument[this]"
-  or
-  index >= 0 and result = "Argument[" + index + "]"
-}
-
 /**
  * By convention, the subtypes property of the MaD declaration should only be
  * true when there _can_ exist any subtypes with a different implementation.