Skip to content

Commit 9615519

Browse files
siddhijainsiddhijain
authored
Merge pull request #516 from AzureAD/SJAIN/managed-identity-token-caching
enable token caching for managed identity
2 parents 7c0f19a + c6fc879 commit 9615519

9 files changed

+278
-11
lines changed

‎src/integrationtest/java/com.microsoft.aad.msal4j/ConfidentialClientApplicationUnitT.java

Lines changed: 80 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,11 @@
2929
import java.security.*;
3030
import java.security.cert.CertificateException;
3131
import java.util.*;
32+
import java.util.concurrent.CompletableFuture;
3233
import java.util.concurrent.Future;
34+
import java.util.function.Function;
3335

36+
import static com.microsoft.aad.msal4j.TestConstants.KEYVAULT_DEFAULT_SCOPE;
3437
import static org.easymock.EasyMock.*;
3538
import static org.testng.Assert.*;
3639

@@ -277,12 +280,10 @@ private ClientCredentialRequest getClientCredentialRequest(ConfidentialClientApp
277280
PublicApi.ACQUIRE_TOKEN_FOR_CLIENT,
278281
clientCredentials);
279282

280-
ClientCredentialRequest clientCredentialRequest =
281-
new ClientCredentialRequest(
283+
return new ClientCredentialRequest(
282284
clientCredentials,
283285
app,
284286
requestContext);
285-
return clientCredentialRequest;
286287
}
287288

288289
@Test(expectedExceptions = MsalClientException.class)
@@ -298,6 +299,82 @@ public void testClientAssertion_throwsException() throws Exception{
298299

299300
}
300301

302+
@Test
303+
public void validateAppTokenProviderAsync() throws Exception{
304+
305+
SignedJWT jwt = createClientAssertion("issuer");
306+
307+
ClientAssertion clientAssertion = new ClientAssertion(jwt.serialize());
308+
309+
IClientCredential iClientCredential = ClientCredentialFactory.createFromClientAssertion(
310+
clientAssertion.assertion());
311+
312+
//builds client with AppTokenProvider
313+
ConfidentialClientApplication cca = ConfidentialClientApplication.
314+
builder(TestConfiguration.AAD_CLIENT_ID, iClientCredential)
315+
.appTokenProvider((parameters) -> {
316+
Assert.assertNotNull(parameters.scopes);
317+
Assert.assertNotNull(parameters.correlationId);
318+
Assert.assertNotNull(parameters.tenantId);
319+
return getAppTokenProviderResult("/default");
320+
})
321+
.build();
322+
323+
IAuthenticationResult result1 = cca.acquireToken(ClientCredentialParameters
324+
.builder(Collections.singleton(KEYVAULT_DEFAULT_SCOPE))
325+
.tenant("tenant1")
326+
.build())
327+
.get();
328+
329+
Assert.assertNotNull(result1.accessToken());
330+
331+
Assert.assertEquals(cca.tokenCache.accessTokens.size(), 1);
332+
333+
//Acquire token from cache
334+
335+
IAuthenticationResult result2 = cca.acquireToken(ClientCredentialParameters
336+
.builder(Collections.singleton(KEYVAULT_DEFAULT_SCOPE))
337+
.build())
338+
.get();
339+
340+
Assert.assertEquals(result1.accessToken(), result2.accessToken());
341+
342+
Assert.assertEquals(cca.tokenCache.accessTokens.size(), 1);
343+
344+
cca = ConfidentialClientApplication.
345+
builder(TestConfiguration.AAD_CLIENT_ID, iClientCredential)
346+
.appTokenProvider((parameters) -> {
347+
Assert.assertNotNull(parameters.scopes);
348+
Assert.assertNotNull(parameters.correlationId);
349+
Assert.assertNotNull(parameters.tenantId);
350+
return getAppTokenProviderResult("/newScope");
351+
})
352+
.build();
353+
354+
IAuthenticationResult result3 = cca.acquireToken(ClientCredentialParameters
355+
.builder(Collections.singleton("/newScope"))
356+
.tenant("tenant1")
357+
// .claims(new ClaimsRequest().formatAsClaimsRequest(TestConstants.CLAIMS))
358+
.build())
359+
.get();
360+
361+
Assert.assertNotEquals(result2.accessToken(), result3.accessToken());
362+
Assert.assertEquals(cca.tokenCache.accessTokens.size(), 1);
363+
364+
}
365+
366+
private CompletableFuture<TokenProviderResult> getAppTokenProviderResult(String differentScopesForAt)
367+
{
368+
long currTimestampSec = new Date().getTime() / 1000;
369+
TokenProviderResult token = new TokenProviderResult();
370+
token.setAccessToken(TestConstants.DEFAULT_ACCESS_TOKEN + differentScopesForAt); //Used to indicate that there is a new access token for a different set of scopes
371+
token.setTenantId("tenantId");
372+
token.setExpiresInSeconds(currTimestampSec + 1000000);
373+
token.setRefreshInSeconds(currTimestampSec + 800000);
374+
375+
return CompletableFuture.completedFuture(token);
376+
}
377+
301378
private SignedJWT createClientAssertion(String issuer) throws KeyStoreException, IOException, NoSuchAlgorithmException, CertificateException, UnrecoverableKeyException, NoSuchProviderException, JOSEException {
302379
IClientCertificate certificate = CertificateHelper.getClientCertificate();
303380

‎src/integrationtest/java/com.microsoft.aad.msal4j/TestConstants.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,4 +61,6 @@ public class TestConstants {
6161
public final static String AUTHORITY_ARLINGTON = "https://login.microsoftonline.us/" + ARLINGTON_AUTHORITY_TENANT;
6262
public final static String AUTHORITY_MOONCAKE = "https://login.chinacloudapi.cn/mncmsidlab1.partner.onmschina.cn";
6363
public final static String AUTHORITY_PUBLIC_TENANT_SPECIFIC = "https://login.microsoftonline.com/" + MICROSOFT_AUTHORITY_TENANT;
64+
65+
public final static String DEFAULT_ACCESS_TOKEN = "defaultAccessToken";
6466
}

‎src/main/java/com/microsoft/aad/msal4j/AcquireTokenByAppProviderSupplier.java

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.microsoft.aad.msal4j;
5+
6+
import java.util.concurrent.CompletableFuture;
7+
import java.util.concurrent.ExecutionException;
8+
9+
class AcquireTokenByAppProviderSupplier extends AuthenticationResultSupplier {
10+
11+
private AppTokenProviderParameters appTokenProviderParameters;
12+
13+
private ClientCredentialRequest clientCredentialRequest;
14+
15+
AcquireTokenByAppProviderSupplier(AbstractClientApplicationBase clientApplication,
16+
ClientCredentialRequest clientCredentialRequest,
17+
AppTokenProviderParameters appTokenProviderParameters) {
18+
super(clientApplication, clientCredentialRequest);
19+
this.clientCredentialRequest = clientCredentialRequest;
20+
this.appTokenProviderParameters = appTokenProviderParameters;
21+
}
22+
23+
private static void validateTokenProviderResult(TokenProviderResult tokenProviderResult) {
24+
if (null == tokenProviderResult.getAccessToken() || tokenProviderResult.getAccessToken().isEmpty()) {
25+
handleInvalidExternalValueError(tokenProviderResult.getAccessToken());
26+
}
27+
28+
if (tokenProviderResult.getExpiresInSeconds() == 0 || tokenProviderResult.getExpiresInSeconds() < 0) {
29+
handleInvalidExternalValueError(Long.valueOf(tokenProviderResult.getExpiresInSeconds()).toString());
30+
}
31+
32+
if (null == tokenProviderResult.getTenantId() || tokenProviderResult.getTenantId().isEmpty()) {
33+
handleInvalidExternalValueError(tokenProviderResult.getTenantId());
34+
}
35+
}
36+
37+
private static void handleInvalidExternalValueError(String nameOfValue) {
38+
throw new MsalClientException("The following token provider result value is invalid" + nameOfValue, "Invalid_TokenProviderResult_Input");
39+
}
40+
41+
@Override
42+
AuthenticationResult execute() throws Exception {
43+
44+
AuthenticationResult authenticationResult = fetchTokenUsingAppTokenProvider(appTokenProviderParameters);
45+
46+
TokenRequestExecutor tokenRequestExecutor = new TokenRequestExecutor(
47+
clientCredentialRequest.application().authenticationAuthority,
48+
msalRequest,
49+
clientApplication.getServiceBundle()
50+
);
51+
52+
clientApplication.tokenCache.saveTokens(tokenRequestExecutor, authenticationResult, clientCredentialRequest.application().authenticationAuthority.host);
53+
54+
return authenticationResult;
55+
}
56+
57+
public AuthenticationResult fetchTokenUsingAppTokenProvider(AppTokenProviderParameters appTokenProviderParameters) throws ExecutionException, InterruptedException {
58+
59+
CompletableFuture<TokenProviderResult> completableFuture = this.clientCredentialRequest.appTokenProvider.apply(appTokenProviderParameters);
60+
61+
TokenProviderResult tokenProviderResult = completableFuture.get();
62+
63+
validateTokenProviderResult(tokenProviderResult);
64+
65+
return AuthenticationResult.builder()
66+
.accessToken(tokenProviderResult.getAccessToken())
67+
.refreshToken(null)
68+
.idToken(null)
69+
.expiresOn(tokenProviderResult.getExpiresInSeconds())
70+
.refreshOn(tokenProviderResult.getRefreshInSeconds())
71+
.build();
72+
73+
}
74+
}

‎src/main/java/com/microsoft/aad/msal4j/AcquireTokenByClientCredentialSupplier.java

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99
class AcquireTokenByClientCredentialSupplier extends AuthenticationResultSupplier {
1010

11-
private final static Logger LOG = LoggerFactory.getLogger(AcquireTokenByClientCredentialSupplier.class);
11+
private static final Logger LOG = LoggerFactory.getLogger(AcquireTokenByClientCredentialSupplier.class);
1212
private ClientCredentialRequest clientCredentialRequest;
1313

1414
AcquireTokenByClientCredentialSupplier(ConfidentialClientApplication clientApplication,
@@ -55,11 +55,36 @@ AuthenticationResult execute() throws Exception {
5555
}
5656

5757
private AuthenticationResult acquireTokenByClientCredential() throws Exception {
58+
59+
if (this.clientCredentialRequest.appTokenProvider != null) {
60+
61+
String claims = "";
62+
if (null != clientCredentialRequest.parameters.claims()) {
63+
claims = clientCredentialRequest.parameters.claims().toString();
64+
}
65+
66+
AppTokenProviderParameters appTokenProviderParameters = new AppTokenProviderParameters(
67+
clientCredentialRequest.parameters.scopes(),
68+
clientCredentialRequest.requestContext().correlationId(),
69+
claims,
70+
clientCredentialRequest.parameters.tenant()
71+
);
72+
73+
AcquireTokenByAppProviderSupplier supplier =
74+
new AcquireTokenByAppProviderSupplier(this.clientApplication,
75+
clientCredentialRequest,
76+
appTokenProviderParameters);
77+
78+
return supplier.execute();
79+
}
80+
5881
AcquireTokenByAuthorizationGrantSupplier supplier = new AcquireTokenByAuthorizationGrantSupplier(
5982
this.clientApplication,
6083
clientCredentialRequest,
6184
null);
6285

6386
return supplier.execute();
6487
}
88+
89+
6590
}

‎src/main/java/com/microsoft/aad/msal4j/AcquireTokenSilentSupplier.java

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,12 @@
33

44
package com.microsoft.aad.msal4j;
55

6+
import lombok.extern.slf4j.Slf4j;
7+
68
import java.net.URL;
79
import java.util.Date;
810

11+
@Slf4j
912
class AcquireTokenSilentSupplier extends AuthenticationResultSupplier {
1013

1114
private SilentRequest silentRequest;
@@ -105,6 +108,8 @@ AuthenticationResult execute() throws Exception {
105108
throw new MsalClientException(AuthenticationErrorMessage.NO_TOKEN_IN_CACHE, AuthenticationErrorCode.CACHE_MISS);
106109
}
107110

111+
log.info("Returning token from cache");
112+
108113
return res;
109114
}
110-
}
115+
}

‎src/main/java/com/microsoft/aad/msal4j/AppTokenProviderParameters.java

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.microsoft.aad.msal4j;
5+
6+
import lombok.AllArgsConstructor;
7+
import lombok.Getter;
8+
import lombok.Setter;
9+
10+
import java.util.Set;
11+
12+
@Getter
13+
@Setter
14+
@AllArgsConstructor
15+
/// The authentication parameters provided to the app token provider callback.
16+
public class AppTokenProviderParameters {
17+
18+
/// Specifies which scopes to request.
19+
public Set<String> scopes;
20+
/// Correlation id of the authentication request.
21+
public String correlationId;
22+
/// A string with one or multiple claims.
23+
public String claims;
24+
/// tenant id
25+
public String tenantId;
26+
}

‎src/main/java/com/microsoft/aad/msal4j/ClientCredentialRequest.java

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,29 @@
55

66
import com.nimbusds.oauth2.sdk.ClientCredentialsGrant;
77

8+
import java.util.concurrent.CompletableFuture;
9+
import java.util.function.Function;
10+
811
class ClientCredentialRequest extends MsalRequest {
912

1013
ClientCredentialParameters parameters;
14+
Function<AppTokenProviderParameters, CompletableFuture<TokenProviderResult>> appTokenProvider;
1115

1216
ClientCredentialRequest(ClientCredentialParameters parameters,
1317
ConfidentialClientApplication application,
1418
RequestContext requestContext) {
1519
super(application, createMsalGrant(parameters), requestContext);
1620
this.parameters = parameters;
21+
appTokenProvider = null;
22+
}
23+
24+
ClientCredentialRequest(ClientCredentialParameters parameters,
25+
ConfidentialClientApplication application,
26+
RequestContext requestContext,
27+
Function<AppTokenProviderParameters, CompletableFuture<TokenProviderResult>> appTokenProvider) {
28+
super(application, createMsalGrant(parameters), requestContext);
29+
this.parameters = parameters;
30+
this.appTokenProvider = appTokenProvider;
1731
}
1832

1933
private static OAuthAuthorizationGrant createMsalGrant(ClientCredentialParameters parameters) {

0 commit comments

Comments
 (0)