diff --git a/sdk/identity/azure-identity-extensions/CHANGELOG.md b/sdk/identity/azure-identity-extensions/CHANGELOG.md index 38751608824ba..45cfff7592319 100644 --- a/sdk/identity/azure-identity-extensions/CHANGELOG.md +++ b/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). diff --git a/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/credential/provider/CachingTokenCredentialProvider.java b/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/credential/provider/CachingTokenCredentialProvider.java new file mode 100644 index 0000000000000..40df900f2ce17 --- /dev/null +++ b/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 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 cache, + TokenCredentialProviderOptions options, TokenCredentialProvider delegate, + Function 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(",")); + } +} diff --git a/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/credential/provider/DefaultTokenCredentialProvider.java b/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/credential/provider/DefaultTokenCredentialProvider.java index 4bc8f29058f2d..0ebfb4c525e9b 100644 --- a/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/credential/provider/DefaultTokenCredentialProvider.java +++ b/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); } diff --git a/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/enums/AuthProperty.java b/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/enums/AuthProperty.java index 198524d909e7d..4670a043b2889 100644 --- a/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/enums/AuthProperty.java +++ b/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; diff --git a/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/template/AzureAuthenticationTemplate.java b/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/template/AzureAuthenticationTemplate.java index 4240aa50313ff..cc671496f4def 100644 --- a/sdk/identity/azure-identity-extensions/src/main/java/com/azure/identity/extensions/implementation/template/AzureAuthenticationTemplate.java +++ b/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; } - } diff --git a/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/CachingTokenCredentialProviderTest.java b/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/CachingTokenCredentialProviderTest.java new file mode 100644 index 0000000000000..d1d658f24ed99 --- /dev/null +++ b/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; + } +} diff --git a/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/SpringTokenCredentialProviderTest.java b/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/TestSpringTokenCredentialProvider.java similarity index 77% rename from sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/SpringTokenCredentialProviderTest.java rename to sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/TestSpringTokenCredentialProvider.java index 3e26afe4d844f..fd17433d2edc5 100644 --- a/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/SpringTokenCredentialProviderTest.java +++ b/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 diff --git a/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/TokenCredentialProvidersTest.java b/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/TokenCredentialProvidersTest.java index 84d160d88c9e1..79ec3f0895446 100644 --- a/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/TokenCredentialProvidersTest.java +++ b/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); } } diff --git a/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/template/AzureAuthenticationTemplateTest.java b/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/template/AzureAuthenticationTemplateTest.java index b7750f5cd4c1c..67280d1584073 100644 --- a/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/template/AzureAuthenticationTemplateTest.java +++ b/sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/template/AzureAuthenticationTemplateTest.java @@ -6,6 +6,8 @@ import com.azure.core.credential.AccessToken; import com.azure.core.credential.TokenCredential; import com.azure.core.util.Configuration; +import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions; +import com.azure.identity.extensions.implementation.credential.provider.CachingTokenCredentialProvider; import com.azure.identity.extensions.implementation.credential.provider.DefaultTokenCredentialProvider; import com.azure.identity.extensions.implementation.enums.AuthProperty; import org.junit.jupiter.api.Test; @@ -16,6 +18,7 @@ import java.time.OffsetDateTime; import java.util.Properties; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import static com.azure.identity.extensions.implementation.enums.AuthProperty.GET_TOKEN_TIMEOUT; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -78,7 +81,43 @@ void testGetTokenAsPasswordAsync() { } @Test - void testGetTokenAsPassword() throws InterruptedException { + void cacheTokenCredential() { + Properties properties = new Properties(); + AzureAuthenticationTemplate template = new AzureAuthenticationTemplate(); + template.init(properties); + + AzureAuthenticationTemplate template2 = new AzureAuthenticationTemplate(); + template2.init(properties); + + TokenCredentialProviderOptions providerOptions = new TokenCredentialProviderOptions(properties); + + TokenCredential tokenCredential = template.getTokenCredentialProvider().get(providerOptions); + TokenCredential tokenCredential2 = template2.getTokenCredentialProvider().get(providerOptions); + assertNotNull(tokenCredential); + assertNotNull(tokenCredential2); + assertTrue(tokenCredential == tokenCredential2); + } + + @Test + void nonCacheTokenCredential() { + Properties properties = new Properties(); + properties.setProperty("azure.tokenCredentialCacheEnabled", "false"); + AzureAuthenticationTemplate template = new AzureAuthenticationTemplate(); + template.init(properties); + + AzureAuthenticationTemplate template2 = new AzureAuthenticationTemplate(); + template2.init(properties); + + TokenCredentialProviderOptions providerOptions = new TokenCredentialProviderOptions(properties); + + assertNotNull(template.getTokenCredentialProvider().get(providerOptions)); + assertNotNull(template2.getTokenCredentialProvider().get(providerOptions)); + assertNotEquals(template.getTokenCredentialProvider().get(providerOptions), + template2.getTokenCredentialProvider().get(providerOptions)); + } + + @Test + void getTokenAsPasswordWithDefaultCredentialProvider() throws InterruptedException { // setup String token1 = "token1"; String token2 = "token2"; @@ -93,11 +132,12 @@ void testGetTokenAsPassword() throws InterruptedException { } }); // mock - try (MockedConstruction identityClientMock + try (MockedConstruction defaultCredentialProviderMock = mockConstruction(DefaultTokenCredentialProvider.class, (defaultTokenCredentialProvider, context) -> { when(defaultTokenCredentialProvider.get()).thenReturn(mockTokenCredential); })) { Properties properties = new Properties(); + properties.setProperty("azure.tokenCredentialCacheEnabled", "false"); AzureAuthenticationTemplate template = new AzureAuthenticationTemplate(); template.init(properties); @@ -107,7 +147,48 @@ void testGetTokenAsPassword() throws InterruptedException { TimeUnit.SECONDS.sleep(tokenExpireSeconds); assertEquals(token2, template.getTokenAsPassword()); - assertNotNull(identityClientMock); + assertNotNull(defaultCredentialProviderMock); + } + } + + @Test + void getTokenAsPasswordWithCachingCredentialProvider() throws InterruptedException { + int tokenExpireSeconds = 2; + AtomicInteger tokenIndex1 = new AtomicInteger(); + AtomicInteger tokenIndex2 = new AtomicInteger(1); + TokenCredential mockTokenCredential = mock(TokenCredential.class); + OffsetDateTime offsetDateTime = OffsetDateTime.now().plusSeconds(tokenExpireSeconds); + when(mockTokenCredential.getToken(any())).thenAnswer(u -> { + if (OffsetDateTime.now().isBefore(offsetDateTime)) { + return Mono.just(new AccessToken("token1-" + (tokenIndex1.getAndIncrement()), offsetDateTime)); + } else { + return Mono.just(new AccessToken("token2-" + (tokenIndex2.getAndIncrement()), + offsetDateTime.plusSeconds(tokenExpireSeconds))); + } + }); + // mock + try (MockedConstruction credentialProviderMock + = mockConstruction(CachingTokenCredentialProvider.class, (defaultTokenCredentialProvider, context) -> { + when(defaultTokenCredentialProvider.get()).thenReturn(mockTokenCredential); + })) { + Properties properties = new Properties(); + + AzureAuthenticationTemplate template = new AzureAuthenticationTemplate(); + template.init(properties); + AzureAuthenticationTemplate template2 = new AzureAuthenticationTemplate(); + template2.init(properties); + + verifyToken("token1-", 0, template); + TimeUnit.SECONDS.sleep(tokenExpireSeconds + 1); + verifyToken("token2-", 1, template); + assertNotNull(credentialProviderMock); + } + } + + private static void verifyToken(String tokenPrefix, int tokenInitialIndexValue, + AzureAuthenticationTemplate template) { + for (int i = 0; i < 5; i++) { + assertEquals(tokenPrefix + (tokenInitialIndexValue + i), template.getTokenAsPassword()); } }