Skip to content

Commit 896ab96

Browse files
neha-bhargavaneha-bhargava
committed
Update code
1 parent b24bd25 commit 896ab96

9 files changed

Lines changed: 84 additions & 97 deletions

‎msal4j-sdk/pom.xml‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -157,12 +157,6 @@
157157
<version>2.14.0</version>
158158
<scope>test</scope>
159159
</dependency>
160-
<dependency>
161-
<groupId>jakarta.xml.bind</groupId>
162-
<artifactId>jakarta.xml.bind-api</artifactId>
163-
<version>2.3.2</version>
164-
<scope>compile</scope>
165-
</dependency>
166160
</dependencies>
167161

168162
<!-- force https -->

‎msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractManagedIdentitySource.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ public ManagedIdentityResponse getManagedIdentityResponse(
4141
ManagedIdentityParameters parameters) {
4242

4343
createManagedIdentityRequest(parameters.resource);
44+
managedIdentityRequest.addTokenRevocationParametersToQuery(parameters);
4445
IHttpResponse response;
4546

4647
try {

‎msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AcquireTokenByManagedIdentitySupplier.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,12 @@ AuthenticationResult execute() throws Exception {
8383
result.metadata().tokenSource(TokenSource.CACHE);
8484
return result;
8585
} else {
86+
if (cacheRefreshReason == CacheRefreshReason.CLAIMS) {
87+
LOG.debug("Claims are passed, creating token hash and refreshing the token");
88+
managedIdentityParameters.revokedTokenHash = StringHelper.createSha256HashHexString(result.accessToken());
89+
return fetchNewAccessTokenAndSaveToCache(tokenRequestExecutor, CacheRefreshReason.CLAIMS);
90+
}
91+
8692
LOG.debug(String.format("Refreshing access token. Cache refresh reason: %s", cacheRefreshReason));
8793
return fetchNewAccessTokenAndSaveToCache(tokenRequestExecutor, cacheRefreshReason);
8894
}

‎msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/Constants.java‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@
33

44
package com.microsoft.aad.msal4j;
55

6+
import java.util.HashSet;
7+
import java.util.Set;
8+
69
final class Constants {
710

811
static final String CACHE_KEY_SEPARATOR = "-";
@@ -24,13 +27,12 @@ final class Constants {
2427
public static final String IDENTITY_SERVER_THUMBPRINT = "IDENTITY_SERVER_THUMBPRINT";
2528

2629
// Constants for token revocation and client capabilities
27-
public static final String CLIENT_CAPABILITY_CP1 = "cp1";
28-
public static final String TOKEN_REVOCATION_REQUEST_PARAM = "x-ms-revoke-token";
29-
public static final String TOKEN_HASH_CLAIM = "x-ms-token-hash";
30+
public static final String TOKEN_HASH_CLAIM = "token_sha256_to_refresh";
31+
public static final String CLIENT_CAPABILITY_REQUEST_PARAM = "xms_cc";
3032

3133
// Only Service Fabric and App Service managed identity environments support token revocation
32-
public static final ManagedIdentitySourceType[] TOKEN_REVOCATION_SUPPORTED_ENVIRONMENTS = {
33-
ManagedIdentitySourceType.SERVICE_FABRIC
34-
};
34+
public static final Set<ManagedIdentitySourceType> TOKEN_REVOCATION_SUPPORTED_ENVIRONMENTS = new HashSet<ManagedIdentitySourceType>() {{
35+
add(ManagedIdentitySourceType.SERVICE_FABRIC);
36+
}};
3537

3638
}

‎msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ManagedIdentityParameters.java‎

Lines changed: 5 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,11 @@ public class ManagedIdentityParameters implements IAcquireTokenParameters {
1515
String resource;
1616
boolean forceRefresh;
1717
String claims;
18-
private String tokenToRevoke;
19-
private ManagedIdentityParameters(String resource, boolean forceRefresh, String claims, String tokenToRevoke) {
18+
String revokedTokenHash;
19+
private ManagedIdentityParameters(String resource, boolean forceRefresh, String claims) {
2020
this.resource = resource;
2121
this.forceRefresh = forceRefresh;
2222
this.claims = claims;
23-
this.tokenToRevoke = tokenToRevoke;
2423
}
2524

2625
@Override
@@ -79,19 +78,14 @@ public String resource() {
7978
return this.resource;
8079
}
8180

82-
/**
83-
* Gets the token to be revoked in supported Managed Identity environments.
84-
* @return the token to be revoked
85-
*/
86-
public String getTokenToRevoke() {
87-
return this.tokenToRevoke;
81+
public String revokedTokenHash() {
82+
return this.revokedTokenHash;
8883
}
8984

9085
public static class ManagedIdentityParametersBuilder {
9186
private String resource;
9287
private boolean forceRefresh;
9388
private String claims;
94-
private String tokenToRevoke;
9589

9690
ManagedIdentityParametersBuilder() {
9791
}
@@ -122,23 +116,9 @@ public ManagedIdentityParametersBuilder claims(String claims) {
122116
this.claims = claims;
123117
return this;
124118
}
125-
126-
/**
127-
* Specifies a token to be revoked in supported Managed Identity environments (App Service, Service Fabric).
128-
* The token will be converted to a SHA256 hash for revocation to avoid transmitting the original token.
129-
*
130-
* @param tokenToRevoke The access token to revoke
131-
* @return this builder instance
132-
*/
133-
public ManagedIdentityParametersBuilder tokenToRevoke(String tokenToRevoke) {
134-
ParameterValidationUtils.validateNotBlank("tokenToRevoke", tokenToRevoke);
135-
136-
this.tokenToRevoke = tokenToRevoke;
137-
return this;
138-
}
139119

140120
public ManagedIdentityParameters build() {
141-
return new ManagedIdentityParameters(this.resource, this.forceRefresh, this.claims, this.tokenToRevoke);
121+
return new ManagedIdentityParameters(this.resource, this.forceRefresh, this.claims);
142122
}
143123

144124
public String toString() {

‎msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ManagedIdentityRequest.java‎

Lines changed: 31 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -32,36 +32,6 @@ class ManagedIdentityRequest extends MsalRequest {
3232

3333
public ManagedIdentityRequest(ManagedIdentityApplication managedIdentityApplication, RequestContext requestContext) {
3434
super(managedIdentityApplication, requestContext);
35-
36-
// Check if the environment supports token revocation
37-
ManagedIdentitySourceType sourceType = ManagedIdentityClient.getManagedIdentitySource();
38-
boolean supportsTokenRevocation = false;
39-
40-
for (ManagedIdentitySourceType type : Constants.TOKEN_REVOCATION_SUPPORTED_ENVIRONMENTS) {
41-
if (type == sourceType) {
42-
supportsTokenRevocation = true;
43-
break;
44-
}
45-
}
46-
47-
// If client capabilities include CP1 (token revocation) and the source type supports it,
48-
// add the token revocation parameter to query parameters
49-
if (supportsTokenRevocation &&
50-
managedIdentityApplication.getClientCapabilities() != null &&
51-
managedIdentityApplication.getClientCapabilities().contains(Constants.CLIENT_CAPABILITY_CP1)) {
52-
53-
if (requestContext.apiParameters() instanceof ManagedIdentityParameters) {
54-
ManagedIdentityParameters parameters = (ManagedIdentityParameters) requestContext.apiParameters();
55-
if (parameters.getTokenToRevoke() != null && !parameters.getTokenToRevoke().isEmpty()) {
56-
LOG.info("[Managed Identity] Adding token revocation parameter to request");
57-
if (queryParameters == null) {
58-
queryParameters = new HashMap<>();
59-
}
60-
String tokenHash = StringHelper.createSha256HashHexString(parameters.getTokenToRevoke());
61-
queryParameters.put(Constants.TOKEN_REVOCATION_REQUEST_PARAM, Collections.singletonList(tokenHash));
62-
}
63-
}
64-
}
6535
}
6636

6737
public String getBodyAsString() {
@@ -106,4 +76,35 @@ void addUserAssignedIdToQuery(ManagedIdentityIdType idType, String userAssignedI
10676
break;
10777
}
10878
}
79+
80+
void addTokenRevocationParametersToQuery(ManagedIdentityParameters parameters) {
81+
// Check if the environment supports token revocation
82+
ManagedIdentitySourceType sourceType = ManagedIdentityClient.getManagedIdentitySource();
83+
boolean supportsTokenRevocation = Constants.TOKEN_REVOCATION_SUPPORTED_ENVIRONMENTS
84+
.contains(sourceType);
85+
86+
// If token revocation is supported, pass the client capabilities and token revocation parameters
87+
if (supportsTokenRevocation) {
88+
ManagedIdentityApplication managedIdentityApplication =
89+
(ManagedIdentityApplication) this.application();
90+
91+
// Pass capabilities if present.
92+
if (managedIdentityApplication.getClientCapabilities() != null &&
93+
!managedIdentityApplication.getClientCapabilities().isEmpty()) {
94+
// Add client capabilities as a comma separated string for all the values in client capabilities
95+
String clientCapabilities = String.join(",", managedIdentityApplication.getClientCapabilities());
96+
97+
queryParameters.put(Constants.CLIENT_CAPABILITY_REQUEST_PARAM, Collections.singletonList(clientCapabilities.toString()));
98+
}
99+
100+
// Pass the token revocation parameter if the claims are present and there is a token to revoke
101+
if (!StringHelper.isNullOrBlank(parameters.claims) && !StringHelper.isNullOrBlank(parameters.revokedTokenHash())) {
102+
LOG.info("[Managed Identity] Adding token revocation parameter to request");
103+
if (queryParameters == null) {
104+
queryParameters = new HashMap<>();
105+
}
106+
queryParameters.put(Constants.TOKEN_HASH_CLAIM, Collections.singletonList(parameters.revokedTokenHash()));
107+
}
108+
}
109+
}
109110
}

‎msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ServiceFabricManagedIdentitySource.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ public ManagedIdentityResponse getManagedIdentityResponse(
5757
ManagedIdentityParameters parameters) {
5858

5959
createManagedIdentityRequest(parameters.resource);
60+
managedIdentityRequest.addTokenRevocationParametersToQuery(parameters);
6061
IHttpResponse response;
6162

6263
try {

‎msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ManagedIdentityTests.java‎

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -70,16 +70,34 @@ private String getMsiErrorResponseNoRetry() {
7070
return "{\"statusCode\":\"123\",\"message\":\"Not one of the retryable error responses\",\"correlationId\":\"7d0c9763-ff1d-4842-a3f3-6d49e64f4513\"}";
7171
}
7272

73+
private HttpRequest expectedRequest(ManagedIdentitySourceType source, String resource, boolean hasClaims, boolean hasCapabilities, String expectedTokenHash) {
74+
return expectedRequest(source, resource, ManagedIdentityId.systemAssigned(), hasClaims, hasCapabilities, expectedTokenHash);
75+
}
76+
77+
private HttpRequest expectedRequest(ManagedIdentitySourceType source, String resource, ManagedIdentityId id) {
78+
return expectedRequest(source, resource, id, false, false, null);
79+
}
80+
7381
private HttpRequest expectedRequest(ManagedIdentitySourceType source, String resource) {
74-
return expectedRequest(source, resource, ManagedIdentityId.systemAssigned());
82+
return expectedRequest(source, resource, ManagedIdentityId.systemAssigned(), false, false, null);
7583
}
7684

7785
private HttpRequest expectedRequest(ManagedIdentitySourceType source, String resource,
78-
ManagedIdentityId id) {
86+
ManagedIdentityId id, boolean hasClaims, boolean hasCapabilities, String expectedTokenHash) {
7987
String endpoint = null;
8088
Map<String, String> headers = new HashMap<>();
8189
Map<String, List<String>> queryParameters = new HashMap<>();
8290

91+
if (Constants.TOKEN_REVOCATION_SUPPORTED_ENVIRONMENTS.contains(source)) {
92+
if (hasCapabilities) {
93+
queryParameters.put(Constants.CLIENT_CAPABILITY_REQUEST_PARAM, Collections.singletonList("cp1"));
94+
}
95+
96+
if (hasClaims) {
97+
queryParameters.put(Constants.TOKEN_HASH_CLAIM, Collections.singletonList(expectedTokenHash));
98+
}
99+
}
100+
83101
switch (source) {
84102
case APP_SERVICE:
85103
endpoint = appServiceEndpoint;
@@ -93,12 +111,6 @@ private HttpRequest expectedRequest(ManagedIdentitySourceType source, String res
93111
headers.put("Metadata", "true");
94112
queryParameters.put("resource", Collections.singletonList(resource));
95113
break;
96-
case IMDS:
97-
endpoint = IMDS_ENDPOINT;
98-
queryParameters.put("api-version", Collections.singletonList("2018-02-01"));
99-
queryParameters.put("resource", Collections.singletonList(resource));
100-
headers.put("Metadata", "true");
101-
break;
102114
case AZURE_ARC:
103115
endpoint = azureArcEndpoint;
104116
queryParameters.put("api-version", Collections.singletonList("2019-11-01"));
@@ -111,6 +123,7 @@ private HttpRequest expectedRequest(ManagedIdentitySourceType source, String res
111123
queryParameters.put("resource", Collections.singletonList(resource));
112124
headers.put("secret", "secret");
113125
break;
126+
case IMDS:
114127
case NONE:
115128
case DEFAULT_TO_IMDS:
116129
endpoint = IMDS_ENDPOINT;
@@ -657,6 +670,9 @@ void managedIdentityTest_WithClaims(ManagedIdentitySourceType source, String end
657670
assertNotNull(result.accessToken());
658671
assertEquals(TokenSource.CACHE, result.metadata().tokenSource());
659672

673+
String expectedTokenHash = StringHelper.createSha256HashHexString(result.accessToken());
674+
when(httpClientMock.send(expectedRequest(source, resource, true, false, expectedTokenHash))).thenReturn(expectedResponse(200, getSuccessfulResponse(resource)));
675+
660676
// Third call, when claims are passed bypass the cache.
661677
result = miApp.acquireTokenForManagedIdentity(
662678
ManagedIdentityParameters.builder(resource)
@@ -679,7 +695,7 @@ void managedIdentity_ClaimsAndCapabilities(ManagedIdentitySourceType source, Str
679695
ServiceFabricManagedIdentitySource.setHttpClient(httpClientMock);
680696
}
681697

682-
when(httpClientMock.send(expectedRequest(source, resource))).thenReturn(expectedResponse(200, getSuccessfulResponse(resource)));
698+
when(httpClientMock.send(expectedRequest(source, resource, false, true, null))).thenReturn(expectedResponse(200, getSuccessfulResponse(resource)));
683699

684700
miApp = ManagedIdentityApplication
685701
.builder(ManagedIdentityId.systemAssigned())
@@ -707,6 +723,9 @@ void managedIdentity_ClaimsAndCapabilities(ManagedIdentitySourceType source, Str
707723
assertNotNull(result.accessToken());
708724
assertEquals(TokenSource.CACHE, result.metadata().tokenSource());
709725

726+
String expectedTokenHash = StringHelper.createSha256HashHexString(result.accessToken());
727+
when(httpClientMock.send(expectedRequest(source, resource, true, true, expectedTokenHash))).thenReturn(expectedResponse(200, getSuccessfulResponse(resource)));
728+
710729
// Third call, when claims are passed bypass the cache.
711730
result = miApp.acquireTokenForManagedIdentity(
712731
ManagedIdentityParameters.builder(resource)
@@ -715,8 +734,6 @@ void managedIdentity_ClaimsAndCapabilities(ManagedIdentitySourceType source, Str
715734

716735
assertNotNull(result.accessToken());
717736
assertEquals(TokenSource.IDENTITY_PROVIDER, result.metadata().tokenSource());
718-
719-
verify(httpClientMock, times(2)).send(any());
720737
}
721738

722739
@ParameterizedTest

‎msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/TokenRevocationTest.java‎

Lines changed: 4 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -33,26 +33,11 @@ public void testTokenToRevokeValidation() {
3333
@Test
3434
public void testManagedIdentityParametersBuilder() {
3535
ManagedIdentityParameters params = ManagedIdentityParameters.builder(TEST_RESOURCE)
36-
.tokenToRevoke(TEST_TOKEN)
3736
.build();
38-
37+
38+
params.revokedTokenHash = StringHelper.createSha256HashHexString(TEST_TOKEN);
39+
3940
assertEquals(TEST_RESOURCE, params.resource());
40-
assertEquals(TEST_TOKEN, params.getTokenToRevoke());
41-
}
42-
43-
@Test
44-
public void testManagedIdentityParametersBuilderValidation() {
45-
// Should throw exception when tokenToRevoke is null or empty
46-
assertThrows(IllegalArgumentException.class, () -> {
47-
ManagedIdentityParameters.builder(TEST_RESOURCE)
48-
.tokenToRevoke(null)
49-
.build();
50-
});
51-
52-
assertThrows(IllegalArgumentException.class, () -> {
53-
ManagedIdentityParameters.builder(TEST_RESOURCE)
54-
.tokenToRevoke("")
55-
.build();
56-
});
41+
assertEquals(EXPECTED_TOKEN_HASH, params.revokedTokenHash());
5742
}
5843
}

0 commit comments

Comments
 (0)