[Change Safety] Az.Accounts: AutoRest AcquirePolicyToken pipeline step (Stage A)

src/Accounts/Accounts.Test/ChangeSafetyParameterContractTests.cs

@@ -0,0 +1,64 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using Microsoft.WindowsAzure.Commands.Common;
+using Microsoft.WindowsAzure.Commands.ScenarioTest;
+using System.Linq;
+using System.Management.Automation;
+using Xunit;
+
+namespace Microsoft.Azure.Commands.Profile.Test
+{
+    /// <summary>
+    /// Pins the Change Safety parameter names and help text that Az.Accounts reads from the cmdlet's
+    /// BoundParameters (via the pipeline step in <see cref="Microsoft.Azure.Commands.Common.ContextAdapter" />).
+    ///
+    /// The AutoRest generator hardcodes these same literal strings when it emits the static
+    /// -AcquirePolicyToken / -ChangeReference parameters on write-verb cmdlets, and there is no
+    /// compile-time link between the generator (powershell/cmdlets/class.ts) and this library. If a name
+    /// or help message changes here, these tests fail as a reminder to update the generator to match and
+    /// regenerate the modules; otherwise the header would silently stop being stamped.
+    /// </summary>
+    public class ChangeSafetyParameterContractTests
+    {
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void ParameterNamesMatchGeneratorLiterals()
+        {
+            Assert.Equal("AcquirePolicyToken", ChangeSafetyParameters.AcquirePolicyTokenParamName);
+            Assert.Equal("ChangeReference", ChangeSafetyParameters.ChangeReferenceParamName);
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void ParameterHelpTextMatchesGeneratorLiterals()
+        {
+            var dict = new RuntimeDefinedParameterDictionary();
+            ChangeSafetyParameters.AddChangeSafetyParameters(dict);
+
+            Assert.Equal(
+                "Acquire an Azure Policy token automatically for this resource operation.",
+                GetHelpMessage(dict, ChangeSafetyParameters.AcquirePolicyTokenParamName));
+            Assert.Equal(
+                "The change reference resource ID for this resource operation.",
+                GetHelpMessage(dict, ChangeSafetyParameters.ChangeReferenceParamName));
+        }
+
+        private static string GetHelpMessage(RuntimeDefinedParameterDictionary dict, string name)
+        {
+            var attribute = dict[name].Attributes.OfType<ParameterAttribute>().First();
+            return attribute.HelpMessage;
+        }
+    }
+}

src/Accounts/Accounts.Test/UnitTest/AcquirePolicyTokenHandlerTests.cs

@@ -0,0 +1,137 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using Microsoft.Azure.Commands.Common;
+using Microsoft.Azure.Commands.Common.Authentication;
+using Microsoft.Azure.Commands.Common.Authentication.Models;
+using Microsoft.WindowsAzure.Commands.Common;
+using System;
+using System.Collections.Generic;
+using System.Management.Automation;
+using System.Net.Http;
+using System.Runtime.CompilerServices;
+using System.Threading;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace Microsoft.Azure.Commands.Profile.Test.UnitTest
+{
+    // Matches the PipelineChangeDelegate alias declared in ContextAdapter.cs so the internal
+    // AddChangeSafetyPolicyTokenHandler overload can be invoked directly from tests.
+    using PipelineStep = Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Task<HttpResponseMessage>>, Task<HttpResponseMessage>>;
+
+    public class AcquirePolicyTokenHandlerTests
+    {
+        public AcquirePolicyTokenHandlerTests()
+        {
+            // ContextAdapter's constructor reads AzureSession.Instance, so ensure a session exists.
+            AzureSessionInitializer.CreateOrReplaceSession(new MemoryDataStore());
+        }
+
+        // InvocationInfo has no public constructor, so build an uninitialized instance and set its
+        // BoundParameters via the non-public setter to exercise the handler's parameter evaluation.
+        private static InvocationInfo CreateInvocationInfo(IDictionary<string, object> boundParameters)
+        {
+            var invocationInfo = (InvocationInfo)RuntimeHelpers.GetUninitializedObject(typeof(InvocationInfo));
+            if (boundParameters != null)
+            {
+                var setter = typeof(InvocationInfo).GetProperty(nameof(InvocationInfo.BoundParameters)).GetSetMethod(nonPublic: true);
+                setter.Invoke(invocationInfo, new object[] { new Dictionary<string, object>(boundParameters) });
+            }
+            return invocationInfo;
+        }
+
+        private static int CountAppendedSteps(IDictionary<string, object> boundParameters)
+        {
+            int appended = 0;
+            Action<PipelineStep> appendStep = _ => appended++;
+            ContextAdapter.Instance.AddChangeSafetyPolicyTokenHandler(CreateInvocationInfo(boundParameters), appendStep);
+            return appended;
+        }
+
+        [Fact]
+        public void NoChangeSafetyParameters_DoesNotAppendStep()
+        {
+            var boundParameters = new Dictionary<string, object>();
+            Assert.Equal(0, CountAppendedSteps(boundParameters));
+        }
+
+        [Fact]
+        public void NullBoundParameters_DoesNotAppendStep()
+        {
+            Assert.Equal(0, CountAppendedSteps(null));
+        }
+
+        [Fact]
+        public void AcquirePolicyTokenSwitchPresent_AppendsSingleStep()
+        {
+            var boundParameters = new Dictionary<string, object>
+            {
+                { ChangeSafetyParameters.AcquirePolicyTokenParamName, new SwitchParameter(true) }
+            };
+            Assert.Equal(1, CountAppendedSteps(boundParameters));
+        }
+
+        [Fact]
+        public void AcquirePolicyTokenSwitchFalse_DoesNotAppendStep()
+        {
+            var boundParameters = new Dictionary<string, object>
+            {
+                { ChangeSafetyParameters.AcquirePolicyTokenParamName, new SwitchParameter(false) }
+            };
+            Assert.Equal(0, CountAppendedSteps(boundParameters));
+        }
+
+        [Fact]
+        public void ChangeReferenceNonEmpty_AppendsSingleStep()
+        {
+            var boundParameters = new Dictionary<string, object>
+            {
+                { ChangeSafetyParameters.ChangeReferenceParamName, "/subscriptions/change-ref" }
+            };
+            Assert.Equal(1, CountAppendedSteps(boundParameters));
+        }
+
+        [Fact]
+        public void ChangeReferenceEmpty_DoesNotAppendStep()
+        {
+            var boundParameters = new Dictionary<string, object>
+            {
+                { ChangeSafetyParameters.ChangeReferenceParamName, string.Empty }
+            };
+            Assert.Equal(0, CountAppendedSteps(boundParameters));
+        }
+
+        [Fact]
+        public void ChangeReferenceWhitespace_DoesNotAppendStep()
+        {
+            var boundParameters = new Dictionary<string, object>
+            {
+                { ChangeSafetyParameters.ChangeReferenceParamName, "   " }
+            };
+            Assert.Equal(0, CountAppendedSteps(boundParameters));
+        }
+
+        [Fact]
+        public void BothAcquireSwitchAndChangeReference_AppendsSingleStep()
+        {
+            var boundParameters = new Dictionary<string, object>
+            {
+                { ChangeSafetyParameters.AcquirePolicyTokenParamName, new SwitchParameter(true) },
+                { ChangeSafetyParameters.ChangeReferenceParamName, "/subscriptions/change-ref" }
+            };
+            Assert.Equal(1, CountAppendedSteps(boundParameters));
+        }
+    }
+}

src/Accounts/Accounts/ChangeLog.md

@@ -23,6 +23,7 @@
 ## Version 5.5.2
 * Upgraded `Azure.Core` dependency from 1.56.0 to 1.57.0.
 * Upgraded `System.ClientModel` dependency from 1.12.0 to 1.13.0.
+* Upgraded common library to `1.3.114-preview`.
 
 ## Version 5.5.1
 * Upgraded `Azure.Core` dependency from 1.50.0 to 1.56.0.

src/Accounts/Accounts/CommonModule/ContextAdapter.cs

@@ -123,6 +123,43 @@ internal void AddAuthorizeRequestHandler(
                 });
         }
 
+        /// <summary>
+        /// Change safety pipeline hook, exposed as its own VTable delegate and invoked by the generated
+        /// module right after OnNewRequest. Conditionally appends a step that acquires an Azure Policy
+        /// token and stamps it onto outgoing write requests, based on the -AcquirePolicyToken /
+        /// -ChangeReference bound parameters. When the feature is off, no step is added (zero added cost).
+        /// The write-verb gate lives inside <c>StampPolicyTokenAsync</c>, so GET sub-requests are skipped
+        /// even though the step is added for the cmdlet.
+        /// </summary>
+        internal void AddChangeSafetyPolicyTokenHandler(InvocationInfo invocationInfo, PipelineChangeDelegate appendStep)
+        {
+            var boundParameters = invocationInfo?.BoundParameters;
+            if (boundParameters == null) { return; }
+
+            bool acquire = boundParameters.TryGetValue(Microsoft.WindowsAzure.Commands.Common.ChangeSafetyParameters.AcquirePolicyTokenParamName, out var acquireVal)
+                && acquireVal is SwitchParameter sp && sp.ToBool();
+            string changeReference = boundParameters.TryGetValue(Microsoft.WindowsAzure.Commands.Common.ChangeSafetyParameters.ChangeReferenceParamName, out var crVal)
+                ? crVal as string : null;
+            // Treat a whitespace-only change reference as not provided, so it doesn't trigger acquisition.
+            if (string.IsNullOrWhiteSpace(changeReference)) { changeReference = null; }
+            bool shouldAcquire = acquire || changeReference != null;
+            if (!shouldAcquire) { return; } // feature off -> no added pipeline step (zero cost)
+
+            var acquirer = new Microsoft.WindowsAzure.Commands.Common.PolicyTokenAcquirer();
+            appendStep(
+                async (request, cancelToken, cancelAction, signal, next) =>
+                {
+                    await acquirer.StampPolicyTokenAsync(
+                        request,
+                        shouldAcquire: shouldAcquire,
+                        changeReference: changeReference,
+                        debugMessages: null,
+                        tokenHttpClient: null,
+                        cancellationToken: cancelToken).ConfigureAwait(false);
+                    return await next(request, cancelToken, cancelAction, signal).ConfigureAwait(false);
+                });
+        }
+
         /// <summary>
         ///  Called for well-known parameters that require argument completers
         ///  </summary>

src/Accounts/Accounts/CommonModule/RegisterAzModule.cs

@@ -74,6 +74,9 @@ protected override void ProcessRecord()
 
                     AddAuthorizeRequestHandler = ContextAdapter.Instance.AddAuthorizeRequestHandler,
 
+                    // change safety policy-token step; the generated module invokes this after OnNewRequest
+                    AddChangeSafetyPolicyTokenHandler = ContextAdapter.Instance.AddChangeSafetyPolicyTokenHandler,
+
                     // Called for well-known parameters that require argument completers
                     ArgumentCompleter = ContextAdapter.Instance.CompleteArgument,
 

src/Accounts/Accounts/CommonModule/VTable.cs

@@ -25,6 +25,7 @@ namespace Microsoft.Azure.Commands.Common
     using GetTelemetryIdDelegate = Func<string>;
     using ModuleLoadPipelineDelegate = Action<string, string, Action<Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Task<HttpResponseMessage>>, Task<HttpResponseMessage>>>, Action<Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Task<HttpResponseMessage>>, Task<HttpResponseMessage>>>>;
     using NewRequestPipelineDelegate = Action<System.Management.Automation.InvocationInfo, string, string, Action<Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Task<HttpResponseMessage>>, Task<HttpResponseMessage>>>, Action<Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Task<HttpResponseMessage>>, Task<HttpResponseMessage>>>>;
+    using ChangeSafetyPolicyTokenDelegate = Action<System.Management.Automation.InvocationInfo, Action<Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Task<HttpResponseMessage>>, Task<HttpResponseMessage>>>>;
     using ArgumentCompleterDelegate = Func<string, System.Management.Automation.InvocationInfo, string, string[], string[], string[]>;
     using AuthorizeRequestDelegate = global::System.Action<System.Management.Automation.InvocationInfo,
                                         string,
@@ -105,6 +106,12 @@ public class VTable
 
         public AuthorizeRequestDelegate AddAuthorizeRequestHandler;
 
+        /// <summary>
+        /// Called by the generated module after OnNewRequest to conditionally add the change safety
+        /// policy-token step, based on the -AcquirePolicyToken / -ChangeReference bound parameters.
+        /// </summary>
+        public ChangeSafetyPolicyTokenDelegate AddChangeSafetyPolicyTokenHandler;
+
         public SanitizerDelegate SanitizerHandler;
 
         public GetTelemetryInfoDelegate GetTelemetryInfo;

tools/Common.Netcore.Dependencies.targets

@@ -3,22 +3,22 @@
   <ItemGroup>
     <PackageReference Include="Microsoft.Rest.ClientRuntime" Version="2.3.24"/>
     <PackageReference Include="Microsoft.Rest.ClientRuntime.Azure" Version="3.3.19"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Aks" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Authentication.Abstractions" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Authorization" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Common" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Compute" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Graph.Rbac" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.KeyVault" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Monitor" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Network" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.PolicyInsights" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.ResourceManager" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Storage" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Storage.Management" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Strategies" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Websites" Version="1.3.113-preview"/>
-    <PackageReference Include="Microsoft.Azure.PowerShell.Common.Share" Version="1.3.113-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Aks" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Authentication.Abstractions" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Authorization" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Common" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Compute" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Graph.Rbac" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.KeyVault" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Monitor" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Network" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.PolicyInsights" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.ResourceManager" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Storage" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Storage.Management" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Strategies" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Clients.Websites" Version="1.3.114-preview"/>
+    <PackageReference Include="Microsoft.Azure.PowerShell.Common.Share" Version="1.3.114-preview"/>
     <PackageReference Include="Microsoft.CSharp" Version="4.7.0" />
   </ItemGroup>
   <ItemGroup>
@@ -37,7 +37,7 @@
     <PackageReference Include="PowerShellStandard.Library" Version="5.1.0" PrivateAssets="All" />
   </ItemGroup>
   <PropertyGroup>
-    <StorageToolsPath>$(NugetPackageRoot)\microsoft.azure.powershell.storage\1.3.113-preview\tools\</StorageToolsPath>
+    <StorageToolsPath>$(NugetPackageRoot)\microsoft.azure.powershell.storage\1.3.114-preview\tools\</StorageToolsPath>
   </PropertyGroup>
   <ItemGroup Condition="'$(OmitJsonPackage)' != 'true'">
     <PackageReference Include="Newtonsoft.Json" Version="13.0.2"/>