Add -StrictMode to Invoke-Command to allow specifying strict mode…

… when invoking command locally (PowerShell#16545)

src/System.Management.Automation/engine/ExperimentalFeature/ExperimentalFeature.cs

@@ -27,6 +27,7 @@ public class ExperimentalFeature
         internal const string PSCleanBlockFeatureName = "PSCleanBlock";
         internal const string PSAMSIMethodInvocationLogging = "PSAMSIMethodInvocationLogging";
         internal const string PSExecFeatureName = "PSExec";
+        internal const string PSStrictModeAssignment = "PSStrictModeAssignment";
 
         #endregion
 
@@ -142,6 +143,9 @@ static ExperimentalFeature()
                 new ExperimentalFeature(
                     name: PSExecFeatureName,
                     description: "Add 'exec' built-in command on Linux and macOS"),
+                new ExperimentalFeature(
+                    name: PSStrictModeAssignment,
+                    description: "Add support of setting Strict-Mode with Invoke-Command"),
             };
 
             EngineExperimentalFeatures = new ReadOnlyCollection<ExperimentalFeature>(engineFeatures);

src/System.Management.Automation/engine/remoting/commands/InvokeCommandCommand.cs

@@ -263,6 +263,72 @@ public override SwitchParameter UseSSL
             }
         }
 
+        private sealed class ArgumentToPSVersionTransformationAttribute : ArgumentToVersionTransformationAttribute
+        {
+            protected override bool TryConvertFromString(string versionString, [NotNullWhen(true)] out Version version)
+            {
+                if (string.Equals("off", versionString, StringComparison.OrdinalIgnoreCase))
+                {
+                    version = new Version(0, 0);
+                    return true;
+                }
+
+                if (string.Equals("latest", versionString, StringComparison.OrdinalIgnoreCase))
+                {
+                    version = PSVersionInfo.PSVersion;
+                    return true;
+                }
+
+                return base.TryConvertFromString(versionString, out version);
+            }
+        }
+
+        private static readonly Version s_OffVersion = new Version(0, 0);
+
+        private sealed class ValidateVersionAttribute : ValidateArgumentsAttribute
+        {
+            protected override void Validate(object arguments, EngineIntrinsics engineIntrinsics)
+            {
+                Version version = arguments as Version;
+                if (version == s_OffVersion)
+                {
+                    return;
+                }
+
+                if (version == null || !PSVersionInfo.IsValidPSVersion(version))
+                {
+                    // No conversion succeeded so throw an exception...
+                    throw new ValidationMetadataException(
+                        "InvalidPSVersion",
+                        null,
+                        Metadata.ValidateVersionFailure,
+                        arguments);
+                }
+            }
+        }
+
+        /// <summary>
+        /// Gets or sets strict mode.
+        /// </summary>
+        [Experimental(ExperimentalFeature.PSStrictModeAssignment, ExperimentAction.Show)]
+        [Parameter(ParameterSetName = InvokeCommandCommand.InProcParameterSet)]
+        [ArgumentToPSVersionTransformation]
+        [ValidateVersion]
+        public Version StrictMode
+        {
+            get
+            {
+                return _strictmodeversion;
+            }
+
+            set
+            {
+                _strictmodeversion = value;
+            }
+        }
+
+        private Version _strictmodeversion = null;
+
         /// <summary>
         /// For WSMan session:
         /// If this parameter is not specified then the value specified in
@@ -823,6 +889,8 @@ public virtual SwitchParameter RemoteDebug
 
         #endregion
 
+        private Version _savedStrictModeVersion;
+
         #endregion Parameters
 
         #region Overrides
@@ -946,6 +1014,12 @@ protected override void BeginProcessing()
                     }
                 }
 
+                if (_strictmodeversion != null)
+                {
+                    _savedStrictModeVersion = Context.EngineSessionState.CurrentScope.StrictModeVersion;
+                    Context.EngineSessionState.CurrentScope.StrictModeVersion = _strictmodeversion;
+                }
+
                 return;
             }
 
@@ -1162,7 +1236,19 @@ protected override void ProcessRecord()
                 }
                 else if (ParameterSetName.Equals(InvokeCommandCommand.InProcParameterSet) && (_steppablePipeline != null))
                 {
-                    _steppablePipeline.Process(InputObject);
+                    try
+                    {
+                        _steppablePipeline.Process(InputObject);
+                    }
+                    catch
+                    {
+                        if (_strictmodeversion != null)
+                        {
+                            Context.EngineSessionState.CurrentScope.StrictModeVersion = _savedStrictModeVersion;
+                        }
+
+                        throw;
+                    }
                 }
                 else
                 {
@@ -1193,20 +1279,30 @@ protected override void EndProcessing()
             {
                 if (ParameterSetName.Equals(InvokeCommandCommand.InProcParameterSet))
                 {
-                    if (_steppablePipeline != null)
-                    {
-                        _steppablePipeline.End();
+                    try
+                    {   
+                        if (_steppablePipeline != null)
+                        {
+                            _steppablePipeline.End();
+                        }
+                        else
+                        {
+                            ScriptBlock.InvokeUsingCmdlet(
+                                contextCmdlet: this,
+                                useLocalScope: !NoNewScope,
+                                errorHandlingBehavior: ScriptBlock.ErrorHandlingBehavior.WriteToCurrentErrorPipe,
+                                dollarUnder: AutomationNull.Value,
+                                input: _input,
+                                scriptThis: AutomationNull.Value,
+                                args: ArgumentList);
+                        }
                     }
-                    else
+                    finally
                     {
-                        ScriptBlock.InvokeUsingCmdlet(
-                            contextCmdlet: this,
-                            useLocalScope: !NoNewScope,
-                            errorHandlingBehavior: ScriptBlock.ErrorHandlingBehavior.WriteToCurrentErrorPipe,
-                            dollarUnder: AutomationNull.Value,
-                            input: _input,
-                            scriptThis: AutomationNull.Value,
-                            args: ArgumentList);
+                        if (_strictmodeversion != null)
+                        {
+                            Context.EngineSessionState.CurrentScope.StrictModeVersion = _savedStrictModeVersion;
+                        }
                     }
                 }
                 else

test/powershell/Modules/Microsoft.PowerShell.Core/Invoke-Command.Tests.ps1

@@ -0,0 +1,41 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+Describe "Invoke-Command" -Tags "CI" {
+    Context "StrictMode tests" {
+        BeforeAll {
+            $skipTest = !($EnabledExperimentalFeatures -contains "PSStrictModeAssignment");
+            If (Test-Path Variable:InvokeCommand__Test) {
+                Remove-Item Variable:InvokeCommand__Test
+            }
+        }
+
+        It "Setting -StrictMode parameter with uninitialized variable throws error"  -skip:$skipTest {
+            { Invoke-Command -StrictMode 3.0 {$InvokeCommand__Test} } | Should -Throw -ErrorId 'VariableIsUndefined'
+        }
+
+        It "Setting -StrictMode parameter with initialized variable does not throw error" -skip:$skipTest {
+            $InvokeCommand__Test = 'Something'
+            Invoke-Command -StrictMode 3.0 {$InvokeCommand__Test} | Should -Be 'Something'
+            Remove-Item Variable:InvokeCommand__Test
+        }
+
+        It "-StrictMode parameter sets StrictMode back to original state after process completes" -skip:$skipTest {
+            { Invoke-Command -StrictMode 3.0 {$InvokeCommand__Test} } | Should -Throw -ErrorId 'VariableIsUndefined'
+            { Invoke-Command {$InvokeCommand__Test} } | Should -Not -Throw
+        }
+
+        It "-StrictMode parameter works on piped input" -skip:$skipTest {
+            "There" | Invoke-Command -ScriptBlock { "Hello $input" } -StrictMode 3.0 | Should -Be 'Hello There'
+            { "There" | Invoke-Command -ScriptBlock { "Hello $InvokeCommand__Test" } -StrictMode 3.0 } | Should -Throw -ErrorId 'VariableIsUndefined'
+        }
+
+        It "-StrictMode latest works" -skip:$skipTest {
+            { Invoke-Command -StrictMode latest {$InvokeCommand__Test} } | Should -Throw -ErrorId 'VariableIsUndefined'
+        }
+
+        It "-StrictMode off works" -skip:$skipTest {
+            { Invoke-Command -StrictMode off {$InvokeCommand__Test} } | Should -Not -Throw
+        }
+    }
+}