Add utility predicates to FunctionModel

Co-authored-by: Chris Smowton <smowton@github.com>

Author: Sauyon Lee

ql/src/Security/CWE-352/ConstantOauth2State.ql

@@ -101,7 +101,7 @@ class PrivateUrlFlowsToAuthCodeUrlCall extends DataFlow::Configuration {
     TaintTracking::referenceStep(pred, succ)
     or
     // Propagate across Sprintf and similar calls
-    TaintTracking::functionModelStep(any(Fmt::Sprinter s), pred, succ)
+    any(Fmt::Sprinter s).taintStep(pred, succ)
   }
 
   predicate isSink(DataFlow::Node sink, DataFlow::CallNode call) {

ql/src/semmle/go/dataflow/internal/DataFlowUtil.qll

@@ -896,6 +896,26 @@ class TypeCastNode extends ExprNode {
 abstract class FunctionModel extends Function {
   /** Holds if data flows through this function from `input` to `output`. */
   abstract predicate hasDataFlow(FunctionInput input, FunctionOutput output);
+
+  /** Gets an input node for this model for the call `c`. */
+  DataFlow::Node getAnInputNode(DataFlow::CallNode c) { this.flowStepForCall(result, _, c) }
+
+  /** Gets an output node for this model for the call `c`. */
+  DataFlow::Node getAnOutputNode(DataFlow::CallNode c) { this.flowStepForCall(_, result, c) }
+
+  /** Holds if this function model causes data to flow from `pred` to `succ` for the call `c`. */
+  predicate flowStepForCall(DataFlow::Node pred, DataFlow::Node succ, DataFlow::CallNode c) {
+    c = this.getACall() and
+    exists(FunctionInput inp, FunctionOutput outp | this.hasDataFlow(inp, outp) |
+      pred = inp.getNode(c) and
+      succ = outp.getNode(c)
+    )
+  }
+
+  /** Holds if this function model causes data to flow from `pred` to `succ`. */
+  predicate flowStep(DataFlow::Node pred, DataFlow::Node succ) {
+    this.flowStepForCall(pred, succ, _)
+  }
 }
 
 /**
@@ -1006,12 +1026,7 @@ predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
   basicLocalFlowStep(nodeFrom, nodeTo)
   or
   // step through function model
-  exists(FunctionModel m, CallNode c, FunctionInput inp, FunctionOutput outp |
-    c = m.getACall() and
-    m.hasDataFlow(inp, outp) and
-    nodeFrom = inp.getNode(c) and
-    nodeTo = outp.getNode(c)
-  )
+  any(FunctionModel m).flowStep(nodeFrom, nodeTo)
 }
 
 /**

ql/src/semmle/go/dataflow/internal/TaintTrackingUtil.qll

@@ -60,7 +60,7 @@ predicate localAdditionalTaintStep(DataFlow::Node pred, DataFlow::Node succ) {
   tupleStep(pred, succ) or
   stringConcatStep(pred, succ) or
   sliceStep(pred, succ) or
-  functionModelStep(_, pred, succ) or
+  any(FunctionModel fm).taintStep(pred, succ) or
   any(AdditionalTaintStep a).step(pred, succ)
 }
 
@@ -140,23 +140,33 @@ predicate sliceStep(DataFlow::Node pred, DataFlow::Node succ) {
   succ.(DataFlow::SliceNode).getBase() = pred
 }
 
-/** Holds if taint flows from `pred` to `succ` via a function model. */
-predicate functionModelStep(FunctionModel fn, DataFlow::Node pred, DataFlow::Node succ) {
-  exists(DataFlow::CallNode c, FunctionInput inp, FunctionOutput outp |
-    c = fn.getACall() and
-    fn.hasTaintFlow(inp, outp) and
-    pred = inp.getNode(c) and
-    succ = outp.getNode(c)
-  )
-}
-
 /**
  * A model of a function specifying that the function propagates taint from
  * a parameter or qualifier to a result.
  */
 abstract class FunctionModel extends Function {
   /** Holds if taint propagates through this function from `input` to `output`. */
   abstract predicate hasTaintFlow(FunctionInput input, FunctionOutput output);
+
+  /** Gets an input node for this model for the call `c`. */
+  DataFlow::Node getAnInputNode(DataFlow::CallNode c) { this.taintStepForCall(result, _, c) }
+
+  /** Gets an output node for this model for the call `c`. */
+  DataFlow::Node getAnOutputNode(DataFlow::CallNode c) { this.taintStepForCall(_, result, c) }
+
+  /** Holds if this function model causes taint to flow from `pred` to `succ` for the call `c`. */
+  predicate taintStepForCall(DataFlow::Node pred, DataFlow::Node succ, DataFlow::CallNode c) {
+    c = this.getACall() and
+    exists(FunctionInput inp, FunctionOutput outp | this.hasTaintFlow(inp, outp) |
+      pred = inp.getNode(c) and
+      succ = outp.getNode(c)
+    )
+  }
+
+  /** Holds if this function model causes taint to flow from `pred` to `succ`. */
+  predicate taintStep(DataFlow::Node pred, DataFlow::Node succ) {
+    this.taintStepForCall(pred, succ, _)
+  }
 }
 
 /**

ql/src/semmle/go/security/SafeUrlFlowCustomizations.qll

@@ -25,6 +25,6 @@ module SafeUrlFlow {
 
   /** A function model step using `UnsafeUrlMethod`, considered as a sanitizer for safe URL flow. */
   private class UnsafeUrlMethodEdge extends SanitizerEdge {
-    UnsafeUrlMethodEdge() { TaintTracking::functionModelStep(any(UnsafeUrlMethod um), this, _) }
+    UnsafeUrlMethodEdge() { this = any(UnsafeUrlMethod um).getAnInputNode(_) }
   }
 }

ql/test/library-tests/semmle/go/dataflow/FunctionInputsAndOutputs/FunctionModelStep.ql

@@ -16,6 +16,6 @@ class ReaderReset extends TaintTracking::FunctionModel, Method {
   }
 }
 
-from Function fn, DataFlow::Node pred, DataFlow::Node succ
-where TaintTracking::functionModelStep(fn, pred, succ)
+from TaintTracking::FunctionModel fn, DataFlow::Node pred, DataFlow::Node succ
+where fn.taintStep(pred, succ)
 select fn, pred, succ