fix(boto3): Fix botocore SigV4 failures caused by post-sign trace propagation

sentry_sdk/integrations/boto3.py

@@ -6,8 +6,13 @@
 from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version
 from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.traces import StreamedSpan
-from sentry_sdk.tracing import Span
-from sentry_sdk.tracing_utils import add_http_breadcrumb, has_span_streaming_enabled
+from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span
+from sentry_sdk.tracing_utils import (
+    add_http_breadcrumb,
+    add_sentry_baggage_to_headers,
+    has_span_streaming_enabled,
+    should_propagate_trace,
+)
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     parse_url,
@@ -49,6 +54,8 @@ def sentry_patched_init(
                 "request-created",
                 partial(_sentry_request_created, service_id=service_id),
             )
+            # run after other `before-sign` handlers, allowing it to see and preserve existing baggage.
+            meta.events.register_last("before-sign", _sentry_before_sign)
             meta.events.register("after-call", _sentry_after_call)
             meta.events.register("after-call-error", _sentry_after_call_error)
 
@@ -143,6 +150,60 @@ def _sentry_request_created(
         request.context["_sentrysdk_span"] = span
 
 
+def _sentry_before_sign(
+    request: "AWSRequest", signature_version: "Any", **kwargs: "Any"
+) -> None:
+    client = sentry_sdk.get_client()
+    if client.get_integration(Boto3Integration) is None:
+        return
+
+    with capture_internal_exceptions():
+        # presigned requests are executed later by another caller. Adding propagation
+        # headers here would make those headers part of the signature, requiring the caller to reproduce the same values.
+        if isinstance(signature_version, str) and signature_version.endswith(
+            ("-query", "-presign-post")
+        ):
+            return
+
+        if request.url is None or not should_propagate_trace(client, request.url):
+            return
+
+        def _replace_header(request: "AWSRequest", key: str, value: str) -> None:
+            """
+            Botocore's `HTTPHeaders` inherits from `email.message.Message`, where:
+                headers["foo"] = "old"
+                headers["foo"] = "new"
+            produces two fields: {"foo": "old", "foo": "new"}. So delete existing
+            fields before assigning replacement.
+            """
+            if key in request.headers:
+                del request.headers[key]
+            request.headers[key] = value
+
+        # use span associated with this botocore request
+        span = request.context.get("_sentrysdk_span")
+
+        headers = sentry_sdk.get_current_scope().iter_trace_propagation_headers(
+            span=span
+        )
+        for header_name, header_value in headers:
+            if header_name != BAGGAGE_HEADER_NAME:
+                # normal headers (e.g. `sentry-trace`) are non-shared, so replace stale values
+                _replace_header(request, header_name, header_value)
+                continue
+
+            # merge existing `baggage` values under single header
+            existing_values = request.headers.get_all(BAGGAGE_HEADER_NAME, [])
+            combined_baggage = {
+                BAGGAGE_HEADER_NAME: ",".join(str(value) for value in existing_values)
+            }
+            # preserve third-party baggage, replace stale `sentry-*` values
+            add_sentry_baggage_to_headers(combined_baggage, header_value)
+            _replace_header(
+                request, BAGGAGE_HEADER_NAME, combined_baggage[BAGGAGE_HEADER_NAME]
+            )
+
+
 def _sentry_after_call(
     context: "Dict[str, Any]", parsed: "Dict[str, Any]", **kwargs: "Any"
 ) -> None:

sentry_sdk/integrations/stdlib.py

@@ -10,7 +10,7 @@
 from sentry_sdk.integrations import Integration
 from sentry_sdk.scope import add_global_event_processor, should_send_default_pii
 from sentry_sdk.traces import StreamedSpan
-from sentry_sdk.tracing import Span
+from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, SENTRY_TRACE_HEADER_NAME, Span
 from sentry_sdk.tracing_utils import (
     EnvironHeaders,
     add_http_breadcrumb,
@@ -20,6 +20,8 @@
 )
 from sentry_sdk.utils import (
     SENSITIVE_DATA_SUBSTITUTE,
+    _get_aws_sigv4_signed_headers_from_authorization_header,
+    _get_aws_sigv4_signed_headers_from_url_query_string,
     capture_internal_exceptions,
     ensure_integration_enabled,
     is_sentry_url,
@@ -29,7 +31,7 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Dict, List, Optional, Union
+    from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
     from sentry_sdk._types import Event, Hint
 
@@ -40,6 +42,14 @@
     "build": sys.version,
 }
 
+_SENTRY_HEADER_NAMES = frozenset((BAGGAGE_HEADER_NAME, SENTRY_TRACE_HEADER_NAME))
+
+try:
+    from botocore.awsrequest import AWSHTTPConnection, AWSHTTPSConnection
+except ImportError:
+    AWSHTTPConnection = None  # type: ignore[misc,assignment]
+    AWSHTTPSConnection = None  # type: ignore[misc,assignment]
+
 
 class StdlibIntegration(Integration):
     identifier = "stdlib"
@@ -73,7 +83,178 @@
             add_http_request_source(span)
 
 
+def _get_wrapped_putheader(
+    original_putheader: "Callable[..., Any]",
+) -> "Callable[..., Any]":
+    """
+    Responsible for tracking which sentry headers are present and whether
+    they are listed in AWS SigV4 `SignedHeaders`.
+    """
+
+    def putheader(self: "HTTPConnection", header: "Any", *values: "Any") -> "Any":
+        rv = original_putheader(self, header, *values)
+
+        request_headers: "Optional[Dict[str, Tuple[bool, bool]]]" = getattr(
+            self, "_sentrysdk_request_headers", None
+        )
+        if request_headers is None:
+            return rv
+
+        if isinstance(header, bytes):
+            normalized_header = header.decode("ascii", "ignore").lower()
+        elif isinstance(header, str):
+            normalized_header = header.lower()
+        else:
+            return rv
+
+        if normalized_header in _SENTRY_HEADER_NAMES:
+            _, is_signed = request_headers.get(normalized_header, (False, False))
+            request_headers[normalized_header] = (True, is_signed)
+
+        if normalized_header == "authorization" and values:
+            with capture_internal_exceptions():
+                authorization = values[0]
+                if isinstance(authorization, bytes):
+                    authorization = authorization.decode("latin-1")
+                for (
+                    signed_header
+                ) in _get_aws_sigv4_signed_headers_from_authorization_header(
+                    authorization
+                ):
+                    if signed_header in _SENTRY_HEADER_NAMES:
+                        is_present, _ = request_headers.get(
+                            signed_header, (False, False)
+                        )
+                        request_headers[signed_header] = (is_present, True)
+
+        return rv
+
+    return putheader
+
+
+def _get_wrapped_endheaders(
+    original_endheaders: "Callable[..., Any]",
+) -> "Callable[..., Any]":
+    """
+    Responsible for injecting trace propagation headers, ensuring that the request is not invalidated
+    by honoring signed headers.
+    """
+
+    def endheaders(self: "HTTPConnection", *args: "Any", **kwargs: "Any") -> "Any":
+        real_url = getattr(self, "_sentrysdk_trace_url", None)
+        span = getattr(self, "_sentrysdk_span", None)
+
+        try:
+            if real_url is not None:
+                with capture_internal_exceptions():
+                    request_headers: "Optional[Dict[str, Tuple[bool, bool]]]" = getattr(
+                        self, "_sentrysdk_request_headers", {}
+                    )
+                    if request_headers is not None:
+                        for (
+                            signed_header
+                        ) in _get_aws_sigv4_signed_headers_from_url_query_string(
+                            real_url
+                        ):
+                            if signed_header in _SENTRY_HEADER_NAMES:
+                                is_present, _ = request_headers.get(
+                                    signed_header, (False, False)
+                                )
+                                request_headers[signed_header] = (is_present, True)
+
+                        for (
+                            header_name,
+                            header_value,
+                        ) in sentry_sdk.get_current_scope().iter_trace_propagation_headers(
+                            span=span
+                        ):
+                            normalized_header = header_name.lower()
+                            # preserve signed headers and avoid duplicate `sentry-trace`.
+                            is_present, is_signed = request_headers.get(
+                                normalized_header, (False, False)
+                            )
+                            if is_signed or (
+                                is_present and normalized_header != BAGGAGE_HEADER_NAME
+                            ):
+                                continue
+
+                            logger.debug(
+                                "[Tracing] Adding `{key}` header {value} to outgoing request to {real_url}.".format(
+                                    key=header_name,
+                                    value=header_value,
+                                    real_url=real_url,
+                                )
+                            )
+                            self.putheader(header_name, header_value)
+            return original_endheaders(self, *args, **kwargs)
+        finally:
+            self._sentrysdk_trace_url = None  # type: ignore[attr-defined]
+
+    return endheaders
+
+
+def _get_wrapped_putrequest(
+    original_putrequest: "Callable[..., Any]",
+) -> "Callable[..., Any]":
+    """
+    Responsible for initializing request and signed header tracking on the instance.
+    """
+
+    def putrequest(
+        self: "HTTPConnection", method: str, url: str, *args: "Any", **kwargs: "Any"
+    ) -> "Any":
+        # track which propagation headers are present and signed e.g. {"sentry-trace": (is_present, is_signed)}
+        request_headers: "Optional[Dict[str, Tuple[bool, bool]]]" = {}
+        self._sentrysdk_request_headers = request_headers  # type: ignore[attr-defined]
+
+        try:
+            rv = original_putrequest(self, method, url, *args, **kwargs)
+        except BaseException:
+            self._sentrysdk_request_headers = None  # type: ignore[attr-defined]
+            raise
+
+        return rv
+
+    return putrequest
+
+
+def _patch_aws_connection() -> None:
+    """
+    Patch AWS connection classes. These classes provide functions to sign HTTP headers, and subsequently
+    injecting trace propagation headers would invalidate the request.
+
+    Detect whether propagation headers are present and signed by patching
+    `putheader()`. Store that state in `_sentrysdk_request_headers`, initialized
+    by the `putrequest()` patch.
+
+    Do not edit signed headers when adding trace propagation headers in the `endheaders()` patch.
+    """
+    if AWSHTTPConnection is not None:
+        AWSHTTPConnection.putheader = _get_wrapped_putheader(  # type: ignore[method-assign]
+            AWSHTTPConnection.putheader
+        )
+        AWSHTTPConnection.endheaders = _get_wrapped_endheaders(  # type: ignore[method-assign]
+            AWSHTTPConnection.endheaders
+        )
+        AWSHTTPConnection.putrequest = _get_wrapped_putrequest(  # type: ignore[method-assign]
+            AWSHTTPConnection.putrequest
+        )
+
+    if AWSHTTPSConnection is not None:
+        AWSHTTPSConnection.putheader = _get_wrapped_putheader(  # type: ignore[method-assign]
+            AWSHTTPSConnection.putheader
+        )
+        AWSHTTPSConnection.endheaders = _get_wrapped_endheaders(  # type: ignore[method-assign]
+            AWSHTTPSConnection.endheaders
+        )
+        AWSHTTPSConnection.putrequest = _get_wrapped_putrequest(  # type: ignore[method-assign]
+            AWSHTTPSConnection.putrequest
+        )
+
+
 def _install_httplib() -> None:
+    _patch_aws_connection()
+
     real_putrequest = HTTPConnection.putrequest
     real_getresponse = HTTPConnection.getresponse
     real_read = HTTPResponse.read
@@ -189,19 +370,25 @@
 
         rv = real_putrequest(self, method, url, *args, **kwargs)
 
-        if should_propagate_trace(client, real_url):
+        # If _sentrysdk_request_headers is present, trace propagation headers should
+        # be injected in an `endheaders()` patch.
+        if should_propagate_trace(client, real_url) and not hasattr(
+            self, "_sentrysdk_request_headers"
+        ):
             for (
                 key,
                 value,
             ) in sentry_sdk.get_current_scope().iter_trace_propagation_headers(
                 span=span
             ):
                 logger.debug(
                     "[Tracing] Adding `{key}` header {value} to outgoing request to {real_url}.".format(
                         key=key, value=value, real_url=real_url
                     )
                 )
                 self.putheader(key, value)
+        elif should_propagate_trace(client, real_url):
+            self._sentrysdk_trace_url = real_url  # type: ignore[attr-defined]
 
         self._sentrysdk_span = span  # type: ignore[attr-defined]
         self._sentrysdk_breadcrumb = breadcrumb  # type: ignore[attr-defined]

sentry_sdk/utils.py

@@ -87,6 +87,9 @@
 
 FALSY_ENV_VALUES = frozenset(("false", "f", "n", "no", "off", "0"))
 TRUTHY_ENV_VALUES = frozenset(("true", "t", "y", "yes", "on", "1"))
+_AWS_SIGV4_SIGNING_ALGORITHMS = frozenset(
+    ("AWS4-HMAC-SHA256", "AWS4-ECDSA-P256-SHA256")
+)
 
 MAX_STACK_FRAMES = 2000
 """Maximum number of stack frames to send to Sentry.
@@ -1697,6 +1700,40 @@ def parse_url(url: str, sanitize: bool = True) -> "ParsedUrl":
     )
 
 
+def _get_aws_sigv4_signed_headers_from_authorization_header(
+    authorization: str,
+) -> "Set[str]":
+    # only AWS SigV4 authorization has the SignedHeaders parameter.
+    value = authorization.lstrip()
+    algorithm, _, parameters = value.partition(" ")
+    if algorithm not in _AWS_SIGV4_SIGNING_ALGORITHMS:
+        return set()
+
+    for part in parameters.split(","):
+        part = part.strip()
+        if part.startswith("SignedHeaders="):
+            _, _, header_names = part.partition("=")
+            return {header.lower() for header in header_names.split(";") if header}
+
+    return set()
+
+
+def _get_aws_sigv4_signed_headers_from_url_query_string(url: str) -> "Set[str]":
+    query = {
+        key.lower(): values for key, values in parse_qs(urlsplit(url).query).items()
+    }
+    algorithm = query.get("x-amz-algorithm", [""])[0]
+    if algorithm not in _AWS_SIGV4_SIGNING_ALGORITHMS:
+        return set()
+
+    # presigned requests have SignedHeaders in the URL query.
+    return {
+        header.lower()
+        for header in query.get("x-amz-signedheaders", [""])[0].split(";")
+        if header
+    }
+
+
 def is_valid_sample_rate(rate: "Any", source: str) -> bool:
     """
     Checks the given sample rate to make sure it is valid type and value (a