feat: configure header extraction for ASGI middleware via constructor params (#2026)

* feat: configure header extraction for ASGI middleware via constructor params * fix django middleware * lint * remove import * Fix lint * Update instrumentation/opentelemetry-instrumentation-asgi/src/opentelemetry/instrumentation/asgi/__init__.py
2025-08-02 11:31:52 +08:00 · 2024-01-31 00:11:00 -03:00
parent a93bd74dc3
commit 4b1a9c75db
7 changed files with 209 additions and 99 deletions
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@ -48,6 +48,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  ([#1948](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1948))
 - Added schema_url (`"https://opentelemetry.io/schemas/1.11.0"`) to all metrics and traces
  ([#1977](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1977))
+- Add support for configuring ASGI middleware header extraction via runtime constructor parameters
+  ([#2026](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2026))

 ### Fixed

--- a/instrumentation/opentelemetry-instrumentation-asgi/src/opentelemetry/instrumentation/asgi/init.py
+++ b/instrumentation/opentelemetry-instrumentation-asgi/src/opentelemetry/instrumentation/asgi/init.py
@ -189,11 +189,13 @@ API
 ---
 """

+from __future__ import annotations
+
 import typing
 import urllib
 from functools import wraps
 from timeit import default_timer
-from typing import Tuple
+from typing import Any, Awaitable, Callable, Tuple, cast

 from asgiref.compatibility import guarantee_single_callable

@ -332,55 +334,28 @@ def collect_request_attributes(scope):
    return result


-def collect_custom_request_headers_attributes(scope):
-    """returns custom HTTP request headers to be added into SERVER span as span attributes
-    Refer specification https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers
+def collect_custom_headers_attributes(
+    scope_or_response_message: dict[str, Any],
+    sanitize: SanitizeValue,
+    header_regexes: list[str],
+    normalize_names: Callable[[str], str],
+) -> dict[str, str]:
    """
+    Returns custom HTTP request or response headers to be added into SERVER span as span attributes.

-    sanitize = SanitizeValue(
-        get_custom_headers(
-            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS
-        )
-    )
-
+    Refer specifications:
+     - https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers
+    """
    # Decode headers before processing.
-    headers = {
+    headers: dict[str, str] = {
        _key.decode("utf8"): _value.decode("utf8")
-        for (_key, _value) in scope.get("headers")
+        for (_key, _value) in scope_or_response_message.get("headers")
+        or cast("list[tuple[bytes, bytes]]", [])
    }
-
    return sanitize.sanitize_header_values(
        headers,
-        get_custom_headers(
-            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST
-        ),
-        normalise_request_header_name,
-    )
-
-
-def collect_custom_response_headers_attributes(message):
-    """returns custom HTTP response headers to be added into SERVER span as span attributes
-    Refer specification https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers
-    """
-
-    sanitize = SanitizeValue(
-        get_custom_headers(
-            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS
-        )
-    )
-
-    # Decode headers before processing.
-    headers = {
-        _key.decode("utf8"): _value.decode("utf8")
-        for (_key, _value) in message.get("headers")
-    }
-
-    return sanitize.sanitize_header_values(
-        headers,
-        get_custom_headers(
-            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE
-        ),
-        normalise_response_header_name,
+        header_regexes,
+        normalize_names,
    )


@ -493,6 +468,9 @@ class OpenTelemetryMiddleware:
        tracer_provider=None,
        meter_provider=None,
        meter=None,
+        http_capture_headers_server_request: list[str] | None = None,
+        http_capture_headers_server_response: list[str] | None = None,
+        http_capture_headers_sanitize_fields: list[str] | None = None,
    ):
        self.app = guarantee_single_callable(app)
        self.tracer = trace.get_tracer(
@ -540,7 +518,41 @@ class OpenTelemetryMiddleware:
        self.client_response_hook = client_response_hook
        self.content_length_header = None

-    async def __call__(self, scope, receive, send):
+        # Environment variables as constructor parameters
+        self.http_capture_headers_server_request = (
+            http_capture_headers_server_request
+            or (
+                get_custom_headers(
+                    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST
+                )
+            )
+            or None
+        )
+        self.http_capture_headers_server_response = (
+            http_capture_headers_server_response
+            or (
+                get_custom_headers(
+                    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE
+                )
+            )
+            or None
+        )
+        self.http_capture_headers_sanitize_fields = SanitizeValue(
+            http_capture_headers_sanitize_fields
+            or (
+                get_custom_headers(
+                    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS
+                )
+            )
+            or []
+        )
+
+    async def __call__(
+        self,
+        scope: dict[str, Any],
+        receive: Callable[[], Awaitable[dict[str, Any]]],
+        send: Callable[[dict[str, Any]], Awaitable[None]],
+    ) -> None:
        """The ASGI application

        Args:
@ -583,7 +595,14 @@ class OpenTelemetryMiddleware:

                    if current_span.kind == trace.SpanKind.SERVER:
                        custom_attributes = (
-                            collect_custom_request_headers_attributes(scope)
+                            collect_custom_headers_attributes(
+                                scope,
+                                self.http_capture_headers_sanitize_fields,
+                                self.http_capture_headers_server_request,
+                                normalise_request_header_name,
+                            )
+                            if self.http_capture_headers_server_request
+                            else {}
                        )
                        if len(custom_attributes) > 0:
                            current_span.set_attributes(custom_attributes)
@ -658,7 +677,7 @@ class OpenTelemetryMiddleware:
        expecting_trailers = False

        @wraps(send)
-        async def otel_send(message):
+        async def otel_send(message: dict[str, Any]):
            nonlocal expecting_trailers
            with self.tracer.start_as_current_span(
                " ".join((server_span_name, scope["type"], "send"))
@ -685,7 +704,14 @@ class OpenTelemetryMiddleware:
                        and "headers" in message
                    ):
                        custom_response_attributes = (
-                            collect_custom_response_headers_attributes(message)
+                            collect_custom_headers_attributes(
+                                message,
+                                self.http_capture_headers_sanitize_fields,
+                                self.http_capture_headers_server_response,
+                                normalise_response_header_name,
+                            )
+                            if self.http_capture_headers_server_response
+                            else {}
                        )
                        if len(custom_response_attributes) > 0:
                            server_span.set_attributes(
--- a/instrumentation/opentelemetry-instrumentation-asgi/tests/test_asgi_custom_headers.py
+++ b/instrumentation/opentelemetry-instrumentation-asgi/tests/test_asgi_custom_headers.py
@ -1,4 +1,4 @@
-from unittest import mock
+import os

 import opentelemetry.instrumentation.asgi as otel_asgi
 from opentelemetry.test.asgitestutil import AsgiTestBase
@ -72,21 +72,22 @@ async def websocket_app_with_custom_headers(scope, receive, send):
            break


-@mock.patch.dict(
-    "os.environ",
-    {
-        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS: ".*my-secret.*",
-        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,Regex-Test-Header-.*,Regex-Invalid-Test-Header-.*,.*my-secret.*",
-        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,my-custom-regex-header-.*,invalid-regex-header-.*,.*my-secret.*",
-    },
-)
 class TestCustomHeaders(AsgiTestBase, TestBase):
+    constructor_params = {}
+    __test__ = False
+
+    def __init_subclass__(cls) -> None:
+        if cls is not TestCustomHeaders:
+            cls.__test__ = True
+
    def setUp(self):
        super().setUp()
        self.tracer_provider, self.exporter = TestBase.create_tracer_provider()
        self.tracer = self.tracer_provider.get_tracer(__name__)
        self.app = otel_asgi.OpenTelemetryMiddleware(
-            simple_asgi, tracer_provider=self.tracer_provider
+            simple_asgi,
+            tracer_provider=self.tracer_provider,
+            **self.constructor_params,
        )

    def test_http_custom_request_headers_in_span_attributes(self):
@ -148,7 +149,9 @@ class TestCustomHeaders(AsgiTestBase, TestBase):

    def test_http_custom_response_headers_in_span_attributes(self):
        self.app = otel_asgi.OpenTelemetryMiddleware(
-            http_app_with_custom_headers, tracer_provider=self.tracer_provider
+            http_app_with_custom_headers,
+            tracer_provider=self.tracer_provider,
+            **self.constructor_params,
        )
        self.seed_app(self.app)
        self.send_default_request()
@ -175,7 +178,9 @@ class TestCustomHeaders(AsgiTestBase, TestBase):

    def test_http_custom_response_headers_not_in_span_attributes(self):
        self.app = otel_asgi.OpenTelemetryMiddleware(
-            http_app_with_custom_headers, tracer_provider=self.tracer_provider
+            http_app_with_custom_headers,
+            tracer_provider=self.tracer_provider,
+            **self.constructor_params,
        )
        self.seed_app(self.app)
        self.send_default_request()
@ -277,6 +282,7 @@ class TestCustomHeaders(AsgiTestBase, TestBase):
        self.app = otel_asgi.OpenTelemetryMiddleware(
            websocket_app_with_custom_headers,
            tracer_provider=self.tracer_provider,
+            **self.constructor_params,
        )
        self.seed_app(self.app)
        self.send_input({"type": "websocket.connect"})
@ -317,6 +323,7 @@ class TestCustomHeaders(AsgiTestBase, TestBase):
        self.app = otel_asgi.OpenTelemetryMiddleware(
            websocket_app_with_custom_headers,
            tracer_provider=self.tracer_provider,
+            **self.constructor_params,
        )
        self.seed_app(self.app)
        self.send_input({"type": "websocket.connect"})
@ -333,3 +340,46 @@ class TestCustomHeaders(AsgiTestBase, TestBase):
            if span.kind == SpanKind.SERVER:
                for key, _ in not_expected.items():
                    self.assertNotIn(key, span.attributes)
+
+
+SANITIZE_FIELDS_TEST_VALUE = ".*my-secret.*"
+SERVER_REQUEST_TEST_VALUE = "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,Regex-Test-Header-.*,Regex-Invalid-Test-Header-.*,.*my-secret.*"
+SERVER_RESPONSE_TEST_VALUE = "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,my-custom-regex-header-.*,invalid-regex-header-.*,.*my-secret.*"
+
+
+class TestCustomHeadersEnv(TestCustomHeaders):
+    def setUp(self):
+        os.environ.update(
+            {
+                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS: SANITIZE_FIELDS_TEST_VALUE,
+                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: SERVER_REQUEST_TEST_VALUE,
+                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: SERVER_RESPONSE_TEST_VALUE,
+            }
+        )
+        super().setUp()
+
+    def tearDown(self):
+        os.environ.pop(
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS, None
+        )
+        os.environ.pop(
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST, None
+        )
+        os.environ.pop(
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE, None
+        )
+        super().tearDown()
+
+
+class TestCustomHeadersConstructor(TestCustomHeaders):
+    constructor_params = {
+        "http_capture_headers_sanitize_fields": SANITIZE_FIELDS_TEST_VALUE.split(
+            ","
+        ),
+        "http_capture_headers_server_request": SERVER_REQUEST_TEST_VALUE.split(
+            ","
+        ),
+        "http_capture_headers_server_response": SERVER_RESPONSE_TEST_VALUE.split(
+            ","
+        ),
+    }
--- a/instrumentation/opentelemetry-instrumentation-asgi/tests/test_asgi_middleware.py
+++ b/instrumentation/opentelemetry-instrumentation-asgi/tests/test_asgi_middleware.py
@ -983,18 +983,16 @@ class TestAsgiApplicationRaisingError(AsgiTestBase):
    def tearDown(self):
        pass

-    @mock.patch(
-        "opentelemetry.instrumentation.asgi.collect_custom_request_headers_attributes",
-        side_effect=ValueError("whatever"),
-    )
-    def test_asgi_issue_1883(
-        self, mock_collect_custom_request_headers_attributes
-    ):
+    def test_asgi_issue_1883(self):
        """
        Test that exception UnboundLocalError local variable 'start' referenced before assignment is not raised
        See https://github.com/open-telemetry/opentelemetry-python-contrib/issues/1883
        """
-        app = otel_asgi.OpenTelemetryMiddleware(simple_asgi)
+
+        async def bad_app(_scope, _receive, _send):
+            raise ValueError("whatever")
+
+        app = otel_asgi.OpenTelemetryMiddleware(bad_app)
        self.seed_app(app)
        self.send_default_request()
        try:
--- a/instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware/otel_middleware.py
+++ b/instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware/otel_middleware.py
@ -43,10 +43,17 @@ from opentelemetry.instrumentation.wsgi import wsgi_getter
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace import Span, SpanKind, use_span
 from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+    SanitizeValue,
    _parse_active_request_count_attrs,
    _parse_duration_attrs,
+    get_custom_headers,
    get_excluded_urls,
    get_traced_request_attrs,
+    normalise_request_header_name,
+    normalise_response_header_name,
 )

 try:
@ -91,10 +98,7 @@ else:
 try:
    from opentelemetry.instrumentation.asgi import asgi_getter, asgi_setter
    from opentelemetry.instrumentation.asgi import (
-        collect_custom_request_headers_attributes as asgi_collect_custom_request_attributes,
-    )
-    from opentelemetry.instrumentation.asgi import (
-        collect_custom_response_headers_attributes as asgi_collect_custom_response_attributes,
+        collect_custom_headers_attributes as asgi_collect_custom_headers_attributes,
    )
    from opentelemetry.instrumentation.asgi import (
        collect_request_attributes as asgi_collect_request_attributes,
@ -108,7 +112,6 @@ except ImportError:
    set_status_code = None
    _is_asgi_supported = False

-
 _logger = getLogger(__name__)
 _attributes_by_preference = [
    [
@ -249,7 +252,18 @@ class _DjangoMiddleware(MiddlewareMixin):
                )
                if span.is_recording() and span.kind == SpanKind.SERVER:
                    attributes.update(
-                        asgi_collect_custom_request_attributes(carrier)
+                        asgi_collect_custom_headers_attributes(
+                            carrier,
+                            SanitizeValue(
+                                get_custom_headers(
+                                    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS
+                                )
+                            ),
+                            get_custom_headers(
+                                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST
+                            ),
+                            normalise_request_header_name,
+                        )
                    )
            else:
                if span.is_recording() and span.kind == SpanKind.SERVER:
@ -336,8 +350,17 @@ class _DjangoMiddleware(MiddlewareMixin):
                    for key, value in response.items():
                        asgi_setter.set(custom_headers, key, value)

-                    custom_res_attributes = (
-                        asgi_collect_custom_response_attributes(custom_headers)
+                    custom_res_attributes = asgi_collect_custom_headers_attributes(
+                        custom_headers,
+                        SanitizeValue(
+                            get_custom_headers(
+                                OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS
+                            )
+                        ),
+                        get_custom_headers(
+                            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE
+                        ),
+                        normalise_response_header_name,
                    )
                    for key, value in custom_res_attributes.items():
                        span.set_attribute(key, value)
--- a/instrumentation/opentelemetry-instrumentation-starlette/tests/test_starlette_instrumentation.py
+++ b/instrumentation/opentelemetry-instrumentation-starlette/tests/test_starlette_instrumentation.py
@ -467,15 +467,18 @@ class TestBaseWithCustomHeaders(TestBase):
        return app


-@patch.dict(
-    "os.environ",
-    {
-        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS: ".*my-secret.*",
-        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,Regex-Test-Header-.*,Regex-Invalid-Test-Header-.*,.*my-secret.*",
-        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,my-custom-regex-header-.*,invalid-regex-header-.*,.*my-secret.*",
-    },
-)
 class TestHTTPAppWithCustomHeaders(TestBaseWithCustomHeaders):
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS: ".*my-secret.*",
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,Regex-Test-Header-.*,Regex-Invalid-Test-Header-.*,.*my-secret.*",
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,my-custom-regex-header-.*,invalid-regex-header-.*,.*my-secret.*",
+        },
+    )
+    def setUp(self) -> None:
+        super().setUp()
+
    def test_custom_request_headers_in_span_attributes(self):
        expected = {
            "http.request.header.custom_test_header_1": (
@ -590,15 +593,18 @@ class TestHTTPAppWithCustomHeaders(TestBaseWithCustomHeaders):
            self.assertNotIn(key, server_span.attributes)


-@patch.dict(
-    "os.environ",
-    {
-        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS: ".*my-secret.*",
-        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,Regex-Test-Header-.*,Regex-Invalid-Test-Header-.*,.*my-secret.*",
-        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,my-custom-regex-header-.*,invalid-regex-header-.*,.*my-secret.*",
-    },
-)
 class TestWebSocketAppWithCustomHeaders(TestBaseWithCustomHeaders):
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS: ".*my-secret.*",
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,Regex-Test-Header-.*,Regex-Invalid-Test-Header-.*,.*my-secret.*",
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3,my-custom-regex-header-.*,invalid-regex-header-.*,.*my-secret.*",
+        },
+    )
+    def setUp(self) -> None:
+        super().setUp()
+
    def test_custom_request_headers_in_span_attributes(self):
        expected = {
            "http.request.header.custom_test_header_1": (
--- a/util/opentelemetry-util-http/src/opentelemetry/util/http/init.py
+++ b/util/opentelemetry-util-http/src/opentelemetry/util/http/init.py
@ -12,11 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

+from __future__ import annotations
+
 from os import environ
 from re import IGNORECASE as RE_IGNORECASE
 from re import compile as re_compile
 from re import search
-from typing import Iterable, List, Optional
+from typing import Callable, Iterable, Optional
 from urllib.parse import urlparse, urlunparse

 from opentelemetry.semconv.trace import SpanAttributes
@ -84,9 +86,12 @@ class SanitizeValue:
        )

    def sanitize_header_values(
-        self, headers: dict, header_regexes: list, normalize_function: callable
-    ) -> dict:
-        values = {}
+        self,
+        headers: dict[str, str],
+        header_regexes: list[str],
+        normalize_function: Callable[[str], str],
+    ) -> dict[str, str]:
+        values: dict[str, str] = {}

        if header_regexes:
            header_regexes_compiled = re_compile(
@ -216,14 +221,14 @@ def sanitize_method(method: Optional[str]) -> Optional[str]:
    return "UNKNOWN"


-def get_custom_headers(env_var: str) -> List[str]:
-    custom_headers = environ.get(env_var, [])
+def get_custom_headers(env_var: str) -> list[str]:
+    custom_headers = environ.get(env_var, None)
    if custom_headers:
-        custom_headers = [
+        return [
            custom_headers.strip()
            for custom_headers in custom_headers.split(",")
        ]
-    return custom_headers
+    return []


 def _parse_active_request_count_attrs(req_attrs):