JS: Only featurize endpoints that are part of a flow path

Author: henrymercer

javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointFeatures.qll

@@ -6,7 +6,7 @@
 
 import javascript
 import CodeToFeatures
-import EndpointScoring
+private import EndpointScoring
 
 /**
  * Gets the value of the token-based feature named `featureName` for the endpoint `endpoint`.
@@ -269,14 +269,22 @@ private string getASupportedFeatureName() {
     ]
 }
 
+/** A configuration that defines which endpoints should be featurized. */
+abstract class FeaturizationConfig extends string {
+  bindingset[this]
+  FeaturizationConfig() { any() }
+
+  abstract DataFlow::Node getAnEndpointToFeaturize();
+}
+
 /**
  * Generic token-based features for ATM.
  *
  * This predicate holds if the generic token-based feature named `featureName` has the value
  * `featureValue` for the endpoint `endpoint`.
  */
 predicate tokenFeatures(DataFlow::Node endpoint, string featureName, string featureValue) {
-  ModelScoring::endpoints(endpoint) and
+  endpoint = any(FeaturizationConfig cfg).getAnEndpointToFeaturize() and
   (
     if strictcount(getTokenFeature(endpoint, featureName)) = 1
     then featureValue = getTokenFeature(endpoint, featureName)

javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointScoring.qll

@@ -80,23 +80,30 @@ DatabaseFeatures::Entity getRepresentativeEntityForEndpoint(DataFlow::Node endpo
 }
 
 module ModelScoring {
-  predicate endpoints(DataFlow::Node endpoint) {
-    getCfg().isEffectiveSource(endpoint) or
-    getCfg().isEffectiveSink(endpoint)
+  /**
+   * A featurization config that only featurizes new candidate endpoints that are part of a flow
+   * path.
+   */
+  class RelevantFeaturizationConfig extends EndpointFeatures::FeaturizationConfig {
+    RelevantFeaturizationConfig() { this = "RelevantFeaturization" }
+
+    override DataFlow::Node getAnEndpointToFeaturize() {
+      getCfg().isEffectiveSource(result) and any(DataFlow::Configuration cfg).hasFlow(result, _)
+      or
+      getCfg().isEffectiveSink(result) and any(DataFlow::Configuration cfg).hasFlow(_, result)
+    }
   }
 
-  private int requestedEndpointTypes() { result = any(EndpointType type).getEncoding() }
-
-  private predicate relevantTokenFeatures(
-    DataFlow::Node endpoint, string featureName, string featureValue
-  ) {
-    endpoints(endpoint) and
-    EndpointFeatures::tokenFeatures(endpoint, featureName, featureValue)
+  DataFlow::Node getARequestedEndpoint() {
+    result = any(EndpointFeatures::FeaturizationConfig cfg).getAnEndpointToFeaturize()
   }
 
+  private int getARequestedEndpointType() { result = any(EndpointType type).getEncoding() }
+
   predicate endpointScores(DataFlow::Node endpoint, int encodedEndpointType, float score) =
-    scoreEndpoints(endpoints/1, requestedEndpointTypes/0, relevantTokenFeatures/3,
-      getACompatibleModelChecksum/0)(endpoint, encodedEndpointType, score)
+    scoreEndpoints(getARequestedEndpoint/0, getARequestedEndpointType/0,
+      EndpointFeatures::tokenFeatures/3, getACompatibleModelChecksum/0)(endpoint,
+      encodedEndpointType, score)
 }
 
 /**
@@ -212,7 +219,9 @@ class EndpointScoringResults extends ScoringResults {
 }
 
 module Debugging {
-  query predicate hopInputEndpoints = ModelScoring::endpoints/1;
+  query predicate hopInputEndpoints(DataFlow::Node endpoint) {
+    endpoint = ModelScoring::getARequestedEndpoint()
+  }
 
   query predicate endpointScores = ModelScoring::endpointScores/3;