Java: support remote sources in automodel positive example extraction

Author: Stephan Brandauer

java/ql/automodel/src/AutomodelApplicationModeCharacteristics.qll

@@ -205,6 +205,13 @@ module ApplicationCandidatesImpl implements SharedCharacteristics::CandidateSig
     isCustomSink(e, kind) and provenance = "custom-sink"
   }
 
+  predicate isSource(Endpoint e, string kind, string provenance) {
+    exists(string package, string type, string name, string signature, string ext, string output |
+      sourceSpec(e, package, type, name, signature, ext, output) and
+      ExternalFlow::sourceModel(package, type, _, name, [signature, ""], ext, output, kind, provenance)
+    )
+  }
+
   predicate isNeutral(Endpoint e) {
     exists(string package, string type, string name, string signature |
       sinkSpec(e, package, type, name, signature, _, _) and
@@ -222,6 +229,15 @@ module ApplicationCandidatesImpl implements SharedCharacteristics::CandidateSig
     input = e.getMaDInput()
   }
 
+  additional predicate sourceSpec(
+    Endpoint e, string package, string type, string name, string signature, string ext, string output
+  ) {
+    ApplicationModeGetCallable::getCallable(e).hasQualifiedName(package, type, name) and
+    signature = ExternalFlow::paramsString(ApplicationModeGetCallable::getCallable(e)) and
+    ext = "" and
+    output = e.getMaDOutput()
+  }
+
   /**
    * Gets the related location for the given endpoint.
    *

java/ql/automodel/src/AutomodelApplicationModeExtractCandidates.ql

@@ -83,7 +83,7 @@ where
   // a non-sink, and we surface only endpoints that have at least one such sink type.
   message =
     strictconcat(AutomodelEndpointTypes::SinkType sinkType |
-      not CharacteristicsImpl::isKnownSink(endpoint, sinkType, _) and
+      not CharacteristicsImpl::isKnownAs(endpoint, sinkType, _) and
       CharacteristicsImpl::isSinkCandidate(endpoint, sinkType)
     |
       sinkType, ", "

java/ql/automodel/src/AutomodelApplicationModeExtractPositiveExamples.ql

@@ -13,7 +13,7 @@ private import AutomodelEndpointTypes
 private import AutomodelJavaUtil
 
 from
-  Endpoint endpoint, SinkType sinkType, ApplicationModeMetadataExtractor meta,
+  Endpoint endpoint, EndpointType endpointType, ApplicationModeMetadataExtractor meta,
   DollarAtString package, DollarAtString type, DollarAtString subtypes, DollarAtString name,
   DollarAtString signature, DollarAtString input, DollarAtString output, DollarAtString isVarargsArray
 where
@@ -22,10 +22,10 @@ where
   not erroneousEndpoints(endpoint, _, _, _, _, false) 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
+  CharacteristicsImpl::isKnownAs(endpoint, endpointType, _) and
   exists(CharacteristicsImpl::getRelatedLocationOrCandidate(endpoint, CallContext()))
 select endpoint.asNode(),
-  sinkType + "\nrelated locations: $@." + "\nmetadata: $@, $@, $@, $@, $@, $@, $@, $@.", //
+  endpointType + "\nrelated locations: $@." + "\nmetadata: $@, $@, $@, $@, $@, $@, $@, $@.", //
   CharacteristicsImpl::getRelatedLocationOrCandidate(endpoint, CallContext()), "CallContext", //
   package, "package", //
   type, "type", //

java/ql/automodel/src/AutomodelEndpointTypes.qll

@@ -28,12 +28,6 @@ abstract class SinkType extends EndpointType {
   SinkType() { any() }
 }
 
-/** A class for source types that can be predicted by a classifier. */
-abstract class SourceType extends EndpointType {
-  bindingset[this]
-  SourceType() { any() }
-}
-
 /** The `Negative` class for non-sinks. */
 class NegativeSinkType extends SinkType {
   NegativeSinkType() { this = "non-sink" }
@@ -58,3 +52,14 @@ class RequestForgerySinkType extends SinkType {
 class CommandInjectionSinkType extends SinkType {
   CommandInjectionSinkType() { this = "command-injection" }
 }
+
+/** A class for source types that can be predicted by a classifier. */
+abstract class SourceType extends EndpointType {
+  bindingset[this]
+  SourceType() { any() }
+}
+
+/** A source of remote data. */
+class RemoteSourceType extends SourceType {
+  RemoteSourceType() { this = "remote" }
+}

java/ql/automodel/src/AutomodelFrameworkModeCharacteristics.qll

@@ -117,6 +117,10 @@ module FrameworkCandidatesImpl implements SharedCharacteristics::CandidateSig {
     )
   }
 
+  predicate isSource(Endpoint e, string kind, string provenance) {
+    none() // TODO: implement
+  }
+
   predicate isNeutral(Endpoint e) {
     exists(string package, string type, string name, string signature |
       sinkSpec(e, package, type, name, signature, _, _) and

java/ql/automodel/src/AutomodelFrameworkModeExtractCandidates.ql

@@ -40,7 +40,7 @@ where
   // a non-sink, and we surface only endpoints that have at least one such sink type.
   message =
     strictconcat(AutomodelEndpointTypes::SinkType sinkType |
-      not CharacteristicsImpl::isKnownSink(endpoint, sinkType, _) and
+      not CharacteristicsImpl::isKnownAs(endpoint, sinkType, _) and
       CharacteristicsImpl::isSinkCandidate(endpoint, sinkType)
     |
       sinkType, ", "

java/ql/automodel/src/AutomodelFrameworkModeExtractPositiveExamples.ql

@@ -22,7 +22,7 @@ where
   not erroneousEndpoints(endpoint, _, _, _, _, false) and
   meta.hasMetadata(endpoint, package, type, subtypes, name, signature, input, parameterName) and
   // Extract positive examples of sinks belonging to the existing ATM query configurations.
-  CharacteristicsImpl::isKnownSink(endpoint, sinkType, _)
+  CharacteristicsImpl::isKnownAs(endpoint, sinkType, _)
 select endpoint,
   sinkType + "\nrelated locations: $@, $@." + "\nmetadata: $@, $@, $@, $@, $@, $@, $@.", //
   CharacteristicsImpl::getRelatedLocationOrCandidate(endpoint, MethodDoc()), "MethodDoc", //

java/ql/automodel/src/AutomodelJavaUtil.qll

@@ -39,6 +39,9 @@ predicate isKnownKind(string kind, AutomodelEndpointTypes::EndpointType type) {
   or
   kind = "command-injection" and
   type instanceof AutomodelEndpointTypes::CommandInjectionSinkType
+  or
+  kind = "remote" and
+  type instanceof AutomodelEndpointTypes::RemoteSourceType
 }
 
 /** Gets the models-as-data description for the method argument with the index `index`. */

java/ql/automodel/src/AutomodelSharedCharacteristics.qll

@@ -62,6 +62,11 @@ signature module CandidateSig {
    */
   predicate isSink(Endpoint e, string kind, string provenance);
 
+  /**
+   * Holds if `e` is a sink with the label `kind`, and provenance `provenance`.
+   */
+  predicate isSource(Endpoint e, string kind, string provenance);
+
   /**
    * Holds if `e` is not a sink of any kind.
    */
@@ -92,16 +97,16 @@ module SharedCharacteristics<CandidateSig Candidate> {
   predicate isNeutral = Candidate::isNeutral/1;
 
   /**
-   * Holds if `sink` is a known sink of type `endpointType`.
+   * Holds if `endpoint` is modeled as `endpointType` (endpoint type must not be negative).
    */
-  predicate isKnownSink(
-    Candidate::Endpoint sink, Candidate::EndpointType endpointType,
+  predicate isKnownAs(
+    Candidate::Endpoint endpoint, Candidate::EndpointType endpointType,
     EndpointCharacteristic characteristic
   ) {
     // If the list of characteristics includes positive indicators with maximal confidence for this class, then it's a
     // known sink for the class.
     not endpointType instanceof Candidate::NegativeEndpointType and
-    characteristic.appliesToEndpoint(sink) and
+    characteristic.appliesToEndpoint(endpoint) and
     characteristic.hasImplications(endpointType, true, maximalConfidence())
   }
 
@@ -209,6 +214,25 @@ module SharedCharacteristics<CandidateSig Candidate> {
     }
   }
 
+  /**
+   * A high-confidence characteristic that indicates that an endpoint is a source of a specified type. These endpoints can
+   * be used as positive samples for training or for a few-shot prompt.
+   */
+  abstract class SourceCharacteristic extends EndpointCharacteristic {
+    bindingset[this]
+    SourceCharacteristic() { any() }
+
+    abstract Candidate::EndpointType getSourceType();
+
+    final override predicate hasImplications(
+      Candidate::EndpointType endpointType, boolean isPositiveIndicator, float confidence
+    ) {
+      endpointType = this.getSourceType() and
+      isPositiveIndicator = true and
+      confidence = maximalConfidence()
+    }
+  }
+
   /**
    * A high-confidence characteristic that indicates that an endpoint is not a sink of any type. These endpoints can be
    * used as negative samples for training or for a few-shot prompt.
@@ -292,6 +316,25 @@ module SharedCharacteristics<CandidateSig Candidate> {
       override Candidate::EndpointType getSinkType() { result = endpointType }
     }
 
+    private class KnownSourceCharacteristic extends SourceCharacteristic {
+      string madKind;
+      Candidate::EndpointType endpointType;
+      string provenance;
+
+      KnownSourceCharacteristic() {
+        Candidate::isKnownKind(madKind, endpointType) and
+        // bind "this" to a unique string differing from that of the SinkType classes
+        this = madKind + "_" + provenance + "_characteristic" and
+        Candidate::isSource(_, madKind, provenance)
+      }
+
+      override predicate appliesToEndpoint(Candidate::Endpoint e) {
+        Candidate::isSource(e, madKind, provenance)
+      }
+
+      override Candidate::EndpointType getSourceType() { result = endpointType }
+    }
+
     /**
      * A negative characteristic that indicates that an endpoint was manually modeled as a neutral model.
      */