Rust: Implement known{Source,Sink}Model

github · Jan 9, 2025 · a7bb952 · a7bb952
1 parent 868caf9
commit a7bb952
Show file tree

Hide file tree

Showing 5 changed files with 115 additions and 53 deletions.
diff --git a/rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll b/rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll
@@ -202,13 +202,13 @@ module Node {
     }
 
     /** Holds is this node is a source node of kind `kind`. */
-    predicate isSource(string kind) {
-      this.getSummaryNode().(FlowSummaryImpl::Private::SourceOutputNode).isEntry(kind)
+    predicate isSource(string kind, string model) {
+      this.getSummaryNode().(FlowSummaryImpl::Private::SourceOutputNode).isEntry(kind, model)
     }
 
     /** Holds is this node is a sink node of kind `kind`. */
-    predicate isSink(string kind) {
-      this.getSummaryNode().(FlowSummaryImpl::Private::SinkInputNode).isExit(kind)
+    predicate isSink(string kind, string model) {
+      this.getSummaryNode().(FlowSummaryImpl::Private::SinkInputNode).isExit(kind, model)
     }
 
     override CfgScope getCfgScope() {
@@ -1305,9 +1305,13 @@ module RustDataFlow implements InputSig<Location> {
   /** Extra data flow steps needed for lambda flow analysis. */
   predicate additionalLambdaFlowStep(Node nodeFrom, Node nodeTo, boolean preservesValue) { none() }
 
-  predicate knownSourceModel(Node source, string model) { none() }
+  predicate knownSourceModel(Node source, string model) {
+    source.(Node::FlowSummaryNode).isSource(_, model)
+  }
 
-  predicate knownSinkModel(Node sink, string model) { none() }
+  predicate knownSinkModel(Node sink, string model) {
+    sink.(Node::FlowSummaryNode).isSink(_, model)
+  }
 
   class DataFlowSecondLevelScope = Void;
 }
@@ -1575,11 +1579,11 @@ private module Cached {
 
   /** Holds if `n` is a flow source of kind `kind`. */
   cached
-  predicate sourceNode(Node n, string kind) { n.(Node::FlowSummaryNode).isSource(kind) }
+  predicate sourceNode(Node n, string kind) { n.(Node::FlowSummaryNode).isSource(kind, _) }
 
   /** Holds if `n` is a flow sink of kind `kind`. */
   cached
-  predicate sinkNode(Node n, string kind) { n.(Node::FlowSummaryNode).isSink(kind) }
+  predicate sinkNode(Node n, string kind) { n.(Node::FlowSummaryNode).isSink(kind, _) }
 }
 
 import Cached
diff --git a/rust/ql/test/library-tests/dataflow/models/main.rs b/rust/ql/test/library-tests/dataflow/models/main.rs
@@ -222,6 +222,24 @@ fn test_enum_method_sink() {
     e.sink(); // $ hasValueFlow=15
 }
 
+// has a source model
+fn simple_source(i: i64) -> i64 {
+    0
+}
+
+fn test_simple_source() {
+    let s = simple_source(16);
+    sink(s) // $ hasValueFlow=16
+}
+
+// has a sink model
+fn simple_sink(i: i64) {}
+
+fn test_simple_sink() {
+    let s = source(17);
+    simple_sink(s); // $ hasValueFlow=17
+}
+
 fn main() {
     test_identify();
     test_get_var_pos();
@@ -238,5 +256,7 @@ fn main() {
     test_enum_method_source();
     test_enum_sink();
     test_enum_method_sink();
+    test_simple_source();
+    test_simple_sink();
     let dummy = Some(0); // ensure that the the `lang:core` crate is extracted
 }