add sync flow to workload identity

eng/code-quality-reports/src/main/resources/spotbugs/spotbugs-exclude.xml

@@ -901,6 +901,13 @@
     <Bug pattern="RCN_REDUNDANT_NULLCHECK_OF_NONNULL_VALUE,REC_CATCH_EXCEPTION"/>
   </Match>
 
+  <!-- Exception needs to be caught here.-->
+  <Match>
+    <Class name="~com\.azure\.identity\.implementation\.IdentitySyncClient"/>
+    <Method name="authenticateWithWorkloadIdentityConfidentialClient"/>
+    <Bug pattern="REC_CATCH_EXCEPTION"/>
+  </Match>
+
   <!-- ClassLoader::getResourceAsStream can return null, ref: https://docs.oracle.com/javase/8/docs/technotes/guides/lang/resources.html -->
   <Match>
     <Class name="com.azure.cosmos.implementation.directconnectivity.RntbdTransportClient$Options$Builder"/>

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

@@ -12,6 +12,7 @@
 import com.azure.identity.implementation.IdentityClient;
 import com.azure.identity.implementation.IdentityClientBuilder;
 import com.azure.identity.implementation.IdentityClientOptions;
+import com.azure.identity.implementation.IdentitySyncClient;
 import com.azure.identity.implementation.util.ValidationUtil;
 import reactor.core.publisher.Mono;
 
@@ -54,6 +55,7 @@
 public class WorkloadIdentityCredential implements TokenCredential {
     private static final ClientLogger LOGGER = new ClientLogger(WorkloadIdentityCredential.class);
     private final IdentityClient identityClient;
+    private final IdentitySyncClient identitySyncClient;
 
     /**
      * WorkloadIdentityCredential supports Azure workload identity on Kubernetes.
@@ -82,14 +84,16 @@ public class WorkloadIdentityCredential implements TokenCredential {
             || CoreUtils.isNullOrEmpty(federatedTokenFilePathInput)
             || CoreUtils.isNullOrEmpty(clientIdInput)
             || CoreUtils.isNullOrEmpty(identityClientOptions.getAuthorityHost()))) {
-            identityClient = new IdentityClientBuilder()
+            IdentityClientBuilder builder = new IdentityClientBuilder()
                 .clientAssertionPath(federatedTokenFilePathInput)
                 .clientId(clientId)
                 .tenantId(tenantIdInput)
-                .identityClientOptions(identityClientOptions)
-                .build();
+                .identityClientOptions(identityClientOptions);
+            identityClient = builder.build();
+            identitySyncClient = builder.buildSyncClient();
         } else {
             identityClient = null;
+            identitySyncClient = null;
         }
     }
 
@@ -101,7 +105,18 @@ public Mono<AccessToken> getToken(TokenRequestContext request) {
                 + " guide for more information."
                 + " https://aka.ms/azsdk/java/identity/workloadidentitycredential/troubleshoot")));
         }
-        return identityClient.authenticateWithExchangeToken(request);
+        return identityClient.authenticateWithWorkloadIdentityConfidentialClient(request);
+    }
+
+    @Override
+    public AccessToken getTokenSync(TokenRequestContext request) {
+        if (identitySyncClient == null) {
+            throw LOGGER.logExceptionAsError(new CredentialUnavailableException("WorkloadIdentityCredential"
+                + " authentication unavailable. The workload options are not fully configured. See the troubleshooting"
+                + " guide for more information."
+                + " https://aka.ms/azsdk/java/identity/workloadidentitycredential/troubleshoot"));
+        }
+        return identitySyncClient.authenticateWithWorkloadIdentityConfidentialClient(request);
     }
 
     String getClientId() {

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

@@ -18,6 +18,7 @@
 import com.azure.identity.implementation.util.ScopeUtil;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.microsoft.aad.msal4j.AuthorizationCodeParameters;
+import com.microsoft.aad.msal4j.AppTokenProviderParameters;
 import com.microsoft.aad.msal4j.ClaimsRequest;
 import com.microsoft.aad.msal4j.ClientCredentialFactory;
 import com.microsoft.aad.msal4j.ClientCredentialParameters;
@@ -30,13 +31,13 @@
 import com.microsoft.aad.msal4j.PublicClientApplication;
 import com.microsoft.aad.msal4j.RefreshTokenParameters;
 import com.microsoft.aad.msal4j.SilentParameters;
+import com.microsoft.aad.msal4j.TokenProviderResult;
 import com.microsoft.aad.msal4j.UserNamePasswordParameters;
 import com.sun.jna.Platform;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import javax.net.ssl.HttpsURLConnection;
-import java.io.DataOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.HttpURLConnection;
@@ -62,6 +63,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 /**
@@ -72,6 +74,7 @@ public class IdentityClient extends IdentityClientBase {
     private final SynchronizedAccessor<PublicClientApplication> publicClientApplicationAccessor;
     private final SynchronizedAccessor<ConfidentialClientApplication> confidentialClientApplicationAccessor;
     private final SynchronizedAccessor<ConfidentialClientApplication> managedIdentityConfidentialClientApplicationAccessor;
+    private final SynchronizedAccessor<ConfidentialClientApplication> workloadIdentityConfidentialClientApplicationAccessor;
     private final SynchronizedAccessor<String> clientAssertionAccessor;
 
 
@@ -106,6 +109,9 @@ public class IdentityClient extends IdentityClientBase {
         this.managedIdentityConfidentialClientApplicationAccessor =
             new SynchronizedAccessor<>(this::getManagedIdentityConfidentialClientApplication);
 
+        this.workloadIdentityConfidentialClientApplicationAccessor =
+            new SynchronizedAccessor<>(this::getWorkloadIdentityConfidentialClientApplication);
+
         Duration cacheTimeout = (clientAssertionTimeout == null) ? Duration.ofMinutes(5) : clientAssertionTimeout;
         this.clientAssertionAccessor = new SynchronizedAccessor<>(this::parseClientAssertion, cacheTimeout);
     }
@@ -130,6 +136,16 @@ private Mono<ConfidentialClientApplication> getManagedIdentityConfidentialClient
         });
     }
 
+    private Mono<ConfidentialClientApplication> getWorkloadIdentityConfidentialClientApplication() {
+        return Mono.defer(() -> {
+            try {
+                return Mono.just(super.getWorkloadIdentityConfidentialClient());
+            } catch (RuntimeException e) {
+                return Mono.error(e);
+            }
+        });
+    }
+
     @Override
     Mono<AccessToken> getTokenFromTargetManagedIdentity(TokenRequestContext tokenRequestContext) {
         ManagedIdentityParameters parameters = options.getManagedIdentityParameters();
@@ -492,10 +508,6 @@ public Mono<AccessToken> authenticateWithConfidentialClient(TokenRequestContext
                     builder.clientCredential(ClientCredentialFactory
                         .createFromClientAssertion(clientAssertionSupplier.get()));
                 }
-                if (request.getClaims() != null) {
-                    ClaimsRequest customClaimRequest = CustomClaimRequest.formatAsClaimsRequest(request.getClaims());
-                    builder.claims(customClaimRequest);
-                }
                 return confidentialClient.acquireToken(builder.build());
             }
         )).map(MsalToken::new);
@@ -514,6 +526,19 @@ public Mono<AccessToken> authenticateWithManagedIdentityConfidentialClient(Token
             .map(MsalToken::new);
     }
 
+    public Mono<AccessToken> authenticateWithWorkloadIdentityConfidentialClient(TokenRequestContext request) {
+        return workloadIdentityConfidentialClientApplicationAccessor.getValue()
+            .flatMap(confidentialClient -> Mono.fromFuture(() -> {
+                    ClientCredentialParameters.ClientCredentialParametersBuilder builder =
+                        ClientCredentialParameters.builder(new HashSet<>(request.getScopes()))
+                            .tenant(IdentityUtil
+                                .resolveTenantId(tenantId, request, options));
+                    return confidentialClient.acquireToken(builder.build());
+                }
+            )).onErrorMap(t -> new CredentialUnavailableException("Managed Identity authentication is not available.", t))
+            .map(MsalToken::new);
+    }
+
     /**
      * Asynchronously acquire a token from Active Directory with a username and a password.
      *
@@ -602,11 +627,6 @@ public Mono<AccessToken> authenticateWithConfidentialClientCache(TokenRequestCon
                 SilentParameters.SilentParametersBuilder parametersBuilder = SilentParameters.builder(
                         new HashSet<>(request.getScopes()))
                     .tenant(IdentityUtil.resolveTenantId(tenantId, request, options));
-                if (request.getClaims() != null) {
-                    ClaimsRequest customClaimRequest = CustomClaimRequest
-                        .formatAsClaimsRequest(request.getClaims());
-                    parametersBuilder.claims(customClaimRequest);
-                }
                 try {
                     return confidentialClient.acquireTokenSilently(parametersBuilder.build());
                 } catch (MalformedURLException e) {
@@ -902,41 +922,7 @@ private Mono<AccessToken> authenticateToArcManagedIdentityEndpoint(String identi
     public Mono<AccessToken> authenticateWithExchangeToken(TokenRequestContext request) {
 
         return clientAssertionAccessor.getValue()
-            .flatMap(assertionToken -> Mono.fromCallable(() -> {
-                String authorityUrl = TRAILING_FORWARD_SLASHES.matcher(options.getAuthorityHost()).replaceAll("")
-                    + "/" + tenantId + "/oauth2/v2.0/token";
-
-                String urlParams = "client_assertion=" + assertionToken
-                    + "&client_assertion_type=urn:ietf:params:oauth:client-assertion-type:jwt-bearer&client_id="
-                    + clientId + "&grant_type=client_credentials&scope=" + urlEncode(request.getScopes().get(0));
-
-                byte[] postData = urlParams.getBytes(StandardCharsets.UTF_8);
-                int postDataLength = postData.length;
-
-                HttpURLConnection connection = null;
-
-                URL url = getUrl(authorityUrl);
-
-                try {
-                    connection = (HttpURLConnection) url.openConnection();
-                    connection.setRequestMethod("POST");
-                    connection.setRequestProperty("Content-Type", "application/x-www-form-urlencoded");
-                    connection.setRequestProperty("Content-Length", Integer.toString(postDataLength));
-                    connection.setRequestProperty("User-Agent", userAgent);
-                    connection.setDoOutput(true);
-                    try (DataOutputStream outputStream = new DataOutputStream(connection.getOutputStream())) {
-                        outputStream.write(postData);
-                    }
-                    connection.connect();
-
-                    return SERIALIZER_ADAPTER.deserialize(connection.getInputStream(), MSIToken.class,
-                        SerializerEncoding.JSON);
-                } finally {
-                    if (connection != null) {
-                        connection.disconnect();
-                    }
-                }
-            }));
+            .flatMap(assertionToken -> Mono.fromCallable(() -> authenticateWithExchangeTokenHelper(request, assertionToken)));
     }
 
     /**
@@ -1088,9 +1074,6 @@ public Mono<AccessToken> authenticateToManagedIdentityEndpoint(String identityEn
         });
     }
 
-    static URL getUrl(String uri) throws MalformedURLException {
-        return new URL(uri);
-    }
     /**
      * Asynchronously acquire a token from the Virtual Machine IMDS endpoint.
      *
@@ -1237,8 +1220,6 @@ private static Proxy proxyOptionsToJavaNetProxy(ProxyOptions options) {
         }
     }
 
-
-
     void openUrl(String url) throws IOException {
         Runtime rt = Runtime.getRuntime();
 
@@ -1273,7 +1254,22 @@ private boolean isADFSTenant() {
         return ADFS_TENANT.equals(this.tenantId);
     }
 
-    private static String urlEncode(String value) throws IOException {
-        return URLEncoder.encode(value, StandardCharsets.UTF_8.name());
+    Function<AppTokenProviderParameters, CompletableFuture<TokenProviderResult>> getWorkloadIdentityTokenProvider() {
+        return appTokenProviderParameters -> {
+            TokenRequestContext trc = new TokenRequestContext()
+                .setScopes(new ArrayList<>(appTokenProviderParameters.scopes))
+                .setClaims(appTokenProviderParameters.claims)
+                .setTenantId(appTokenProviderParameters.tenantId);
+
+            Mono<AccessToken> accessTokenAsync = authenticateWithExchangeToken(trc);
+
+            return accessTokenAsync.map(accessToken -> {
+                TokenProviderResult result = new TokenProviderResult();
+                result.setAccessToken(accessToken.getToken());
+                result.setTenantId(trc.getTenantId());
+                result.setExpiresInSeconds(accessToken.getExpiresAt().toEpochSecond());
+                return result;
+            }).toFuture();
+        };
     }
 }