Skip to content

Commit 2cbcdf6

Browse files
moarychanmoarychan
authored
Support token credential cache for azure-identity-extension (#43659)
1 parent f1f5e29 commit 2cbcdf6

File tree

9 files changed

+306
-13
lines changed

9 files changed

+306
-13
lines changed

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22

33
## 1.2.0-beta.2 (Unreleased)
44

5+
#### Features Added
6+
- Support cache for token credential object. [#39393](https://github.com/Azure/azure-sdk-for-java/issues/39393).
7+
58
#### Bugs Fixed
69
- 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).
710

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

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.azure.identity.extensions.implementation.credential.provider;
5+
6+
import com.azure.core.credential.TokenCredential;
7+
import com.azure.core.util.logging.ClientLogger;
8+
import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions;
9+
10+
import java.util.Arrays;
11+
import java.util.Map;
12+
import java.util.concurrent.ConcurrentHashMap;
13+
import java.util.function.Function;
14+
import java.util.stream.Collectors;
15+
16+
/**
17+
* Caching tokenCredentialProvider implementation that provides tokenCredential instance.
18+
*/
19+
public class CachingTokenCredentialProvider implements TokenCredentialProvider {
20+
21+
private static final ClientLogger LOGGER = new ClientLogger(CachingTokenCredentialProvider.class);
22+
23+
private static final ConcurrentHashMap<String, TokenCredential> CACHE = new ConcurrentHashMap<>();
24+
25+
private final TokenCredentialProviderOptions defaultOptions;
26+
27+
private final TokenCredentialProvider delegate;
28+
29+
/**
30+
* CachingTokenCredentialProvider constructor.
31+
* @param defaultOptions the {@link TokenCredentialProviderOptions} for the delegate {@link TokenCredentialProvider} initialization.
32+
* @param tokenCredentialProvider the delegate {@link TokenCredentialProvider}.
33+
*/
34+
public CachingTokenCredentialProvider(TokenCredentialProviderOptions defaultOptions,
35+
TokenCredentialProvider tokenCredentialProvider) {
36+
this.defaultOptions = defaultOptions;
37+
this.delegate = tokenCredentialProvider;
38+
}
39+
40+
@Override
41+
public TokenCredential get() {
42+
return getOrCreate(CACHE, this.defaultOptions, this.delegate,
43+
tokenCredentialProvider -> tokenCredentialProvider.get());
44+
}
45+
46+
@Override
47+
public TokenCredential get(TokenCredentialProviderOptions options) {
48+
return getOrCreate(CACHE, options, this.delegate,
49+
tokenCredentialProvider -> tokenCredentialProvider.get(options));
50+
}
51+
52+
private static TokenCredential getOrCreate(Map<String, TokenCredential> cache,
53+
TokenCredentialProviderOptions options, TokenCredentialProvider delegate,
54+
Function<TokenCredentialProvider, TokenCredential> fn) {
55+
String tokenCredentialCacheKey = convertToTokenCredentialCacheKey(options);
56+
57+
if (cache.containsKey(tokenCredentialCacheKey)) {
58+
LOGGER.verbose("Retrieving token credential from cache.");
59+
} else {
60+
LOGGER.verbose("Caching token credential.");
61+
cache.put(tokenCredentialCacheKey, fn.apply(delegate));
62+
}
63+
64+
return cache.get(tokenCredentialCacheKey);
65+
}
66+
67+
private static String convertToTokenCredentialCacheKey(TokenCredentialProviderOptions options) {
68+
if (options == null) {
69+
return CachingTokenCredentialProvider.class.getSimpleName();
70+
}
71+
72+
return Arrays
73+
.stream(new String[] {
74+
options.getTenantId(),
75+
options.getClientId(),
76+
options.getClientCertificatePath(),
77+
options.getUsername(),
78+
String.valueOf(options.isManagedIdentityEnabled()),
79+
options.getTokenCredentialProviderClassName(),
80+
options.getTokenCredentialBeanName() })
81+
.map(option -> option == null ? "" : option)
82+
.collect(Collectors.joining(","));
83+
}
84+
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,8 @@ private TokenCredential resolveTokenCredential(TokenCredentialProviderOptions op
6767
.clientId(clientId);
6868

6969
if (hasText(options.getClientCertificatePassword())) {
70-
builder.pfxCertificate(clientCertificatePath, options.getClientCertificatePassword());
70+
builder.pfxCertificate(clientCertificatePath)
71+
.clientCertificatePassword(options.getClientCertificatePassword());
7172
} else {
7273
builder.pemCertificate(clientCertificatePath);
7374
}

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,12 @@ public enum AuthProperty {
7777
* The given bean name of a TokenCredential bean in the Spring context.
7878
*/
7979
TOKEN_CREDENTIAL_BEAN_NAME("azure.tokenCredentialBeanName", "springCloudAzureDefaultCredential",
80-
"The given bean name of a TokenCredential bean in the Spring context.", false);
80+
"The given bean name of a TokenCredential bean in the Spring context.", false),
81+
/**
82+
* Whether to enable token credential cache.
83+
*/
84+
TOKEN_CREDENTIAL_CACHE_ENABLED("azure.tokenCredentialCacheEnabled", "true",
85+
"Whether to enable the token credential cache.", false);
8186

8287
String propertyKey;
8388
String defaultValue;

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

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,16 @@
55

66
import com.azure.core.credential.AccessToken;
77
import com.azure.core.util.logging.ClientLogger;
8+
import com.azure.identity.extensions.implementation.credential.provider.CachingTokenCredentialProvider;
89
import com.azure.identity.extensions.implementation.credential.provider.TokenCredentialProvider;
910
import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions;
11+
import com.azure.identity.extensions.implementation.enums.AuthProperty;
1012
import com.azure.identity.extensions.implementation.token.AccessTokenResolver;
1113
import com.azure.identity.extensions.implementation.token.AccessTokenResolverOptions;
1214
import java.time.Duration;
1315
import java.util.Properties;
1416
import java.util.concurrent.atomic.AtomicBoolean;
17+
1518
import reactor.core.publisher.Mono;
1619
import static com.azure.identity.extensions.implementation.enums.AuthProperty.GET_TOKEN_TIMEOUT;
1720

@@ -41,7 +44,7 @@ public AzureAuthenticationTemplate() {
4144
/**
4245
* AzureAuthenticationTemplate constructor.
4346
*
44-
* @param tokenCredentialProvider An TokenCredentialProvider class instance.
47+
* @param tokenCredentialProvider A TokenCredentialProvider class instance.
4548
* @param accessTokenResolver An AccessTokenResolver class instance.
4649
*/
4750
public AzureAuthenticationTemplate(TokenCredentialProvider tokenCredentialProvider,
@@ -60,8 +63,13 @@ public void init(Properties properties) {
6063
LOGGER.verbose("Initializing AzureAuthenticationTemplate.");
6164

6265
if (getTokenCredentialProvider() == null) {
63-
this.tokenCredentialProvider
64-
= TokenCredentialProvider.createDefault(new TokenCredentialProviderOptions(properties));
66+
TokenCredentialProviderOptions options = new TokenCredentialProviderOptions(properties);
67+
this.tokenCredentialProvider = TokenCredentialProvider.createDefault(options);
68+
69+
if (Boolean.TRUE.equals(AuthProperty.TOKEN_CREDENTIAL_CACHE_ENABLED.getBoolean(properties))) {
70+
this.tokenCredentialProvider
71+
= new CachingTokenCredentialProvider(options, this.tokenCredentialProvider);
72+
}
6573
}
6674

6775
if (getAccessTokenResolver() == null) {
@@ -125,5 +133,4 @@ Duration getBlockTimeout() {
125133
AtomicBoolean getIsInitialized() {
126134
return isInitialized;
127135
}
128-
129136
}

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

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.azure.identity.extensions.implementation.credential.provider;
5+
6+
import com.azure.core.credential.TokenCredential;
7+
import com.azure.identity.DefaultAzureCredential;
8+
import com.azure.identity.ManagedIdentityCredential;
9+
import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions;
10+
import org.junit.jupiter.api.Test;
11+
12+
import static org.junit.jupiter.api.Assertions.assertTrue;
13+
14+
class CachingTokenCredentialProviderTest {
15+
16+
@Test
17+
void returnCacheUsingDefaultAuthMethodViaDifferentProviderInstances() {
18+
DefaultTokenCredentialProvider defaultTokenCredentialProvider1 = new DefaultTokenCredentialProvider(null);
19+
CachingTokenCredentialProvider provider1
20+
= new CachingTokenCredentialProvider(null, defaultTokenCredentialProvider1);
21+
TokenCredential tokenCredential1 = provider1.get();
22+
23+
DefaultTokenCredentialProvider defaultTokenCredentialProvider2 = new DefaultTokenCredentialProvider(null);
24+
CachingTokenCredentialProvider provider2
25+
= new CachingTokenCredentialProvider(null, defaultTokenCredentialProvider2);
26+
27+
TokenCredential tokenCredential2 = provider2.get();
28+
assertTrue(tokenCredential1 instanceof DefaultAzureCredential);
29+
assertTrue(tokenCredential2 instanceof DefaultAzureCredential);
30+
assertTrue(tokenCredential1 == tokenCredential2);
31+
}
32+
33+
@Test
34+
void returnCacheUsingSameAuthMethodViaDifferentProviderInstances() {
35+
TokenCredentialProviderOptions customOptions = getSystemManagedIdentityCredentialProviderOptions();
36+
37+
DefaultTokenCredentialProvider defaultTokenCredentialProvider1
38+
= new DefaultTokenCredentialProvider(customOptions);
39+
CachingTokenCredentialProvider cachingProvider1
40+
= new CachingTokenCredentialProvider(customOptions, defaultTokenCredentialProvider1);
41+
TokenCredential tokenCredential1 = cachingProvider1.get();
42+
43+
TokenCredentialProviderOptions customOptions2 = getSystemManagedIdentityCredentialProviderOptions();
44+
DefaultTokenCredentialProvider defaultTokenCredentialProvider2
45+
= new DefaultTokenCredentialProvider(customOptions2);
46+
CachingTokenCredentialProvider cachingProvider2
47+
= new CachingTokenCredentialProvider(customOptions2, defaultTokenCredentialProvider2);
48+
49+
TokenCredential tokenCredential2 = cachingProvider2.get();
50+
assertTrue(tokenCredential1 instanceof ManagedIdentityCredential);
51+
assertTrue(tokenCredential2 instanceof ManagedIdentityCredential);
52+
assertTrue(tokenCredential1 == tokenCredential2);
53+
}
54+
55+
@Test
56+
void returnCacheUsingSameAuthMethodAndInvokingDifferentGetMethods() {
57+
TokenCredentialProviderOptions customOptions = getSystemManagedIdentityCredentialProviderOptions();
58+
59+
DefaultTokenCredentialProvider defaultTokenCredentialProvider1
60+
= new DefaultTokenCredentialProvider(customOptions);
61+
CachingTokenCredentialProvider cachingProvider1
62+
= new CachingTokenCredentialProvider(customOptions, defaultTokenCredentialProvider1);
63+
TokenCredential tokenCredential1 = cachingProvider1.get();
64+
65+
TokenCredentialProviderOptions customOptions2 = getSystemManagedIdentityCredentialProviderOptions();
66+
DefaultTokenCredentialProvider defaultTokenCredentialProvider2
67+
= new DefaultTokenCredentialProvider(customOptions2);
68+
CachingTokenCredentialProvider cachingProvider2
69+
= new CachingTokenCredentialProvider(customOptions2, defaultTokenCredentialProvider2);
70+
71+
TokenCredential tokenCredential2 = cachingProvider2.get(customOptions2);
72+
assertTrue(tokenCredential1 instanceof ManagedIdentityCredential);
73+
assertTrue(tokenCredential2 instanceof ManagedIdentityCredential);
74+
assertTrue(tokenCredential1 == tokenCredential2);
75+
}
76+
77+
@Test
78+
void returnDifferentCachesUsingDifferentAuthenticationMethods() {
79+
TokenCredentialProviderOptions customOptions = getSystemManagedIdentityCredentialProviderOptions();
80+
81+
DefaultTokenCredentialProvider defaultTokenCredentialProvider1
82+
= new DefaultTokenCredentialProvider(customOptions);
83+
CachingTokenCredentialProvider cachingProvider1
84+
= new CachingTokenCredentialProvider(customOptions, defaultTokenCredentialProvider1);
85+
TokenCredential tokenCredential1 = cachingProvider1.get();
86+
87+
TokenCredentialProviderOptions customOptions2
88+
= getUserManagedIdentityCredentialProviderOptions("test-client-id");
89+
DefaultTokenCredentialProvider defaultTokenCredentialProvider2
90+
= new DefaultTokenCredentialProvider(customOptions2);
91+
CachingTokenCredentialProvider cachingProvider2
92+
= new CachingTokenCredentialProvider(customOptions2, defaultTokenCredentialProvider2);
93+
94+
TokenCredential tokenCredential2 = cachingProvider2.get(customOptions2);
95+
assertTrue(tokenCredential1 instanceof ManagedIdentityCredential);
96+
assertTrue(tokenCredential2 instanceof ManagedIdentityCredential);
97+
assertTrue(tokenCredential1 != tokenCredential2);
98+
}
99+
100+
private static TokenCredentialProviderOptions getUserManagedIdentityCredentialProviderOptions(String clientId) {
101+
TokenCredentialProviderOptions customOptions = new TokenCredentialProviderOptions();
102+
customOptions.setManagedIdentityEnabled(true);
103+
customOptions.setClientId(clientId);
104+
return customOptions;
105+
}
106+
107+
private static TokenCredentialProviderOptions getSystemManagedIdentityCredentialProviderOptions() {
108+
TokenCredentialProviderOptions customOptions = new TokenCredentialProviderOptions();
109+
customOptions.setManagedIdentityEnabled(true);
110+
return customOptions;
111+
}
112+
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,9 @@
66
import com.azure.core.credential.TokenCredential;
77
import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions;
88

9-
class SpringTokenCredentialProviderTest implements TokenCredentialProvider {
9+
class TestSpringTokenCredentialProvider implements TokenCredentialProvider {
1010

11-
SpringTokenCredentialProviderTest(TokenCredentialProviderOptions options) {
11+
TestSpringTokenCredentialProvider(TokenCredentialProviderOptions options) {
1212
}
1313

1414
@Override

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
class TokenCredentialProvidersTest {
1111

1212
private static final String SPRING_TOKEN_CREDENTIAL_PROVIDER_CLASS_NAME
13-
= SpringTokenCredentialProviderTest.class.getName();
13+
= TestSpringTokenCredentialProvider.class.getName();
1414

1515
@Test
1616
void testOptionsIsNull() {
@@ -29,7 +29,7 @@ void testCreateSpringTokenCredentialProvider() {
2929
TokenCredentialProviderOptions option = new TokenCredentialProviderOptions();
3030
option.setTokenCredentialProviderClassName(SPRING_TOKEN_CREDENTIAL_PROVIDER_CLASS_NAME);
3131
TokenCredentialProvider credentialProvider = TokenCredentialProviders.createInstance(option);
32-
Assertions.assertTrue(credentialProvider instanceof SpringTokenCredentialProviderTest);
32+
Assertions.assertTrue(credentialProvider instanceof TestSpringTokenCredentialProvider);
3333
}
3434

3535
}

0 commit comments

Comments
 (0)