Merge pull request #6211 from aschackmull/dataflow/refactor-call-context-check

Dataflow: Refactor call context check

python/ql/src/semmle/python/dataflow/new/internal/DataFlowImpl.qll

@@ -943,13 +943,8 @@ private module Stage2 {
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) { any() }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) { any() }
-
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    any()
-  }
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) { any() }
 
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) { any() }
@@ -1122,8 +1117,7 @@ private module Stage2 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )
@@ -1615,13 +1609,8 @@ private module Stage3 {
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) { any() }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) { any() }
-
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    any()
-  }
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) { any() }
 
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) { any() }
@@ -1816,8 +1805,7 @@ private module Stage3 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )
@@ -2364,20 +2352,16 @@ private module Stage4 {
 
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) {
-    c = resolveCall(call, outercc) and
+    checkCallContextCall(outercc, call, c) and
     if recordDataFlowCallSite(call, c) then result = TSpecificCall(call) else result = TSomeCall()
   }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) {
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) {
+    checkCallContextReturn(innercc, c, call) and
     if reducedViableImplInReturn(c, call) then result = TReturn(c, call) else result = ccNone()
   }
 
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    resolveReturn(innercc, inner, call)
-  }
-
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) {
     localFlowEntry(node, config) and
@@ -2579,8 +2563,7 @@ private module Stage4 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )

python/ql/src/semmle/python/dataflow/new/internal/DataFlowImpl2.qll

@@ -943,13 +943,8 @@ private module Stage2 {
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) { any() }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) { any() }
-
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    any()
-  }
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) { any() }
 
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) { any() }
@@ -1122,8 +1117,7 @@ private module Stage2 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )
@@ -1615,13 +1609,8 @@ private module Stage3 {
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) { any() }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) { any() }
-
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    any()
-  }
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) { any() }
 
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) { any() }
@@ -1816,8 +1805,7 @@ private module Stage3 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )
@@ -2364,20 +2352,16 @@ private module Stage4 {
 
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) {
-    c = resolveCall(call, outercc) and
+    checkCallContextCall(outercc, call, c) and
     if recordDataFlowCallSite(call, c) then result = TSpecificCall(call) else result = TSomeCall()
   }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) {
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) {
+    checkCallContextReturn(innercc, c, call) and
     if reducedViableImplInReturn(c, call) then result = TReturn(c, call) else result = ccNone()
   }
 
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    resolveReturn(innercc, inner, call)
-  }
-
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) {
     localFlowEntry(node, config) and
@@ -2579,8 +2563,7 @@ private module Stage4 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )

python/ql/src/semmle/python/dataflow/new/internal/DataFlowImpl3.qll

@@ -943,13 +943,8 @@ private module Stage2 {
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) { any() }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) { any() }
-
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    any()
-  }
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) { any() }
 
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) { any() }
@@ -1122,8 +1117,7 @@ private module Stage2 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )
@@ -1615,13 +1609,8 @@ private module Stage3 {
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) { any() }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) { any() }
-
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    any()
-  }
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) { any() }
 
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) { any() }
@@ -1816,8 +1805,7 @@ private module Stage3 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )
@@ -2364,20 +2352,16 @@ private module Stage4 {
 
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) {
-    c = resolveCall(call, outercc) and
+    checkCallContextCall(outercc, call, c) and
     if recordDataFlowCallSite(call, c) then result = TSpecificCall(call) else result = TSomeCall()
   }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) {
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) {
+    checkCallContextReturn(innercc, c, call) and
     if reducedViableImplInReturn(c, call) then result = TReturn(c, call) else result = ccNone()
   }
 
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    resolveReturn(innercc, inner, call)
-  }
-
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) {
     localFlowEntry(node, config) and
@@ -2579,8 +2563,7 @@ private module Stage4 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )

python/ql/src/semmle/python/dataflow/new/internal/DataFlowImpl4.qll

@@ -943,13 +943,8 @@ private module Stage2 {
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) { any() }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) { any() }
-
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    any()
-  }
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) { any() }
 
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) { any() }
@@ -1122,8 +1117,7 @@ private module Stage2 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )
@@ -1615,13 +1609,8 @@ private module Stage3 {
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) { any() }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) { any() }
-
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    any()
-  }
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) { any() }
 
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) { any() }
@@ -1816,8 +1805,7 @@ private module Stage3 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )
@@ -2364,20 +2352,16 @@ private module Stage4 {
 
   bindingset[call, c, outercc]
   private CcCall getCallContextCall(DataFlowCall call, DataFlowCallable c, Cc outercc) {
-    c = resolveCall(call, outercc) and
+    checkCallContextCall(outercc, call, c) and
     if recordDataFlowCallSite(call, c) then result = TSpecificCall(call) else result = TSomeCall()
   }
 
-  bindingset[call, c]
-  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call) {
+  bindingset[call, c, innercc]
+  private CcNoCall getCallContextReturn(DataFlowCallable c, DataFlowCall call, Cc innercc) {
+    checkCallContextReturn(innercc, c, call) and
     if reducedViableImplInReturn(c, call) then result = TReturn(c, call) else result = ccNone()
   }
 
-  bindingset[innercc, inner, call]
-  private predicate checkCallContextReturn(Cc innercc, DataFlowCallable inner, DataFlowCall call) {
-    resolveReturn(innercc, inner, call)
-  }
-
   bindingset[node, cc, config]
   private LocalCc getLocalCc(NodeEx node, Cc cc, Configuration config) {
     localFlowEntry(node, config) and
@@ -2579,8 +2563,7 @@ private module Stage4 {
       fwdFlow(ret, innercc, argAp, ap, config) and
       flowOutOfCall(call, ret, out, allowsFieldFlow, config) and
       inner = ret.getEnclosingCallable() and
-      checkCallContextReturn(innercc, inner, call) and
-      ccOut = getCallContextReturn(inner, call)
+      ccOut = getCallContextReturn(inner, call, innercc)
     |
       ap instanceof ApNil or allowsFieldFlow = true
     )

python/ql/src/semmle/python/dataflow/new/internal/DataFlowImplCommon.qll

@@ -1117,6 +1117,44 @@ ReturnPosition getReturnPosition(ReturnNodeExt ret) {
   result = getReturnPosition0(ret, ret.getKind())
 }
 
+/**
+ * Checks whether `inner` can return to `call` in the call context `innercc`.
+ * Assumes a context of `inner = viableCallableExt(call)`.
+ */
+bindingset[innercc, inner, call]
+predicate checkCallContextReturn(CallContext innercc, DataFlowCallable inner, DataFlowCall call) {
+  innercc instanceof CallContextAny
+  or
+  exists(DataFlowCallable c0, DataFlowCall call0 |
+    callEnclosingCallable(call0, inner) and
+    innercc = TReturn(c0, call0) and
+    c0 = prunedViableImplInCallContextReverse(call0, call)
+  )
+}
+
+/**
+ * Checks whether `call` can resolve to `calltarget` in the call context `cc`.
+ * Assumes a context of `calltarget = viableCallableExt(call)`.
+ */
+bindingset[cc, call, calltarget]
+predicate checkCallContextCall(CallContext cc, DataFlowCall call, DataFlowCallable calltarget) {
+  exists(DataFlowCall ctx | cc = TSpecificCall(ctx) |
+    if reducedViableImplInCallContext(call, _, ctx)
+    then calltarget = prunedViableImplInCallContext(call, ctx)
+    else any()
+  )
+  or
+  cc instanceof CallContextSomeCall
+  or
+  cc instanceof CallContextAny
+  or
+  cc instanceof CallContextReturn
+}
+
+/**
+ * Resolves a return from `callable` in `cc` to `call`. This is equivalent to
+ * `callable = viableCallableExt(call) and checkCallContextReturn(cc, callable, call)`.
+ */
 bindingset[cc, callable]
 predicate resolveReturn(CallContext cc, DataFlowCallable callable, DataFlowCall call) {
   cc instanceof CallContextAny and callable = viableCallableExt(call)
@@ -1128,6 +1166,10 @@ predicate resolveReturn(CallContext cc, DataFlowCallable callable, DataFlowCall
   )
 }
 
+/**
+ * Resolves a call from `call` in `cc` to `result`. This is equivalent to
+ * `result = viableCallableExt(call) and checkCallContextCall(cc, call, result)`.
+ */
 bindingset[call, cc]
 DataFlowCallable resolveCall(DataFlowCall call, CallContext cc) {
   exists(DataFlowCall ctx | cc = TSpecificCall(ctx) |