PS: Generalize the current ReturnContainer computation.

Author: MathiasVP

powershell/ql/lib/semmle/code/powershell/StatementBlock.qll

@@ -1,4 +1,12 @@
 import powershell
+private import semmle.code.powershell.internal.AstEscape::Private
+
+private module ReturnContainerInterpreter implements InterpretAstInputSig {
+  class T = Ast;
+
+  pragma[inline]
+  T interpret(Ast a) { result = a }
+}
 
 class StmtBlock extends @statement_block, Ast {
   override SourceLocation getLocation() { statement_block_location(this, result) }
@@ -16,4 +24,9 @@ class StmtBlock extends @statement_block, Ast {
   TrapStmt getATrapStmt() { result = this.getTrapStmt(_) }
 
   override string toString() { result = "{...}" }
+
+  /** Gets an element that may escape this `StmtBlock`. */
+  Ast getAnElement() {
+    result = this.(AstEscape<ReturnContainerInterpreter>::Element).getAnEscapingElement()
+  }
 }

powershell/ql/lib/semmle/code/powershell/dataflow/internal/DataFlowPrivate.qll

@@ -550,100 +550,45 @@ abstract class ReturnNode extends Node {
   abstract ReturnKind getKind();
 }
 
-private module ReturnNodes {
-  /** An AST element that may produce return values when evaluated. */
-  abstract private class ReturnContainer extends Ast {
-    /**
-     * Gets a direct node that will may be returned when evaluating this node.
-     */
-    CfgNode getANode() { none() }
-
-    /** Gets a child that may produce more nodes that may be returned. */
-    abstract ReturnContainer getAChild();
-
-    /**
-     * Gets a (possibly transitive) node that may be returned when evaluating
-     * this node.
-     */
-    final CfgNode getAReturnedNode() {
-      result = this.getANode()
+private module EscapeContainer {
+  private import semmle.code.powershell.internal.AstEscape::Private
+
+  private module ReturnContainerInterpreter implements InterpretAstInputSig {
+    class T = CfgNodes::AstCfgNode;
+
+    T interpret(Ast a) {
+      result.(CfgNodes::ExprCfgNode).getExpr() = a
       or
-      result = this.getAChild().getAReturnedNode()
+      result.(CfgNodes::StmtCfgNode).getStmt() = a.(Cmd)
     }
+  }
 
+  class EscapeContainer extends AstEscape<ReturnContainerInterpreter>::Element {
     /** Holds if `n` may be returned multiples times. */
     predicate mayBeMultiReturned(CfgNode n) {
       n = this.getANode() and
       n.getASuccessor+() = n
       or
-      this.getAChild().mayBeMultiReturned(n)
+      this.getAChild().(EscapeContainer).mayBeMultiReturned(n)
     }
   }
+}
 
-  class ScriptBlockReturnContainer extends ReturnContainer, ScriptBlock {
-    final override ReturnContainer getAChild() { result = this.getEndBlock() }
-  }
-
-  class NamedBlockReturnContainer extends ReturnContainer, NamedBlock {
-    final override ReturnContainer getAChild() { result = this.getAStmt() }
-  }
-
-  class CmdExprReturnContainer extends ReturnContainer, CmdExpr {
-    final override CfgNodes::ExprCfgNode getANode() { result.getExpr() = this.getExpr() }
-
-    final override ReturnContainer getAChild() { none() }
-  }
-
-  class LoopStmtReturnContainer extends ReturnContainer, LoopStmt {
-    final override ReturnContainer getAChild() { result = this.getBody() }
-  }
-
-  class StmtBlockReturnConainer extends ReturnContainer, StmtBlock {
-    final override ReturnContainer getAChild() { result = this.getAStmt() }
-  }
-
-  class TryStmtReturnContainer extends ReturnContainer, TryStmt {
-    final override ReturnContainer getAChild() {
-      result = this.getBody() or result = this.getACatchClause() or result = this.getFinally()
-    }
-  }
-
-  class ReturnStmtReturnContainer extends ReturnContainer, ReturnStmt {
-    final override ReturnContainer getAChild() { result = this.getPipeline() }
-  }
-
-  class CatchClausReturnContainer extends ReturnContainer, CatchClause {
-    final override ReturnContainer getAChild() { result = this.getBody() }
-  }
-
-  class SwitchStmtReturnContainer extends ReturnContainer, SwitchStmt {
-    final override ReturnContainer getAChild() { result = this.getACase() }
-  }
-
-  class CmdBaseReturnContainer extends ReturnContainer, CmdExpr {
-    final override CfgNodes::ExprCfgNode getANode() { result.getExpr() = this.getExpr() }
-
-    final override ReturnContainer getAChild() { none() }
-  }
-
-  class CmdReturnContainer extends ReturnContainer, Cmd {
-    final override CfgNodes::StmtCfgNode getANode() { result.getStmt() = this }
-
-    final override ReturnContainer getAChild() { none() }
-  }
+private module ReturnNodes {
+  private import EscapeContainer
 
-  private predicate isReturnedImpl(CfgNodes::AstCfgNode n, ReturnContainer container) {
+  private predicate isReturnedImpl(CfgNodes::AstCfgNode n, EscapeContainer container) {
     container = n.getScope() and
-    n = container.getAReturnedNode()
+    n = container.getAnEscapingElement()
   }
 
   /**
    * Holds if `n` may be returned, and there are possibly
    * more than one return value from the function.
    */
   predicate isMultiReturned(CfgNodes::AstCfgNode n) {
-    exists(ReturnContainer container | isReturnedImpl(n, container) |
-      strictcount(container.getAReturnedNode()) > 1
+    exists(EscapeContainer container | isReturnedImpl(n, container) |
+      strictcount(container.getAnEscapingElement()) > 1
       or
       container.mayBeMultiReturned(n)
     )

powershell/ql/lib/semmle/code/powershell/internal/AstEscape.qll

@@ -0,0 +1,96 @@
+private import powershell as PS
+
+/**
+ * TODO: This whole computation cab be sped up by providing a set of "root"s and doing
+ * a forward/backwards traversal first.
+ */
+module Private {
+  signature module InterpretAstInputSig {
+    /** The type on which to translate `Ast` elements during escape calculations */
+    class T;
+
+    /** Interpret `a` into a `T` */
+    T interpret(PS::Ast a);
+  }
+
+  module AstEscape<InterpretAstInputSig Interpret> {
+    private import Interpret
+
+    /** An AST element that may produce a value which can escape from this `Ast` when evaluated. */
+    abstract private class ElementImpl instanceof PS::Ast {
+      string toString() { result = super.toString() }
+
+      /** Gets a direct node that will may escape when evaluating this element. */
+      T getANode() { none() }
+
+      /** Gets a child that may produce more elements that may escape. */
+      abstract Element getAChild();
+
+      /**
+       * Gets a (possibly transitive) element that may escape when evaluating
+       * this element.
+       */
+      final T getAnEscapingElement() {
+        result = this.getANode()
+        or
+        result = this.getAChild().getAnEscapingElement()
+      }
+    }
+
+    final class Element = ElementImpl;
+
+    private class ScriptBlockElement extends ElementImpl instanceof PS::ScriptBlock {
+      final override Element getAChild() { result = super.getEndBlock() }
+    }
+
+    private class NamedBlockElement extends ElementImpl instanceof PS::NamedBlock {
+      final override Element getAChild() { result = super.getAStmt() }
+    }
+
+    private class CmdExprElement extends ElementImpl instanceof PS::CmdExpr {
+      final override T getANode() { result = interpret(super.getExpr()) }
+
+      final override Element getAChild() { none() }
+    }
+
+    private class LoopStmtElement extends ElementImpl instanceof PS::LoopStmt {
+      final override Element getAChild() { result = super.getBody() }
+    }
+
+    private class StmtBlockElement extends ElementImpl instanceof PS::StmtBlock {
+      final override Element getAChild() { result = super.getAStmt() }
+    }
+
+    private class TryStmtElement extends ElementImpl instanceof PS::TryStmt {
+      final override Element getAChild() {
+        result = super.getBody() or result = super.getACatchClause() or result = super.getFinally()
+      }
+    }
+
+    private class ReturnStmtElement extends ElementImpl instanceof PS::ReturnStmt {
+      final override Element getAChild() { result = super.getPipeline() }
+    }
+
+    private class CatchClausElement extends ElementImpl instanceof PS::CatchClause {
+      final override Element getAChild() { result = super.getBody() }
+    }
+
+    private class SwitchStmtElement extends ElementImpl instanceof PS::SwitchStmt {
+      final override Element getAChild() { result = super.getACase() }
+    }
+
+    private class CmdBaseElement extends ElementImpl instanceof PS::CmdExpr {
+      final override T getANode() { result = interpret(super.getExpr()) }
+
+      final override Element getAChild() { none() }
+    }
+
+    private class CmdElement extends ElementImpl instanceof PS::Cmd {
+      final override T getANode() { result = interpret(this) }
+
+      final override Element getAChild() { none() }
+    }
+  }
+}
+
+module Public { }