feat: Implement strict_trace_continuation

sentry_sdk/consts.py

@@ -1024,6 +1024,7 @@ def __init__(
         enable_metrics=True,  # type: bool
         before_send_metric=None,  # type: Optional[Callable[[Metric, Hint], Optional[Metric]]]
         org_id=None,  # type: Optional[str]
+        strict_trace_continuation=False,  # type: bool
     ):
         # type: (...) -> None
         """Initialize the Sentry SDK with the given parameters. All parameters described here can be used in a call to `sentry_sdk.init()`.
@@ -1427,6 +1428,15 @@ def __init__(
             If `trace_ignore_status_codes` is not provided, requests with any status code
             may be traced.
 
+        :param strict_trace_continuation: If set to `True`, the SDK will only continue a trace if the `org_id` of the incoming trace found in the
+           `baggage` header matches the `org_id` of the current Sentry client and only if BOTH are present.
+
+            If set to `False`, consistency of `org_id` will only be enforced if both are present. If either are missing, the trace will be continued.
+
+            The client's organization ID is extracted from the DSN or can be set with the `org_id` option.
+            If the organization IDs do not match, the SDK will start a new trace instead of continuing the incoming one.
+            This is useful to prevent traces of unknown third-party services from being continued in your application.
+
         :param org_id: An optional organization ID. The SDK will try to extract if from the DSN in most cases
             but you can provide it explicitly for self-hosted and Relay setups. This value is used for
             trace propagation and for features like `strict_trace_continuation`.

sentry_sdk/tracing_utils.py

@@ -15,7 +15,6 @@
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     filename_for_module,
-    Dsn,
     logger,
     match_regex_list,
     qualname_from_function,
@@ -453,15 +452,23 @@ def from_incoming_data(cls, incoming_data):
 
         sentry_trace_header = normalized_data.get(SENTRY_TRACE_HEADER_NAME)
         sentrytrace_data = extract_sentrytrace_data(sentry_trace_header)
+
+        # nothing to propagate if no sentry-trace
         if sentrytrace_data is None:
             return None
 
+        baggage_header = normalized_data.get(BAGGAGE_HEADER_NAME)
+        baggage = (
+            Baggage.from_incoming_header(baggage_header) if baggage_header else None
+        )
+
+        if not _should_continue_trace(baggage):
+            return None
+
         propagation_context = PropagationContext()
         propagation_context.update(sentrytrace_data)
-
-        baggage_header = normalized_data.get(BAGGAGE_HEADER_NAME)
-        if baggage_header:
-            propagation_context.baggage = Baggage.from_incoming_header(baggage_header)
+        if baggage:
+            propagation_context.baggage = baggage
 
         propagation_context._fill_sample_rand()
 
@@ -1230,6 +1237,41 @@ def _set_output_attributes(span, template, send_pii, result):
     span.update_data(_get_output_attributes(template, send_pii, result) or {})
 
 
+def _should_continue_trace(baggage):
+    # type: (Optional[Baggage]) -> bool
+    """
+    Check if we should continue the incoming trace according to the strict_trace_continuation spec.
+    https://develop.sentry.dev/sdk/telemetry/traces/#stricttracecontinuation
+    """
+
+    client = sentry_sdk.get_client()
+    parsed_dsn = client.parsed_dsn
+    client_org_id = parsed_dsn.org_id if parsed_dsn else None
+    baggage_org_id = baggage.sentry_items.get("org_id") if baggage else None
+
+    if (
+        client_org_id is not None
+        and baggage_org_id is not None
+        and client_org_id != baggage_org_id
+    ):
+        logger.debug(
+            f"Starting a new trace because org IDs don't match (incoming baggage org_id: {baggage_org_id}, SDK org_id: {client_org_id})"
+        )
+        return False
+
+    strict_trace_continuation = client.options.get("strict_trace_continuation", False)  # type: bool
+    if strict_trace_continuation:
+        if (baggage_org_id is not None and client_org_id is None) or (
+            baggage_org_id is None and client_org_id is not None
+        ):
+            logger.debug(
+                f"Starting a new trace because strict trace continuation is enabled and one org ID is missing (incoming baggage org_id: {baggage_org_id}, SDK org_id: {client_org_id})"
+            )
+            return False
+
+    return True
+
+
 # Circular imports
 from sentry_sdk.tracing import (
     BAGGAGE_HEADER_NAME,

tests/conftest.py

@@ -221,6 +221,19 @@ def capture_envelope(self, _: Envelope) -> None:
         pass
 
 
+class TestTransportWithOptions(Transport):
+    """TestTransport above does not pass in the options and for some tests we need them"""
+
+    __test__ = False
+
+    def __init__(self, options=None):
+        Transport.__init__(self, options)
+
+    def capture_envelope(self, _: Envelope) -> None:
+        """No-op capture_envelope for tests"""
+        pass
+
+
 @pytest.fixture
 def capture_events(monkeypatch):
     def inner():

tests/test_dsc.py

@@ -13,19 +13,7 @@
 import pytest
 
 import sentry_sdk
-from sentry_sdk.transport import Transport
-from sentry_sdk.envelope import Envelope
-
-
-class TransportWithOptions(Transport):
-    """conftest.TestTransport does not pass in the options so we need this here"""
-
-    def __init__(self, options=None):
-        Transport.__init__(self, options)
-
-    def capture_envelope(self, _: Envelope) -> None:
-        """No-op capture_envelope for tests"""
-        pass
+from tests.conftest import TestTransportWithOptions
 
 
 def test_dsc_head_of_trace(sentry_init, capture_envelopes):
@@ -38,7 +26,7 @@ def test_dsc_head_of_trace(sentry_init, capture_envelopes):
         release="[email protected]",
         environment="canary",
         traces_sample_rate=1.0,
-        transport=TransportWithOptions,
+        transport=TestTransportWithOptions,
     )
     envelopes = capture_envelopes()
 
@@ -94,7 +82,7 @@ def test_dsc_head_of_trace_uses_custom_org_id(sentry_init, capture_envelopes):
         release="[email protected]",
         environment="canary",
         traces_sample_rate=1.0,
-        transport=TransportWithOptions,
+        transport=TestTransportWithOptions,
     )
     envelopes = capture_envelopes()
 
@@ -122,7 +110,7 @@ def test_dsc_continuation_of_trace(sentry_init, capture_envelopes):
         release="[email protected]",
         environment="canary",
         traces_sample_rate=1.0,
-        transport=TransportWithOptions,
+        transport=TestTransportWithOptions,
     )
     envelopes = capture_envelopes()
 
@@ -200,7 +188,7 @@ def my_traces_sampler(sampling_context):
         release="[email protected]",
         environment="canary",
         traces_sampler=my_traces_sampler,
-        transport=TransportWithOptions,
+        transport=TestTransportWithOptions,
     )
     envelopes = capture_envelopes()
 
@@ -270,7 +258,7 @@ def test_dsc_issue(sentry_init, capture_envelopes):
         dsn="https://[email protected]/12312012",
         release="[email protected]",
         environment="canary",
-        transport=TransportWithOptions,
+        transport=TestTransportWithOptions,
     )
     envelopes = capture_envelopes()
 
@@ -322,7 +310,7 @@ def test_dsc_issue_with_tracing(sentry_init, capture_envelopes):
         release="[email protected]",
         environment="canary",
         traces_sample_rate=1.0,
-        transport=TransportWithOptions,
+        transport=TestTransportWithOptions,
     )
     envelopes = capture_envelopes()
 
@@ -394,7 +382,7 @@ def test_dsc_issue_twp(sentry_init, capture_envelopes, traces_sample_rate):
         release="[email protected]",
         environment="canary",
         traces_sample_rate=traces_sample_rate,
-        transport=TransportWithOptions,
+        transport=TestTransportWithOptions,
     )
     envelopes = capture_envelopes()
 

tests/test_tracing_utils.py

@@ -1,8 +1,13 @@
+import pytest
 from dataclasses import asdict, dataclass
 from typing import Optional, List
 
-from sentry_sdk.tracing_utils import _should_be_included, Baggage
-import pytest
+from sentry_sdk.tracing_utils import (
+    _should_be_included,
+    _should_continue_trace,
+    Baggage,
+)
+from tests.conftest import TestTransportWithOptions
 
 
 def id_function(val):
@@ -146,3 +151,130 @@ def test_strip_sentry_baggage(header, expected):
 )
 def test_baggage_repr(baggage, expected_repr):
     assert repr(baggage) == expected_repr
+
+
+@pytest.mark.parametrize(
+    (
+        "baggage_header",
+        "dsn",
+        "explicit_org_id",
+        "strict_trace_continuation",
+        "should_continue_trace",
+    ),
+    (
+        # continue cases when strict_trace_continuation=False
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700, sentry-org_id=1234",
+            "https://[email protected]/12312012",
+            None,
+            False,
+            True,
+        ),
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700",
+            "https://[email protected]/12312012",
+            None,
+            False,
+            True,
+        ),
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700, sentry-org_id=1234",
+            None,
+            None,
+            False,
+            True,
+        ),
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700, sentry-org_id=1234",
+            None,
+            "1234",
+            False,
+            True,
+        ),
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700, sentry-org_id=1234",
+            "https://mysecret@not_org_id.ingest.sentry.io/12312012",
+            None,
+            False,
+            True,
+        ),
+        # start new cases when strict_trace_continuation=False
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700, sentry-org_id=1234",
+            "https://[email protected]/12312012",
+            None,
+            False,
+            False,
+        ),
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700, sentry-org_id=1234",
+            "https://[email protected]/12312012",
+            "9999",
+            False,
+            False,
+        ),
+        # continue cases when strict_trace_continuation=True
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700, sentry-org_id=1234",
+            "https://[email protected]/12312012",
+            None,
+            True,
+            True,
+        ),
+        ("sentry-trace_id=771a43a4192642f0b136d5159a501700", None, None, True, True),
+        # start new cases when strict_trace_continuation=True
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700",
+            "https://[email protected]/12312012",
+            None,
+            True,
+            False,
+        ),
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700, sentry-org_id=1234",
+            None,
+            None,
+            True,
+            False,
+        ),
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700, sentry-org_id=1234",
+            "https://mysecret@not_org_id.ingest.sentry.io/12312012",
+            None,
+            True,
+            False,
+        ),
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700, sentry-org_id=1234",
+            "https://[email protected]/12312012",
+            None,
+            True,
+            False,
+        ),
+        (
+            "sentry-trace_id=771a43a4192642f0b136d5159a501700, sentry-org_id=1234",
+            "https://[email protected]/12312012",
+            "9999",
+            True,
+            False,
+        ),
+    ),
+)
+def test_should_continue_trace(
+    sentry_init,
+    baggage_header,
+    dsn,
+    explicit_org_id,
+    strict_trace_continuation,
+    should_continue_trace,
+):
+    sentry_init(
+        dsn=dsn,
+        org_id=explicit_org_id,
+        strict_trace_continuation=strict_trace_continuation,
+        traces_sample_rate=1.0,
+        transport=TestTransportWithOptions,
+    )
+
+    baggage = Baggage.from_incoming_header(baggage_header) if baggage_header else None
+    assert _should_continue_trace(baggage) == should_continue_trace