Adding code to remove dependency with static value specified in AzureKeyVaultServiceEndpointResourceId of Get-AzureRmEnvironment. The keyvault resource Id is now obtained from server using 401 challenge. The Azure PowerShell common infrastructure currently do not support this in an intuitive fashion, so tweaking its usage by overriding the environment endpoint value before acquiring the token.

Author: Kapil Jain

src/ResourceManager/KeyVault/Commands.KeyVault/Models/DataServiceCredential.cs

@@ -24,26 +24,23 @@ namespace Microsoft.Azure.Commands.KeyVault.Models
 {
     internal class DataServiceCredential
     {
+        private readonly IAuthenticationFactory _authenticationFactory;
+        private readonly AzureContext _context;
+        private readonly AzureEnvironment.Endpoint _endpointName;
+
         public DataServiceCredential(IAuthenticationFactory authFactory, AzureContext context, AzureEnvironment.Endpoint resourceIdEndpoint)
         {
             if (authFactory == null)
                 throw new ArgumentNullException("authFactory");
             if (context == null)
                 throw new ArgumentNullException("context");
-
-            var bundle = GetToken(authFactory, context, resourceIdEndpoint);
-            this.token = bundle.Item1;
+            _authenticationFactory = authFactory;
+            _context = context;
+            _endpointName = resourceIdEndpoint;
+            this.TenantId = GetTenantId(context);
         }
 
-        public string AccessToken
-        {
-            get
-            {
-                return token.AccessToken;
-            }
-        }
-
-        public string TenantId { get; set; }
+        public string TenantId { get; private set; }
 
         /// <summary>
         /// Authentication callback method required by KeyVaultClient
@@ -56,16 +53,29 @@ public Task<string> OnAuthentication(string authority, string resource, string s
         {
             // TODO: Add trace to log tokenType, resource, authority, scope etc
             string tokenStr = string.Empty;
-            this.token.AuthorizeRequest((tokenType, tokenValue) =>
+
+            // overriding the cached resourceId value to resource returned from the server
+            if (!string.IsNullOrEmpty(resource))
+            {
+                _context.Environment.Endpoints[_endpointName] = resource;
+            }
+
+            var bundle = GetTokenInternal(this.TenantId, this._authenticationFactory, this._context, this._endpointName);
+            bundle.Item1.AuthorizeRequest((tokenType, tokenValue) =>
             {
                 tokenStr = tokenValue;
             });
-
             return Task.FromResult<string>(tokenStr);
         }
 
-        private Tuple<IAccessToken, string> GetToken(IAuthenticationFactory authFactory, AzureContext context, AzureEnvironment.Endpoint resourceIdEndpoint)
+        public string GetToken()
         {
+            return GetTokenInternal(this.TenantId, this._authenticationFactory, this._context, this._endpointName).Item1.AccessToken;
+        }
+
+        private static string GetTenantId(AzureContext context)
+        {
+            var tenantId = string.Empty;
             if (context.Account == null)
                 throw new ArgumentException(KeyVaultProperties.Resources.ArmAccountNotFound);
 
@@ -74,14 +84,18 @@ private Tuple<IAccessToken, string> GetToken(IAuthenticationFactory authFactory,
                 throw new ArgumentException(string.Format(KeyVaultProperties.Resources.UnsupportedAccountType, context.Account.Type));
 
             if (context.Subscription != null && context.Account != null)
-                TenantId = context.Subscription.GetPropertyAsArray(AzureSubscription.Property.Tenants)
+                tenantId = context.Subscription.GetPropertyAsArray(AzureSubscription.Property.Tenants)
                        .Intersect(context.Account.GetPropertyAsArray(AzureAccount.Property.Tenants))
                        .FirstOrDefault();
 
-            if (string.IsNullOrWhiteSpace(TenantId) && context.Tenant != null && context.Tenant.Id != Guid.Empty)
-                TenantId = context.Tenant.Id.ToString();
+            if (string.IsNullOrWhiteSpace(tenantId) && context.Tenant != null && context.Tenant.Id != Guid.Empty)
+                tenantId = context.Tenant.Id.ToString();
+            return tenantId;
+        }
 
-            if (string.IsNullOrWhiteSpace(TenantId))
+        private static Tuple<IAccessToken, string> GetTokenInternal(string tenantId, IAuthenticationFactory authFactory, AzureContext context, AzureEnvironment.Endpoint resourceIdEndpoint)
+        {
+            if (string.IsNullOrWhiteSpace(tenantId))
                 throw new ArgumentException(KeyVaultProperties.Resources.NoTenantInContext);
 
             try
@@ -92,7 +106,7 @@ private Tuple<IAccessToken, string> GetToken(IAuthenticationFactory authFactory,
                     tokenCache = new TokenCache(context.TokenCache);
                 }
 
-                var accesstoken = authFactory.Authenticate(context.Account, context.Environment, TenantId, null, ShowDialog.Never,
+                var accesstoken = authFactory.Authenticate(context.Account, context.Environment, tenantId, null, ShowDialog.Never,
                     tokenCache, resourceIdEndpoint);
 
                 if (context.TokenCache != null && context.TokenCache.Length > 0)
@@ -107,7 +121,5 @@ private Tuple<IAccessToken, string> GetToken(IAuthenticationFactory authFactory,
                 throw new ArgumentException(KeyVaultProperties.Resources.InvalidSubscriptionState, ex);
             }
         }
-
-        private IAccessToken token;
     }
 }

src/ResourceManager/KeyVault/Commands.KeyVault/Models/KeyVaultManagementCmdletBase.cs

@@ -61,7 +61,7 @@ public ActiveDirectoryClient ActiveDirectoryClient
                     _dataServiceCredential = new DataServiceCredential(AzureSession.AuthenticationFactory, DefaultProfile.Context, AzureEnvironment.Endpoint.Graph);
                     _activeDirectoryClient = new ActiveDirectoryClient(new Uri(string.Format("{0}/{1}",
                         DefaultProfile.Context.Environment.Endpoints[AzureEnvironment.Endpoint.Graph], _dataServiceCredential.TenantId)),
-                        () => Task.FromResult(_dataServiceCredential.AccessToken));
+                        () => Task.FromResult(_dataServiceCredential.GetToken()));
                 }
                 return this._activeDirectoryClient;
             }