Support token credential cache for azure-identity-extension

sdk/identity/azure-identity-extensions/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## 1.2.0-beta.2 (Unreleased)
 
+#### Features Added
+- Support cache for token credential object. [#39393](https://github.com/Azure/azure-sdk-for-java/issues/39393).
+
 #### Bugs Fixed
 - Fix the issue where the token acquisition timeout is not set via the property `azure.accessTokenTimeoutInSeconds`. [#43512](https://github.com/Azure/azure-sdk-for-java/issues/43512).
 

sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/credential/provider/CachingTokenCredentialProvider.java

@@ -0,0 +1,84 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.identity.extensions.implementation.credential.provider;
+
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.util.logging.ClientLogger;
+import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions;
+
+import java.util.Arrays;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * Caching tokenCredentialProvider implementation that provides tokenCredential instance.
+ */
+public class CachingTokenCredentialProvider implements TokenCredentialProvider {
+
+    private static final ClientLogger LOGGER = new ClientLogger(CachingTokenCredentialProvider.class);
+
+    private static final ConcurrentHashMap<String, TokenCredential> CACHE = new ConcurrentHashMap<>();
+
+    private final TokenCredentialProviderOptions defaultOptions;
+
+    private final TokenCredentialProvider delegate;
+
+    /**
+     * CachingTokenCredentialProvider constructor.
+     * @param defaultOptions the {@link TokenCredentialProviderOptions} for the delegate {@link TokenCredentialProvider} initialization.
+     * @param tokenCredentialProvider the delegate {@link TokenCredentialProvider}.
+     */
+    public CachingTokenCredentialProvider(TokenCredentialProviderOptions defaultOptions,
+        TokenCredentialProvider tokenCredentialProvider) {
+        this.defaultOptions = defaultOptions;
+        this.delegate = tokenCredentialProvider;
+    }
+
+    @Override
+    public TokenCredential get() {
+        return getOrCreate(CACHE, this.defaultOptions, this.delegate,
+            tokenCredentialProvider -> tokenCredentialProvider.get());
+    }
+
+    @Override
+    public TokenCredential get(TokenCredentialProviderOptions options) {
+        return getOrCreate(CACHE, options, this.delegate,
+            tokenCredentialProvider -> tokenCredentialProvider.get(options));
+    }
+
+    private static TokenCredential getOrCreate(Map<String, TokenCredential> cache,
+        TokenCredentialProviderOptions options, TokenCredentialProvider delegate,
+        Function<TokenCredentialProvider, TokenCredential> fn) {
+        String tokenCredentialCacheKey = convertToTokenCredentialCacheKey(options);
+
+        if (cache.containsKey(tokenCredentialCacheKey)) {
+            LOGGER.verbose("Retrieving token credential from cache.");
+        } else {
+            LOGGER.verbose("Caching token credential.");
+            cache.put(tokenCredentialCacheKey, fn.apply(delegate));
+        }
+
+        return cache.get(tokenCredentialCacheKey);
+    }
+
+    private static String convertToTokenCredentialCacheKey(TokenCredentialProviderOptions options) {
+        if (options == null) {
+            return CachingTokenCredentialProvider.class.getSimpleName();
+        }
+
+        return Arrays
+            .stream(new String[] {
+                options.getTenantId(),
+                options.getClientId(),
+                options.getClientCertificatePath(),
+                options.getUsername(),
+                String.valueOf(options.isManagedIdentityEnabled()),
+                options.getTokenCredentialProviderClassName(),
+                options.getTokenCredentialBeanName() })
+            .map(option -> option == null ? "" : option)
+            .collect(Collectors.joining(","));
+    }
+}

sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/credential/provider/DefaultTokenCredentialProvider.java

@@ -67,7 +67,8 @@ private TokenCredential resolveTokenCredential(TokenCredentialProviderOptions op
                         .clientId(clientId);
 
                 if (hasText(options.getClientCertificatePassword())) {
-                    builder.pfxCertificate(clientCertificatePath, options.getClientCertificatePassword());
+                    builder.pfxCertificate(clientCertificatePath)
+                        .clientCertificatePassword(options.getClientCertificatePassword());
                 } else {
                     builder.pemCertificate(clientCertificatePath);
                 }

sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/enums/AuthProperty.java

@@ -77,7 +77,12 @@ public enum AuthProperty {
      * The given bean name of a TokenCredential bean in the Spring context.
      */
     TOKEN_CREDENTIAL_BEAN_NAME("azure.tokenCredentialBeanName", "springCloudAzureDefaultCredential",
-        "The given bean name of a TokenCredential bean in the Spring context.", false);
+        "The given bean name of a TokenCredential bean in the Spring context.", false),
+    /**
+     *  Whether to enable token credential cache.
+     */
+    TOKEN_CREDENTIAL_CACHE_ENABLED("azure.tokenCredentialCacheEnabled", "true",
+        "Whether to enable the token credential cache.", false);
 
     String propertyKey;
     String defaultValue;

sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/template/AzureAuthenticationTemplate.java

@@ -5,13 +5,16 @@
 
 import com.azure.core.credential.AccessToken;
 import com.azure.core.util.logging.ClientLogger;
+import com.azure.identity.extensions.implementation.credential.provider.CachingTokenCredentialProvider;
 import com.azure.identity.extensions.implementation.credential.provider.TokenCredentialProvider;
 import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions;
+import com.azure.identity.extensions.implementation.enums.AuthProperty;
 import com.azure.identity.extensions.implementation.token.AccessTokenResolver;
 import com.azure.identity.extensions.implementation.token.AccessTokenResolverOptions;
 import java.time.Duration;
 import java.util.Properties;
 import java.util.concurrent.atomic.AtomicBoolean;
+
 import reactor.core.publisher.Mono;
 import static com.azure.identity.extensions.implementation.enums.AuthProperty.GET_TOKEN_TIMEOUT;
 
@@ -41,7 +44,7 @@ public AzureAuthenticationTemplate() {
     /**
      * AzureAuthenticationTemplate constructor.
      *
-     * @param tokenCredentialProvider An TokenCredentialProvider class instance.
+     * @param tokenCredentialProvider A TokenCredentialProvider class instance.
      * @param accessTokenResolver An AccessTokenResolver class instance.
      */
     public AzureAuthenticationTemplate(TokenCredentialProvider tokenCredentialProvider,
@@ -60,8 +63,13 @@ public void init(Properties properties) {
             LOGGER.verbose("Initializing AzureAuthenticationTemplate.");
 
             if (getTokenCredentialProvider() == null) {
-                this.tokenCredentialProvider
-                    = TokenCredentialProvider.createDefault(new TokenCredentialProviderOptions(properties));
+                TokenCredentialProviderOptions options = new TokenCredentialProviderOptions(properties);
+                this.tokenCredentialProvider = TokenCredentialProvider.createDefault(options);
+
+                if (Boolean.TRUE.equals(AuthProperty.TOKEN_CREDENTIAL_CACHE_ENABLED.getBoolean(properties))) {
+                    this.tokenCredentialProvider
+                        = new CachingTokenCredentialProvider(options, this.tokenCredentialProvider);
+                }
             }
 
             if (getAccessTokenResolver() == null) {
@@ -125,5 +133,4 @@ Duration getBlockTimeout() {
     AtomicBoolean getIsInitialized() {
         return isInitialized;
     }
-
 }

sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/CachingTokenCredentialProviderTest.java

@@ -0,0 +1,112 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.identity.extensions.implementation.credential.provider;
+
+import com.azure.core.credential.TokenCredential;
+import com.azure.identity.DefaultAzureCredential;
+import com.azure.identity.ManagedIdentityCredential;
+import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class CachingTokenCredentialProviderTest {
+
+    @Test
+    void returnCacheUsingDefaultAuthMethodViaDifferentProviderInstances() {
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider1 = new DefaultTokenCredentialProvider(null);
+        CachingTokenCredentialProvider provider1
+            = new CachingTokenCredentialProvider(null, defaultTokenCredentialProvider1);
+        TokenCredential tokenCredential1 = provider1.get();
+
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider2 = new DefaultTokenCredentialProvider(null);
+        CachingTokenCredentialProvider provider2
+            = new CachingTokenCredentialProvider(null, defaultTokenCredentialProvider2);
+
+        TokenCredential tokenCredential2 = provider2.get();
+        assertTrue(tokenCredential1 instanceof DefaultAzureCredential);
+        assertTrue(tokenCredential2 instanceof DefaultAzureCredential);
+        assertTrue(tokenCredential1 == tokenCredential2);
+    }
+
+    @Test
+    void returnCacheUsingSameAuthMethodViaDifferentProviderInstances() {
+        TokenCredentialProviderOptions customOptions = getSystemManagedIdentityCredentialProviderOptions();
+
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider1
+            = new DefaultTokenCredentialProvider(customOptions);
+        CachingTokenCredentialProvider cachingProvider1
+            = new CachingTokenCredentialProvider(customOptions, defaultTokenCredentialProvider1);
+        TokenCredential tokenCredential1 = cachingProvider1.get();
+
+        TokenCredentialProviderOptions customOptions2 = getSystemManagedIdentityCredentialProviderOptions();
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider2
+            = new DefaultTokenCredentialProvider(customOptions2);
+        CachingTokenCredentialProvider cachingProvider2
+            = new CachingTokenCredentialProvider(customOptions2, defaultTokenCredentialProvider2);
+
+        TokenCredential tokenCredential2 = cachingProvider2.get();
+        assertTrue(tokenCredential1 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential2 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential1 == tokenCredential2);
+    }
+
+    @Test
+    void returnCacheUsingSameAuthMethodAndInvokingDifferentGetMethods() {
+        TokenCredentialProviderOptions customOptions = getSystemManagedIdentityCredentialProviderOptions();
+
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider1
+            = new DefaultTokenCredentialProvider(customOptions);
+        CachingTokenCredentialProvider cachingProvider1
+            = new CachingTokenCredentialProvider(customOptions, defaultTokenCredentialProvider1);
+        TokenCredential tokenCredential1 = cachingProvider1.get();
+
+        TokenCredentialProviderOptions customOptions2 = getSystemManagedIdentityCredentialProviderOptions();
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider2
+            = new DefaultTokenCredentialProvider(customOptions2);
+        CachingTokenCredentialProvider cachingProvider2
+            = new CachingTokenCredentialProvider(customOptions2, defaultTokenCredentialProvider2);
+
+        TokenCredential tokenCredential2 = cachingProvider2.get(customOptions2);
+        assertTrue(tokenCredential1 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential2 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential1 == tokenCredential2);
+    }
+
+    @Test
+    void returnDifferentCachesUsingDifferentAuthenticationMethods() {
+        TokenCredentialProviderOptions customOptions = getSystemManagedIdentityCredentialProviderOptions();
+
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider1
+            = new DefaultTokenCredentialProvider(customOptions);
+        CachingTokenCredentialProvider cachingProvider1
+            = new CachingTokenCredentialProvider(customOptions, defaultTokenCredentialProvider1);
+        TokenCredential tokenCredential1 = cachingProvider1.get();
+
+        TokenCredentialProviderOptions customOptions2
+            = getUserManagedIdentityCredentialProviderOptions("test-client-id");
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider2
+            = new DefaultTokenCredentialProvider(customOptions2);
+        CachingTokenCredentialProvider cachingProvider2
+            = new CachingTokenCredentialProvider(customOptions2, defaultTokenCredentialProvider2);
+
+        TokenCredential tokenCredential2 = cachingProvider2.get(customOptions2);
+        assertTrue(tokenCredential1 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential2 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential1 != tokenCredential2);
+    }
+
+    private static TokenCredentialProviderOptions getUserManagedIdentityCredentialProviderOptions(String clientId) {
+        TokenCredentialProviderOptions customOptions = new TokenCredentialProviderOptions();
+        customOptions.setManagedIdentityEnabled(true);
+        customOptions.setClientId(clientId);
+        return customOptions;
+    }
+
+    private static TokenCredentialProviderOptions getSystemManagedIdentityCredentialProviderOptions() {
+        TokenCredentialProviderOptions customOptions = new TokenCredentialProviderOptions();
+        customOptions.setManagedIdentityEnabled(true);
+        return customOptions;
+    }
+}

sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/SpringTokenCredentialProviderTest.java → sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/TestSpringTokenCredentialProvider.java

@@ -6,9 +6,9 @@
 import com.azure.core.credential.TokenCredential;
 import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions;
 
-class SpringTokenCredentialProviderTest implements TokenCredentialProvider {
+class TestSpringTokenCredentialProvider implements TokenCredentialProvider {
 
-    SpringTokenCredentialProviderTest(TokenCredentialProviderOptions options) {
+    TestSpringTokenCredentialProvider(TokenCredentialProviderOptions options) {
     }
 
     @Override

sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/TokenCredentialProvidersTest.java

@@ -10,7 +10,7 @@
 class TokenCredentialProvidersTest {
 
     private static final String SPRING_TOKEN_CREDENTIAL_PROVIDER_CLASS_NAME
-        = SpringTokenCredentialProviderTest.class.getName();
+        = TestSpringTokenCredentialProvider.class.getName();
 
     @Test
     void testOptionsIsNull() {
@@ -29,7 +29,7 @@ void testCreateSpringTokenCredentialProvider() {
         TokenCredentialProviderOptions option = new TokenCredentialProviderOptions();
         option.setTokenCredentialProviderClassName(SPRING_TOKEN_CREDENTIAL_PROVIDER_CLASS_NAME);
         TokenCredentialProvider credentialProvider = TokenCredentialProviders.createInstance(option);
-        Assertions.assertTrue(credentialProvider instanceof SpringTokenCredentialProviderTest);
+        Assertions.assertTrue(credentialProvider instanceof TestSpringTokenCredentialProvider);
     }
 
 }