Use tdbr claim to route telemetry traffic to EU region, Fixes AB#3200872

changelog.txt

@@ -1,5 +1,6 @@
 vNext
 ----------
+- [MINOR] Use tdbr claim to route telemetry traffic to EU region (#2679)
 
 Version 21.3.0
 ----------

common/src/main/java/com/microsoft/identity/common/internal/controllers/BrokerMsalController.java

@@ -101,6 +101,7 @@
 import com.microsoft.identity.common.java.exception.ServiceException;
 import com.microsoft.identity.common.java.exception.UnsupportedBrokerException;
 import com.microsoft.identity.common.java.interfaces.IPlatformComponents;
+import com.microsoft.identity.common.java.opentelemetry.EUClaimStorageUtil;
 import com.microsoft.identity.common.java.providers.microsoft.MicrosoftRefreshToken;
 import com.microsoft.identity.common.java.providers.microsoft.azureactivedirectory.ClientInfo;
 import com.microsoft.identity.common.java.providers.microsoft.microsoftsts.MicrosoftStsAccount;
@@ -420,17 +421,21 @@ public void onReceive(@NonNull PropertyBag propertyBag) {
             //Wait to be notified of the result being returned... we could add a timeout here if we want to
             final Bundle resultBundle = mBrokerResultFuture.get();
 
+            final BrokerResult brokerResult = new MsalBrokerResultAdapter().brokerResultFromBundle(resultBundle);
+            final boolean isSuccess = resultBundle.getBoolean(AuthenticationConstants.Broker.BROKER_REQUEST_V2_SUCCESS);
+
             final String negotiatedBrokerProtocolVersion = interactiveRequestIntent.getStringExtra(NEGOTIATED_BP_VERSION_KEY);
             // For MSA Accounts Broker doesn't save the accounts, instead it just passes the result along,
             // MSAL needs to save this account locally for future token calls.
             // parameters.getOAuth2TokenCache() will be non-null only in case of MSAL native
             // If the request is from MSALCPP , OAuth2TokenCache will be null.
             if (parameters.getOAuth2TokenCache() != null && !BrokerProtocolVersionUtil.canSupportMsaAccountsInBroker(negotiatedBrokerProtocolVersion)) {
-                saveMsaAccountToCache(resultBundle, (MsalOAuth2TokenCache) parameters.getOAuth2TokenCache());
+                saveMsaAccountToCache(brokerResult, isSuccess, (MsalOAuth2TokenCache) parameters.getOAuth2TokenCache());
             }
 
             verifyBrokerVersionIsSupported(resultBundle, parameters.getRequiredBrokerProtocolVersion());
             result = mResultAdapter.getAcquireTokenResultFromResultBundle(resultBundle);
+            trackTelemetryRegionFromResultBundle(brokerResult, isSuccess);
         } catch (final BaseException | ExecutionException e) {
             Telemetry.emit(
                     new ApiEndEvent()
@@ -606,16 +611,18 @@ public AcquireTokenResult extractResultBundle(final @Nullable Bundle resultBundl
                             throw mResultAdapter.getExceptionForEmptyResultBundle();
                         }
 
-                        verifyBrokerVersionIsSupported(resultBundle, parameters.getRequiredBrokerProtocolVersion());
+                        final BrokerResult brokerResult = mResultAdapter.brokerResultFromBundle(resultBundle);
+                        verifyBrokerVersionIsSupportedFromBrokerResult(brokerResult, parameters.getRequiredBrokerProtocolVersion());
 
-                        AcquireTokenResult acquireTokenResult = mResultAdapter.getDeviceCodeFlowTokenResultFromResultBundle(resultBundle);
+                        AcquireTokenResult acquireTokenResult = mResultAdapter.getDeviceCodeFlowTokenResultFromBrokerResultAndBundle(brokerResult, resultBundle);
                         // If authorization_pending continue polling for token
                         if (acquireTokenResult == null) {
                             // Wait between polls for 5 secs
                             ThreadUtils.sleepSafely((int) WAIT_BETWEEN_DCF_POLLING_MILLISECONDS, TAG,
                                     "Attempting to sleep thread during Device Code Flow token polling...");
                             return acquireDeviceCodeFlowToken(authorizationResult, parameters);
                         } else {
+                            trackTelemetryRegionFromResultBundle(brokerResult, resultBundle.getBoolean(AuthenticationConstants.Broker.BROKER_REQUEST_V2_SUCCESS));
                             return acquireTokenResult;
                         }
                     }
@@ -1305,14 +1312,11 @@ public void putValueInSuccessEvent(@NonNull final ApiEndEvent event,
     /**
      * Checks if the account returns is a MSA Account and sets single on state in cache
      */
-    private void saveMsaAccountToCache(final @NonNull Bundle resultBundle,
+    private void saveMsaAccountToCache(final @NonNull BrokerResult brokerResult, final boolean isSuccess,
                                        @SuppressWarnings(WarningType.rawtype_warning) final @NonNull MsalOAuth2TokenCache msalOAuth2TokenCache) throws BaseException {
         final String methodTag = TAG + ":saveMsaAccountToCache";
 
-        final BrokerResult brokerResult = new MsalBrokerResultAdapter().brokerResultFromBundle(resultBundle);
-
-        if (resultBundle.getBoolean(AuthenticationConstants.Broker.BROKER_REQUEST_V2_SUCCESS) &&
-                AzureActiveDirectoryAudience.MSA_MEGA_TENANT_ID.equalsIgnoreCase(brokerResult.getTenantId())) {
+        if (isSuccess && AzureActiveDirectoryAudience.MSA_MEGA_TENANT_ID.equalsIgnoreCase(brokerResult.getTenantId())) {
             Logger.info(methodTag, "Result returned for MSA Account, saving to cache");
 
             if (StringUtil.isNullOrEmpty(brokerResult.getClientInfo())) {
@@ -1322,6 +1326,9 @@ private void saveMsaAccountToCache(final @NonNull Bundle resultBundle,
 
             try {
                 final ClientInfo clientInfo = new ClientInfo(brokerResult.getClientInfo());
+                // Store the telemetry region info in shared preferences
+                EUClaimStorageUtil.Companion.storeTelemetryRegionByTenant(mComponents.getStorageSupplier(), clientInfo);
+
                 final MicrosoftStsAccount microsoftStsAccount = new MicrosoftStsAccount(
                         new IDToken(brokerResult.getIdToken()),
                         clientInfo
@@ -1421,4 +1428,40 @@ private void verifyBrokerVersionIsSupported(@Nullable final Bundle resultBundle,
                     "So, this is not likely a broker version supported issue. Continuing.");
         }
     }
+
+    /**
+     * Check if have received broker version not supported error.
+     *
+     * @param brokerResult                  Broker result from a broker operation.
+     * @param requiredBrokerProtocolVersion Required broker protocol version sent in request.
+     * @throws UnsupportedBrokerException if result contains broker not supported error.
+     */
+    private void verifyBrokerVersionIsSupportedFromBrokerResult(@Nullable final BrokerResult brokerResult, @Nullable final String requiredBrokerProtocolVersion) throws UnsupportedBrokerException {
+        final String methodTag = TAG + ":verifyBrokerVersionIsSupported";
+
+        // check if result bundle contains unsupported broker version exception
+        if (!brokerResult.isSuccess()
+                && ErrorStrings.UNSUPPORTED_BROKER_VERSION_ERROR_CODE.equals(brokerResult.getErrorCode())) {
+            mHelloCache.saveHandshakeError(requiredBrokerProtocolVersion, CLIENT_MAX_PROTOCOL_VERSION);
+            throw new UnsupportedBrokerException(mActiveBrokerPackageName);
+        }
+    }
+
+    private void trackTelemetryRegionFromResultBundle(@NonNull final BrokerResult brokerResult, final boolean isSuccess) throws BaseException {
+        final String methodTag = TAG + ":trackTelemetryRegionFromResultBundle";
+        if (isSuccess) {
+            if (StringUtil.isNullOrEmpty(brokerResult.getClientInfo())) {
+                return;
+            }
+
+            try {
+                final ClientInfo clientInfo = new ClientInfo(brokerResult.getClientInfo());
+
+                // Store the telemetry region info in shared preferences
+                EUClaimStorageUtil.Companion.storeTelemetryRegionByTenant(mComponents.getStorageSupplier(), clientInfo);
+            } catch (Exception e) {
+                Logger.error(methodTag, "Exception while trying to store telemetry region from client info in broker result.", e);
+            }
+        }
+    }
 }

common/src/main/java/com/microsoft/identity/common/internal/result/MsalBrokerResultAdapter.java

@@ -802,9 +802,7 @@ public AuthorizationResult getDeviceCodeFlowAuthResultFromResultBundle(@NonNull
      * @throws BaseException
      * @throws ClientException
      */
-    public AcquireTokenResult getDeviceCodeFlowTokenResultFromResultBundle(@NonNull final Bundle resultBundle) throws BaseException, ClientException {
-
-        BrokerResult brokerResult = brokerResultFromBundle(resultBundle);
+    public AcquireTokenResult getDeviceCodeFlowTokenResultFromBrokerResultAndBundle(@NonNull final BrokerResult brokerResult, @NonNull final Bundle resultBundle) throws BaseException, ClientException {
         final Span span = OTelUtility.createSpan(SpanName.AcquireTokenDcfFetchToken.name());
 
         span.setAttribute(AttributeName.correlation_id.name(), brokerResult.getCorrelationId());

common4j/src/main/com/microsoft/identity/common/java/opentelemetry/AttributeName.java

@@ -27,6 +27,10 @@
  * NOTE : Any changes to this enum should also be made in the corresponding enum in Broker.
  */
 public enum AttributeName {
+    /**
+     * The tenant id for the home tenant of the account for which PRT is required.
+     */
+    tenant_id,
     /**
      * The length of the response body returned from network request.
      */

common4j/src/main/com/microsoft/identity/common/java/opentelemetry/EUClaimStorageUtil.kt

@@ -0,0 +1,81 @@
+//  Copyright (c) Microsoft Corporation.
+//  All rights reserved.
+//
+//  This code is licensed under the MIT License.
+//
+//  Permission is hereby granted, free of charge, to any person obtaining a copy
+//  of this software and associated documentation files(the "Software"), to deal
+//  in the Software without restriction, including without limitation the rights
+//  to use, copy, modify, merge, publish, distribute, sublicense, and / or sell
+//  copies of the Software, and to permit persons to whom the Software is
+//  furnished to do so, subject to the following conditions :
+//
+//  The above copyright notice and this permission notice shall be included in
+//  all copies or substantial portions of the Software.
+//
+//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+//  THE SOFTWARE.
+package com.microsoft.identity.common.java.opentelemetry
+
+import com.microsoft.identity.common.java.interfaces.IStorageSupplier
+import com.microsoft.identity.common.java.logging.Logger
+import com.microsoft.identity.common.java.providers.microsoft.azureactivedirectory.ClientInfo
+import com.microsoft.identity.common.java.util.StringUtil
+
+/**
+ * Utility class for storing telemetry region TDBR claims by tenant.
+ */
+class EUClaimStorageUtil {
+    companion object {
+        private val TAG = EUClaimStorageUtil::class.java.simpleName
+
+        /**
+         * Store telemetry region by tenant.
+         *
+         * @param clientInfo the client info containing tenant information and tdbr claim
+         */
+        fun storeTelemetryRegionByTenant(
+            supplier: IStorageSupplier,
+            clientInfo: ClientInfo
+        ) {
+            val methodTag = "$TAG:storeTelemetryRegionByTenant"
+            val tenantId = clientInfo.utid
+            val tdbrClaim = clientInfo.tdbrClaim
+
+            if (StringUtil.isNullOrEmpty(tenantId)) {
+                Logger.warn(
+                    methodTag,
+                    "tenantId is null or empty. Not storing telemetry region by tenant."
+                )
+                return
+            }
+
+            if (StringUtil.isNullOrEmpty(tdbrClaim)) {
+                Logger.warn(
+                    methodTag,
+                    "Received no tdbr claim, not storing anything in shared preferences.."
+                )
+                return
+            }
+
+            // Store the tdbr claim for a specific tenant ID
+            Logger.info(
+                methodTag,
+                "Storing telemetry region by tenant: $tenantId, TDBR Claim: $tdbrClaim"
+            )
+            val tdbrValueStore = supplier.getUnencryptedNameValueStore(
+                ClientInfo.TDBR_CLAIM,
+                String::class.java
+            )
+            tdbrValueStore.put(tenantId, tdbrClaim)
+
+            // Attach tenant id to the current span
+            SpanExtension.current().setAttribute(AttributeName.tenant_id.name, tenantId)
+        }
+    }
+}

common4j/src/main/com/microsoft/identity/common/java/providers/microsoft/azureactivedirectory/ClientInfo.java

@@ -44,6 +44,7 @@ public class ClientInfo implements Serializable {
 
     private static final String UNIQUE_IDENTIFIER = "uid";
     private static final String UNIQUE_TENANT_IDENTIFIER = "utid";
+    public static final String TDBR_CLAIM = "xms_tdbr";
     private static final long serialVersionUID = 3326461566190095403L;
 
     /**
@@ -56,6 +57,11 @@ public class ClientInfo implements Serializable {
      */
     private String mUtid;
 
+    /**
+     * TDBR Claim, denotes what region the user belongs to.
+     */
+    private String mTdbrClaim;
+
     private String mRawClientInfo;
 
     /**
@@ -80,6 +86,7 @@ public ClientInfo(@NonNull String rawClientInfo) throws ServiceException {
 
         mUid = clientInfoItems.get(ClientInfo.UNIQUE_IDENTIFIER);
         mUtid = clientInfoItems.get(ClientInfo.UNIQUE_TENANT_IDENTIFIER);
+        mTdbrClaim = clientInfoItems.get(ClientInfo.TDBR_CLAIM);
         mRawClientInfo = rawClientInfo;
     }
 
@@ -101,6 +108,15 @@ public String getUtid() {
         return mUtid;
     }
 
+    /**
+     * Gets the TDBR claim.
+     *
+     * @return The TDBR claim to get.
+     */
+    public String getTdbrClaim() {
+        return mTdbrClaim;
+    }
+
     /**
      * Returns the raw String underlying this object.
      *