C++ WIP: Count return dispatch based on 2nd level scopes.

Author: aschackmull

cpp/ql/lib/semmle/code/cpp/ir/dataflow/internal/DataFlowImplSpecific.qll

@@ -21,6 +21,8 @@ module CppDataFlow implements InputSig {
 
   predicate getAdditionalFlowIntoCallNodeTerm = Private::getAdditionalFlowIntoCallNodeTerm/2;
 
+  predicate getSecondLevelScope = Private::getSecondLevelScope/1;
+
   predicate validParameterAliasStep = Private::validParameterAliasStep/2;
 
   predicate mayBenefitFromCallContext = Private::mayBenefitFromCallContext/1;

cpp/ql/lib/semmle/code/cpp/ir/dataflow/internal/DataFlowPrivate.qll

@@ -1259,3 +1259,54 @@ predicate validParameterAliasStep(Node node1, Node node2) {
     )
   )
 }
+
+private predicate isTopLevel(Cpp::Stmt s) { any(Function f).getBlock().getAStmt() = s }
+
+private Cpp::Stmt getAChainedBranch(Cpp::IfStmt s) {
+  result = s.getThen()
+  or
+  exists(Cpp::Stmt elseBranch | s.getElse() = elseBranch |
+    result = getAChainedBranch(elseBranch)
+    or
+    result = elseBranch and not elseBranch instanceof Cpp::IfStmt
+  )
+}
+
+private Instruction getInstruction(Node n) {
+  result = n.asInstruction() or
+  result = n.asOperand().getUse() or
+  result = n.(SsaPhiNode).getPhiNode().getBasicBlock().getFirstInstruction() or
+  n.(IndirectInstruction).hasInstructionAndIndirectionIndex(result, _) or
+  result = getInstruction(n.(PostUpdateNode).getPreUpdateNode())
+}
+
+private newtype TDataFlowSecondLevelScope =
+  TTopLevelIfBranch(Cpp::Stmt s) {
+    exists(Cpp::IfStmt ifstmt | s = getAChainedBranch(ifstmt) and isTopLevel(ifstmt))
+  } or
+  TTopLevelSwitchCase(Cpp::SwitchCase s) {
+    exists(Cpp::SwitchStmt switchstmt | s = switchstmt.getASwitchCase() and isTopLevel(switchstmt))
+  }
+
+class DataFlowSecondLevelScope extends TDataFlowSecondLevelScope {
+  string toString() {
+    exists(Cpp::Stmt s | this = TTopLevelIfBranch(s) | result = s.toString()) or
+    exists(Cpp::SwitchCase s | this = TTopLevelSwitchCase(s) | result = s.toString())
+  }
+
+  Cpp::Location getLocation() {
+    exists(Cpp::Stmt s | this = TTopLevelIfBranch(s) | result = s.getLocation()) or
+    exists(Cpp::SwitchCase s | this = TTopLevelSwitchCase(s) | result = s.getLocation())
+  }
+
+  private Cpp::Stmt getAStmt() {
+    exists(Cpp::Stmt s | this = TTopLevelIfBranch(s) | result = s) or
+    exists(Cpp::SwitchCase s | this = TTopLevelSwitchCase(s) | result = s.getAStmt())
+  }
+
+  Node getANode() {
+    getInstruction(result).getAst().(Cpp::ControlFlowNode).getEnclosingStmt().getParentStmt*() = this.getAStmt()
+  }
+}
+
+DataFlowSecondLevelScope getSecondLevelScope(Node n) { result.getANode() = n }

shared/dataflow/codeql/dataflow/DataFlow.qll

@@ -267,6 +267,13 @@ signature module InputSig {
    */
   default int getAdditionalFlowIntoCallNodeTerm(ArgumentNode arg, ParameterNode p) { none() }
 
+  class DataFlowSecondLevelScope {
+    /** Gets a textual representation of this element. */
+    string toString();
+  }
+
+  default DataFlowSecondLevelScope getSecondLevelScope(Node n) { none() }
+
   bindingset[call, p, arg]
   default predicate golangSpecificParamArgFilter(
     DataFlowCall call, ParameterNode p, ArgumentNode arg

shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll

@@ -1076,28 +1076,46 @@ module MakeImpl<InputSig Lang> {
       result = getAdditionalFlowIntoCallNodeTerm(arg.projectToNode(), p.projectToNode())
     }
 
+    private module SndLevelScopeOption = Option<DataFlowSecondLevelScope>;
+
+    private class SndLevelScopeOption = SndLevelScopeOption::Option;
+
+    pragma[nomagic]
+    private SndLevelScopeOption getScope(RetNodeEx ret) {
+      result = SndLevelScopeOption::some(getSecondLevelScope(ret.asNode()))
+      or
+      result instanceof SndLevelScopeOption::None and not exists(getSecondLevelScope(ret.asNode()))
+    }
+
     pragma[nomagic]
-    private predicate returnCallEdge1(DataFlowCallable c, DataFlowCall call, NodeEx out) {
+    private predicate returnCallEdge1(
+      DataFlowCallable c, SndLevelScopeOption scope, DataFlowCall call, NodeEx out
+    ) {
       exists(RetNodeEx ret |
-        flowOutOfCallNodeCand1(call, ret, _, out) and c = ret.getEnclosingCallable()
+        flowOutOfCallNodeCand1(call, ret, _, out) and
+        c = ret.getEnclosingCallable() and
+        scope = getScope(ret)
       )
     }
 
     private int simpleDispatchFanoutOnReturn(DataFlowCall call, NodeEx out) {
-      result = strictcount(DataFlowCallable c | returnCallEdge1(c, call, out))
+      result =
+        strictcount(DataFlowCallable c, SndLevelScopeOption scope |
+          returnCallEdge1(c, scope, call, out)
+        )
     }
 
     private int ctxDispatchFanoutOnReturn(NodeEx out, DataFlowCall ctx) {
       exists(DataFlowCall call, DataFlowCallable c |
         simpleDispatchFanoutOnReturn(call, out) > 1 and
         not Stage1::revFlow(out, false) and
         call.getEnclosingCallable() = c and
-        returnCallEdge1(c, ctx, _) and
+        returnCallEdge1(c, _, ctx, _) and
         mayBenefitFromCallContextExt(call, _) and
         result =
-          count(DataFlowCallable tgt |
+          count(DataFlowCallable tgt, SndLevelScopeOption scope |
             tgt = viableImplInCallContextExt(call, ctx) and
-            returnCallEdge1(tgt, call, out)
+            returnCallEdge1(tgt, scope, call, out)
           )
       )
     }