Light up CAE support for ManagedIdentityCredential (Azure#48549)

Author: christothes

eng/Packages.Data.props

@@ -168,9 +168,9 @@
     <!-- Other approved packages -->
     <PackageReference Update="Microsoft.Azure.Amqp" Version="2.6.9" />
     <PackageReference Update="Microsoft.Azure.WebPubSub.Common" Version="1.4.0" />
-    <PackageReference Update="Microsoft.Identity.Client" Version="4.67.2" />
-    <PackageReference Update="Microsoft.Identity.Client.Extensions.Msal" Version="4.67.2" />
-    <PackageReference Update="Microsoft.Identity.Client.Broker" Version="4.67.2" />
+    <PackageReference Update="Microsoft.Identity.Client" Version="4.69.1" />
+    <PackageReference Update="Microsoft.Identity.Client.Extensions.Msal" Version="4.69.1" />
+    <PackageReference Update="Microsoft.Identity.Client.Broker" Version="4.69.1" />
 
     <!-- TODO: Make sure this package is arch-board approved -->
     <PackageReference Update="Microsoft.IdentityModel.Protocols.OpenIdConnect" Version="6.35.0" />

sdk/identity/Azure.Identity/CHANGELOG.md

@@ -10,6 +10,8 @@
 - `VisualStudioCredential` will now correctly fall through to the next credential in the chain when no account is found by Visual Studio. ([#48464](https://github.com/Azure/azure-sdk-for-net/issues/48464))
 
 ### Other Changes
+- Updated Microsoft.Identity.Client dependency to version 4.69.1.
+- ManagedIdentityCredential now properly supports CAE.
 - An event is now logged when the `ManagedIdentityCredential` is used directly or indirectly via a credential chain indicating which managed identity source was selected and which `ManagedIdentityId` was specified.
 - Marked `UsernamePasswordCredential` as obsolete because Resource Owner Password Credentials (ROPC) token grant flow is incompatible with multifactor authentication (MFA), which Microsoft Entra ID requires for all tenants. See https://aka.ms/azsdk/identity/mfa for details about MFA enforcement and migration guidance.
 

sdk/identity/Azure.Identity/src/MsalManagedIdentityClient.cs

@@ -15,13 +15,15 @@ namespace Azure.Identity
     internal class MsalManagedIdentityClient
     {
         private readonly AsyncLockWithValue<IManagedIdentityApplication> _clientAsyncLock;
+        private readonly AsyncLockWithValue<IManagedIdentityApplication> _clientCaeAsyncLock;
         private bool _isForceRefreshEnabled { get; }
 
         internal bool IsSupportLoggingEnabled { get; }
         internal Microsoft.Identity.Client.AppConfig.ManagedIdentityId ManagedIdentityId { get; }
         internal bool DisableInstanceDiscovery { get; }
         internal CredentialPipeline Pipeline { get; }
         internal Uri AuthorityHost { get; }
+        protected string[] cp1Capabilities = ["CP1"];
 
         protected MsalManagedIdentityClient()
         { }
@@ -49,34 +51,45 @@ public MsalManagedIdentityClient(ManagedIdentityClientOptions clientOptions)
 
             Pipeline = clientOptions.Pipeline;
             _clientAsyncLock = new AsyncLockWithValue<IManagedIdentityApplication>();
+            _clientCaeAsyncLock = new AsyncLockWithValue<IManagedIdentityApplication>();
             _isForceRefreshEnabled = clientOptions.IsForceRefreshEnabled;
         }
 
-        protected ValueTask<IManagedIdentityApplication> CreateClientAsync(bool async, CancellationToken cancellationToken)
+        protected ValueTask<IManagedIdentityApplication> CreateClientAsync(bool async, bool enableCae, CancellationToken cancellationToken)
         {
-            return CreateClientCoreAsync(async, cancellationToken);
+            return CreateClientCoreAsync(async, enableCae, cancellationToken);
         }
 
-        protected virtual ValueTask<IManagedIdentityApplication> CreateClientCoreAsync(bool async, CancellationToken cancellationToken)
+        protected virtual ValueTask<IManagedIdentityApplication> CreateClientCoreAsync(bool async, bool enableCae, CancellationToken cancellationToken)
         {
+            string[] clientCapabilities =
+                enableCae ? cp1Capabilities : Array.Empty<string>();
+
             ManagedIdentityApplicationBuilder miAppBuilder = ManagedIdentityApplicationBuilder
                 .Create(ManagedIdentityId)
                 .WithHttpClientFactory(new HttpPipelineClientFactory(Pipeline.HttpPipeline), false)
                 .WithLogging(AzureIdentityEventSource.Singleton, enablePiiLogging: IsSupportLoggingEnabled);
 
+            if (clientCapabilities.Length > 0)
+            {
+                miAppBuilder.WithClientCapabilities(clientCapabilities);
+            }
+
             return new ValueTask<IManagedIdentityApplication>(miAppBuilder.Build());
         }
 
-        protected async ValueTask<IManagedIdentityApplication> GetClientAsync(bool async, CancellationToken cancellationToken)
+        protected async ValueTask<IManagedIdentityApplication> GetClientAsync(bool async, bool enableCae, CancellationToken cancellationToken)
         {
-            using var asyncLock = await _clientAsyncLock.GetLockOrValueAsync(async, cancellationToken).ConfigureAwait(false);
+            using var asyncLock = enableCae ?
+                await _clientCaeAsyncLock.GetLockOrValueAsync(async, cancellationToken).ConfigureAwait(false) :
+                await _clientAsyncLock.GetLockOrValueAsync(async, cancellationToken).ConfigureAwait(false);
 
             if (asyncLock.HasValue)
             {
                 return asyncLock.Value;
             }
 
-            var client = await CreateClientAsync(async, cancellationToken).ConfigureAwait(false);
+            var client = await CreateClientAsync(async, enableCae, cancellationToken).ConfigureAwait(false);
             asyncLock.SetValue(client);
             return client;
         }
@@ -89,10 +102,15 @@ public virtual AuthenticationResult AcquireTokenForManagedIdentity(TokenRequestC
 
         public virtual async ValueTask<AuthenticationResult> AcquireTokenForManagedIdentityAsyncCore(bool async, TokenRequestContext requestContext, CancellationToken cancellationToken)
         {
-            IManagedIdentityApplication client = await GetClientAsync(async, cancellationToken).ConfigureAwait(false);
+            IManagedIdentityApplication client = await GetClientAsync(async, requestContext.IsCaeEnabled, cancellationToken).ConfigureAwait(false);
 
             var builder = client.AcquireTokenForManagedIdentity(requestContext.Scopes.FirstOrDefault());
 
+            if (!string.IsNullOrEmpty(requestContext.Claims))
+            {
+                builder.WithClaims(requestContext.Claims);
+            }
+
             if (_isForceRefreshEnabled)
             {
                 builder.WithForceRefresh(true);

sdk/identity/Azure.Identity/tests/ManagedIdentityCredentialTests.cs

@@ -551,7 +551,16 @@ public async Task VerifyIMDSRequestWithPodIdentityEnvVarResourceIdMockAsync()
             var mockTransport = new MockTransport(response);
             var options = new TokenCredentialOptions() { Transport = mockTransport };
 
-            ManagedIdentityCredential credential = InstrumentClient(new ManagedIdentityCredential(new ResourceIdentifier(_expectedResourceId), options));
+            ManagedIdentityCredential credential = InstrumentClient(new ManagedIdentityCredential(
+                new ManagedIdentityClient(
+                    new ManagedIdentityClientOptions()
+                    {
+                        Pipeline = CredentialPipeline.GetInstance(options),
+                        ManagedIdentityId = ManagedIdentityId.FromUserAssignedResourceId(new ResourceIdentifier(_expectedResourceId)),
+                        IsForceRefreshEnabled = true,
+                        Options = options
+                    })
+            ));
 
             AccessToken actualToken = await credential.GetTokenAsync(new TokenRequestContext(MockScopes.Default));
 
@@ -1182,6 +1191,40 @@ public async Task VerifyManagedIdentityIdIsLogged(ManagedIdentityId managedIdent
             Assert.That(messages, Does.Contain(expectedMessage));
         }
 
+        [NonParallelizable]
+        [Test]
+        public async Task VerifyImdsRequestWithCAEMockAsync()
+        {
+            string caeClaims = """{"access_token":{"nbf":{"essential":true, "value":"1724337680"}}}""";
+            using var environment = new TestEnvVar(new() { { "MSI_ENDPOINT", null }, { "MSI_SECRET", null }, { "IDENTITY_ENDPOINT", null }, { "IDENTITY_HEADER", null }, { "AZURE_POD_IDENTITY_AUTHORITY_HOST", null } });
+
+            var mockTransport = new MockTransport(_ => CreateMockResponse(200, Guid.NewGuid().ToString()));
+            var options = new TokenCredentialOptions { Transport = mockTransport };
+
+            ManagedIdentityCredential credential = InstrumentClient(new ManagedIdentityCredential(
+                new ManagedIdentityClient(
+                    new ManagedIdentityClientOptions()
+                    {
+                        Pipeline = CredentialPipeline.GetInstance(options),
+                        ManagedIdentityId = ManagedIdentityId.FromUserAssignedResourceId(new ResourceIdentifier(_expectedResourceId)),
+                        IsForceRefreshEnabled = false,
+                        Options = options
+                    })
+            ));
+
+            AccessToken tokenWithoutCAE1 = await credential.GetTokenAsync(new TokenRequestContext(MockScopes.Default));
+            AccessToken tokenWithoutCAE2 = await credential.GetTokenAsync(new TokenRequestContext(MockScopes.Default));
+            AccessToken tokenWithCAE1 = await credential.GetTokenAsync(new TokenRequestContext(MockScopes.Default, isCaeEnabled: true, claims: caeClaims));
+            AccessToken tokenWithoutCAE3 = await credential.GetTokenAsync(new TokenRequestContext(MockScopes.Default));
+
+            // TokenWithoutCAE1 and TokenWithoutCAE2 should be the same since subsequent calls to GetTokenAsync should return the same token if no claims were provided
+            Assert.AreEqual(tokenWithoutCAE2.Token, tokenWithoutCAE1.Token);
+            // TokenWithCAE1 and TokenWithoutCAE3 should be the same since subsequent calls to GetTokenAsync should return the same token if no claims were provided.
+            Assert.AreEqual(tokenWithoutCAE3.Token, tokenWithCAE1.Token);
+            // TokenWithCAE1 and TokenWithoutCAE1 should be different since the first token was requested with claims.
+            Assert.AreNotEqual(tokenWithCAE1.Token, tokenWithoutCAE1.Token);
+        }
+
         private static IEnumerable<TestCaseData> ResourceAndClientIds()
         {
             yield return new TestCaseData(new object[] { null, false });

sdk/identity/Azure.Identity/tests/Mock/MockMsalManagedIdentityClient.cs

@@ -24,11 +24,11 @@ public MockMsalManagedIdentityClient(AuthenticationResult result)
         public MockMsalManagedIdentityClient(ManagedIdentityClientOptions options)
             : base(options) { }
 
-        protected override ValueTask<IManagedIdentityApplication> CreateClientCoreAsync(bool async, CancellationToken cancellationToken)
+        protected override ValueTask<IManagedIdentityApplication> CreateClientCoreAsync(bool async, bool enableCae, CancellationToken cancellationToken)
         {
             if (ClientAppFactory == null)
             {
-                return base.CreateClientCoreAsync(async, cancellationToken);
+                return base.CreateClientCoreAsync(async, enableCae, cancellationToken);
             }
 
             return new ValueTask<IManagedIdentityApplication>(ClientAppFactory(cancellationToken));