Add Client Assertion Credential (#26900)

sdk/identity/azure-identity/src/main/java/com/azure/identity/AksExchangeTokenCredential.java

@@ -0,0 +1,38 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.identity;
+
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenRequestContext;
+import com.azure.core.util.logging.ClientLogger;
+import com.azure.identity.implementation.IdentityClient;
+
+import reactor.core.publisher.Mono;
+
+/**
+ * Authenticates a service principal with AAD using a client assertion.
+ */
+class AksExchangeTokenCredential extends ManagedIdentityServiceCredential {
+    private final ClientLogger logger = new ClientLogger(AksExchangeTokenCredential.class);
+
+    /**
+     * Creates an instance of AksExchangeTokenCredential.
+     *
+     * @param clientId the client id of user assigned or system assigned identity.
+     * @param identityClient the identity client to acquire a token with.
+     */
+    AksExchangeTokenCredential(String clientId, IdentityClient identityClient) {
+        super(clientId, identityClient, "AZURE AKS TOKEN EXCHANGE");
+    }
+
+    @Override
+    public Mono<AccessToken> authenticate(TokenRequestContext request) {
+        if (this.getClientId() == null) {
+            return Mono.error(logger.logExceptionAsError(new IllegalStateException("The client id is not configured via"
+                + " 'AZURE_CLIENT_ID' environment variable or through the credential builder."
+                + " Please ensure client id is provided to authenticate via token exchange in AKS environment.")));
+        }
+        return identityClient.authenticatewithExchangeToken(request);
+    }
+}

sdk/identity/azure-identity/src/main/java/com/azure/identity/ClientAssertionCredential.java

@@ -4,35 +4,47 @@
 package com.azure.identity;
 
 import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenCredential;
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.identity.implementation.IdentityClient;
-
+import com.azure.identity.implementation.IdentityClientBuilder;
+import com.azure.identity.implementation.IdentityClientOptions;
+import com.azure.identity.implementation.util.LoggingUtil;
 import reactor.core.publisher.Mono;
 
+import java.util.function.Supplier;
+
 /**
  * Authenticates a service principal with AAD using a client assertion.
  */
-class ClientAssertionCredential extends ManagedIdentityServiceCredential {
+public class ClientAssertionCredential implements TokenCredential {
     private final ClientLogger logger = new ClientLogger(ClientAssertionCredential.class);
-
+    private final IdentityClient identityClient;
     /**
      * Creates an instance of ClientAssertionCredential.
      *
      * @param clientId the client id of user assigned or system assigned identity.
-     * @param identityClient the identity client to acquire a token with.
+     * @param tenantId the tenant ID of the application
+     * @param clientId the client ID of the application
+     * @param identityClientOptions the options to configure the identity client
      */
-    ClientAssertionCredential(String clientId, IdentityClient identityClient) {
-        super(clientId, identityClient, "AZURE AKS TOKEN EXCHANGE");
+    ClientAssertionCredential(String clientId, String tenantId, Supplier<String> clientAssertion,
+                              IdentityClientOptions identityClientOptions) {
+        identityClient = new IdentityClientBuilder()
+            .tenantId(tenantId)
+            .clientId(clientId)
+            .clientAssertionSupplier(clientAssertion)
+            .identityClientOptions(identityClientOptions)
+            .build();
     }
 
     @Override
-    public Mono<AccessToken> authenticate(TokenRequestContext request) {
-        if (this.getClientId() == null) {
-            return Mono.error(logger.logExceptionAsError(new IllegalStateException("The client id is not configured via"
-                + " 'AZURE_CLIENT_ID' environment variable or through the credential builder."
-                + " Please ensure client id is provided to authenticate via token exchange in AKS environment.")));
-        }
-        return identityClient.authenticatewithExchangeToken(request);
+    public Mono<AccessToken> getToken(TokenRequestContext request) {
+        return identityClient.authenticateWithConfidentialClientCache(request)
+            .onErrorResume(t -> Mono.empty())
+            .switchIfEmpty(Mono.defer(() -> identityClient.authenticateWithConfidentialClient(request)))
+            .doOnNext(token -> LoggingUtil.logTokenSuccess(logger, request))
+            .doOnError(error -> LoggingUtil.logTokenError(logger, request, error));
     }
 }

sdk/identity/azure-identity/src/main/java/com/azure/identity/ClientAssertionCredentialBuilder.java

@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.identity;
+
+import com.azure.core.util.logging.ClientLogger;
+import com.azure.identity.implementation.RegionalAuthority;
+import com.azure.identity.implementation.util.ValidationUtil;
+
+import java.util.HashMap;
+import java.util.function.Supplier;
+
+/**
+ * Fluent credential builder for instantiating a {@link ClientAssertionCredential}.
+ *
+ * @see ClientAssertionCredential
+ */
+public class ClientAssertionCredentialBuilder extends AadCredentialBuilderBase<ClientAssertionCredentialBuilder> {
+    private Supplier<String> clientAssertionSupplier;
+    private final ClientLogger logger = new ClientLogger(ClientAssertionCredentialBuilder.class);
+
+    /**
+     * Sets the supplier containing the logic to supply the client assertion when invoked.
+     *
+     * @param clientAssertionSupplier the supplier supplying client assertion.
+     * @return An updated instance of this builder.
+     */
+    public ClientAssertionCredentialBuilder clientAssertion(Supplier<String> clientAssertionSupplier) {
+        this.clientAssertionSupplier = clientAssertionSupplier;
+        return this;
+    }
+
+    /**
+     * Configures the persistent shared token cache options and enables the persistent token cache which is disabled
+     * by default. If configured, the credential will store tokens in a cache persisted to the machine, protected to
+     * the current user, which can be shared by other credentials and processes.
+     *
+     * @param tokenCachePersistenceOptions the token cache configuration options
+     * @return An updated instance of this builder with the token cache options configured.
+     */
+    public ClientAssertionCredentialBuilder tokenCachePersistenceOptions(TokenCachePersistenceOptions
+                                                                          tokenCachePersistenceOptions) {
+        this.identityClientOptions.setTokenCacheOptions(tokenCachePersistenceOptions);
+        return this;
+    }
+
+    /**
+     * Specifies either the specific regional authority, or use {@link RegionalAuthority#AUTO_DISCOVER_REGION} to
+     * attempt to auto-detect the region. If unset, a non-regional authority will be used. This argument should be used
+     * only by applications deployed to Azure VMs.
+     *
+     * @param regionalAuthority the regional authority
+     * @return An updated instance of this builder with the regional authority configured.
+     */
+    ClientAssertionCredentialBuilder regionalAuthority(RegionalAuthority regionalAuthority) {
+        this.identityClientOptions.setRegionalAuthority(regionalAuthority);
+        return this;
+    }
+
+    /**
+     * Creates a new {@link ClientAssertionCredential} with the current configurations.
+     *
+     * @return a {@link ClientAssertionCredential} with the current configurations.
+     * @throws IllegalArgumentException if either of clientId, tenantId or clientAssertion is not present.
+     */
+    public ClientAssertionCredential build() {
+        ValidationUtil.validate(getClass().getSimpleName(), new HashMap<String, Object>() {{
+                put("clientId", clientId);
+                put("tenantId", tenantId);
+                put("clientAssertion", clientAssertionSupplier);
+            }});
+
+        return new ClientAssertionCredential(clientId, tenantId, clientAssertionSupplier, identityClientOptions);
+    }
+}

sdk/identity/azure-identity/src/main/java/com/azure/identity/ManagedIdentityCredential.java

@@ -65,7 +65,7 @@ public final class ManagedIdentityCredential implements TokenCredential {
             clientBuilder.tenantId(configuration.get(Configuration.PROPERTY_AZURE_TENANT_ID));
             clientBuilder.clientAssertionPath(configuration.get(AZURE_FEDERATED_TOKEN_FILE));
             clientBuilder.clientAssertionTimeout(Duration.ofMinutes(5));
-            managedIdentityServiceCredential = new ClientAssertionCredential(clientIdentifier, clientBuilder.build());
+            managedIdentityServiceCredential = new AksExchangeTokenCredential(clientIdentifier, clientBuilder.build());
         } else {
             managedIdentityServiceCredential = new VirtualMachineMsiCredential(clientId, clientBuilder.build());
         }

sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClient.java

@@ -91,6 +91,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 /**
  * The identity client that contains APIs to retrieve access tokens
@@ -125,6 +126,7 @@ public class IdentityClient {
     private final String clientAssertionFilePath;
     private final InputStream certificate;
     private final String certificatePath;
+    private final Supplier<String> clientAssertionSupplier;
     private final String certificatePassword;
     private HttpPipelineAdapter httpPipelineAdapter;
     private final SynchronizedAccessor<PublicClientApplication> publicClientApplicationAccessor;
@@ -147,8 +149,8 @@ public class IdentityClient {
      * @param options the options configuring the client.
      */
     IdentityClient(String tenantId, String clientId, String clientSecret, String certificatePath,
-                   String clientAssertionFilePath, InputStream certificate, String certificatePassword,
-                   boolean isSharedTokenCacheCredential, Duration clientAssertionTimeout,
+                   String clientAssertionFilePath, Supplier<String> clientAssertionSupplier, InputStream certificate,
+                   String certificatePassword, boolean isSharedTokenCacheCredential, Duration clientAssertionTimeout,
                    IdentityClientOptions options) {
         if (tenantId == null) {
             tenantId = "organizations";
@@ -163,6 +165,7 @@ public class IdentityClient {
         this.certificatePath = certificatePath;
         this.certificate = certificate;
         this.certificatePassword = certificatePassword;
+        this.clientAssertionSupplier = clientAssertionSupplier;
         this.options = options;
 
         this.publicClientApplicationAccessor = new SynchronizedAccessor<>(() ->
@@ -215,6 +218,8 @@ private Mono<ConfidentialClientApplication> getConfidentialClientApplication() {
                     return Mono.error(logger.logExceptionAsError(new RuntimeException(
                         "Failed to parse the certificate for the credential: " + e.getMessage(), e)));
                 }
+            } else if (clientAssertionSupplier != null) {
+                credential = ClientCredentialFactory.createFromClientAssertion(clientAssertionSupplier.get());
             } else {
                 return Mono.error(logger.logExceptionAsError(
                     new IllegalArgumentException("Must provide client secret or client certificate path."

sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClientBuilder.java

@@ -7,6 +7,7 @@
 
 import java.io.InputStream;
 import java.time.Duration;
+import java.util.function.Supplier;
 
 /**
  * Fluent client builder for instantiating an {@link IdentityClient}.
@@ -24,6 +25,7 @@ public final class IdentityClientBuilder {
     private String certificatePassword;
     private boolean sharedTokenCacheCred;
     private Duration clientAssertionTimeout;
+    private Supplier<String> clientAssertionSupplier;
 
     /**
      * Sets the tenant ID for the client.
@@ -66,6 +68,17 @@ public IdentityClientBuilder certificatePath(String certificatePath) {
         return this;
     }
 
+    /**
+     * Sets the supplier for client assertion.
+     *
+     * @param clientAssertionSupplier the supplier of client assertion.
+     * @return the IdentityClientBuilder itself
+     */
+    public IdentityClientBuilder clientAssertionSupplier(Supplier<String> clientAssertionSupplier) {
+        this.clientAssertionSupplier = clientAssertionSupplier;
+        return this;
+    }
+
     /**
      * Sets the client certificate for the client.
      *
@@ -136,7 +149,8 @@ public IdentityClientBuilder clientAssertionTimeout(Duration clientAssertionTime
      * @return a {@link IdentityClient} with the current configurations.
      */
     public IdentityClient build() {
-        return new IdentityClient(tenantId, clientId, clientSecret, certificatePath, clientAssertionPath, certificate,
-            certificatePassword, sharedTokenCacheCred, clientAssertionTimeout, identityClientOptions);
+        return new IdentityClient(tenantId, clientId, clientSecret, certificatePath, clientAssertionPath,
+            clientAssertionSupplier, certificate, certificatePassword, sharedTokenCacheCred, clientAssertionTimeout,
+            identityClientOptions);
     }
 }

sdk/identity/azure-identity/src/test/java/com/azure/identity/AzureApplicationCredentialTest.java

@@ -50,7 +50,7 @@ public void testUseEnvironmentCredential() throws Exception {
             IdentityClient identityClient = PowerMockito.mock(IdentityClient.class);
             when(identityClient.authenticateWithConfidentialClientCache(any())).thenReturn(Mono.empty());
             when(identityClient.authenticateWithConfidentialClient(request1)).thenReturn(TestUtils.getMockAccessToken(token1, expiresOn));
-            PowerMockito.whenNew(IdentityClient.class).withArguments(eq(TENANT_ID), eq(CLIENT_ID), eq(secret), isNull(), isNull(), isNull(), isNull(), eq(false), isNull(), any()).thenReturn(identityClient);
+            PowerMockito.whenNew(IdentityClient.class).withArguments(eq(TENANT_ID), eq(CLIENT_ID), eq(secret), isNull(), isNull(), isNull(), isNull(), isNull(), eq(false), isNull(), any()).thenReturn(identityClient);
 
             // test
             AzureApplicationCredential credential = new AzureApplicationCredentialBuilder().build();