From bbbb4ee9e6de820cc38579eb02e88b6ee52bcb97 Mon Sep 17 00:00:00 2001
From: Joel Sallow <32407840+vexx32@users.noreply.github.com>
Date: Fri, 22 May 2020 17:01:19 -0400
Subject: [PATCH 1/2] :bug: Fix PipelineVariable not being set correctly

Previously, for script functions -PipelineVariable would only take
effect for the duration of the first output, and then stopped working.

Borrowed some code from the compiled cmdlets' implementation
which has always worked just fine. This appears to resolve the issue.

:bug: Fix NRE with dynamicparam

Dynamicparam blocks also do enter/exit scope, but seem to lack a state.
Workaround is to check for null state and fallback to previous
behaviour for dynamicparam blocks.

This should not cause any issues since output can't be passed from
dynamicparam, so -PipelineVariable can't be set from this block.
---
 .../engine/MshCommandRuntime.cs                       | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/src/System.Management.Automation/engine/MshCommandRuntime.cs b/src/System.Management.Automation/engine/MshCommandRuntime.cs
index 2816dd42f13..12f0904f289 100644
--- a/src/System.Management.Automation/engine/MshCommandRuntime.cs
+++ b/src/System.Management.Automation/engine/MshCommandRuntime.cs
@@ -3744,6 +3744,17 @@ internal void SetVariableListsInPipe()
 
             if (this.PipelineVariable != null)
             {
+                // _state can be null if the current script block is dynamicparam, etc.
+                if (_state != null)
+                {
+                    // Create the pipeline variable
+                    _state.PSVariable.Set(_pipelineVarReference);
+
+                    // Get the reference again in case we re-used one from the
+                    // same scope.
+                    _pipelineVarReference = _state.PSVariable.Get(this.PipelineVariable);
+                }
+
                 this.OutputPipe.SetPipelineVariable(_pipelineVarReference);
             }
         }

From 7b0a58b47b88538c56846e778fa7ae2d79dcc1e5 Mon Sep 17 00:00:00 2001
From: Joel Sallow <32407840+vexx32@users.noreply.github.com>
Date: Sat, 23 May 2020 02:20:30 -0400
Subject: [PATCH 2/2] :white_check_mark: Add test for pipelinevariable

---
 .../ParameterBinding.Tests.ps1                | 56 +++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/test/powershell/engine/ParameterBinding/ParameterBinding.Tests.ps1 b/test/powershell/engine/ParameterBinding/ParameterBinding.Tests.ps1
index 07d132012a0..a0fe5973c8e 100644
--- a/test/powershell/engine/ParameterBinding/ParameterBinding.Tests.ps1
+++ b/test/powershell/engine/ParameterBinding/ParameterBinding.Tests.ps1
@@ -256,6 +256,62 @@ Describe "Parameter Binding Tests" -Tags "CI" {
         DynamicParamTest -PipelineVariable bar | ForEach-Object { $bar } | Should -Be "hi"
     }
 
+    Context "PipelineVariable Behaviour" {
+
+        BeforeAll {
+            function Write-PipelineVariable {
+                [CmdletBinding()]
+                [OutputType([int])]
+                param(
+                    [Parameter(ValueFromPipeline)]
+                    $a
+                )
+                begin { 1 }
+                process { 2 }
+                end { 3 }
+            }
+
+            $testScripts = @(
+                @{
+                    CmdletType = 'Script Cmdlet'
+                    Script     = {
+                        1..3 |
+                            Write-PipelineVariable -PipelineVariable pipe |
+                            Select-Object -Property @(
+                                @{ Name = "PipelineVariableSet"; Expression = { $null -ne $pipe ? $true : $false } }
+                                @{ Name = "PipelineVariable"; Expression = { $pipe } }
+                            )
+                    }
+                }
+                @{
+                    CmdletType = 'Compiled Cmdlet'
+                    Script     = {
+                        1..3 |
+                            Write-PipelineVariable |
+                            ForEach-Object { $_ } -PipelineVariable pipe |
+                            Select-Object -Property @(
+                                @{ Name = "PipelineVariableSet"; Expression = { $null -ne $pipe ? $true : $false } }
+                                @{ Name = "PipelineVariable"; Expression = { $pipe } }
+                            )
+                    }
+                }
+            )
+        }
+
+        AfterAll {
+            Remove-Item -Path 'function:Write-PipelineVariable'
+        }
+
+        It 'should set the pipeline variable every time for a <CmdletType>' -TestCases $testScripts {
+            param($Script, $CmdletType)
+
+            $result = & $Script
+            $result.Count | Should -Be 5
+            $result.PipelineVariableSet | Should -Not -Contain $false
+            $result.PipelineVariable | Should -Be 1, 2, 2, 2, 3
+        }
+    }
+
     Context "Use automatic variables as default value for parameters" {
         BeforeAll {
             ## Explicit use of 'CmdletBinding' make it a script cmdlet