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

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,
@@ -28,7 +28,7 @@
 )
 
 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 +61,41 @@
             return event
 
 
+def _aws_sigv4_signed_headers(buffer: "Optional[List[bytes]]") -> "Set[str]":
+    if buffer is None:
+        return set()
+    for line in buffer:
+        name, separator, value = line.partition(b":")
+        if not separator or name.lower() != b"authorization":
+            continue
+
+        value = value.lstrip()
+        if not value.startswith((b"AWS4-HMAC-SHA256", b"AWS4-ECDSA-P256-SHA256")):
+            continue
+
+        for part in value.split(b","):
+            part = part.strip()
+            if part.startswith(b"SignedHeaders="):
+                _, _, header_names = part.partition(b"=")
+                return {
+                    header.decode("ascii", "ignore").lower()
+                    for header in header_names.split(b";")
+                    if header
+                }
+    return set()
+
+
+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 +109,7 @@
 
 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 +193,56 @@
             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:
+                request_buffer = getattr(self, "_buffer", None)
+                existing_headers = _request_header_names(request_buffer)
+                signed_headers = _aws_sigv4_signed_headers(request_buffer)
+
+                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 +299,7 @@
                 _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]