Fix '-PipelineVariable' to set variable in the right scope (#16199)

This commit is contained in:
Dongbo Wang 2021-10-07 10:26:21 -07:00 committed by GitHub
parent 53ac646cc0
commit e98a8c8601
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 78 additions and 24 deletions

View file

@ -928,7 +928,8 @@ namespace System.Management.Automation
/// </summary>
internal string PipelineVariable { get; set; }
private PSVariable _pipelineVarReference = null;
private PSVariable _pipelineVarReference;
private bool _shouldRemovePipelineVariable;
internal void SetupOutVariable()
{
@ -972,7 +973,7 @@ namespace System.Management.Automation
// This can't use the common SetupVariable implementation, as this needs to persist for an entire
// pipeline.
if (string.IsNullOrEmpty(this.PipelineVariable))
if (string.IsNullOrEmpty(PipelineVariable))
{
return;
}
@ -983,12 +984,24 @@ namespace System.Management.Automation
_state = new SessionState(Context.EngineSessionState);
// Create the pipeline variable
_pipelineVarReference = new PSVariable(this.PipelineVariable);
_state.PSVariable.Set(_pipelineVarReference);
_pipelineVarReference = new PSVariable(PipelineVariable);
object varToUse = _state.Internal.SetVariable(
_pipelineVarReference,
force: false,
CommandOrigin.Internal);
// Get the reference again in case we re-used one from the
// same scope.
_pipelineVarReference = _state.PSVariable.Get(this.PipelineVariable);
if (ReferenceEquals(_pipelineVarReference, varToUse))
{
// The returned variable is the exact same instance, which means we set a new variable.
// In this case, we will try removing the pipeline variable in the end.
_shouldRemovePipelineVariable = true;
}
else
{
// A variable with the same name already exists in the same scope and it was returned.
// In this case, we update the reference and don't remove the variable in the end.
_pipelineVarReference = (PSVariable)varToUse;
}
if (_thisCommand is not PSScriptCmdlet)
{
@ -996,6 +1009,15 @@ namespace System.Management.Automation
}
}
internal void RemovePipelineVariable()
{
if (_shouldRemovePipelineVariable)
{
// Remove pipeline variable when a pipeline is being torn down.
_state.PSVariable.Remove(PipelineVariable);
}
}
/// <summary>
/// Configures the number of objects to buffer before calling the downstream Cmdlet.
/// </summary>
@ -3765,25 +3787,12 @@ namespace System.Management.Automation
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);
}
}
internal void RemoveVariableListsInPipe()
{
// Diagnostics.Assert(thisCommand is PSScriptCmdlet, "this is only done for script cmdlets");
if (_outVarList != null)
{
this.OutputPipe.RemoveVariableList(VariableStreamKind.Output, _outVarList);
@ -3807,9 +3816,6 @@ namespace System.Management.Automation
if (this.PipelineVariable != null)
{
this.OutputPipe.RemovePipelineVariable();
// '_state' could be null when a 'DynamicParam' block runs because the 'DynamicParam' block runs in 'DoPrepare',
// before 'PipelineProcessor.SetupParameterVariables' is called, where '_state' is initialized.
_state?.PSVariable.Remove(this.PipelineVariable);
}
}
}

View file

@ -1298,7 +1298,21 @@ namespace System.Management.Automation.Internal
// pipeline failure and continue disposing cmdlets.
try
{
commandProcessor.CommandRuntime.RemoveVariableListsInPipe();
// Only cmdlets can have variables defined via the common parameters.
// We handle the cleanup of those variables only if we need to.
if (commandProcessor is CommandProcessor)
{
if (commandProcessor.Command is not PSScriptCmdlet)
{
// For script cmdlets, the variable lists were already removed when exiting a scope.
// So we only need to take care of binary cmdlets here.
commandProcessor.CommandRuntime.RemoveVariableListsInPipe();
}
// Remove the pipeline variable if we need to.
commandProcessor.CommandRuntime.RemovePipelineVariable();
}
commandProcessor.Dispose();
}
// 2005/04/13-JonN: The only vaguely plausible reason

View file

@ -264,6 +264,40 @@ Describe "Parameter Binding Tests" -Tags "CI" {
{ Copy-Item "~\$guid*" -Destination ~ -ToSession $null } | Should -Throw -ErrorId 'ParameterArgumentValidationError'
}
It 'PipelineVariable should not cause variable-removal exception (issue #16155)' {
function Invoke-AddOne {
param (
[Parameter(ValueFromPipeline)]
[int]$Number
)
Begin {
$testValue = 'prefix-'
}
Process {
$testValue + $Number
}
}
1,2,3 | Invoke-AddOne -PipelineVariable testValue | ForEach-Object { $testValue } | Should -Be @('prefix-1', 'prefix-2', 'prefix-3')
{ Get-Variable -Name testValue -ErrorAction Stop } | Should -Throw -ErrorId 'VariableNotFound,Microsoft.PowerShell.Commands.GetVariableCommand'
$results = & { $test = 'str'; 1,2,3 | Invoke-AddOne -PipelineVariable test | ForEach-Object { $test }; Get-Variable test }
$results | Should -HaveCount 4
$results[0] | Should -BeExactly 'prefix-1'
$results[1] | Should -BeExactly 'prefix-2'
$results[2] | Should -BeExactly 'prefix-3'
$results[3] -is [psvariable] | Should -BeTrue
$results = & { Set-Variable -Name test -Value 'str'; 1,2,3 | Invoke-AddOne -PipelineVariable test | ForEach-Object { $test }; Get-Variable test }
$results | Should -HaveCount 4
$results[0] | Should -BeExactly 'prefix-1'
$results[1] | Should -BeExactly 'prefix-2'
$results[2] | Should -BeExactly 'prefix-3'
$results[3] -is [psvariable] | Should -BeTrue
}
Context "PipelineVariable Behaviour" {
BeforeAll {