Take reference on new TokenRequestContext (Azure#18440)

jongio · Feb 9, 2021 · db94912 · db94912
1 parent a80563a
commit db94912
Show file tree

Hide file tree

Showing 9 changed files with 59 additions and 32 deletions.
diff --git a/sdk/identity/Azure.Identity/CHANGELOG.md b/sdk/identity/Azure.Identity/CHANGELOG.md
@@ -6,6 +6,10 @@
 
 - The `IDisposable` interface has been removed from `TokenCache`.
 
+### New Features
+
+- All credentials added support to handle the `Claims` property on `TokenRequestContext`
+
 ## 1.4.0-beta.2 (2021-01-29)
 
 ### Fixes and improvements

diff --git a/sdk/identity/Azure.Identity/src/Azure.Identity.csproj b/sdk/identity/Azure.Identity/src/Azure.Identity.csproj
@@ -10,6 +10,7 @@
     <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
   </PropertyGroup>
   <ItemGroup>
+   <ProjectReference Include="$(MSBuildThisFileDirectory)..\..\..\core\Azure.Core\src\Azure.Core.csproj" />
     <PackageReference Include="System.Memory" />
     <PackageReference Include="System.Text.Json" />
     <PackageReference Include="System.Threading.Tasks.Extensions" />
@@ -30,5 +31,6 @@
     <Compile Include="$(AzureCoreSharedSources)Base64Url.cs" />
   </ItemGroup>
   <!-- Import the Azure.Base project -->
-  <Import Project="$(MSBuildThisFileDirectory)..\..\..\core\Azure.Core\src\Azure.Core.props" />
-</Project>
+  <!-- TODO: Revert after TokenRequestContext changes ship in Azure.Core -->
+  <!-- <Import Project="$(MSBuildThisFileDirectory)..\..\..\core\Azure.Core\src\Azure.Core.props" /> -->
+</Project>
diff --git a/sdk/identity/Azure.Identity/src/DeviceCodeCredential.cs b/sdk/identity/Azure.Identity/src/DeviceCodeCredential.cs
@@ -172,7 +172,7 @@ private async Task<AuthenticationRecord> AuthenticateImplAsync(bool async, Token
 
             try
             {
-                AccessToken token = await GetTokenViaDeviceCodeAsync(requestContext.Scopes, async, cancellationToken).ConfigureAwait(false);
+                AccessToken token = await GetTokenViaDeviceCodeAsync(requestContext, async, cancellationToken).ConfigureAwait(false);
 
                 scope.Succeeded(token);
 
@@ -196,7 +196,7 @@ private async ValueTask<AccessToken> GetTokenImplAsync(bool async, TokenRequestC
                 {
                     try
                     {
-                        AuthenticationResult result = await Client.AcquireTokenSilentAsync(requestContext.Scopes, Record, async, cancellationToken).ConfigureAwait(false);
+                        AuthenticationResult result = await Client.AcquireTokenSilentAsync(requestContext.Scopes, requestContext.Claims, Record, async, cancellationToken).ConfigureAwait(false);
 
                         return scope.Succeeded(new AccessToken(result.AccessToken, result.ExpiresOn));
                     }
@@ -211,17 +211,17 @@ private async ValueTask<AccessToken> GetTokenImplAsync(bool async, TokenRequestC
                     throw new AuthenticationRequiredException(AuthenticationRequiredMessage, requestContext, inner);
                 }
 
-                return scope.Succeeded(await GetTokenViaDeviceCodeAsync(requestContext.Scopes, async, cancellationToken).ConfigureAwait(false));
+                return scope.Succeeded(await GetTokenViaDeviceCodeAsync(requestContext, async, cancellationToken).ConfigureAwait(false));
             }
             catch (Exception e)
             {
                 throw scope.FailWrapAndThrow(e);
             }
         }
 
-        private async Task<AccessToken> GetTokenViaDeviceCodeAsync(string[] scopes, bool async, CancellationToken cancellationToken)
+        private async Task<AccessToken> GetTokenViaDeviceCodeAsync(TokenRequestContext context, bool async, CancellationToken cancellationToken)
         {
-            AuthenticationResult result = await Client.AcquireTokenWithDeviceCodeAsync(scopes, code => DeviceCodeCallbackImpl(code, cancellationToken), async, cancellationToken).ConfigureAwait(false);
+            AuthenticationResult result = await Client.AcquireTokenWithDeviceCodeAsync(context.Scopes, context.Claims, code => DeviceCodeCallbackImpl(code, cancellationToken), async, cancellationToken).ConfigureAwait(false);
 
             Record = new AuthenticationRecord(result, ClientId);
 

diff --git a/sdk/identity/Azure.Identity/src/InteractiveBrowserCredential.cs b/sdk/identity/Azure.Identity/src/InteractiveBrowserCredential.cs
@@ -160,7 +160,7 @@ private async Task<AuthenticationRecord> AuthenticateImplAsync(bool async, Token
 
             try
             {
-                scope.Succeeded(await GetTokenViaBrowserLoginAsync(requestContext.Scopes, async, cancellationToken).ConfigureAwait(false));
+                scope.Succeeded(await GetTokenViaBrowserLoginAsync(requestContext, async, cancellationToken).ConfigureAwait(false));
 
                 return Record;
             }
@@ -182,7 +182,7 @@ private async ValueTask<AccessToken> GetTokenImplAsync(bool async, TokenRequestC
                 {
                     try
                     {
-                        AuthenticationResult result = await Client.AcquireTokenSilentAsync(requestContext.Scopes, Record, async, cancellationToken).ConfigureAwait(false);
+                        AuthenticationResult result = await Client.AcquireTokenSilentAsync(requestContext.Scopes, requestContext.Claims, Record, async, cancellationToken).ConfigureAwait(false);
 
                         return scope.Succeeded(new AccessToken(result.AccessToken, result.ExpiresOn));
                     }
@@ -197,17 +197,17 @@ private async ValueTask<AccessToken> GetTokenImplAsync(bool async, TokenRequestC
                     throw new AuthenticationRequiredException(AuthenticationRequiredMessage, requestContext, inner);
                 }
 
-                return scope.Succeeded(await GetTokenViaBrowserLoginAsync(requestContext.Scopes, async, cancellationToken).ConfigureAwait(false));
+                return scope.Succeeded(await GetTokenViaBrowserLoginAsync(requestContext, async, cancellationToken).ConfigureAwait(false));
             }
             catch (Exception e)
             {
                 throw scope.FailWrapAndThrow(e);
             }
         }
 
-        private async Task<AccessToken> GetTokenViaBrowserLoginAsync(string[] scopes, bool async, CancellationToken cancellationToken)
+        private async Task<AccessToken> GetTokenViaBrowserLoginAsync(TokenRequestContext context, bool async, CancellationToken cancellationToken)
         {
-            AuthenticationResult result = await Client.AcquireTokenInteractiveAsync(scopes, Prompt.SelectAccount, async, cancellationToken).ConfigureAwait(false);
+            AuthenticationResult result = await Client.AcquireTokenInteractiveAsync(context.Scopes, context.Claims, Prompt.SelectAccount, async, cancellationToken).ConfigureAwait(false);
 
             Record = new AuthenticationRecord(result, ClientId);
 

diff --git a/sdk/identity/Azure.Identity/src/MsalPublicClient.cs b/sdk/identity/Azure.Identity/src/MsalPublicClient.cs
@@ -38,6 +38,8 @@ protected override ValueTask<IPublicClientApplication> CreateClientAsync(bool as
                 pubAppBuilder = pubAppBuilder.WithRedirectUri(RedirectUrl);
             }
 
+            pubAppBuilder.WithClientCapabilities(new string[] { "CP1" });
+
             return new ValueTask<IPublicClientApplication>(pubAppBuilder.Build());
         }
 
@@ -47,12 +49,15 @@ public virtual async ValueTask<List<IAccount>> GetAccountsAsync(bool async, Canc
             return await GetAccountsAsync(client, async).ConfigureAwait(false);
         }
 
-        public virtual async ValueTask<AuthenticationResult> AcquireTokenSilentAsync(string[] scopes, IAccount account, bool async, CancellationToken cancellationToken)
+        public virtual async ValueTask<AuthenticationResult> AcquireTokenSilentAsync(string[] scopes, string claims, IAccount account, bool async, CancellationToken cancellationToken)
         {
             IPublicClientApplication client = await GetClientAsync(async, cancellationToken).ConfigureAwait(false);
-            return await client.AcquireTokenSilent(scopes, account).ExecuteAsync(async, cancellationToken).ConfigureAwait(false);
+            return await client.AcquireTokenSilent(scopes, account)
+                .WithClaims(claims)
+                .ExecuteAsync(async, cancellationToken)
+                .ConfigureAwait(false);
         }
-        public virtual async ValueTask<AuthenticationResult> AcquireTokenSilentAsync(string[] scopes, AuthenticationRecord record, bool async, CancellationToken cancellationToken)
+        public virtual async ValueTask<AuthenticationResult> AcquireTokenSilentAsync(string[] scopes, string claims, AuthenticationRecord record, bool async, CancellationToken cancellationToken)
         {
             IPublicClientApplication client = await GetClientAsync(async, cancellationToken).ConfigureAwait(false);
 
@@ -61,31 +66,47 @@ public virtual async ValueTask<AuthenticationResult> AcquireTokenSilentAsync(str
             // user authenticated to originally.
             return await client.AcquireTokenSilent(scopes, (AuthenticationAccount)record)
                 .WithAuthority(Pipeline.AuthorityHost.AbsoluteUri, TenantId ?? record.TenantId)
-                .ExecuteAsync(async, cancellationToken).ConfigureAwait(false);
+                .WithClaims(claims)
+                .ExecuteAsync(async, cancellationToken)
+                .ConfigureAwait(false);
         }
 
-        public virtual async ValueTask<AuthenticationResult> AcquireTokenInteractiveAsync(string[] scopes, Prompt prompt, bool async, CancellationToken cancellationToken)
+        public virtual async ValueTask<AuthenticationResult> AcquireTokenInteractiveAsync(string[] scopes, string claims, Prompt prompt, bool async, CancellationToken cancellationToken)
         {
             IPublicClientApplication client = await GetClientAsync(async, cancellationToken).ConfigureAwait(false);
-            return await client.AcquireTokenInteractive(scopes).WithPrompt(prompt).ExecuteAsync(async, cancellationToken).ConfigureAwait(false);
+            return await client.AcquireTokenInteractive(scopes)
+                .WithPrompt(prompt)
+                .WithClaims(claims)
+                .ExecuteAsync(async, cancellationToken)
+                .ConfigureAwait(false);
         }
 
-        public virtual async ValueTask<AuthenticationResult> AcquireTokenByUsernamePasswordAsync(string[] scopes, string username, SecureString password, bool async, CancellationToken cancellationToken)
+        public virtual async ValueTask<AuthenticationResult> AcquireTokenByUsernamePasswordAsync(string[] scopes, string claims, string username, SecureString password, bool async, CancellationToken cancellationToken)
         {
             IPublicClientApplication client = await GetClientAsync(async, cancellationToken).ConfigureAwait(false);
-            return await client.AcquireTokenByUsernamePassword(scopes, username, password).ExecuteAsync(async, cancellationToken).ConfigureAwait(false);
+            return await client.AcquireTokenByUsernamePassword(scopes, username, password)
+                .WithClaims(claims)
+                .ExecuteAsync(async, cancellationToken)
+                .ConfigureAwait(false);
         }
 
-        public virtual async ValueTask<AuthenticationResult> AcquireTokenWithDeviceCodeAsync(string[] scopes, Func<DeviceCodeResult, Task> deviceCodeCallback, bool async, CancellationToken cancellationToken)
+        public virtual async ValueTask<AuthenticationResult> AcquireTokenWithDeviceCodeAsync(string[] scopes, string claims, Func<DeviceCodeResult, Task> deviceCodeCallback, bool async, CancellationToken cancellationToken)
         {
             IPublicClientApplication client = await GetClientAsync(async, cancellationToken).ConfigureAwait(false);
-            return await client.AcquireTokenWithDeviceCode(scopes, deviceCodeCallback).ExecuteAsync(async, cancellationToken).ConfigureAwait(false);
+            return await client.AcquireTokenWithDeviceCode(scopes, deviceCodeCallback)
+                .WithClaims(claims)
+                .ExecuteAsync(async, cancellationToken)
+                .ConfigureAwait(false);
         }
 
-        public virtual async ValueTask<AuthenticationResult> AcquireTokenByRefreshToken(string[] scopes, string refreshToken, AzureCloudInstance azureCloudInstance, string tenant, bool async, CancellationToken cancellationToken)
+        public virtual async ValueTask<AuthenticationResult> AcquireTokenByRefreshToken(string[] scopes, string claims, string refreshToken, AzureCloudInstance azureCloudInstance, string tenant, bool async, CancellationToken cancellationToken)
         {
             IPublicClientApplication client = await GetClientAsync(async, cancellationToken).ConfigureAwait(false);
-            return await ((IByRefreshToken)client).AcquireTokenByRefreshToken(scopes, refreshToken).WithAuthority(azureCloudInstance, tenant).ExecuteAsync(async, cancellationToken).ConfigureAwait(false);
+            return await ((IByRefreshToken)client).AcquireTokenByRefreshToken(scopes, refreshToken)
+                .WithAuthority(azureCloudInstance, tenant)
+                .WithClaims(claims)
+                .ExecuteAsync(async, cancellationToken)
+                .ConfigureAwait(false);
         }
 
         private static async ValueTask<List<IAccount>> GetAccountsAsync(IPublicClientApplication client, bool async)

diff --git a/sdk/identity/Azure.Identity/src/SharedTokenCacheCredential.cs b/sdk/identity/Azure.Identity/src/SharedTokenCacheCredential.cs
@@ -113,7 +113,7 @@ private async ValueTask<AccessToken> GetTokenImplAsync(bool async, TokenRequestC
             try
             {
                 IAccount account = await GetAccountAsync(async, cancellationToken).ConfigureAwait(false);
-                AuthenticationResult result = await Client.AcquireTokenSilentAsync(requestContext.Scopes, account, async, cancellationToken).ConfigureAwait(false);
+                AuthenticationResult result = await Client.AcquireTokenSilentAsync(requestContext.Scopes, requestContext.Claims, account, async, cancellationToken).ConfigureAwait(false);
                 return scope.Succeeded(new AccessToken(result.AccessToken, result.ExpiresOn));
             }
             catch (MsalUiRequiredException)

diff --git a/sdk/identity/Azure.Identity/src/UsernamePasswordCredential.cs b/sdk/identity/Azure.Identity/src/UsernamePasswordCredential.cs
@@ -185,7 +185,7 @@ private async Task<AccessToken> GetTokenImplAsync(bool async, TokenRequestContex
             try
             {
                 AuthenticationResult result = await _client
-                    .AcquireTokenByUsernamePasswordAsync(requestContext.Scopes, _username, _password, async, cancellationToken)
+                    .AcquireTokenByUsernamePasswordAsync(requestContext.Scopes, requestContext.Claims, _username, _password, async, cancellationToken)
                     .ConfigureAwait(false);
 
                 _record = new AuthenticationRecord(result, _clientId);

diff --git a/sdk/identity/Azure.Identity/src/VisualStudioCodeCredential.cs b/sdk/identity/Azure.Identity/src/VisualStudioCodeCredential.cs
@@ -71,7 +71,7 @@ private async ValueTask<AccessToken> GetTokenImplAsync(TokenRequestContext reque
                 var cloudInstance = GetAzureCloudInstance(environmentName);
                 string storedCredentials = GetStoredCredentials(environmentName);
 
-                var result = await _client.AcquireTokenByRefreshToken(requestContext.Scopes, storedCredentials, cloudInstance, tenant, async, cancellationToken).ConfigureAwait(false);
+                var result = await _client.AcquireTokenByRefreshToken(requestContext.Scopes, requestContext.Claims, storedCredentials, cloudInstance, tenant, async, cancellationToken).ConfigureAwait(false);
                 return scope.Succeeded(new AccessToken(result.AccessToken, result.ExpiresOn));
             }
             catch (MsalUiRequiredException e)

diff --git a/sdk/identity/Azure.Identity/tests/Mock/MockMsalPublicClient.cs b/sdk/identity/Azure.Identity/tests/Mock/MockMsalPublicClient.cs
@@ -32,7 +32,7 @@ public override ValueTask<List<IAccount>> GetAccountsAsync(bool async, Cancellat
             return new ValueTask<List<IAccount>>(Accounts);
         }
 
-        public override ValueTask<AuthenticationResult> AcquireTokenByUsernamePasswordAsync(string[] scopes, string username, SecureString password, bool async, CancellationToken cancellationToken)
+        public override ValueTask<AuthenticationResult> AcquireTokenByUsernamePasswordAsync(string[] scopes, string claims, string username, SecureString password, bool async, CancellationToken cancellationToken)
         {
             Func<string[], AuthenticationResult> factory = UserPassAuthFactory ?? AuthFactory;
 
@@ -44,7 +44,7 @@ public override ValueTask<AuthenticationResult> AcquireTokenByUsernamePasswordAs
             throw new NotImplementedException();
         }
 
-        public override ValueTask<AuthenticationResult> AcquireTokenInteractiveAsync(string[] scopes, Prompt prompt, bool async, CancellationToken cancellationToken)
+        public override ValueTask<AuthenticationResult> AcquireTokenInteractiveAsync(string[] scopes, string claims, Prompt prompt, bool async, CancellationToken cancellationToken)
         {
             Func<string[], AuthenticationResult> factory = InteractiveAuthFactory ?? AuthFactory;
 
@@ -56,7 +56,7 @@ public override ValueTask<AuthenticationResult> AcquireTokenInteractiveAsync(str
             throw new NotImplementedException();
         }
 
-        public override ValueTask<AuthenticationResult> AcquireTokenSilentAsync(string[] scopes, IAccount account, bool async, CancellationToken cancellationToken)
+        public override ValueTask<AuthenticationResult> AcquireTokenSilentAsync(string[] scopes, string claims, IAccount account, bool async, CancellationToken cancellationToken)
         {
             if (ExtendedSilentAuthFactory != null)
             {
@@ -73,7 +73,7 @@ public override ValueTask<AuthenticationResult> AcquireTokenSilentAsync(string[]
             throw new NotImplementedException();
         }
 
-        public override ValueTask<AuthenticationResult> AcquireTokenSilentAsync(string[] scopes, AuthenticationRecord record, bool async, CancellationToken cancellationToken)
+        public override ValueTask<AuthenticationResult> AcquireTokenSilentAsync(string[] scopes, string claims, AuthenticationRecord record, bool async, CancellationToken cancellationToken)
         {
             Func<string[], AuthenticationResult> factory = SilentAuthFactory ?? AuthFactory;
 
@@ -85,7 +85,7 @@ public override ValueTask<AuthenticationResult> AcquireTokenSilentAsync(string[]
             throw new NotImplementedException();
         }
 
-        public override ValueTask<AuthenticationResult> AcquireTokenWithDeviceCodeAsync(string[] scopes, Func<DeviceCodeResult, Task> deviceCodeCallback, bool async, CancellationToken cancellationToken)
+        public override ValueTask<AuthenticationResult> AcquireTokenWithDeviceCodeAsync(string[] scopes, string claims, Func<DeviceCodeResult, Task> deviceCodeCallback, bool async, CancellationToken cancellationToken)
         {
             Func<string[], AuthenticationResult> factory = DeviceCodeAuthFactory ?? AuthFactory;