Ensure correlation ID is never null

Author: Avery-Dunn

msal4j-brokers/src/main/java/com/microsoft/aad/msal4jbrokers/MsalRuntimeBroker.java

@@ -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;
 
@@ -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);
             }
@@ -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(),
@@ -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(),
@@ -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(),
@@ -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(),
@@ -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(),
@@ -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(),
@@ -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);
@@ -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();
+    }
 }