Java: Export MaD output in application mode extraction queries

Author: Stephan Brandauer

java/ql/automodel/src/AutomodelApplicationModeCharacteristics.qll

@@ -46,11 +46,21 @@ abstract private class ApplicationModeEndpoint extends TApplicationModeEndpoint
 
   abstract string getMaDInput();
 
+  abstract string getMaDOutput();
+
   abstract Top asTop();
 
   abstract DataFlow::Node asNode();
 
-  abstract string getExtensibleType();
+  string getExtensibleType() {
+    // XXX the sourceModel still implements a bogus getMaDInput() method, so we can't use this yet
+    if /* not exists(this.getMaDInput()) and */ exists(this.getMaDOutput())
+    then result = "sourceModel"
+    else
+      if exists(this.getMaDInput()) and not exists(this.getMaDOutput())
+      then result = "sinkModel"
+      else none() // if both exist, it would be a summaryModel (not yet supported)
+  }
 
   abstract string toString();
 }
@@ -68,16 +78,14 @@ class ExplicitArgument extends ApplicationModeEndpoint, TExplicitArgument {
 
   private int getArgIndex() { this.asTop() = call.getArgument(result) }
 
-  override string getMaDInput() {
-    result = "Argument[" + this.getArgIndex() + "]"
-  }
+  override string getMaDInput() { result = "Argument[" + this.getArgIndex() + "]" }
+
+  override string getMaDOutput() { none() }
 
   override Top asTop() { result = arg.asExpr() }
 
   override DataFlow::Node asNode() { result = arg }
 
-  override string getExtensibleType() { result = "sinkModel" }
-
   override string toString() { result = arg.toString() }
 }
 
@@ -91,12 +99,12 @@ class InstanceArgument extends ApplicationModeEndpoint, TInstanceArgument {
 
   override string getMaDInput() { result = "Argument[this]" }
 
+  override string getMaDOutput() { none() }
+
   override Top asTop() { if exists(arg.asExpr()) then result = arg.asExpr() else result = call }
 
   override DataFlow::Node asNode() { result = arg }
 
-  override string getExtensibleType() { result = "sinkModel" }
-
   override string toString() { result = arg.toString() }
 }
 
@@ -118,16 +126,14 @@ class ImplicitVarargsArray extends ApplicationModeEndpoint, TImplicitVarargsArra
 
   override Call getCall() { result = call }
 
-  override string getMaDInput() {
-    result = "Argument[" + idx + "]"
-  }
+  override string getMaDInput() { result = "Argument[" + idx + "]" }
+
+  override string getMaDOutput() { none() }
 
   override Top asTop() { result = this.getCall() }
 
   override DataFlow::Node asNode() { result = vararg }
 
-  override string getExtensibleType() { result = "sinkModel" }
-
   override string toString() { result = vararg.toString() }
 }
 
@@ -143,12 +149,12 @@ class MethodCall extends ApplicationModeEndpoint, TMethodCall {
 
   override string getMaDInput() { result = "Argument[this]" }
 
+  override string getMaDOutput() { result = "ReturnValue" }
+
   override Top asTop() { result = call }
 
   override DataFlow::Node asNode() { result.asExpr() = call }
 
-  override string getExtensibleType() { result = "sourceModel" }
-
   override string toString() { result = call.toString() }
 }
 
@@ -210,7 +216,6 @@ module ApplicationCandidatesImpl implements SharedCharacteristics::CandidateSig
   additional predicate sinkSpec(
     Endpoint e, string package, string type, string name, string signature, string ext, string input
   ) {
-    e.getExtensibleType() = "sinkModel" and
     ApplicationModeGetCallable::getCallable(e).hasQualifiedName(package, type, name) and
     signature = ExternalFlow::paramsString(ApplicationModeGetCallable::getCallable(e)) and
     ext = "" and
@@ -269,11 +274,12 @@ class ApplicationModeMetadataExtractor extends string {
 
   predicate hasMetadata(
     Endpoint e, string package, string type, string subtypes, string name, string signature,
-    string input, string isVarargsArray
+    string input, string output, string isVarargsArray
   ) {
     exists(Callable callable |
       e.getCall().getCallee() = callable and
-      input = e.getMaDInput() and
+      (if exists(e.getMaDInput()) then input = e.getMaDInput() else input = "") and
+      (if exists(e.getMaDOutput()) then output = e.getMaDOutput() else output = "") and
       package = callable.getDeclaringType().getPackage().getName() and
       // we're using the erased types because the MaD convention is to not specify type parameters.
       // Whether something is or isn't a sink doesn't usually depend on the type parameters.

java/ql/automodel/src/AutomodelApplicationModeExtractCandidates.ql

@@ -25,20 +25,20 @@ private import AutomodelJavaUtil
 bindingset[limit]
 private Endpoint getSampleForSignature(
   int limit, string package, string type, string subtypes, string name, string signature,
-  string input, string isVarargs, string extensibleType
+  string input, string output, string isVarargs, string extensibleType
 ) {
   exists(int n, int num_endpoints, ApplicationModeMetadataExtractor meta |
     num_endpoints =
       count(Endpoint e |
         e.getExtensibleType() = extensibleType and
-        meta.hasMetadata(e, package, type, subtypes, name, signature, input, isVarargs)
+        meta.hasMetadata(e, package, type, subtypes, name, signature, input, output, isVarargs)
       )
   |
     result =
       rank[n](Endpoint e, Location loc |
         loc = e.asTop().getLocation() and
         e.getExtensibleType() = extensibleType and
-        meta.hasMetadata(e, package, type, subtypes, name, signature, input, isVarargs)
+        meta.hasMetadata(e, package, type, subtypes, name, signature, input, output, isVarargs)
       |
         e
         order by
@@ -57,13 +57,15 @@ private Endpoint getSampleForSignature(
 from
   Endpoint endpoint, string message, ApplicationModeMetadataExtractor meta, DollarAtString package,
   DollarAtString type, DollarAtString subtypes, DollarAtString name, DollarAtString signature,
-  DollarAtString input, DollarAtString isVarargsArray, DollarAtString alreadyAiModeled, DollarAtString extensibleType
+  DollarAtString input, DollarAtString output, DollarAtString isVarargsArray,
+  DollarAtString alreadyAiModeled, DollarAtString extensibleType
 where
   not exists(CharacteristicsImpl::UninterestingToModelCharacteristic u |
     u.appliesToEndpoint(endpoint)
   ) and
   endpoint =
-    getSampleForSignature(9, package, type, subtypes, name, signature, input, isVarargsArray, extensibleType) and
+    getSampleForSignature(9, package, type, subtypes, name, signature, input, output,
+      isVarargsArray, extensibleType) and
   // If a node is already a known sink for any of our existing ATM queries and is already modeled as a MaD sink, we
   // don't include it as a candidate. Otherwise, we might include it as a candidate for query A, but the model will
   // label it as a sink for one of the sink types of query B, for which it's already a known sink. This would result in
@@ -75,7 +77,7 @@ where
     alreadyAiModeled.matches("%ai-%") and
     CharacteristicsImpl::isSink(endpoint, _, alreadyAiModeled)
   ) and
-  meta.hasMetadata(endpoint, package, type, subtypes, name, signature, input, isVarargsArray) and
+  meta.hasMetadata(endpoint, package, type, subtypes, name, signature, input, output, isVarargsArray) 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.
@@ -87,14 +89,15 @@ where
       sinkType, ", "
     )
 select endpoint.asNode(),
-  message + "\nrelated locations: $@." + "\nmetadata: $@, $@, $@, $@, $@, $@, $@, $@, $@.", //
+  message + "\nrelated locations: $@." + "\nmetadata: $@, $@, $@, $@, $@, $@, $@, $@, $@, $@.", //
   CharacteristicsImpl::getRelatedLocationOrCandidate(endpoint, CallContext()), "CallContext", //
   package, "package", //
   type, "type", //
   subtypes, "subtypes", //
   name, "name", // method name
   signature, "signature", //
   input, "input", //
+  output, "output", //
   isVarargsArray, "isVarargsArray", //
   alreadyAiModeled, "alreadyAiModeled", //
   extensibleType, "extensibleType"

java/ql/automodel/src/AutomodelApplicationModeExtractNegativeExamples.ql

@@ -44,15 +44,15 @@ from
   Endpoint endpoint, EndpointCharacteristic characteristic, float confidence, string message,
   ApplicationModeMetadataExtractor meta, DollarAtString package, DollarAtString type,
   DollarAtString subtypes, DollarAtString name, DollarAtString signature, DollarAtString input,
-  DollarAtString isVarargsArray
+  DollarAtString output, DollarAtString isVarargsArray
 where
   endpoint = getSampleForCharacteristic(characteristic, 100) and
   confidence >= SharedCharacteristics::highConfidence() and
   characteristic.hasImplications(any(NegativeSinkType negative), true, confidence) 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, isVarargsArray) and
+  meta.hasMetadata(endpoint, package, type, subtypes, name, signature, input, output, isVarargsArray) 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
@@ -65,12 +65,13 @@ where
   ) and
   message = characteristic
 select endpoint.asNode(),
-  message + "\nrelated locations: $@." + "\nmetadata: $@, $@, $@, $@, $@, $@, $@.", //
+  message + "\nrelated locations: $@." + "\nmetadata: $@, $@, $@, $@, $@, $@, $@, $@.", //
   CharacteristicsImpl::getRelatedLocationOrCandidate(endpoint, CallContext()), "CallContext", //
   package, "package", //
   type, "type", //
   subtypes, "subtypes", //
   name, "name", //
   signature, "signature", //
   input, "input", //
+  output, "output", //
   isVarargsArray, "isVarargsArray" //

java/ql/automodel/src/AutomodelApplicationModeExtractPositiveExamples.ql

@@ -15,22 +15,23 @@ private import AutomodelJavaUtil
 from
   Endpoint endpoint, SinkType sinkType, ApplicationModeMetadataExtractor meta,
   DollarAtString package, DollarAtString type, DollarAtString subtypes, DollarAtString name,
-  DollarAtString signature, DollarAtString input, DollarAtString isVarargsArray
+  DollarAtString signature, DollarAtString input, DollarAtString output, DollarAtString isVarargsArray
 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, isVarargsArray) and
+  meta.hasMetadata(endpoint, package, type, subtypes, name, signature, input, output, isVarargsArray) and
   // Extract positive examples of sinks belonging to the existing ATM query configurations.
   CharacteristicsImpl::isKnownSink(endpoint, sinkType, _) and
   exists(CharacteristicsImpl::getRelatedLocationOrCandidate(endpoint, CallContext()))
 select endpoint.asNode(),
-  sinkType + "\nrelated locations: $@." + "\nmetadata: $@, $@, $@, $@, $@, $@, $@.", //
+  sinkType + "\nrelated locations: $@." + "\nmetadata: $@, $@, $@, $@, $@, $@, $@, $@.", //
   CharacteristicsImpl::getRelatedLocationOrCandidate(endpoint, CallContext()), "CallContext", //
   package, "package", //
   type, "type", //
   subtypes, "subtypes", //
   name, "name", //
   signature, "signature", //
   input, "input", //
+  output, "output", //
   isVarargsArray, "isVarargsArray"