Merge pull request #18673 from paldepind/rust-higher-order-function-model-generation

Rust: Higher order function model generation

Author: paldepind

rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll

@@ -871,6 +871,36 @@ final class TuplePositionContent extends Content, TTuplePositionContent {
   override Location getLocation() { result instanceof EmptyLocation }
 }
 
+/**
+ * A content for the index of an argument to at function call.
+ *
+ * Used by the model generator to create flow summaries for higher-order
+ * functions.
+ */
+final class FunctionCallArgumentContent extends Content, TFunctionCallArgumentContent {
+  private int pos;
+
+  FunctionCallArgumentContent() { this = TFunctionCallArgumentContent(pos) }
+
+  int getPosition() { result = pos }
+
+  override string toString() { result = "function argument at " + pos }
+
+  override Location getLocation() { result instanceof EmptyLocation }
+}
+
+/**
+ * A content for the return value of function call.
+ *
+ * Used by the model generator to create flow summaries for higher-order
+ * functions.
+ */
+final class FunctionCallReturnContent extends Content, TFunctionCallReturnContent {
+  override string toString() { result = "function return" }
+
+  override Location getLocation() { result instanceof EmptyLocation }
+}
+
 /** Holds if `access` indexes a tuple at an index corresponding to `c`. */
 private predicate fieldTuplePositionContent(FieldExprCfgNode access, TuplePositionContent c) {
   access.getNameRef().getText().toInt() = c.getPosition()
@@ -1192,6 +1222,13 @@ module RustDataFlow implements InputSig<Location> {
         node2.asExpr() = deref
       )
       or
+      // Read from function return
+      exists(DataFlowCall call |
+        lambdaCall(call, _, node1) and
+        call = node2.(OutNode).getCall(TNormalReturnKind()) and
+        c instanceof FunctionCallReturnContent
+      )
+      or
       VariableCapture::readStep(node1, c, node2)
     )
     or
@@ -1273,6 +1310,13 @@ module RustDataFlow implements InputSig<Location> {
       node2.asExpr() = ref
     )
     or
+    // Store in function argument
+    exists(DataFlowCall call, int i |
+      isArgumentNode(node1, call, TPositionalParameterPosition(i)) and
+      lambdaCall(call, _, node2.(PostUpdateNode).getPreUpdateNode()) and
+      c.(FunctionCallArgumentContent).getPosition() = i
+    )
+    or
     VariableCapture::storeStep(node1, c, node2)
   }
 
@@ -1615,6 +1659,10 @@ private module Cached {
     TStructFieldContent(Struct s, string field) {
       field = s.getFieldList().(RecordFieldList).getAField().getName().getText()
     } or
+    TFunctionCallReturnContent() or
+    TFunctionCallArgumentContent(int pos) {
+      pos in [0 .. any(CallExpr c).getArgList().getNumberOfArgs() - 1]
+    } or
     TCapturedVariableContent(VariableCapture::CapturedVariable v) or
     TReferenceContent()
 

rust/ql/lib/codeql/rust/dataflow/internal/FlowSummaryImpl.qll

@@ -51,7 +51,7 @@ module Input implements InputSig<Location, RustDataFlow> {
     override MethodCallExpr getCall() { result = call }
   }
 
-  RustDataFlow::ArgumentPosition callbackSelfParameterPosition() { none() }
+  RustDataFlow::ArgumentPosition callbackSelfParameterPosition() { result.isClosureSelf() }
 
   ReturnKind getStandardReturnValueKind() { result = TNormalReturnKind() }
 

rust/ql/src/utils/modelgenerator/internal/CaptureModels.qll

@@ -88,8 +88,9 @@ module ModelGeneratorInput implements ModelGeneratorInputSig<Location, RustDataF
 
   bindingset[c]
   string paramReturnNodeAsContentOutput(Callable c, ParameterPosition pos) {
-    // TODO: Implement this to support returning through parameters.
-    result = "paramReturnNodeAsContentOutput(" + c + ", " + pos + ")"
+    result = parameterContentAccess(c.getParamList().getParam(pos.getPosition()))
+    or
+    pos.isSelf() and result = qualifierString()
   }
 
   Callable returnNodeEnclosingCallable(DataFlow::Node ret) {
@@ -131,12 +132,34 @@ module ModelGeneratorInput implements ModelGeneratorInputSig<Location, RustDataF
     c.(SingletonContentSet).getContent() instanceof StructFieldContent
   }
 
-  predicate isCallback(DataFlow::ContentSet c) { none() }
+  predicate isCallback(DataFlow::ContentSet cs) {
+    exists(Content c | c = cs.(SingletonContentSet).getContent() |
+      c instanceof FunctionCallReturnContent or
+      c instanceof FunctionCallArgumentContent
+    )
+  }
 
   string getSyntheticName(DataFlow::ContentSet c) { none() }
 
+  private string encodeContent(ContentSet cs, string arg) {
+    result = FlowSummary::Input::encodeContent(cs, arg)
+    or
+    exists(Content c | cs = TSingletonContentSet(c) |
+      exists(int pos |
+        pos = c.(FunctionCallArgumentContent).getPosition() and
+        result = "Parameter" and
+        arg = pos.toString()
+      )
+      or
+      c instanceof FunctionCallReturnContent and result = "ReturnValue" and arg = ""
+    )
+  }
+
   string printContent(DataFlow::ContentSet cs) {
-    exists(string arg | result = FlowSummary::Input::encodeContent(cs, arg) + "[" + arg + "]")
+    exists(string name, string arg |
+      name = encodeContent(cs, arg) and
+      if arg = "" then result = name else result = name + "[" + arg + "]"
+    )
   }
 
   string partialModelRow(Callable api, int i) {

rust/ql/test/library-tests/dataflow/local/DataFlowStep.expected

@@ -603,6 +603,8 @@ storeStep
 | main.rs:414:41:414:67 | default_name | captured default_name | main.rs:414:41:414:67 | \|...\| ... |
 | main.rs:436:27:436:27 | 0 | Some | main.rs:436:22:436:28 | Some(...) |
 readStep
+| file://:0:0:0:0 | [summary param] 0 in lang:core::_::<crate::option::Option>::unwrap_or_else | function return | file://:0:0:0:0 | [summary] read: Argument[0].ReturnValue in lang:core::_::<crate::option::Option>::unwrap_or_else |
+| file://:0:0:0:0 | [summary param] 0 in lang:core::_::<crate::result::Result>::unwrap_or_else | function return | file://:0:0:0:0 | [summary] read: Argument[0].ReturnValue in lang:core::_::<crate::result::Result>::unwrap_or_else |
 | file://:0:0:0:0 | [summary param] self in lang:core::_::<crate::option::Option>::expect | Some | file://:0:0:0:0 | [summary] read: Argument[self].Variant[crate::option::Option::Some(0)] in lang:core::_::<crate::option::Option>::expect |
 | file://:0:0:0:0 | [summary param] self in lang:core::_::<crate::option::Option>::unwrap | Some | file://:0:0:0:0 | [summary] read: Argument[self].Variant[crate::option::Option::Some(0)] in lang:core::_::<crate::option::Option>::unwrap |
 | file://:0:0:0:0 | [summary param] self in lang:core::_::<crate::option::Option>::unwrap_or | Some | file://:0:0:0:0 | [summary] read: Argument[self].Variant[crate::option::Option::Some(0)] in lang:core::_::<crate::option::Option>::unwrap_or |

rust/ql/test/library-tests/dataflow/models/main.rs

@@ -175,6 +175,34 @@ fn test_set_tuple_element() {
     sink(t.1); // $ hasValueFlow=11
 }
 
+// has a flow model
+pub fn apply<F>(n: i64, f: F) -> i64 where F : FnOnce(i64) -> i64 {
+    0
+}
+
+fn test_apply_flow_in() {
+    let s = source(83);
+    let f = |n| {
+        sink(n); // $ hasValueFlow=83
+        n + 3
+    };
+    apply(s, f);
+}
+
+fn test_apply_flow_out() {
+    let s = source(86);
+    let f = |n| if n != 0 { n } else { s };
+    let t = apply(34, f);
+    sink(t); // $ hasValueFlow=86
+}
+
+fn test_apply_flow_through() {
+    let s = source(33);
+    let f = |n| if n != 0 { n } else { 0 };
+    let t = apply(s, f);
+    sink(t); // $ hasValueFlow=33
+}
+
 impl MyFieldEnum {
     // has a source model
     fn source(&self, i: i64) -> MyFieldEnum {