Python: Let the user help us identifying callbacks

python/ql/lib/semmle/python/dataflow/new/TestSummaries.qll

@@ -12,6 +12,8 @@ private class SummarizedCallableIdentity extends SummarizedCallable {
 
   override Call getACall() { result.getFunc().(Name).getId() = this }
 
+  override DataFlow::ArgumentNode getACallback() { result.asExpr().(Name).getId() = this }
+
   override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
     input = "Argument[0]" and
     output = "ReturnValue" and
@@ -25,6 +27,8 @@ private class SummarizedCallableApplyLambda extends SummarizedCallable {
 
   override Call getACall() { result.getFunc().(Name).getId() = this }
 
+  override DataFlow::ArgumentNode getACallback() { result.asExpr().(Name).getId() = this }
+
   override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
     input = "Argument[1]" and
     output = "Argument[0].Parameter[0]" and
@@ -41,6 +45,8 @@ private class SummarizedCallableReversed extends SummarizedCallable {
 
   override Call getACall() { result.getFunc().(Name).getId() = this }
 
+  override DataFlow::ArgumentNode getACallback() { result.asExpr().(Name).getId() = this }
+
   override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
     input = "Argument[0].ListElement" and
     output = "ReturnValue.ListElement" and
@@ -53,6 +59,8 @@ private class SummarizedCallableMap extends SummarizedCallable {
 
   override Call getACall() { result.getFunc().(Name).getId() = this }
 
+  override DataFlow::ArgumentNode getACallback() { result.asExpr().(Name).getId() = this }
+
   override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
     input = "Argument[1].ListElement" and
     output = "Argument[0].Parameter[0]" and
@@ -73,6 +81,10 @@ private class SummarizedCallableJsonLoads extends SummarizedCallable {
     result = API::moduleImport("json").getMember("loads").getACall().asExpr()
   }
 
+  override DataFlow::ArgumentNode getACallback() {
+    result = API::moduleImport("json").getMember("loads").getAUse()
+  }
+
   override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
     input = "Argument[0]" and
     output = "ReturnValue.ListElement" and

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatchPointsTo.qll

@@ -283,6 +283,9 @@ abstract class LibraryCallable extends string {
 
   /** Gets a call to this library callable. */
   abstract Call getACall();
+
+  /** Gets a data-flow node, where this library callable is used as a call-back. */
+  abstract ArgumentNode getACallback();
 }
 
 /**
@@ -404,6 +407,9 @@ class LibraryCallableValue extends DataFlowCallable, TLibraryCallable {
 
   override CallNode getACall() { result.getNode() = callable.getACall() }
 
+  /** Gets a data-flow node, where this library callable is used as a call-back. */
+  ArgumentNode getACallback() { result = callable.getACallback() }
+
   override Scope getScope() { none() }
 
   override NameNode getParameter(int n) { none() }

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll

@@ -942,10 +942,9 @@ predicate lambdaCreation(Node creation, LambdaCallKind kind, DataFlowCallable c)
   )
   or
   // summarized function
-  exists(Call call, Name arg |
+  exists(Call call |
-    arg = call.getAnArg() and
+    creation.asExpr() = call.getAnArg() and
-    c.(LibraryCallableValue).getName() = arg.getId() and
+    creation = c.(LibraryCallableValue).getACallback()
-    creation.asExpr() = arg
   )
 }
 

python/ql/test/experimental/dataflow/summaries/TestSummaries.qll

@@ -7,6 +7,8 @@ private class SummarizedCallableIdentity extends SummarizedCallable {
 
   override Call getACall() { result.getFunc().(Name).getId() = this }
 
+  override DataFlow::ArgumentNode getACallback() { result.asExpr().(Name).getId() = this }
+
   override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
     input = "Argument[0]" and
     output = "ReturnValue" and
@@ -20,6 +22,8 @@ private class SummarizedCallableApplyLambda extends SummarizedCallable {
 
   override Call getACall() { result.getFunc().(Name).getId() = this }
 
+  override DataFlow::ArgumentNode getACallback() { result.asExpr().(Name).getId() = this }
+
   override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
     input = "Argument[1]" and
     output = "Argument[0].Parameter[0]" and
@@ -36,6 +40,8 @@ private class SummarizedCallableReversed extends SummarizedCallable {
 
   override Call getACall() { result.getFunc().(Name).getId() = this }
 
+  override DataFlow::ArgumentNode getACallback() { result.asExpr().(Name).getId() = this }
+
   override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
     input = "Argument[0].ListElement" and
     output = "ReturnValue.ListElement" and
@@ -48,6 +54,8 @@ private class SummarizedCallableMap extends SummarizedCallable {
 
   override Call getACall() { result.getFunc().(Name).getId() = this }
 
+  override DataFlow::ArgumentNode getACallback() { result.asExpr().(Name).getId() = this }
+
   override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
     input = "Argument[1].ListElement" and
     output = "Argument[0].Parameter[0]" and
@@ -68,6 +76,10 @@ private class SummarizedCallableJsonLoads extends SummarizedCallable {
     result = API::moduleImport("json").getMember("loads").getACall().asExpr()
   }
 
+  override DataFlow::ArgumentNode getACallback() {
+    result = API::moduleImport("json").getMember("loads").getAUse()
+  }
+
   override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
     input = "Argument[0]" and
     output = "ReturnValue.ListElement" and