From f2d89a24f8a8f3b06e3e39b674c114db2a4cbe3c Mon Sep 17 00:00:00 2001
From: Mathias Vorreiter Pedersen <mapeder@microsoft.com>
Date: Mon, 23 Sep 2024 16:41:20 +0100
Subject: [PATCH] PS: Add flow steps from SSA and from right-hand of assignment
 to assignment.

---
 .../code/powershell/controlflow/CfgNodes.qll  | 23 +++++++++++++++++++
 .../dataflow/internal/DataFlowPrivate.qll     | 23 +++++++++++++++----
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/powershell/ql/lib/semmle/code/powershell/controlflow/CfgNodes.qll b/powershell/ql/lib/semmle/code/powershell/controlflow/CfgNodes.qll
index a94c43ff0571..3638723a39e5 100644
--- a/powershell/ql/lib/semmle/code/powershell/controlflow/CfgNodes.qll
+++ b/powershell/ql/lib/semmle/code/powershell/controlflow/CfgNodes.qll
@@ -197,6 +197,29 @@ module ExprNodes {
 
     InvokeMemberCfgNode getInvokeMember() { this = result.getQualifier() }
   }
+
+  class ConditionalChildMapping extends ExprChildMapping, ConditionalExpr {
+    override predicate relevantChild(Ast n) { n = this.getCondition() or n = this.getABranch() }
+  }
+
+  /** A control-flow node that wraps a `ConditionalExpr` expression. */
+  class ConditionalCfgNode extends ExprCfgNode {
+    override string getAPrimaryQlClass() { result = "ConditionalCfgNode" }
+
+    override ConditionalChildMapping e;
+
+    final override ConditionalExpr getExpr() { result = super.getExpr() }
+
+    final ExprCfgNode getCondition() { e.hasCfgChild(e.getCondition(), this, result) }
+
+    final ExprCfgNode getBranch(boolean value) { e.hasCfgChild(e.getBranch(value), this, result) }
+
+    final ExprCfgNode getABranch() { result = this.getBranch(_) }
+
+    final ExprCfgNode getIfTrue() { e.hasCfgChild(e.getIfTrue(), this, result) }
+
+    final ExprCfgNode getIfFalse() { e.hasCfgChild(e.getIfFalse(), this, result) }
+  }
 }
 
 module StmtNodes {
diff --git a/powershell/ql/lib/semmle/code/powershell/dataflow/internal/DataFlowPrivate.qll b/powershell/ql/lib/semmle/code/powershell/dataflow/internal/DataFlowPrivate.qll
index 9da8705332e4..90ceeca944fa 100644
--- a/powershell/ql/lib/semmle/code/powershell/dataflow/internal/DataFlowPrivate.qll
+++ b/powershell/ql/lib/semmle/code/powershell/dataflow/internal/DataFlowPrivate.qll
@@ -89,10 +89,14 @@ module SsaFlow {
 module LocalFlow {
   pragma[nomagic]
   predicate localFlowStepCommon(Node nodeFrom, Node nodeTo) {
-    none() // TODO
+    nodeFrom.asExpr() = nodeTo.asExpr().(CfgNodes::ExprNodes::ConditionalCfgNode).getABranch()
+    or
+    nodeFrom.asStmt() = nodeTo.asStmt().(CfgNodes::StmtNodes::AssignStmtCfgNode).getRightHandSide()
   }
 
-  predicate localMustFlowStep(Node node1, Node node2) { none() }
+  predicate localMustFlowStep(Node nodeFrom, Node nodeTo) {
+    nodeFrom.asStmt() = nodeTo.asStmt().(CfgNodes::StmtNodes::AssignStmtCfgNode).getRightHandSide()
+  }
 }
 
 /** Provides logic related to captured variables. */
@@ -125,11 +129,22 @@ private module Cached {
    * data flow.
    */
   cached
-  predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo, string model) { none() }
+  predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo, string model) {
+    (
+      LocalFlow::localFlowStepCommon(nodeFrom, nodeTo)
+      or
+      SsaFlow::localFlowStep(_, nodeFrom, nodeTo, _)
+    ) and
+    model = ""
+  }
 
   /** This is the local flow predicate that is exposed. */
   cached
-  predicate localFlowStepImpl(Node nodeFrom, Node nodeTo) { none() }
+  predicate localFlowStepImpl(Node nodeFrom, Node nodeTo) {
+    LocalFlow::localFlowStepCommon(nodeFrom, nodeTo)
+    or
+    SsaFlow::localFlowStep(_, nodeFrom, nodeTo, _)
+  }
 
   cached
   newtype TContentSet = TSingletonContent(Content c)