Using Baggage to propagate attributes from parent Span, Fixes AB#3262849 (#2671)

### Summary
The passkey telemetry dashboard (and maybe some others) has had issues
with latency due to joins being done between the Fido span and the
parent span related to acquire token interactive. These join statements
by trace id allowed us to associate attributes (such as tenant id) to
the Fido span that weren't available in the context of the Fido classes.
In order to directly save the attributes onto the Fido span itself, we
can make use of [Otel's
Baggage](https://opentelemetry.io/docs/specs/otel/baggage/api/), which
is meant to help hold data across multiple spans (from parent to
children). We can then use an implementation of [Otel's
TextMapPropagator](https://opentelemetry.io/docs/specs/otel/context/api-propagators/)
to serialize the Otel Context and pass it to the AuthenticationActivity
as an Extra (similar to what we are doing for spanContext). This Context
can then be accessed by the WebViewClient and Fido related classes,
where they will retrieve the Baggage.

This PR has the changes to implement this behavior and adds a few static
helper classes for TextMapPropagator and Baggage (the baggage class is
more relevant for the corresponding Broker PR:
AzureAD/ad-accounts-for-android#3119).

While this PR tackles attribute propagation for AuthorizationActivity
specifically, Baggage and the helper classes can be used in other areas
of the code where propagation is needed as well.


[AB#3262849](https://identitydivision.visualstudio.com/Engineering/_workitems/edit/3262849)

Author: melissaahn

changelog.txt

@@ -3,6 +3,7 @@ vNext
 - [PATCH] Fix caching of secret key and add retries for InvalidKeyException during unwrap (#2659)
 - [MINOR] Replace AbstractSecretKeyLoader with ISecretKeyProvider (#2666)
 - [MINOR] Update IP phone app teams signature constants to use SHA-512 format (#2700)
+- [MINOR] Using Baggage to propagate attributes from parent Span (#2671)
 - [PATCH] Fix a few small switch browser bugs (#2710)
 
 Version 21.4.0

common/src/main/java/com/microsoft/identity/common/adal/internal/AuthenticationConstants.java

@@ -2026,6 +2026,8 @@ public static final class AuthorizationIntentKey {
         public static final String WEB_VIEW_ZOOM_CONTROLS_ENABLED = "com.microsoft.identity.web.view.zoom.controls.enabled";
 
         public static final String WEB_VIEW_ZOOM_ENABLED = "com.microsoft.identity.web.view.zoom.enabled";
+
+        public static final String OTEL_CONTEXT_CARRIER = "otel_context_carrier";
     }
 
     public static final class AuthorizationIntentAction {

common/src/main/java/com/microsoft/identity/common/internal/fido/AuthFidoChallengeHandler.kt

@@ -30,12 +30,14 @@ import com.microsoft.identity.common.internal.ui.webview.challengehandlers.IChal
 import com.microsoft.identity.common.java.constants.FidoConstants
 import com.microsoft.identity.common.java.constants.FidoConstants.Companion.PASSKEY_PROTOCOL_ERROR_PREFIX_STRING
 import com.microsoft.identity.common.java.opentelemetry.AttributeName
+import com.microsoft.identity.common.java.opentelemetry.BaggageExtension
 import com.microsoft.identity.common.java.opentelemetry.OTelUtility
+import com.microsoft.identity.common.java.opentelemetry.SpanExtension
 import com.microsoft.identity.common.java.opentelemetry.SpanName
 import com.microsoft.identity.common.logging.Logger
 import io.opentelemetry.api.trace.Span
-import io.opentelemetry.api.trace.SpanContext
 import io.opentelemetry.api.trace.StatusCode
+import io.opentelemetry.context.Context
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.launch
 import java.util.*
@@ -51,18 +53,35 @@ import java.util.*
 class AuthFidoChallengeHandler (
     private val fidoManager: IFidoManager,
     private val webView: WebView,
-    private val spanContext : SpanContext?,
+    private val oTelContext: Context?,
     private val lifecycleOwner: LifecycleOwner?
 ) : IChallengeHandler<FidoChallenge, Void> {
     val TAG = AuthFidoChallengeHandler::class.simpleName.toString()
+    companion object {
+        private val parentAttributeNames = arrayListOf(
+            AttributeName.correlation_id,
+            AttributeName.tenant_id,
+            AttributeName.account_type,
+            AttributeName.calling_package_name
+        )
+    }
 
     override fun processChallenge(fidoChallenge: FidoChallenge): Void? {
         val methodTag = "$TAG:processChallenge"
         Logger.info(methodTag, "Processing FIDO challenge.")
-        val span = if (spanContext != null) {
-            OTelUtility.createSpanFromParent(SpanName.Fido.name, spanContext)
+        val span : Span
+        if (oTelContext != null) {
+            val parentSpan = SpanExtension.fromContext(oTelContext)
+            val spanContext = parentSpan.spanContext
+            span = OTelUtility.createSpanFromParent(SpanName.Fido.name, spanContext)
+            val baggage = BaggageExtension.fromContext(oTelContext)
+            parentAttributeNames.forEach { attributeName ->
+                baggage.getEntryValue(attributeName.name)?.let { value ->
+                    span.setAttribute(attributeName.name, value)
+                }
+            }
         } else {
-            OTelUtility.createSpan(SpanName.Fido.name)
+            span = OTelUtility.createSpan(SpanName.Fido.name)
         }
         span.setAttribute(
             AttributeName.fido_challenge_handler.name,

common/src/main/java/com/microsoft/identity/common/internal/providers/oauth2/AuthorizationActivity.java

@@ -22,6 +22,8 @@
 // THE SOFTWARE.
 package com.microsoft.identity.common.internal.providers.oauth2;
 
+import static com.microsoft.identity.common.adal.internal.AuthenticationConstants.AuthorizationIntentKey.OTEL_CONTEXT_CARRIER;
+
 import android.os.Bundle;
 
 import androidx.annotation.Nullable;
@@ -31,9 +33,14 @@
 import com.microsoft.identity.common.internal.util.CommonMoshiJsonAdapter;
 import com.microsoft.identity.common.java.exception.TerminalException;
 import com.microsoft.identity.common.java.opentelemetry.SerializableSpanContext;
+import com.microsoft.identity.common.java.opentelemetry.TextMapPropagatorExtension;
+import com.microsoft.identity.common.java.util.StringUtil;
 import com.microsoft.identity.common.logging.Logger;
 
+import java.util.HashMap;
+
 import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
 import lombok.Getter;
 import lombok.experimental.Accessors;
 
@@ -43,6 +50,11 @@ public class AuthorizationActivity extends DualScreenActivity {
     @Getter
     @Accessors(prefix = "m")
     private SpanContext mSpanContext;
+
+    @Getter
+    @Accessors(prefix = "m")
+    private Context mOtelContext;
+
     private AuthorizationFragment mFragment;
 
     public AuthorizationFragment getFragment() {
@@ -54,21 +66,32 @@ public void onCreate(@Nullable Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
 
         final String methodTag = TAG + ":onCreate";
-        if (getIntent().getExtras() != null) {
+        final Bundle bundle = getIntent().getExtras();
+        if (bundle != null) {
             try {
-                mSpanContext = new CommonMoshiJsonAdapter().fromJson(
-                        getIntent().getExtras().getString(SerializableSpanContext.SERIALIZABLE_SPAN_CONTEXT),
+                String spanContextJson = bundle.getString(SerializableSpanContext.SERIALIZABLE_SPAN_CONTEXT);
+                mSpanContext = StringUtil.isNullOrEmpty(spanContextJson) ? null : new CommonMoshiJsonAdapter().fromJson(
+                        spanContextJson,
                         SerializableSpanContext.class
                 );
-            } catch (final TerminalException e) {
-                // Don't want to block any features if an error occurs during deserialization of the span context.
-                mSpanContext = null;
+
+                final HashMap<String, String> carrier;
+                // For Android Tiramisu and above, we use the new Serializable interface, which is a bit safer because it performs type checking at the framework level.
+                if (android.os.Build.VERSION.SDK_INT >= android.os.Build.VERSION_CODES.TIRAMISU) {
+                    carrier = bundle.getSerializable(OTEL_CONTEXT_CARRIER, HashMap.class);
+                } else {
+                    carrier = (HashMap<String, String>) bundle.getSerializable(OTEL_CONTEXT_CARRIER);
+                }
+                mOtelContext = TextMapPropagatorExtension.extract(carrier);
+            } catch (final TerminalException | ClassCastException | NullPointerException e) {
+                // Don't want to block any features if an error occurs during deserialization.
+                Logger.error(methodTag, "Exception thrown during extraction: " + e.getMessage(), e);
             }
         }
         final Fragment fragment = AuthorizationActivityFactory.getAuthorizationFragmentFromStartIntent(getIntent());
         if (fragment instanceof AuthorizationFragment) {
             mFragment = (AuthorizationFragment) fragment;
-            mFragment.setInstanceState(getIntent().getExtras());
+            mFragment.setInstanceState(bundle);
         } else {
             final IllegalStateException ex = new IllegalStateException("Unexpected fragment type.");
             Logger.error(methodTag, "Did not receive AuthorizationFragment from factory", ex);

common/src/main/java/com/microsoft/identity/common/internal/providers/oauth2/AuthorizationActivityFactory.kt

@@ -26,6 +26,7 @@ import android.content.Intent
 import android.os.Bundle
 import androidx.fragment.app.Fragment
 import com.microsoft.identity.common.adal.internal.AuthenticationConstants
+import com.microsoft.identity.common.adal.internal.AuthenticationConstants.AuthorizationIntentKey.OTEL_CONTEXT_CARRIER
 import com.microsoft.identity.common.internal.msafederation.getIdProviderExtraQueryParamForAuthorization
 import com.microsoft.identity.common.internal.msafederation.getIdProviderHeadersForAuthorization
 import com.microsoft.identity.common.internal.msafederation.google.SignInWithGoogleApi.Companion.getInstance
@@ -39,8 +40,10 @@ import com.microsoft.identity.common.java.AuthenticationConstants.SdkPlatformFie
 import com.microsoft.identity.common.java.configuration.LibraryConfiguration
 import com.microsoft.identity.common.java.exception.ClientException
 import com.microsoft.identity.common.java.logging.DiagnosticContext
+import com.microsoft.identity.common.java.opentelemetry.OtelContextExtension
 import com.microsoft.identity.common.java.opentelemetry.SerializableSpanContext
 import com.microsoft.identity.common.java.opentelemetry.SpanExtension
+import com.microsoft.identity.common.java.opentelemetry.TextMapPropagatorExtension
 import com.microsoft.identity.common.java.ui.AuthorizationAgent
 import com.microsoft.identity.common.java.util.CommonURIBuilder
 import java.net.URISyntaxException
@@ -119,6 +122,10 @@ object AuthorizationActivityFactory {
                         .build()
                 )
             )
+            putExtra(
+                OTEL_CONTEXT_CARRIER,
+                TextMapPropagatorExtension.inject(OtelContextExtension.current())
+            )
             if (parameters.sourceLibraryName != null) {
                 putExtra(PRODUCT, parameters.sourceLibraryName)
             }

common/src/main/java/com/microsoft/identity/common/internal/ui/webview/AzureActiveDirectoryWebViewClient.java

@@ -104,6 +104,7 @@
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 
 /**
@@ -269,7 +270,7 @@ private boolean handleUrl(final WebView view, final String url) {
                 Logger.info(methodTag,"WebView detected request for passkey protocol.");
                 final FidoChallenge challenge = FidoChallenge.createFromRedirectUri(url);
                 final Activity currentActivity = getActivity();
-                final SpanContext spanContext = currentActivity instanceof AuthorizationActivity ? ((AuthorizationActivity)currentActivity).getSpanContext() : null;
+                final Context oTelContext = currentActivity instanceof AuthorizationActivity ? ((AuthorizationActivity)currentActivity).getOtelContext() : null;
                 // The legacyManager should only be getting added if the device is on Android 13 or lower and the library is MSAL/OneAuth with fragment or dialog mode.
                 // The legacyManager logic should be removed once a larger majority of users are on Android 14+.
                 final IFidoManager legacyManager =
@@ -285,7 +286,7 @@ private boolean handleUrl(final WebView view, final String url) {
                                 legacyManager
                         ),
                         view,
-                        spanContext,
+                        oTelContext,
                         ViewTreeLifecycleOwner.get(view));
                 challengeHandler.processChallenge(challenge);
             } else if (CommonFlightsManager.INSTANCE.getFlightsProvider().isFlightEnabled(CommonFlight.ENABLE_ATTACH_NEW_PRT_HEADER_WHEN_NONCE_EXPIRED) && isNonceRedirect(formattedURL)) {

common/src/test/java/com/microsoft/identity/common/internal/fido/AuthFidoChallengeHandlerTest.kt

@@ -66,7 +66,7 @@ class AuthFidoChallengeHandlerTest {
         authFidoChallengeHandler = AuthFidoChallengeHandler(
             fidoManager = testFidoManager,
             webView = webView,
-            spanContext = null,
+            oTelContext = null,
             lifecycleOwner = testLifecycleOwner
         )
     }
@@ -126,7 +126,7 @@ class AuthFidoChallengeHandlerTest {
         authFidoChallengeHandler = AuthFidoChallengeHandler(
             fidoManager = testFidoManager,
             webView = webView,
-            spanContext = null,
+            oTelContext = null,
             lifecycleOwner = null
         )
         assertFalse(webView.urlLoaded)

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

@@ -365,6 +365,29 @@ public enum AttributeName {
      */
     is_switch_browser_resume_handled,
 
+    /**
+     * The tenant id for the home tenant of the account for which PRT is required.
+     */
+    tenant_id,
+
+    /**
+     * Indicates the type of account such as AAD or MSA.
+     */
+    account_type,
+
+    /**
+     * Indicates the broker app that emits the event.
+     * The broker is not necessarily the active broker.
+     * e.g. An inactive broker app might be invoked during OnUpgrade.
+     * (It should be renamed, but that would mess up the dashboard)
+     */
+    active_broker_package_name,
+
+    /**
+     * Indicates the current broker package name processing the request.
+     */
+    current_broker_package_name,
+
     /**
      * Records if the request is a webcp authorize request.
      */

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

@@ -0,0 +1,77 @@
+// 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.logging.Logger;
+
+import io.opentelemetry.api.baggage.Baggage;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+
+/**
+ * Utility class for working with OpenTelemetry Baggage objects.
+ */
+public class BaggageExtension {
+
+    private static final String TAG = BaggageExtension.class.getSimpleName();
+
+    /**
+     * Default constructor.
+     */
+    private BaggageExtension() {
+        // Utility class, private constructor to prevent instantiation
+    }
+
+    /**
+     * Makes the provided Baggage current in the context, catching any exceptions silently.
+     * This is useful in scenarios where Baggage propagation should not interrupt normal operation flow.
+     *
+     * @param baggage The Baggage to make current.
+     * @return the resulting scope.
+     */
+    public static Scope makeBaggageCurrent(final Baggage baggage) {
+        try {
+            if (baggage == null) {
+                return SpanExtension.NoopScope.INSTANCE;
+            }
+            return baggage.storeInContext(Context.current()).makeCurrent();
+        } catch (final Exception e) {
+            Logger.error(TAG + ":makeBaggageCurrent", "Failed to make baggage current", e);
+            return SpanExtension.NoopScope.INSTANCE;
+        }
+    }
+
+    /**
+     * Returns the current Baggage from the context, or a NoopBaggage if an error occurs.
+     *
+     * @return the current Baggage.
+     */
+    public static Baggage fromContext(final Context context) {
+        try {
+            return Baggage.fromContext(context);
+        } catch (final Exception e) {
+            Logger.error(TAG + ":fromContext", "Failed to get baggage from context", e);
+            return new NoopBaggage();
+        }
+    }
+}