Add AuthorizeRequestDelegate to allow service module to adjust token audience (#15989)

wyunchi-ms · dingmeng-xue · web-flow · commit a75f056aad0e · 2021-10-03T19:47:38.000+08:00
* Add Erich's DataPlane code

* Add Erich's DataPlane code

* Add Erich's DataPlane code

* Update the code

Co-authored-by: wyunchi-ms &lt;yunwang@microsoft.com&gt;
Co-authored-by: dingmeng-xue &lt;dixue@microsoft.com&gt;
diff --git a/src/Accounts/Accounts/ChangeLog.md b/src/Accounts/Accounts/ChangeLog.md
@@ -19,6 +19,7 @@
 -->
 
 ## Upcoming Release
+* Added AuthorizeRequestDelegate to allow service module to adjust token audience.
 * Utilized [AssemblyLoadContext](https://docs.microsoft.com/en-us/dotnet/api/system.runtime.loader.assemblyloadcontext) to resolve assembly conflict issues in PowerShell.
 * Updated Azure.Core from 1.16.0 to 1.19.0.
 
diff --git a/src/Accounts/Accounts/CommonModule/ContextAdapter.cs b/src/Accounts/Accounts/CommonModule/ContextAdapter.cs
@@ -32,6 +32,7 @@ namespace Microsoft.Azure.Commands.Common
     using NextDelegate = Func<HttpRequestMessage, CancellationToken, Action, Func<string, CancellationToken, Func<EventArgs>, Task>, Task<HttpResponseMessage>>;
     using SignalDelegate = Func<string, CancellationToken, Func<EventArgs>, Task>;
     using PipelineChangeDelegate = 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 TokenAudienceConverterDelegate = Func<string, string, string, string, Uri, string>;
 
     /// <summary>
     /// Perform authentication and parameter completion based on the value of the context
@@ -74,6 +75,51 @@ public void OnNewRequest(InvocationInfo invocationInfo, string correlationId, st
             appendStep(this.SendHandler(GetDefaultContext(_provider, invocationInfo), AzureEnvironment.Endpoint.ResourceManager));
         }
 
+        internal void AddRequestUserAgentHandler(
+            InvocationInfo invocationInfo,
+            string correlationId,
+            string processRecordId,
+            PipelineChangeDelegate prependStep,
+            PipelineChangeDelegate appendStep)
+        {
+            appendStep(new UserAgent(invocationInfo).SendAsync);
+        }
+
+        internal void AddPatchRequestUriHandler(
+            InvocationInfo invocationInfo,
+            string correlationId,
+            string processRecordId,
+            PipelineChangeDelegate prependStep,
+            PipelineChangeDelegate appendStep)
+        {
+            appendStep(
+                async (request, cancelToken, cancelAction, signal, next) =>
+                {
+                    var context = GetDefaultContext(_provider, invocationInfo);
+                    PatchRequestUri(context, request);
+                    return await next(request, cancelToken, cancelAction, signal);
+                });
+        }
+
+        internal void AddAuthorizeRequestHandler(
+            InvocationInfo invocationInfo,
+            string endpointResourceIdKey,
+            string endpointSuffixKey,
+            PipelineChangeDelegate prependStep,
+            PipelineChangeDelegate appendStep,
+            TokenAudienceConverterDelegate tokenAudienceConverter,
+            IDictionary<string, object> extensibleParameters = null)
+        {
+            appendStep(
+                async (request, cancelToken, cancelAction, signal, next) =>
+                {
+                    endpointResourceIdKey = endpointResourceIdKey ?? AzureEnvironment.Endpoint.ResourceManager;
+                    var context = GetDefaultContext(_provider, invocationInfo);
+                    await AuthorizeRequest(context, request, cancelToken, endpointResourceIdKey, endpointSuffixKey, tokenAudienceConverter);
+                    return await next(request, cancelToken, cancelAction, signal);
+                });
+        }
+
         /// <summary>
         ///  Called for well-known parameters that require argument completers
         ///  </summary>
@@ -156,7 +202,7 @@ internal Func<HttpRequestMessage, CancellationToken, Action, SignalDelegate, Nex
             return async (request, cancelToken, cancelAction, signal, next) =>
             {
                 PatchRequestUri(context, request);
-                await AuthorizeRequest(context, resourceId, request, cancelToken);
+                await AuthorizeRequest(context, request, cancelToken, resourceId, resourceId);
                 return await next(request, cancelToken, cancelAction, signal);
             };
         }
@@ -165,11 +211,12 @@ internal Func<HttpRequestMessage, CancellationToken, Action, SignalDelegate, Nex
         /// Pipeline step for authenticating requests
         /// </summary>
         /// <param name="context"></param>
-        /// <param name="resourceId"></param>
+        /// <param name="endpointResourceIdKey"></param>
         /// <param name="request"></param>
         /// <param name="outerToken"></param>
         /// <returns></returns>
-        internal async Task AuthorizeRequest(IAzureContext context, string resourceId, HttpRequestMessage request, CancellationToken outerToken)
+        internal async Task AuthorizeRequest(IAzureContext context, HttpRequestMessage request, CancellationToken outerToken, string endpointResourceIdKey,
+                        string endpointSuffixKey, TokenAudienceConverterDelegate tokenAudienceConverter = null, IDictionary<string, object> extensibleParamters = null)
         {
             if (context == null || context.Account == null || context.Environment == null)
             {
@@ -178,12 +225,29 @@ internal async Task AuthorizeRequest(IAzureContext context, string resourceId, H
 
             await Task.Run(() =>
             {
-                resourceId = context?.Environment?.GetAudienceFromRequestUri(request.RequestUri) ?? resourceId;
-                var authToken = _authenticator.Authenticate(context.Account, context.Environment, context.Tenant.Id, null, "Never", null, resourceId);
+                if (tokenAudienceConverter != null)
+                {
+                    var info = GetEndpointInfo(context.Environment, endpointResourceIdKey, endpointSuffixKey);
+                    var tokenAudience = tokenAudienceConverter.Invoke(info.CurEnvEndpointResourceId, info.CurEnvEndpointSuffix, info.BaseEnvEndpointResourceId, info.BaseEnvEndpointSuffix, request.RequestUri);
+                    endpointResourceIdKey = tokenAudience ?? endpointResourceIdKey;
+                }
+                var authToken = _authenticator.Authenticate(context.Account, context.Environment, context.Tenant.Id, null, "Never", null, endpointResourceIdKey);
                 authToken.AuthorizeRequest((type, token) => request.Headers.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue(type, token));
             }, outerToken);
         }
 
+        private (string CurEnvEndpointResourceId, string CurEnvEndpointSuffix, string BaseEnvEndpointResourceId, string BaseEnvEndpointSuffix) GetEndpointInfo(IAzureEnvironment environment, string endpointResourceIdKey, string endpointSuffixKey)
+        {
+            var baseEnvironment = AzureEnvironment.PublicEnvironments[EnvironmentName.AzureCloud];
+
+            string curEnvEndpointResourceId = environment?.GetEndpoint(endpointResourceIdKey);
+            string curEnvEndpointSuffix = environment?.GetEndpoint(endpointSuffixKey);
+            string baseEnvEndpointResourceId = baseEnvironment?.GetEndpoint(endpointResourceIdKey);
+            string baseEnvEndpointSuffix = baseEnvironment?.GetEndpoint(endpointSuffixKey);
+
+            return (curEnvEndpointResourceId, curEnvEndpointSuffix, baseEnvEndpointResourceId, baseEnvEndpointSuffix); ;
+        }
+
         internal void PatchRequestUri(IAzureContext context, HttpRequestMessage request)
         {
             var requestUri = context?.Environment?.GetUriFromBaseRequestUri(request.RequestUri);
diff --git a/src/Accounts/Accounts/CommonModule/EnvironmentExtensions.cs b/src/Accounts/Accounts/CommonModule/EnvironmentExtensions.cs
@@ -103,6 +103,7 @@ public static Uri GetUriFromBaseRequestUri(this IAzureEnvironment environment, U
             return baseEndpoint;
         }
 
+        ////TODO: Update to support all data plane audience
         /// <summary>
         /// Determien the inteneded audience of a request
         /// </summary>
diff --git a/src/Accounts/Accounts/CommonModule/RegisterAzModule.cs b/src/Accounts/Accounts/CommonModule/RegisterAzModule.cs
@@ -54,6 +54,13 @@ protected override void ProcessRecord()
                     // this gets called before the generated cmdlet makes a call across the wire (allows you to change the HTTP pipeline)
                     OnNewRequest = ContextAdapter.Instance.OnNewRequest,
 
+                    //OnNewRequest = AddRequestUserAgentHandler + AddPatchRequestUriHandler + AddAuthorizeRequestHandler
+                    AddRequestUserAgentHandler = ContextAdapter.Instance.AddRequestUserAgentHandler,
+
+                    AddPatchRequestUriHandler = ContextAdapter.Instance.AddPatchRequestUriHandler,
+
+                    AddAuthorizeRequestHandler = ContextAdapter.Instance.AddAuthorizeRequestHandler,
+
                     // Called for well-known parameters that require argument completers
                     ArgumentCompleter = ContextAdapter.Instance.CompleteArgument,
 
diff --git a/src/Accounts/Accounts/CommonModule/VTable.cs b/src/Accounts/Accounts/CommonModule/VTable.cs
@@ -24,6 +24,13 @@ namespace Microsoft.Azure.Commands.Common
     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 ArgumentCompleterDelegate = Func<string, System.Management.Automation.InvocationInfo, string, string[], string[], string[]>;
+    using AuthorizeRequestDelegate = global::System.Action<System.Management.Automation.InvocationInfo,
+                                        string,
+                                        string,
+                                        global::System.Action<global::System.Func<global::System.Net.Http.HttpRequestMessage, global::System.Threading.CancellationToken, global::System.Action, global::System.Func<string, global::System.Threading.CancellationToken, global::System.Func<global::System.EventArgs>, global::System.Threading.Tasks.Task>, global::System.Func<global::System.Net.Http.HttpRequestMessage, global::System.Threading.CancellationToken, global::System.Action, global::System.Func<string, global::System.Threading.CancellationToken, global::System.Func<global::System.EventArgs>, global::System.Threading.Tasks.Task>, global::System.Threading.Tasks.Task<global::System.Net.Http.HttpResponseMessage>>, global::System.Threading.Tasks.Task<global::System.Net.Http.HttpResponseMessage>>>,
+                                        global::System.Action<global::System.Func<global::System.Net.Http.HttpRequestMessage, global::System.Threading.CancellationToken, global::System.Action, global::System.Func<string, global::System.Threading.CancellationToken, global::System.Func<global::System.EventArgs>, global::System.Threading.Tasks.Task>, global::System.Func<global::System.Net.Http.HttpRequestMessage, global::System.Threading.CancellationToken, global::System.Action, global::System.Func<string, global::System.Threading.CancellationToken, global::System.Func<global::System.EventArgs>, global::System.Threading.Tasks.Task>, global::System.Threading.Tasks.Task<global::System.Net.Http.HttpResponseMessage>>, global::System.Threading.Tasks.Task<global::System.Net.Http.HttpResponseMessage>>>,
+                                        global::System.Func<string, string, string, string, global::System.Uri, string>,
+                                        System.Collections.Generic.IDictionary<string, object>>;
 
     /// <summary>
     /// The Virtual Call table of the functions to be exported to the generated module
@@ -74,6 +81,12 @@ public class VTable
         /// <param name="appendStep">a delegate which allows the module to append a step in the HTTP Pipeline</param>
         public NewRequestPipelineDelegate OnNewRequest;
 
+        public NewRequestPipelineDelegate AddRequestUserAgentHandler;
+
+        public NewRequestPipelineDelegate AddPatchRequestUriHandler;
+
+        public AuthorizeRequestDelegate AddAuthorizeRequestHandler;
+
         /// <summary>
         /// Called for well-known parameters that require argument completers
         /// </summary>

Original file line number	Diff line number	Diff line change
`@@ -103,6 +103,7 @@ public static Uri GetUriFromBaseRequestUri(this IAzureEnvironment environment, U`
`103`	`103`	`return baseEndpoint;`
`104`	`104`	`}`
`105`	`105`
	`106`	`+ ////TODO: Update to support all data plane audience`
`106`	`107`	`/// <summary>`
`107`	`108`	`/// Determien the inteneded audience of a request`
`108`	`109`	`/// </summary>`