Rust: Track closure types in data flow

Author: hvitved

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

@@ -263,10 +263,75 @@ predicate lambdaCallExpr(CallExprImpl::DynamicCallExpr call, LambdaCallKind kind
   exists(kind)
 }
 
+// NOTE: We do not yet track type information, except for closures where
+// we use the closure itself to represent the unique type.
+final class DataFlowType extends TDataFlowType {
+  Expr asClosureExpr() { this = TClosureExprType(result) }
+
+  predicate isUnknownType() { this = TUnknownType() }
+
+  predicate isSourceContextParameterType() { this = TSourceContextParameterType() }
+
+  string toString() {
+    exists(this.asClosureExpr()) and
+    result = "... => .."
+    or
+    this.isUnknownType() and
+    result = ""
+    or
+    this.isSourceContextParameterType() and
+    result = "<source context parameter type>"
+  }
+}
+
+pragma[nomagic]
+private predicate compatibleTypesLeft(DataFlowType t1, DataFlowType t2) {
+  t1.isUnknownType()
+  or
+  t1.asClosureExpr() = t2.asClosureExpr()
+  or
+  t1.isSourceContextParameterType() and not exists(t2.asClosureExpr())
+}
+
+predicate compatibleTypes(DataFlowType t1, DataFlowType t2) {
+  compatibleTypesLeft(t1, t2)
+  or
+  compatibleTypesLeft(t2, t1)
+}
+
+pragma[nomagic]
+predicate typeStrongerThan(DataFlowType t1, DataFlowType t2) { exists(t1) and t2.isUnknownType() }
+
+DataFlowType getNodeType(NodePublic node) {
+  result.asClosureExpr() = node.asExpr()
+  or
+  result.asClosureExpr() = node.(ClosureParameterNode).getCfgScope()
+  or
+  exists(VariableCapture::Flow::SynthesizedCaptureNode scn |
+    scn = node.(CaptureNode).getSynthesizedCaptureNode() and
+    if scn.isInstanceAccess()
+    then result.asClosureExpr() = scn.getEnclosingCallable()
+    else result.isUnknownType()
+  )
+  or
+  not lambdaCreationExpr(node.asExpr()) and
+  not node instanceof ClosureParameterNode and
+  not node instanceof CaptureNode and
+  result.isUnknownType()
+}
+
 // Defines a set of aliases needed for the `RustDataFlow` module
 private module Aliases {
   class DataFlowCallableAlias = DataFlowCallable;
 
+  class DataFlowTypeAlias = DataFlowType;
+
+  predicate compatibleTypesAlias = compatibleTypes/2;
+
+  predicate typeStrongerThanAlias = typeStrongerThan/2;
+
+  predicate getNodeTypeAlias = getNodeType/1;
+
   class ReturnKindAlias = ReturnKind;
 
   class DataFlowCallAlias = DataFlowCall;
@@ -398,8 +463,6 @@ module RustDataFlowGen<RustDataFlowInputSig Input> implements InputSig<Location>
     result = node.(Node::Node).getEnclosingCallable()
   }
 
-  DataFlowType getNodeType(Node node) { any() }
-
   predicate nodeIsHidden(Node node) {
     node instanceof SsaNode or
     node.(FlowSummaryNode).getSummaryNode().isHidden() or
@@ -486,15 +549,15 @@ module RustDataFlowGen<RustDataFlowInputSig Input> implements InputSig<Location>
    */
   OutNode getAnOutNode(DataFlowCall call, ReturnKind kind) { call = result.getCall(kind) }
 
-  // NOTE: For now we use the type `Unit` and do not benefit from type
-  // information in the data flow analysis.
-  final class DataFlowType extends Unit {
-    string toString() { result = "" }
-  }
+  class DataFlowType = DataFlowTypeAlias;
 
-  predicate compatibleTypes(DataFlowType t1, DataFlowType t2) { any() }
+  predicate compatibleTypes = compatibleTypesAlias/2;
 
-  predicate typeStrongerThan(DataFlowType t1, DataFlowType t2) { none() }
+  predicate typeStrongerThan = typeStrongerThanAlias/2;
+
+  DataFlowType getSourceContextParameterNodeType() { result.isSourceContextParameterType() }
+
+  predicate getNodeType = getNodeTypeAlias/1;
 
   class Content = ContentAlias;
 
@@ -1110,6 +1173,12 @@ private module Cached {
     TCfgScope(CfgScope scope) or
     TSummarizedCallable(SummarizedCallable c)
 
+  cached
+  newtype TDataFlowType =
+    TClosureExprType(Expr e) { lambdaCreationExpr(e) } or
+    TUnknownType() or
+    TSourceContextParameterType()
+
   /** This is the local flow predicate that is exposed. */
   cached
   predicate localFlowStepImpl(Node nodeFrom, Node nodeTo) {

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

@@ -280,7 +280,7 @@ private module Debug {
   private import Content
   private import codeql.rust.dataflow.internal.DataFlowImpl
   private import codeql.rust.internal.typeinference.TypeMention
-  private import codeql.rust.internal.typeinference.Type
+  private import codeql.rust.internal.typeinference.Type as Type
 
   private predicate relevantManualModel(SummarizedCallableImpl sc, string can) {
     exists(Provenance manual |
@@ -298,7 +298,7 @@ private module Debug {
       sc.propagatesFlow(input, _, _, _, _, _) and
       input.head() = SummaryComponent::argument(pos) and
       p = pos.getParameterIn(sc.getParamList()) and
-      tm.getType() instanceof RefType and
+      tm.getType() instanceof Type::RefType and
       not input.tail().head() = SummaryComponent::content(TSingletonContentSet(TReferenceContent()))
     |
       tm = p.getTypeRepr()
@@ -313,7 +313,7 @@ private module Debug {
     exists(TypeMention tm |
       relevantManualModel(sc, can) and
       sc.propagatesFlow(_, output, _, _, _, _) and
-      tm.getType() instanceof RefType and
+      tm.getType() instanceof Type::RefType and
       output.head() = SummaryComponent::return(_) and
       not output.tail().head() =
         SummaryComponent::content(TSingletonContentSet(TReferenceContent())) and