From b2cf155ff576b80482b184673f91a68d3a07321e Mon Sep 17 00:00:00 2001
From: Mathias Vorreiter Pedersen <mapeder@microsoft.com>
Date: Fri, 28 Mar 2025 19:35:41 +0000
Subject: [PATCH] PS: Define pre-return node and implicit-wrapping nodes using
 the script block instead of the individual elements.

---
 .../dataflow/internal/DataFlowPrivate.qll      | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

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 1dc66c573c5..f71ec194d29 100644
--- a/powershell/ql/lib/semmle/code/powershell/dataflow/internal/DataFlowPrivate.qll
+++ b/powershell/ql/lib/semmle/code/powershell/dataflow/internal/DataFlowPrivate.qll
@@ -131,10 +131,12 @@ module LocalFlow {
       nodeTo = TReturnNodeImpl(cfgNode.getScope())
     )
     or
-    exists(CfgNode cfgNode |
-      cfgNode = nodeFrom.(AstNode).getCfgNode() and
-      isUniqueReturned(cfgNode) and
-      nodeTo.(ReturnNodeImpl).getCfgScope() = cfgNode.getScope()
+    exists(CfgNodes::ExprCfgNode e, CfgNodes::ScriptBlockCfgNode scriptBlock |
+      e = nodeFrom.(AstNode).getCfgNode() and
+      isReturned(e) and
+      e.getScope() = scriptBlock.getAstNode() and
+      not blockMayReturnMultipleValues(scriptBlock) and
+      nodeTo.(ReturnNodeImpl).getCfgScope() = scriptBlock.getAstNode()
     )
   }
 
@@ -1082,12 +1084,12 @@ private import PostUpdateNodes
  * (or statement) is being returned from a function.
  */
 private class ImplicitWrapNode extends TImplicitWrapNode, NodeImpl {
-  private CfgNodes::AstCfgNode n;
+  private CfgNodes::ScriptBlockCfgNode n;
   private boolean shouldWrap;
 
   ImplicitWrapNode() { this = TImplicitWrapNode(n, shouldWrap) }
 
-  CfgNodes::AstCfgNode getReturnedNode() { result = n }
+  CfgNodes::ScriptBlockCfgNode getScriptBlock() { result = n }
 
   predicate shouldWrap() { shouldWrap = true }
 
@@ -1105,12 +1107,12 @@ private class ImplicitWrapNode extends TImplicitWrapNode, NodeImpl {
  * has been performed.
  */
 private class PreReturNodeImpl extends TPreReturnNodeImpl, NodeImpl {
-  private CfgNodes::AstCfgNode n;
+  private CfgNodes::ScriptBlockCfgNode n;
   private boolean isArray;
 
   PreReturNodeImpl() { this = TPreReturnNodeImpl(n, isArray) }
 
-  CfgNodes::AstCfgNode getReturnedNode() { result = n }
+  CfgNodes::AstCfgNode getScriptBlock() { result = n }
 
   override CfgScope getCfgScope() { result = n.getScope() }