Adding code to remove dependency with AzureKeyVaultServiceEndpointResourceId

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;
             }