diff --git a/sdk/identity/azure-identity/src/main/java/com/azure/identity/DeviceCodeCredentialBuilder.java b/sdk/identity/azure-identity/src/main/java/com/azure/identity/DeviceCodeCredentialBuilder.java index 40a6fcc38842..6dc62951b8a5 100644 --- a/sdk/identity/azure-identity/src/main/java/com/azure/identity/DeviceCodeCredentialBuilder.java +++ b/sdk/identity/azure-identity/src/main/java/com/azure/identity/DeviceCodeCredentialBuilder.java @@ -56,9 +56,9 @@ public DeviceCodeCredentialBuilder enablePersistentCache() { /** * Sets the {@link AuthenticationRecord} captured from a previous authentication. * - * @param authenticationRecord the authentication record to ser. + * @param authenticationRecord the authentication record to be configured. * - * @return An updated instance of this builder with if the shared token cache enabled specified. + * @return An updated instance of this builder with the configured authentication record. */ public DeviceCodeCredentialBuilder authenticationRecord(AuthenticationRecord authenticationRecord) { this.identityClientOptions.setAuthenticationRecord(authenticationRecord); diff --git a/sdk/identity/azure-identity/src/main/java/com/azure/identity/SharedTokenCacheCredential.java b/sdk/identity/azure-identity/src/main/java/com/azure/identity/SharedTokenCacheCredential.java index 41bc18e14360..12a303b156a5 100644 --- a/sdk/identity/azure-identity/src/main/java/com/azure/identity/SharedTokenCacheCredential.java +++ b/sdk/identity/azure-identity/src/main/java/com/azure/identity/SharedTokenCacheCredential.java @@ -11,6 +11,7 @@ import com.azure.identity.implementation.IdentityClient; import com.azure.identity.implementation.IdentityClientBuilder; import com.azure.identity.implementation.IdentityClientOptions; +import com.azure.identity.implementation.MsalAuthenticationAccount; import com.azure.identity.implementation.MsalToken; import com.azure.identity.implementation.util.LoggingUtil; import reactor.core.publisher.Mono; @@ -26,7 +27,8 @@ public class SharedTokenCacheCredential implements TokenCredential { private final String username; private final String clientId; private final String tenantId; - private final AtomicReference cachedToken; + private final AtomicReference cachedToken; + private final IdentityClient identityClient; private final ClientLogger logger = new ClientLogger(SharedTokenCacheCredential.class); @@ -65,6 +67,9 @@ public class SharedTokenCacheCredential implements TokenCredential { .identityClientOptions(identityClientOptions) .build(); this.cachedToken = new AtomicReference<>(); + if (identityClientOptions.getAuthenticationRecord() != null) { + cachedToken.set(new MsalAuthenticationAccount(identityClientOptions.getAuthenticationRecord())); + } LoggingUtil.logAvailableEnvironmentVariables(logger, configuration); } @@ -75,18 +80,23 @@ public class SharedTokenCacheCredential implements TokenCredential { public Mono getToken(TokenRequestContext request) { return Mono.defer(() -> { if (cachedToken.get() != null) { - return identityClient.authenticateWithPublicClientCache(request, cachedToken.get().getAccount()) + return identityClient.authenticateWithPublicClientCache(request, cachedToken.get()) .onErrorResume(t -> Mono.empty()); } else { return Mono.empty(); } }).switchIfEmpty( Mono.defer(() -> identityClient.authenticateWithSharedTokenCache(request, username))) - .map(msalToken -> { - cachedToken.set(msalToken); - return (AccessToken) msalToken; - }) + .map(this::updateCache) .doOnNext(token -> LoggingUtil.logTokenSuccess(logger, request)) .doOnError(error -> LoggingUtil.logTokenError(logger, request, error)); } + + private AccessToken updateCache(MsalToken msalToken) { + cachedToken.set( + new MsalAuthenticationAccount( + new AuthenticationRecord(msalToken.getAuthenticationResult(), + identityClient.getTenantId(), identityClient.getClientId()))); + return msalToken; + } } diff --git a/sdk/identity/azure-identity/src/main/java/com/azure/identity/SharedTokenCacheCredentialBuilder.java b/sdk/identity/azure-identity/src/main/java/com/azure/identity/SharedTokenCacheCredentialBuilder.java index 69ffa263a444..5c2f372b327a 100644 --- a/sdk/identity/azure-identity/src/main/java/com/azure/identity/SharedTokenCacheCredentialBuilder.java +++ b/sdk/identity/azure-identity/src/main/java/com/azure/identity/SharedTokenCacheCredentialBuilder.java @@ -34,6 +34,18 @@ public SharedTokenCacheCredentialBuilder allowUnencryptedCache() { return this; } + /** + * Sets the {@link AuthenticationRecord} captured from a previous authentication. + * + * @param authenticationRecord the authentication record to be configured. + * + * @return An updated instance of this builder with the configured authentication record. + */ + public SharedTokenCacheCredentialBuilder authenticationRecord(AuthenticationRecord authenticationRecord) { + this.identityClientOptions.setAuthenticationRecord(authenticationRecord); + return this; + } + /** * Creates a new {@link SharedTokenCacheCredentialBuilder} with the current configurations. *