Support more CFG node types in optional binding flow

Author: atorralba

swift/ql/lib/codeql/swift/controlflow/CfgNodes.qll

@@ -91,20 +91,21 @@ class CfgNode extends ControlFlowNode, TElementNode {
 
   /** Gets a split for this control flow node, if any. */
   final Split getASplit() { result = splits.getASplit() }
-}
 
-private Expr getAst(ControlFlowElement n) {
-  result = n.asAstNode()
-  or
-  result = n.(PropertyGetterElement).getRef()
-  or
-  result = n.(PropertySetterElement).getAssignExpr()
-  or
-  result = n.(PropertyObserverElement).getAssignExpr()
-  or
-  result = n.(ClosureElement).getAst()
-  or
-  result = n.(KeyPathElement).getAst()
+  /** Gets the AST representation of this control flow node, if any. */
+  Expr getAst() {
+    result = n.asAstNode()
+    or
+    result = n.(PropertyGetterElement).getRef()
+    or
+    result = n.(PropertySetterElement).getAssignExpr()
+    or
+    result = n.(PropertyObserverElement).getAssignExpr()
+    or
+    result = n.(ClosureElement).getAst()
+    or
+    result = n.(KeyPathElement).getAst()
+  }
 }
 
 /** A control-flow node that wraps an AST expression. */
@@ -123,7 +124,7 @@ class PropertyGetterCfgNode extends CfgNode {
 
   Expr getRef() { result = n.getRef() }
 
-  CfgNode getBase() { getAst(result.getNode()) = n.getBase() }
+  CfgNode getBase() { result.getAst() = n.getBase() }
 
   AccessorDecl getAccessorDecl() { result = n.getAccessorDecl() }
 }
@@ -134,9 +135,9 @@ class PropertySetterCfgNode extends CfgNode {
 
   AssignExpr getAssignExpr() { result = n.getAssignExpr() }
 
-  CfgNode getBase() { getAst(result.getNode()) = n.getBase() }
+  CfgNode getBase() { result.getAst() = n.getBase() }
 
-  CfgNode getSource() { getAst(result.getNode()) = n.getAssignExpr().getSource() }
+  CfgNode getSource() { result.getAst() = n.getAssignExpr().getSource() }
 
   AccessorDecl getAccessorDecl() { result = n.getAccessorDecl() }
 }
@@ -146,19 +147,19 @@ class PropertyObserverCfgNode extends CfgNode {
 
   AssignExpr getAssignExpr() { result = n.getAssignExpr() }
 
-  CfgNode getBase() { getAst(result.getNode()) = n.getBase() }
+  CfgNode getBase() { result.getAst() = n.getBase() }
 
-  CfgNode getSource() { getAst(result.getNode()) = n.getAssignExpr().getSource() }
+  CfgNode getSource() { result.getAst() = n.getAssignExpr().getSource() }
 
   AccessorDecl getAccessorDecl() { result = n.getObserver() }
 }
 
 class ApplyExprCfgNode extends ExprCfgNode {
   override ApplyExpr e;
 
-  CfgNode getArgument(int index) { getAst(result.getNode()) = e.getArgument(index).getExpr() }
+  CfgNode getArgument(int index) { result.getAst() = e.getArgument(index).getExpr() }
 
-  CfgNode getQualifier() { getAst(result.getNode()) = e.getQualifier() }
+  CfgNode getQualifier() { result.getAst() = e.getQualifier() }
 
   AbstractFunctionDecl getStaticTarget() { result = e.getStaticTarget() }
 

swift/ql/lib/codeql/swift/dataflow/Ssa.qll

@@ -160,10 +160,7 @@ module Ssa {
         pbd.getAPattern() = bb.getNode(blockIndex).getNode().asAstNode() and
         init = var.getParentInitializer()
       |
-        value.getNode().asAstNode() = init
-        or
-        // TODO: We should probably enumerate more cfg nodes here.
-        value.(PropertyGetterCfgNode).getRef() = init
+        value.getAst() = init
       )
       or
       exists(SsaInput::BasicBlock bb, int blockIndex, ConditionElement ce, Expr init |
@@ -172,7 +169,7 @@ module Ssa {
         init = ce.getInitializer() and
         strictcount(Ssa::WriteDefinition alt | alt.definesAt(_, bb, blockIndex)) = 1 // exclude cases where there are multiple writes from the same pattern, this is at best taint flow.
       |
-        value.getNode().asAstNode() = init
+        value.getAst() = init
       )
     }
   }