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

sentry_sdk/integrations/aiohttp.py

@@ -49,6 +49,7 @@
     capture_internal_exceptions,
     ensure_integration_enabled,
     event_from_exception,
+    get_aws_sigv4_signed_headers,
     has_data_collection_enabled,
     logger,
     parse_url,
@@ -456,12 +457,19 @@ async def on_request_start(
             span = legacy_span
 
         if should_propagate_trace(client, str(params.url)):
+            signed_headers = get_aws_sigv4_signed_headers(
+                params.headers, str(params.url)
+            )
             for (
                 key,
                 value,
             ) in sentry_sdk.get_current_scope().iter_trace_propagation_headers(
                 span=span
             ):
+                # do not modify a header whose value is already signed.
+                if key.lower() in signed_headers:
+                    continue
+
                 logger.debug(
                     "[Tracing] Adding `{key}` header {value} to outgoing request to {url}.".format(
                         key=key, value=value, url=params.url

sentry_sdk/integrations/boto3.py

@@ -6,8 +6,12 @@
 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 has_span_streaming_enabled
+from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span
+from sentry_sdk.tracing_utils import (
+    add_sentry_baggage_to_headers,
+    has_span_streaming_enabled,
+    should_propagate_trace,
+)
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     parse_url,
@@ -49,6 +53,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)
 
@@ -114,6 +120,53 @@ 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:
+            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, Span
 from sentry_sdk.tracing_utils import (
     EnvironHeaders,
     add_http_request_source,
@@ -21,14 +21,15 @@
     SENSITIVE_DATA_SUBSTITUTE,
     capture_internal_exceptions,
     ensure_integration_enabled,
+    get_aws_sigv4_signed_headers,
     is_sentry_url,
     logger,
     parse_url,
     safe_repr,
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Dict, List, Optional, Union
+    from typing import Any, Callable, Dict, List, Optional, Set, Union
 
     from sentry_sdk._types import Event, Hint
 
@@ -61,6 +62,17 @@ def add_python_runtime_context(
             return event
 
 
+def _request_header_names(buffer: "Optional[List[bytes]]") -> "Set[str]":
+    if buffer is None:
+        return set()
+    names = set()
+    for line in buffer:
+        name, separator, _ = line.partition(b":")
+        if separator:
+            names.add(name.decode("ascii", "ignore").lower())
+    return names
+
+
 def _complete_span(span: "Union[Span, StreamedSpan]") -> None:
     if isinstance(span, StreamedSpan):
         with capture_internal_exceptions():
@@ -74,6 +86,7 @@ def _complete_span(span: "Union[Span, StreamedSpan]") -> None:
 
 def _install_httplib() -> None:
     real_putrequest = HTTPConnection.putrequest
+    real_endheaders = HTTPConnection.endheaders
     real_getresponse = HTTPConnection.getresponse
     real_read = HTTPResponse.read
     real_close = HTTPResponse.close
@@ -157,26 +170,58 @@ def putrequest(
             set_on_span(SPANDATA.NETWORK_PEER_ADDRESS, self.host)
             set_on_span(SPANDATA.NETWORK_PEER_PORT, self.port)
 
-        rv = real_putrequest(self, method, url, *args, **kwargs)
+        try:
+            rv = real_putrequest(self, method, url, *args, **kwargs)
+        except BaseException:
+            self._sentrysdk_trace_url = None  # type: ignore[attr-defined]
+            raise
 
         if should_propagate_trace(client, real_url):
-            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)
+            self._sentrysdk_trace_url = real_url  # type: ignore[attr-defined]
+        else:
+            self._sentrysdk_trace_url = None  # type: ignore[attr-defined]
 
         self._sentrysdk_span = span  # type: ignore[attr-defined]
 
         return rv
 
+    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_buffer = getattr(self, "_buffer", None)
+                    existing_headers = _request_header_names(request_buffer)
+                    signed_headers = get_aws_sigv4_signed_headers(
+                        request_buffer, real_url
+                    )
+
+                    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`.
+                        if normalized_header in existing_headers and (
+                            normalized_header != BAGGAGE_HEADER_NAME
+                            or normalized_header in signed_headers
+                        ):
+                            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 real_endheaders(self, *args, **kwargs)
+        finally:
+            self._sentrysdk_trace_url = None  # type: ignore[attr-defined]
+
     def getresponse(self: "HTTPConnection", *args: "Any", **kwargs: "Any") -> "Any":
         span = getattr(self, "_sentrysdk_span", None)
 
@@ -233,6 +278,7 @@ def close(self: "HTTPResponse") -> None:
                 _complete_span(span)
 
     HTTPConnection.putrequest = putrequest  # type: ignore[method-assign]
+    HTTPConnection.endheaders = endheaders  # type: ignore[method-assign]
     HTTPConnection.getresponse = getresponse  # type: ignore[method-assign]
     HTTPResponse.read = read  # type: ignore[method-assign]
     HTTPResponse.close = close  # type: ignore[assignment,method-assign]

sentry_sdk/utils.py

@@ -1697,6 +1697,60 @@ def parse_url(url: str, sanitize: bool = True) -> "ParsedUrl":
     )
 
 
+def get_aws_sigv4_signed_headers(
+    headers: "Any", url: "Optional[str]" = None
+) -> "Set[str]":
+    # httpConnection exposes buffer, aiohttp uses header mapping.
+    if isinstance(headers, (str, bytes)):
+        authorization = headers
+    elif headers is None:
+        authorization = ""
+    elif hasattr(headers, "get"):
+        authorization = headers.get("Authorization", "")
+    else:
+        authorization = ""
+        for line in headers:
+            name, separator, value = line.partition(b":")
+            if separator and name.lower() == b"authorization":
+                authorization = value
+                break
+
+    if isinstance(authorization, bytes):
+        authorization = authorization.decode("ascii", "ignore")
+
+    signed_headers: "Set[str]" = set()
+    if isinstance(authorization, str):
+        # only AWS SigV4 authorization has the SignedHeaders parameter.
+        value = authorization.lstrip()
+        if value.startswith(("AWS4-HMAC-SHA256", "AWS4-ECDSA-P256-SHA256")):
+            for part in value.split(","):
+                part = part.strip()
+                if part.startswith("SignedHeaders="):
+                    _, _, header_names = part.partition("=")
+                    signed_headers.update(
+                        header.lower() for header in header_names.split(";") if header
+                    )
+                    break
+
+    if url is None:
+        return signed_headers
+
+    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 ("AWS4-HMAC-SHA256", "AWS4-ECDSA-P256-SHA256"):
+        return signed_headers
+
+    # presigned requests have SignedHeaders in the URL query.
+    signed_headers.update(
+        header.lower()
+        for header in query.get("x-amz-signedheaders", [""])[0].split(";")
+        if header
+    )
+    return signed_headers
+
+
 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