Match enclosing unit without casting to specific nodes

java/ql/src/utils/model-generator/CaptureSinkModels.ql

@@ -16,8 +16,8 @@ class PropagateToSinkConfiguration extends TaintTracking::Configuration {
 
   override predicate isSource(DataFlow::Node source) {
     source instanceof DataFlow::ParameterNode and
-    source.asParameter().getCallable().isPublic() and
-    source.asParameter().getCallable().getDeclaringType().isPublic() and
+    source.getEnclosingCallable().isPublic() and
+    source.getEnclosingCallable().getDeclaringType().isPublic() and
     isRelevantForModels(source.getEnclosingCallable())
   }
 
@@ -25,7 +25,10 @@ class PropagateToSinkConfiguration extends TaintTracking::Configuration {
 }
 
 string asInputArgument(DataFlow::Node source) {
-  result = "Argument[" + source.asParameter().getPosition() + "]"
+  exists(int pos |
+    source.(DataFlow::ParameterNode).isParameterOf(_, pos) and
+    result = "Argument[" + pos + "]"
+  )
 }
 
 string captureSink(Callable api) {

java/ql/src/utils/model-generator/CaptureSourceModels.ql

@@ -22,7 +22,7 @@ class FromSourceConfiguration extends TaintTracking::Configuration {
   override predicate isSink(DataFlow::Node sink) {
     exists(Callable c |
       sink instanceof ReturnNodeExt and
-      sink.asExpr().getEnclosingCallable() = c and
+      sink.getEnclosingCallable() = c and
       c.isPublic() and
       c.fromSource()
     )
@@ -42,7 +42,7 @@ string captureSource(Callable api) {
   |
     config.hasFlow(src, sink) and
     specificSourceNode(sink, output, kind) and
-    api = src.asExpr().getEnclosingCallable() and
+    api = src.getEnclosingCallable() and
     result = asSourceModel(api, output, kind)
   )
 }

java/ql/src/utils/model-generator/CaptureSummaryModels.ql

@@ -28,14 +28,15 @@ string captureQualifierFlow(Callable api) {
 }
 
 string captureFieldFlow(Callable api) {
-  exists(FieldAccess fa, ReturnNodeExt postUpdate |
+  exists(FieldAccess fa, ReturnNodeExt returnNode |
     not (fa.getField().isStatic() and fa.getField().isFinal()) and
-    postUpdate.getEnclosingCallable() = api and
+    returnNode.getEnclosingCallable() = api and
+    fa.getCompilationUnit() = api.getCompilationUnit() and
     isRelevantType(api.getReturnType()) and
     not api.getDeclaringType() instanceof EnumType and
-    TaintTracking::localTaint(DataFlow::exprNode(fa), postUpdate)
+    TaintTracking::localTaint(DataFlow::exprNode(fa), returnNode)
   |
-    result = asTaintModel(api, "Argument[-1]", asOutput(api, postUpdate))
+    result = asTaintModel(api, "Argument[-1]", asOutput(api, returnNode))
   )
 }
 
@@ -59,7 +60,11 @@ class ParameterToFieldConfig extends TaintTracking::Configuration {
   }
 
   override predicate isSink(DataFlow::Node sink) {
-    exists(FieldAssignment a | a.getSource() = sink.asExpr())
+    exists(FieldAssignment a |
+      a.getSource() = sink.asExpr() and
+      a.getDest().(VarAccess).getVariable().getCompilationUnit() =
+        sink.getEnclosingCallable().getCompilationUnit()
+    )
   }
 }