Rust: Restrict type propagation into arguments

Author: hvitved

rust/ql/lib/codeql/rust/internal/Type.qll

@@ -51,6 +51,7 @@ newtype TType =
   TSliceType() or
   TNeverType() or
   TPtrType() or
+  TContextType() or
   TTupleTypeParameter(int arity, int i) { exists(TTuple(arity)) and i in [0 .. arity - 1] } or
   TTypeParamTypeParameter(TypeParam t) or
   TAssociatedTypeTypeParameter(TypeAlias t) { any(TraitItemNode trait).getAnAssocItem() = t } or
@@ -371,6 +372,14 @@ class PtrType extends Type, TPtrType {
   override Location getLocation() { result instanceof EmptyLocation }
 }
 
+class ContextType extends Type, TContextType {
+  override TypeParameter getPositionalTypeParameter(int i) { none() }
+
+  override string toString() { result = "(context typed)" }
+
+  override Location getLocation() { result instanceof EmptyLocation }
+}
+
 /** A type parameter. */
 abstract class TypeParameter extends Type {
   override TypeParameter getPositionalTypeParameter(int i) { none() }

rust/ql/lib/codeql/rust/internal/TypeInference.qll

@@ -1961,7 +1961,30 @@ private Type inferMethodCallType0(
 ) {
   exists(TypePath path0 |
     n = a.getNodeAt(apos) and
-    result = MethodCallMatching::inferAccessType(a, derefChainBorrow, apos, path0)
+    (
+      result = MethodCallMatching::inferAccessType(a, derefChainBorrow, apos, path0)
+      or
+      apos.isReturn() and
+      exists(Method target, TypeParameter tp |
+        target = a.getTarget(derefChainBorrow) and
+        assocFunctionContextTypedAt(target, _, path0, tp) and
+        not exists(TypeArgumentPosition tapos |
+          exists(int i |
+            i = tapos.asMethodTypeArgumentPosition() and
+            tp = TTypeParamTypeParameter(target.getGenericParamList().getTypeParam(i))
+          )
+          or
+          TTypeParamTypeParameter(tapos.asTypeParam()) = tp
+        |
+          exists(a.getTypeArgument(tapos, _))
+        ) and
+        not (
+          tp instanceof TSelfTypeParameter and
+          exists(getCallExprTypeQualifier(a, _))
+        )
+      ) and
+      result = TContextType()
+    )
   |
     if
       // index expression `x[i]` desugars to `*x.index(i)`, so we must account for
@@ -1973,6 +1996,9 @@ private Type inferMethodCallType0(
   )
 }
 
+pragma[nomagic]
+private TypePath getContextTypePath(AstNode n) { inferType(n, result) = TContextType() }
+
 /**
  * Gets the type of `n` at `path`, where `n` is either a method call or an
  * argument/receiver of a method call.
@@ -1983,7 +2009,8 @@ private Type inferMethodCallType(AstNode n, TypePath path) {
     MethodCallMatchingInput::Access a, MethodCallMatchingInput::AccessPosition apos,
     string derefChainBorrow, TypePath path0
   |
-    result = inferMethodCallType0(a, apos, n, derefChainBorrow, path0)
+    result = inferMethodCallType0(a, apos, n, derefChainBorrow, path0) and
+    if not apos.isReturn() then path.startsWith(getContextTypePath(n)) else any()
   |
     (
       not apos.isSelf()
@@ -2171,6 +2198,12 @@ private module NonMethodResolution {
       or
       result = this.resolveCallTargetRec()
     }
+
+    pragma[nomagic]
+    Function resolveTraitFunction() {
+      this.(Call).hasTrait() and
+      result = this.getPathResolutionResolved()
+    }
   }
 
   private newtype TCallAndBlanketPos =
@@ -2431,7 +2464,30 @@ pragma[nomagic]
 private Type inferNonMethodCallType(AstNode n, TypePath path) {
   exists(NonMethodCallMatchingInput::Access a, NonMethodCallMatchingInput::AccessPosition apos |
     n = a.getNodeAt(apos) and
+    if not apos.isReturn() then path.startsWith(getContextTypePath(n)) else any()
+  |
     result = NonMethodCallMatching::inferAccessType(a, apos, path)
+    or
+    apos.isReturn() and
+    exists(Function target, TypeParameter tp |
+      target = [a.getTarget().(Function), a.resolveTraitFunction()] and
+      assocFunctionContextTypedAt(target, _, path, tp) and
+      not exists(TypeArgumentPosition tapos |
+        // exists(int i |
+        //   i = tapos.asMethodTypeArgumentPosition() and
+        //   tp = TTypeParamTypeParameter(target.getGenericParamList().getTypeParam(i))
+        // )
+        // or
+        TTypeParamTypeParameter(tapos.asTypeParam()) = tp
+      |
+        exists(a.getTypeArgument(tapos, _))
+      ) and
+      not (
+        tp instanceof TSelfTypeParameter and
+        exists(getCallExprTypeQualifier(a, _))
+      )
+    ) and
+    result = TContextType()
   )
 }
 
@@ -2510,7 +2566,8 @@ pragma[nomagic]
 private Type inferOperationType(AstNode n, TypePath path) {
   exists(OperationMatchingInput::Access a, OperationMatchingInput::AccessPosition apos |
     n = a.getNodeAt(apos) and
-    result = OperationMatching::inferAccessType(a, apos, path)
+    result = OperationMatching::inferAccessType(a, apos, path) and
+    if not apos.isReturn() then path.startsWith(getContextTypePath(n)) else any()
   )
 }
 
@@ -3291,8 +3348,10 @@ private module Debug {
   Locatable getRelevantLocatable() {
     exists(string filepath, int startline, int startcolumn, int endline, int endcolumn |
       result.getLocation().hasLocationInfo(filepath, startline, startcolumn, endline, endcolumn) and
-      filepath.matches("%/sqlx.rs") and
-      startline = [56 .. 60]
+      filepath.matches("%/crate/data_derive/src/lib.rs") and
+      startline = [48, 74]
+      // filepath.matches("%/main.rs") and
+      // startline = [2525]
     )
   }
 
@@ -3355,7 +3414,7 @@ private module Debug {
   }
 
   predicate countTypesForNodeAtLimit(AstNode n, int c) {
-    n = getRelevantLocatable() and
+    // n = getRelevantLocatable() and
     c = strictcount(Type t, TypePath path | t = debugInferTypeForNodeAtLimit(n, path))
   }
 

rust/ql/lib/codeql/rust/internal/typeinference/FunctionType.qll

@@ -418,3 +418,18 @@ module ArgsAreInstantiationsOf<ArgsAreInstantiationsOfInputSig Input> {
     )
   }
 }
+
+pragma[nomagic]
+predicate assocFunctionContextTypedAt(
+  Function f, ImplOrTraitItemNode i, TypePath path, TypeParameter tp
+) {
+  // f.getName().getText() = "default" and
+  exists(FunctionPosition resPos |
+    resPos.isReturn() and
+    assocFunctionTypeAt(f, i, resPos, path, tp)
+  ) and
+  not exists(FunctionPosition nonResPos |
+    not nonResPos.isReturn() and
+    assocFunctionTypeAt(f, i, nonResPos, _, tp)
+  )
+}

rust/ql/test/library-tests/dataflow/sources/CONSISTENCY/PathResolutionConsistency.expected

@@ -2,6 +2,8 @@ multipleCallTargets
 | test.rs:113:62:113:77 | ...::from(...) |
 | test.rs:120:58:120:73 | ...::from(...) |
 | test.rs:229:22:229:72 | ... .read_to_string(...) |
+| test.rs:911:24:911:34 | row.take(...) |
+| test.rs:998:24:998:34 | row.take(...) |
 | test.rs:1096:50:1096:66 | ...::from(...) |
 | test.rs:1096:50:1096:66 | ...::from(...) |
 | test_futures_io.rs:35:26:35:63 | pinned.poll_read(...) |

rust/ql/test/library-tests/sensitivedata/CONSISTENCY/PathResolutionConsistency.expected

@@ -0,0 +1,2 @@
+multipleCallTargets
+| test.rs:288:7:288:36 | ... .as_str() |