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()
+    )
   }
 }
 

java/ql/test/utils/model-generator/CaptureSummaryModels.expected

@@ -1,5 +1,6 @@
 | p;Factory;false;create;(String);;Argument[0];Argument[-1];taint; |
 | p;Factory;false;create;(String,int);;Argument[0];Argument[-1];taint; |
+| p;Factory;false;getValue;();;Argument[-1];ReturnValue;taint; |
 | p;FinalClass;false;returnsInput;(String);;Argument[0];ReturnValue;taint; |
 | p;FluentAPI;false;returnsThis;(String);;Argument[-1];ReturnValue;value; |
 | p;ImmutablePojo;false;ImmutablePojo;(String,int);;Argument[0];Argument[-1];taint; |

java/ql/test/utils/model-generator/p/Factory.java

@@ -19,4 +19,12 @@ public final class Factory {
         this.intValue = intValue;
     }
 
+    public String getValue() {
+        return value;
+    }
+
+    public int getIntValue() {
+        return intValue;
+    }
+
 }

java/ql/test/utils/model-generator/p/ImmutablePojo.java

@@ -15,6 +15,10 @@ public final class ImmutablePojo {
         return value;
     }
 
+    public long getX() {
+        return x;
+    }
+
     public String or(String defaultValue) {
         return value != null ? value : defaultValue;
     }

java/ql/test/utils/model-generator/p/MultipleImpls.java

@@ -1,9 +1,5 @@
 package p;
 
-import java.io.File;
-import java.io.FileFilter;
-import java.io.IOException;
-import java.nio.file.Files;
 import java.util.concurrent.Callable;
 
 public class MultipleImpls {