Skip to content

Commit

Permalink
Ensure correlation ID is never null
Browse files Browse the repository at this point in the history
  • Loading branch information
Avery-Dunn committed Oct 8, 2023
1 parent f638f61 commit 6fc2c6c
Showing 1 changed file with 23 additions and 10 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;

Expand All @@ -34,12 +35,13 @@ public class MsalRuntimeBroker implements IBroker {

@Override
public CompletableFuture<IAuthenticationResult> acquireToken(PublicClientApplication application, SilentParameters parameters) {
String correlationID = application.correlationId() == null ? generateCorrelationID() : application.correlationId();
Account accountResult = null;

//If request has an account ID, MSALRuntime likely has data cached for that account that we can retrieve
if (parameters.account() != null) {
try {
accountResult = ((ReadAccountResult) interop.readAccountById(parameters.account().homeAccountId().split("\\.")[0], application.correlationId()).get()).getAccount();
accountResult = ((ReadAccountResult) interop.readAccountById(parameters.account().homeAccountId().split("\\.")[0], correlationID).get()).getAccount();
} catch (InterruptedException | ExecutionException ex) {
throw new MsalClientException(String.format("MSALRuntime async operation interrupted when waiting for result: %s", ex.getMessage()), AuthenticationErrorCode.MSALRUNTIME_INTEROP_ERROR);
}
Expand All @@ -62,8 +64,8 @@ public CompletableFuture<IAuthenticationResult> acquireToken(PublicClientApplica
AuthParameters authParameters = authParamsBuilder.build();

if (accountResult == null) {
return interop.signInSilently(authParameters, application.correlationId())
.thenCompose(acctResult -> interop.acquireTokenSilently(authParameters, application.correlationId(), ((AuthResult) acctResult).getAccount()))
return interop.signInSilently(authParameters, correlationID)
.thenCompose(acctResult -> interop.acquireTokenSilently(authParameters, correlationID, ((AuthResult) acctResult).getAccount()))
.thenApply(authResult -> parseBrokerAuthResult(
application.authority(),
((AuthResult) authResult).getIdToken(),
Expand All @@ -73,7 +75,7 @@ public CompletableFuture<IAuthenticationResult> acquireToken(PublicClientApplica
((AuthResult) authResult).getAccessTokenExpirationTime(),
((AuthResult) authResult).isPopAuthorization()));
} else {
return interop.acquireTokenSilently(authParameters, application.correlationId(), accountResult)
return interop.acquireTokenSilently(authParameters, correlationID, accountResult)
.thenApply(authResult -> parseBrokerAuthResult(application.authority(),
((AuthResult) authResult).getIdToken(),
((AuthResult) authResult).getAccessToken(),
Expand All @@ -89,6 +91,8 @@ public CompletableFuture<IAuthenticationResult> acquireToken(PublicClientApplica

@Override
public CompletableFuture<IAuthenticationResult> acquireToken(PublicClientApplication application, InteractiveRequestParameters parameters) {
String correlationID = application.correlationId() == null ? generateCorrelationID() : application.correlationId();

try {
AuthParameters.AuthParametersBuilder authParamsBuilder = new AuthParameters.
AuthParametersBuilder(application.clientId(),
Expand All @@ -106,8 +110,8 @@ public CompletableFuture<IAuthenticationResult> acquireToken(PublicClientApplica

AuthParameters authParameters = authParamsBuilder.build();

return interop.signInInteractively(parameters.windowHandle(), authParameters, application.correlationId(), parameters.loginHint())
.thenCompose(acctResult -> interop.acquireTokenInteractively(parameters.windowHandle(), authParameters, application.correlationId(), ((AuthResult) acctResult).getAccount()))
return interop.signInInteractively(parameters.windowHandle(), authParameters, correlationID, parameters.loginHint())
.thenCompose(acctResult -> interop.acquireTokenInteractively(parameters.windowHandle(), authParameters, correlationID, ((AuthResult) acctResult).getAccount()))
.thenApply(authResult -> parseBrokerAuthResult(
application.authority(),
((AuthResult) authResult).getIdToken(),
Expand All @@ -127,6 +131,8 @@ public CompletableFuture<IAuthenticationResult> acquireToken(PublicClientApplica
@Deprecated
@Override
public CompletableFuture<IAuthenticationResult> acquireToken(PublicClientApplication application, UserNamePasswordParameters parameters) {
String correlationID = application.correlationId() == null ? generateCorrelationID() : application.correlationId();

try {
AuthParameters.AuthParametersBuilder authParamsBuilder = new AuthParameters.
AuthParametersBuilder(application.clientId(),
Expand All @@ -143,8 +149,8 @@ public CompletableFuture<IAuthenticationResult> acquireToken(PublicClientApplica

AuthParameters authParameters = authParamsBuilder.build();

return interop.signInSilently(authParameters, application.correlationId())
.thenCompose(acctResult -> interop.acquireTokenSilently(authParameters, application.correlationId(), ((AuthResult) acctResult).getAccount()))
return interop.signInSilently(authParameters, correlationID)
.thenCompose(acctResult -> interop.acquireTokenSilently(authParameters, correlationID, ((AuthResult) acctResult).getAccount()))
.thenApply(authResult -> parseBrokerAuthResult(
application.authority(),
((AuthResult) authResult).getIdToken(),
Expand All @@ -160,11 +166,13 @@ public CompletableFuture<IAuthenticationResult> acquireToken(PublicClientApplica

@Override
public void removeAccount(PublicClientApplication application, IAccount msalJavaAccount) {
String correlationID = application.correlationId() == null ? generateCorrelationID() : application.correlationId();

try {
Account msalRuntimeAccount = ((ReadAccountResult) interop.readAccountById(msalJavaAccount.homeAccountId().split("\\.")[0], application.correlationId()).get()).getAccount();
Account msalRuntimeAccount = ((ReadAccountResult) interop.readAccountById(msalJavaAccount.homeAccountId().split("\\.")[0], correlationID).get()).getAccount();

if (msalRuntimeAccount != null) {
interop.signOutSilently(application.clientId(), application.correlationId(), msalRuntimeAccount);
interop.signOutSilently(application.clientId(), correlationID, msalRuntimeAccount);
}
} catch (MsalInteropException interopException) {
throw new MsalClientException(interopException.getErrorMessage(), AuthenticationErrorCode.MSALRUNTIME_INTEROP_ERROR);
Expand Down Expand Up @@ -233,4 +241,9 @@ public void enableBrokerPIILogging(boolean enablePII) {
throw new MsalClientException(String.format("Error occurred when calling MSALRuntime PII logging API: %s", ex.getMessage()), AuthenticationErrorCode.MSALRUNTIME_INTEROP_ERROR);
}
}

//Generates a random correlation ID, used when a correlation ID was not set at the application level
private String generateCorrelationID() {
return UUID.randomUUID().toString();
}
}

0 comments on commit 6fc2c6c

Please sign in to comment.