chore(tracing): Store incoming baggage org ID

sentry-core/src/performance/headers.rs

@@ -3,17 +3,23 @@
 use std::error::Error;
 use std::fmt::{Display, Formatter, Result as FmtResult};
 
+use sentry_types::protocol::v7::OrganizationId;
+
 use crate::protocol::{SpanId, TraceId};
 
 /// A key-value header pair.
 type Header<'h> = (&'h str, &'h str);
 
+/// The baggage key for the Sentry org ID.
+const SENTRY_ORG_ID: &str = "sentry-org_id";
+
 /// The [trace propagation] context.
 ///
 /// Contains the information necessary for propagating Sentry traces and continuing traces from
 /// incoming requests.
 ///
-/// The data stored in this struct can be parsed from and transmitted as `sentry-trace` headers.
+/// The data stored in this struct can be parsed from and transmitted as `sentry-trace` and Sentry
+/// baggage headers.
 ///
 /// Note that the Rust SDK only partially supports trace propagation, certain features such as
 /// [dynamic sampling] may be missing or incomplete.
@@ -25,6 +31,7 @@ pub struct TracePropagationContext {
     pub(crate) trace_id: TraceId,
     pub(crate) span_id: SpanId,
     pub(super) sampled: Option<bool>,
+    pub(super) org_id: Option<OrganizationId>,
 }
 
 #[derive(Debug, Clone)]
@@ -53,6 +60,7 @@ impl TracePropagationContext {
             trace_id,
             span_id,
             sampled: None,
+            org_id: None,
         }
     }
 
@@ -68,6 +76,7 @@ impl TracePropagationContext {
             trace_id,
             span_id,
             sampled,
+            org_id: _,
         } = self;
 
         let sampled_suffix = sampled
@@ -85,6 +94,7 @@ impl TracePropagationContext {
         I: IntoIterator<Item = Header<'a>>,
     {
         let mut context_result = Err(HeaderParseError::Missing);
+        let mut baggage = SentryBaggage::default();
 
         for (header, value) in headers {
             if header.eq_ignore_ascii_case("sentry-trace") {
@@ -93,10 +103,15 @@ impl TracePropagationContext {
                 context_result = TracePropagationContext::from_sentry_trace(value)
                     .map_or(context_result, Ok)
                     .map_err(|_| HeaderParseError::Invalid);
+            } else if header.eq_ignore_ascii_case("baggage") {
+                baggage.update_from_header(value);
             }
         }
 
-        context_result
+        let context = context_result?;
+
+        let SentryBaggage { org_id } = baggage;
+        Ok(TracePropagationContext { org_id, ..context })
     }
 
     /// Attempts to construct a [`TracePropagationContext`] from the given Sentry trace header.
@@ -118,6 +133,7 @@ impl TracePropagationContext {
             trace_id,
             span_id,
             sampled,
+            org_id: None,
         })
     }
 }
@@ -134,6 +150,7 @@ where
         trace_id,
         span_id,
         sampled,
+        org_id: _,
     } = TracePropagationContext::try_from_headers(headers).ok()?;
 
     Some(SentryTrace {
@@ -175,6 +192,7 @@ impl From<SentryTrace> for TracePropagationContext {
             trace_id: trace.trace_id,
             span_id: trace.span_id,
             sampled: trace.sampled,
+            org_id: None,
         }
     }
 }
@@ -190,6 +208,42 @@ impl std::fmt::Display for SentryTrace {
     }
 }
 
+/// A struct containing known Sentry baggage values.
+///
+/// For now, this only includes the `org_id`, but we can add more values as we support them.
+#[derive(Debug, Default)]
+struct SentryBaggage {
+    org_id: Option<OrganizationId>,
+}
+
+impl SentryBaggage {
+    /// Update `self` with the known Sentry baggage values in the provided [baggage header].
+    ///
+    /// The header is parsed according to the W3C baggage format: entries are separated by
+    /// commas, each entry is a key-value pair separated by `=`, and optional properties after
+    /// a semicolon are ignored.
+    ///
+    /// [baggage header]: https://www.w3.org/TR/baggage/
+    fn update_from_header(&mut self, value: &str) {
+        value
+            .split(',')
+            .flat_map(|s| s.split_once('='))
+            // Discard optional values after semicolon.
+            .map(|(key, value)| (key, value.split_once(';').map_or(value, |(v, _)| v)))
+            .map(|(key, value)| (key.trim(), value.trim()))
+            .for_each(|(key, value)| self.update_from_value(key, value))
+    }
+
+    /// Update `self` with a key-value pair from the baggage header.
+    ///
+    /// The value is only updated if it is valid, otherwise the old value is kept.
+    fn update_from_value(&mut self, key: &str, value: &str) {
+        if key == SENTRY_ORG_ID {
+            self.org_id = value.parse().ok().or(self.org_id);
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -210,6 +264,7 @@ mod tests {
                 trace_id,
                 span_id: parent_trace_id,
                 sampled: Some(false),
+                org_id: None,
             }
         );
 
@@ -221,4 +276,63 @@ mod tests {
         .expect("should parse successfully");
         assert_eq!(parsed, trace);
     }
+
+    #[test]
+    fn parses_baggage_org_id() {
+        let trace = TracePropagationContext::try_from_headers([
+            (
+                "sentry-trace",
+                "09e04486820349518ac7b5d2adbf6ba5-9cf635fa5b870b3a-0",
+            ),
+            ("baggage", "sentry-org_id=123"),
+        ])
+        .expect("should parse successfully");
+
+        assert_eq!(trace.org_id, Some("123".parse().unwrap()));
+    }
+
+    #[test]
+    fn parses_baggage_org_id_with_unrelated_fields() {
+        let trace = TracePropagationContext::try_from_headers([
+            (
+                "sentry-trace",
+                "09e04486820349518ac7b5d2adbf6ba5-9cf635fa5b870b3a-0",
+            ),
+            (
+                "baggage",
+                "other=value, sentry-org_id=123, another=value;property",
+            ),
+        ])
+        .expect("should parse successfully");
+
+        assert_eq!(trace.org_id, Some("123".parse().unwrap()));
+    }
+
+    #[test]
+    fn accepts_mixed_case_baggage_header_name() {
+        let trace = TracePropagationContext::try_from_headers([
+            (
+                "sentry-trace",
+                "09e04486820349518ac7b5d2adbf6ba5-9cf635fa5b870b3a-0",
+            ),
+            ("BagGaGe", "sentry-org_id=123"),
+        ])
+        .expect("should parse successfully");
+
+        assert_eq!(trace.org_id, Some("123".parse().unwrap()));
+    }
+
+    #[test]
+    fn treats_malformed_baggage_org_id_as_absent() {
+        let trace = TracePropagationContext::try_from_headers([
+            (
+                "sentry-trace",
+                "09e04486820349518ac7b5d2adbf6ba5-9cf635fa5b870b3a-0",
+            ),
+            ("baggage", "sentry-org_id=not-an-org-id"),
+        ])
+        .expect("should parse successfully");
+
+        assert_eq!(trace.org_id, None);
+    }
 }

sentry-core/src/performance/mod.rs

@@ -6,7 +6,7 @@ use std::time::SystemTime;
 
 #[cfg(feature = "client")]
 use sentry_types::protocol::v7::client_report::Reason as ClientReportReason;
-use sentry_types::protocol::v7::SpanId;
+use sentry_types::protocol::v7::{OrganizationId, SpanId};
 
 #[cfg(feature = "client")]
 use crate::clientoptions::TracesSamplingStrategy;
@@ -116,6 +116,11 @@ pub struct TransactionContext {
     parent_span_id: Option<protocol::SpanId>,
     span_id: protocol::SpanId,
     sampled: Option<bool>,
+    #[cfg_attr(
+        not(test),
+        expect(dead_code, reason = "used by future strict trace continuation")
+    )]
+    incoming_org_id: Option<OrganizationId>,
     custom: Option<CustomTransactionContext>,
 }
 
@@ -151,6 +156,7 @@ impl TransactionContext {
             parent_span_id: None,
             span_id: Default::default(),
             sampled: None,
+            incoming_org_id: None,
             custom: None,
         }
     }
@@ -197,6 +203,7 @@ impl TransactionContext {
                 parent_span_id: None,
                 span_id: Default::default(),
                 sampled: None,
+                incoming_org_id: None,
                 custom: None,
             })
     }
@@ -227,6 +234,7 @@ impl TransactionContext {
             trace_id,
             span_id: context_span_id,
             sampled,
+            org_id,
         } = context;
 
         Self {
@@ -235,6 +243,7 @@ impl TransactionContext {
             trace_id,
             parent_span_id: Some(context_span_id),
             sampled,
+            incoming_org_id: org_id,
             span_id: span_id.unwrap_or_default(),
             custom: None,
         }
@@ -274,6 +283,7 @@ impl TransactionContext {
             parent_span_id: Some(parent_span_id),
             span_id: protocol::SpanId::default(),
             sampled,
+            incoming_org_id: None,
             custom: None,
         }
     }
@@ -1313,6 +1323,35 @@ mod tests {
         assert_eq!(ctx.sampled(), Some(true));
     }
 
+    #[test]
+    fn continue_from_headers_stores_incoming_org_id() {
+        let ctx = TransactionContext::continue_from_headers(
+            "noop",
+            "noop",
+            [
+                (
+                    "sentry-trace",
+                    "09e04486820349518ac7b5d2adbf6ba5-9cf635fa5b870b3a-1",
+                ),
+                ("baggage", "sentry-org_id=123"),
+            ],
+        );
+
+        assert_eq!(ctx.incoming_org_id, Some("123".parse().unwrap()));
+    }
+
+    #[test]
+    fn continue_from_headers_does_not_keep_org_id_without_sentry_trace() {
+        let ctx = TransactionContext::continue_from_headers(
+            "noop",
+            "noop",
+            [("baggage", "sentry-org_id=123")],
+        );
+
+        assert_eq!(ctx.incoming_org_id, None);
+        assert_eq!(ctx.parent_span_id, None);
+    }
+
     #[cfg(feature = "client")]
     #[test]
     fn compute_transaction_sample_rate() {