Skip to content

Commit

Permalink
Add configuration for refresh token before expiry (Azure#8715)
Browse files Browse the repository at this point in the history 
* Add configuration for refresh token before expiry

* Fix test utils in identity

* Address feedback

* Builder style setter

* Add tests and address feedback
  • Loading branch information
@jianghaolu
jianghaolu authored Mar 6, 2020
1 parent 4c53746 commit d408073
Show file tree
Hide file tree
Showing 8 changed files with 219 additions and 48 deletions.
17 changes: 17 additions & 0 deletions sdk/identity/azure-identity/src/main/java/com/azure/identity/CredentialBuilderBase.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,4 +84,21 @@ public T httpClient(HttpClient client) {
this.identityClientOptions.setHttpClient(client);
return (T) this;
}

/**
* Sets how long before the actual token expiry to refresh the token. The
* token will be considered expired at and after the time of (actual
* expiry - token refresh offset). The default offset is 2 minutes.
*
* This is useful when network is congested and a request containing the
* token takes longer than normal to get to the server.
*
* @param tokenRefreshOffset the duration before the actual expiry of a token to refresh it
* @return An updated instance of this builder with the token refresh offset set as specified.
*/
@SuppressWarnings("unchecked")
public T tokenRefreshOffset(Duration tokenRefreshOffset) {
this.identityClientOptions.setTokenRefreshOffset(tokenRefreshOffset);
return (T) this;
}
}
74 changes: 35 additions & 39 deletions ...entity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,6 @@
import java.nio.file.Paths;
import java.time.Duration;
import java.time.OffsetDateTime;
import java.time.ZoneOffset;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
Expand Down Expand Up @@ -147,10 +146,8 @@ public Mono<AccessToken> authenticateWithClientSecret(String clientSecret, Token
ConfidentialClientApplication application = applicationBuilder.build();
return Mono.fromFuture(application.acquireToken(
ClientCredentialParameters.builder(new HashSet<>(request.getScopes()))
.build()))
.map(ar -> new AccessToken(ar.accessToken(), OffsetDateTime.ofInstant(ar.expiresOnDate().toInstant(),
ZoneOffset.UTC)));

.build()))
.map(ar -> new MsalToken(ar, options));
} catch (MalformedURLException e) {
return Mono.error(e);
}
Expand Down Expand Up @@ -194,8 +191,7 @@ public Mono<AccessToken> authenticateWithPfxCertificate(String pfxCertificatePat
return applicationBuilder.build();
}).flatMap(application -> Mono.fromFuture(application.acquireToken(
ClientCredentialParameters.builder(new HashSet<>(request.getScopes())).build())))
.map(ar -> new AccessToken(ar.accessToken(), OffsetDateTime.ofInstant(ar.expiresOnDate().toInstant(),
ZoneOffset.UTC)));
.map(ar -> new MsalToken(ar, options));
}

/**
Expand Down Expand Up @@ -226,8 +222,7 @@ public Mono<AccessToken> authenticateWithPemCertificate(String pemCertificatePat
return Mono.fromFuture(application.acquireToken(
ClientCredentialParameters.builder(new HashSet<>(request.getScopes()))
.build()))
.map(ar -> new AccessToken(ar.accessToken(), OffsetDateTime.ofInstant(ar.expiresOnDate().toInstant(),
ZoneOffset.UTC)));
.map(ar -> new MsalToken(ar, options));
} catch (IOException e) {
return Mono.error(e);
}
Expand All @@ -246,7 +241,7 @@ public Mono<MsalToken> authenticateWithUsernamePassword(TokenRequestContext requ
return Mono.fromFuture(publicClientApplication.acquireToken(
UserNamePasswordParameters.builder(new HashSet<>(request.getScopes()), username, password.toCharArray())
.build()))
.map(MsalToken::new);
.map(ar -> new MsalToken(ar, options));
}

/**
Expand All @@ -264,7 +259,8 @@ public Mono<MsalToken> authenticateWithUserRefreshToken(TokenRequestContext requ
}
return Mono.defer(() -> {
try {
return Mono.fromFuture(publicClientApplication.acquireTokenSilently(parameters)).map(MsalToken::new);
return Mono.fromFuture(publicClientApplication.acquireTokenSilently(parameters))
.map(ar -> new MsalToken(ar, options));
} catch (MalformedURLException e) {
return Mono.error(e);
}
Expand All @@ -288,7 +284,7 @@ public Mono<MsalToken> authenticateWithDeviceCode(TokenRequestContext request,
dc -> deviceCodeConsumer.accept(new DeviceCodeInfo(dc.userCode(), dc.deviceCode(),
dc.verificationUri(), OffsetDateTime.now().plusSeconds(dc.expiresIn()), dc.message()))).build();
return publicClientApplication.acquireToken(parameters);
}).map(MsalToken::new);
}).map(ar -> new MsalToken(ar, options));
}

/**
Expand All @@ -305,7 +301,7 @@ public Mono<MsalToken> authenticateWithAuthorizationCode(TokenRequestContext req
AuthorizationCodeParameters.builder(authorizationCode, redirectUrl)
.scopes(new HashSet<>(request.getScopes()))
.build()))
.map(MsalToken::new);
.map(ar -> new MsalToken(ar, options));
}

/**
Expand Down Expand Up @@ -362,11 +358,11 @@ public Mono<MsalToken> authenticateWithBrowserInteraction(TokenRequestContext re
*/
public Mono<AccessToken> authenticateToManagedIdentityEndpoint(String msiEndpoint, String msiSecret,
TokenRequestContext request) {
String resource = ScopeUtil.scopesToResource(request.getScopes());
HttpURLConnection connection = null;
StringBuilder payload = new StringBuilder();
return Mono.fromCallable(() -> {
String resource = ScopeUtil.scopesToResource(request.getScopes());
HttpURLConnection connection = null;
StringBuilder payload = new StringBuilder();

try {
payload.append("resource=");
payload.append(URLEncoder.encode(resource, "UTF-8"));
payload.append("&api-version=");
Expand All @@ -375,32 +371,30 @@ public Mono<AccessToken> authenticateToManagedIdentityEndpoint(String msiEndpoin
payload.append("&clientid=");
payload.append(URLEncoder.encode(clientId, "UTF-8"));
}
} catch (IOException exception) {
return Mono.error(exception);
}
try {
URL url = new URL(String.format("%s?%s", msiEndpoint, payload));
connection = (HttpURLConnection) url.openConnection();
try {
URL url = new URL(String.format("%s?%s", msiEndpoint, payload));
connection = (HttpURLConnection) url.openConnection();

connection.setRequestMethod("GET");
if (msiSecret != null) {
connection.setRequestProperty("Secret", msiSecret);
}
connection.setRequestProperty("Metadata", "true");
connection.setRequestMethod("GET");
if (msiSecret != null) {
connection.setRequestProperty("Secret", msiSecret);
}
connection.setRequestProperty("Metadata", "true");

connection.connect();
connection.connect();

Scanner s = new Scanner(connection.getInputStream(), StandardCharsets.UTF_8.name()).useDelimiter("\\A");
String result = s.hasNext() ? s.next() : "";
Scanner s = new Scanner(connection.getInputStream(), StandardCharsets.UTF_8.name())
.useDelimiter("\\A");
String result = s.hasNext() ? s.next() : "";

return Mono.just(SERIALIZER_ADAPTER.deserialize(result, MSIToken.class, SerializerEncoding.JSON));
} catch (IOException e) {
return Mono.error(e);
} finally {
if (connection != null) {
connection.disconnect();
MSIToken msiToken = SERIALIZER_ADAPTER.deserialize(result, MSIToken.class, SerializerEncoding.JSON);
return new IdentityToken(msiToken.getToken(), msiToken.getExpiresAt(), options);
} finally {
if (connection != null) {
connection.disconnect();
}
}
}
});
}

/**
Expand Down Expand Up @@ -446,7 +440,9 @@ public Mono<AccessToken> authenticateToIMDSEndpoint(TokenRequestContext request)
.useDelimiter("\\A");
String result = s.hasNext() ? s.next() : "";

return SERIALIZER_ADAPTER.<MSIToken>deserialize(result, MSIToken.class, SerializerEncoding.JSON);
MSIToken msiToken = SERIALIZER_ADAPTER.deserialize(result,
MSIToken.class, SerializerEncoding.JSON);
return new IdentityToken(msiToken.getToken(), msiToken.getExpiresAt(), options);
} catch (IOException exception) {
if (connection == null) {
throw logger.logExceptionAsError(new RuntimeException(
Expand Down
25 changes: 25 additions & 0 deletions ...azure-identity/src/main/java/com/azure/identity/implementation/IdentityClientOptions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ public final class IdentityClientOptions {
private Function<Duration, Duration> retryTimeout;
private ProxyOptions proxyOptions;
private HttpPipeline httpPipeline;
private Duration tokenRefreshOffset = Duration.ofMinutes(2);
private HttpClient httpClient;

/**
Expand Down Expand Up @@ -125,6 +126,30 @@ public IdentityClientOptions setHttpPipeline(HttpPipeline httpPipeline) {
return this;
}

/**
* @return how long before the actual token expiry to refresh the token.
*/
public Duration getTokenRefreshOffset() {
return tokenRefreshOffset;
}

/**
* Sets how long before the actual token expiry to refresh the token. The
* token will be considered expired at and after the time of (actual
* expiry - token refresh offset). The default offset is 2 minutes.
*
* This is useful when network is congested and a request containing the
* token takes longer than normal to get to the server.
*
* @param tokenRefreshOffset the duration before the actual expiry of a token to refresh it
*/
public IdentityClientOptions setTokenRefreshOffset(Duration tokenRefreshOffset) {
if (tokenRefreshOffset != null) {
this.tokenRefreshOffset = tokenRefreshOffset;
}
return this;
}

/**
* Specifies the HttpClient to send use for requests.
* @param httpClient the http client to use for requests
Expand Down
24 changes: 24 additions & 0 deletions ...dentity/azure-identity/src/main/java/com/azure/identity/implementation/IdentityToken.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package com.azure.identity.implementation;

import com.azure.core.credential.AccessToken;

import java.time.OffsetDateTime;

/**
* Type representing authentication result from the azure-identity client.
*/
public class IdentityToken extends AccessToken {
/**
* Creates an identity token instance.
*
* @param token the token string.
* @param expiresAt the expiration time.
* @param options the identity client options.
*/
public IdentityToken(String token, OffsetDateTime expiresAt, IdentityClientOptions options) {
super(token, expiresAt.plusMinutes(2).minus(options.getTokenRefreshOffset()));
}
}
10 changes: 5 additions & 5 deletions sdk/identity/azure-identity/src/main/java/com/azure/identity/implementation/MsalToken.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@

package com.azure.identity.implementation;

import com.azure.core.credential.AccessToken;
import com.microsoft.aad.msal4j.IAccount;
import com.microsoft.aad.msal4j.IAuthenticationResult;

Expand All @@ -13,7 +12,7 @@
/**
* Type representing authentication result from the MSAL (Microsoft Authentication Library).
*/
public final class MsalToken extends AccessToken {
public final class MsalToken extends IdentityToken {

private IAccount account;

Expand All @@ -22,9 +21,10 @@ public final class MsalToken extends AccessToken {
*
* @param msalResult the raw authentication result returned by MSAL
*/
public MsalToken(IAuthenticationResult msalResult) {
super(msalResult.accessToken(), OffsetDateTime.ofInstant(msalResult.expiresOnDate().toInstant(),
ZoneOffset.UTC));
public MsalToken(IAuthenticationResult msalResult, IdentityClientOptions options) {
super(msalResult.accessToken(),
OffsetDateTime.ofInstant(msalResult.expiresOnDate().toInstant(), ZoneOffset.UTC),
options);
this.account = msalResult.account();
}

Expand Down
35 changes: 35 additions & 0 deletions sdk/identity/azure-identity/src/test/java/com/azure/identity/ClientSecretCredentialTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;

import java.time.Duration;
import java.time.OffsetDateTime;
import java.time.ZoneOffset;
import java.util.UUID;
Expand Down Expand Up @@ -61,6 +62,40 @@ public void testValidSecrets() throws Exception {
.verifyComplete();
}

@Test
public void testValidSecretsWithTokenRefreshOffset() throws Exception {
// setup
String secret = "secret";
String token1 = "token1";
String token2 = "token2";
TokenRequestContext request1 = new TokenRequestContext().addScopes("https://management.azure.com");
TokenRequestContext request2 = new TokenRequestContext().addScopes("https://vault.azure.net");
OffsetDateTime expiresAt = OffsetDateTime.now(ZoneOffset.UTC).plusHours(1);
Duration offset = Duration.ofMinutes(10);

// mock
IdentityClient identityClient = PowerMockito.mock(IdentityClient.class);
when(identityClient.authenticateWithClientSecret(secret, request1)).thenReturn(TestUtils.getMockAccessToken(token1, expiresAt, offset));
when(identityClient.authenticateWithClientSecret(secret, request2)).thenReturn(TestUtils.getMockAccessToken(token2, expiresAt, offset));
PowerMockito.whenNew(IdentityClient.class).withAnyArguments().thenReturn(identityClient);

// test
ClientSecretCredential credential = new ClientSecretCredentialBuilder()
.tenantId(tenantId)
.clientId(clientId)
.clientSecret(secret)
.tokenRefreshOffset(offset)
.build();
StepVerifier.create(credential.getToken(request1))
.expectNextMatches(accessToken -> token1.equals(accessToken.getToken())
&& expiresAt.getSecond() == accessToken.getExpiresAt().getSecond())
.verifyComplete();
StepVerifier.create(credential.getToken(request2))
.expectNextMatches(accessToken -> token2.equals(accessToken.getToken())
&& expiresAt.getSecond() == accessToken.getExpiresAt().getSecond())
.verifyComplete();
}

@Test
public void testInvalidSecrets() throws Exception {
// setup
Expand Down
Loading

0 comments on commit d408073

Please sign in to comment.