Capture custom request/response headers for wsgi and change in passing response_headers in django, pyramid (#925)

ashu658 · web-flow · commit 2f5bbc416ab3 · 2022-03-07T09:41:46.000-08:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.9.1-0.28b1...HEAD)
 
+- `opentelemetry-instrumentation-wsgi` Capture custom request/response headers in span attributes
+  ([#925])(https://github.com/open-telemetry/opentelemetry-python-contrib/pull/925)
+
 ### Added
 
 - `opentelemetry-instrumentation-dbapi` add experimental sql commenter capability
diff --git a/instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware.py b/instrumentation/opentelemetry-instrumentation-django/src/opentelemetry/instrumentation/django/middleware.py
@@ -275,7 +275,7 @@ def process_response(self, request, response):
                 add_response_attributes(
                     span,
                     f"{response.status_code} {response.reason_phrase}",
-                    response,
+                    response.items(),
                 )
 
             propagator = get_global_response_propagator()
diff --git a/instrumentation/opentelemetry-instrumentation-pyramid/src/opentelemetry/instrumentation/pyramid/callbacks.py b/instrumentation/opentelemetry-instrumentation-pyramid/src/opentelemetry/instrumentation/pyramid/callbacks.py
@@ -161,7 +161,7 @@ def trace_tween(request):
                 otel_wsgi.add_response_attributes(
                     span,
                     response_or_exception.status,
-                    response_or_exception.headers,
+                    response_or_exception.headerlist,
                 )
 
                 propagator = get_global_response_propagator()
diff --git a/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py b/instrumentation/opentelemetry-instrumentation-wsgi/src/opentelemetry/instrumentation/wsgi/__init__.py
@@ -117,7 +117,14 @@ def response_hook(span: Span, environ: WSGIEnvironment, status: str, response_he
 from opentelemetry.propagators.textmap import Getter
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace.status import Status, StatusCode
-from opentelemetry.util.http import remove_url_credentials
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+    get_custom_headers,
+    normalise_request_header_name,
+    normalise_response_header_name,
+    remove_url_credentials,
+)
 
 _HTTP_VERSION_PREFIX = "HTTP/"
 _CARRIER_KEY_PREFIX = "HTTP_"
@@ -208,6 +215,44 @@ def collect_request_attributes(environ):
     return result
 
 
+def add_custom_request_headers(span, environ):
+    """Adds custom HTTP request headers into the span which are configured by the user
+    from the PEP3333-conforming WSGI environ to be used as span creation attributes as described
+    in the specification https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers"""
+    attributes = {}
+    custom_request_headers_name = get_custom_headers(
+        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST
+    )
+    for header_name in custom_request_headers_name:
+        wsgi_env_var = header_name.upper().replace("-", "_")
+        header_values = environ.get(f"HTTP_{wsgi_env_var}")
+        if header_values:
+            key = normalise_request_header_name(header_name)
+            attributes[key] = [header_values]
+    span.set_attributes(attributes)
+
+
+def add_custom_response_headers(span, response_headers):
+    """Adds custom HTTP response headers into the sapn which are configured by the user from the
+    PEP3333-conforming WSGI environ as described in the specification
+    https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-request-and-response-headers"""
+    attributes = {}
+    custom_response_headers_name = get_custom_headers(
+        OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE
+    )
+    response_headers_dict = {}
+    if response_headers:
+        for header_name, header_value in response_headers:
+            response_headers_dict[header_name.lower()] = header_value
+
+    for header_name in custom_response_headers_name:
+        header_values = response_headers_dict.get(header_name.lower())
+        if header_values:
+            key = normalise_response_header_name(header_name)
+            attributes[key] = [header_values]
+    span.set_attributes(attributes)
+
+
 def add_response_attributes(
     span, start_response_status, response_headers
 ):  # pylint: disable=unused-argument
@@ -268,6 +313,8 @@ def _create_start_response(span, start_response, response_hook):
         @functools.wraps(start_response)
         def _start_response(status, response_headers, *args, **kwargs):
             add_response_attributes(span, status, response_headers)
+            if span.kind == trace.SpanKind.SERVER:
+                add_custom_response_headers(span, response_headers)
             if response_hook:
                 response_hook(status, response_headers)
             return start_response(status, response_headers, *args, **kwargs)
@@ -289,6 +336,8 @@ def __call__(self, environ, start_response):
             context_getter=wsgi_getter,
             attributes=collect_request_attributes(environ),
         )
+        if span.kind == trace.SpanKind.SERVER:
+            add_custom_request_headers(span, environ)
 
         if self.request_hook:
             self.request_hook(span, environ)
diff --git a/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py b/instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py
@@ -25,6 +25,10 @@
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.test.wsgitestutil import WsgiTestBase
 from opentelemetry.trace import StatusCode
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+)
 
 
 class Response:
@@ -82,6 +86,19 @@ def error_wsgi_unhandled(environ, start_response):
     raise ValueError
 
 
+def wsgi_with_custom_response_headers(environ, start_response):
+    assert isinstance(environ, dict)
+    start_response(
+        "200 OK",
+        [
+            ("content-type", "text/plain; charset=utf-8"),
+            ("content-length", "100"),
+            ("my-custom-header", "my-custom-value-1,my-custom-header-2"),
+        ],
+    )
+    return [b"*"]
+
+
 class TestWsgiApplication(WsgiTestBase):
     def validate_response(
         self,
@@ -444,5 +461,119 @@ def test_mark_span_internal_in_presence_of_span_from_other_framework(self):
             )
 
 
+class TestAdditionOfCustomRequestResponseHeaders(WsgiTestBase, TestBase):
+    def setUp(self):
+        super().setUp()
+        tracer_provider, _ = TestBase.create_tracer_provider()
+        self.tracer = tracer_provider.get_tracer(__name__)
+
+    def iterate_response(self, response):
+        while True:
+            try:
+                value = next(response)
+                self.assertEqual(value, b"*")
+            except StopIteration:
+                break
+
+    @mock.patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1,Custom-Test-Header-2,Custom-Test-Header-3"
+        },
+    )
+    def test_custom_request_headers_added_in_server_span(self):
+        self.environ.update(
+            {
+                "HTTP_CUSTOM_TEST_HEADER_1": "Test Value 1",
+                "HTTP_CUSTOM_TEST_HEADER_2": "TestValue2,TestValue3",
+            }
+        )
+        app = otel_wsgi.OpenTelemetryMiddleware(simple_wsgi)
+        response = app(self.environ, self.start_response)
+        self.iterate_response(response)
+        span = self.memory_exporter.get_finished_spans()[0]
+        expected = {
+            "http.request.header.custom_test_header_1": ("Test Value 1",),
+            "http.request.header.custom_test_header_2": (
+                "TestValue2,TestValue3",
+            ),
+        }
+        self.assertSpanHasAttributes(span, expected)
+
+    @mock.patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "Custom-Test-Header-1"
+        },
+    )
+    def test_custom_request_headers_not_added_in_internal_span(self):
+        self.environ.update(
+            {
+                "HTTP_CUSTOM_TEST_HEADER_1": "Test Value 1",
+            }
+        )
+
+        with self.tracer.start_as_current_span(
+            "test", kind=trace_api.SpanKind.SERVER
+        ):
+            app = otel_wsgi.OpenTelemetryMiddleware(simple_wsgi)
+            response = app(self.environ, self.start_response)
+            self.iterate_response(response)
+            span = self.memory_exporter.get_finished_spans()[0]
+            not_expected = {
+                "http.request.header.custom_test_header_1": ("Test Value 1",),
+            }
+            for key, _ in not_expected.items():
+                self.assertNotIn(key, span.attributes)
+
+    @mock.patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "content-type,content-length,my-custom-header,invalid-header"
+        },
+    )
+    def test_custom_response_headers_added_in_server_span(self):
+        app = otel_wsgi.OpenTelemetryMiddleware(
+            wsgi_with_custom_response_headers
+        )
+        response = app(self.environ, self.start_response)
+        self.iterate_response(response)
+        span = self.memory_exporter.get_finished_spans()[0]
+        expected = {
+            "http.response.header.content_type": (
+                "text/plain; charset=utf-8",
+            ),
+            "http.response.header.content_length": ("100",),
+            "http.response.header.my_custom_header": (
+                "my-custom-value-1,my-custom-header-2",
+            ),
+        }
+        self.assertSpanHasAttributes(span, expected)
+
+    @mock.patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "my-custom-header"
+        },
+    )
+    def test_custom_response_headers_not_added_in_internal_span(self):
+        with self.tracer.start_as_current_span(
+            "test", kind=trace_api.SpanKind.INTERNAL
+        ):
+            app = otel_wsgi.OpenTelemetryMiddleware(
+                wsgi_with_custom_response_headers
+            )
+            response = app(self.environ, self.start_response)
+            self.iterate_response(response)
+            span = self.memory_exporter.get_finished_spans()[0]
+            not_expected = {
+                "http.response.header.my_custom_header": (
+                    "my-custom-value-1,my-custom-header-2",
+                ),
+            }
+            for key, _ in not_expected.items():
+                self.assertNotIn(key, span.attributes)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/util/opentelemetry-util-http/src/opentelemetry/util/http/__init__.py b/util/opentelemetry-util-http/src/opentelemetry/util/http/__init__.py
@@ -15,9 +15,16 @@
 from os import environ
 from re import compile as re_compile
 from re import search
-from typing import Iterable
+from typing import Iterable, List
 from urllib.parse import urlparse, urlunparse
 
+OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST = (
+    "OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST"
+)
+OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE = (
+    "OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE"
+)
+
 
 class ExcludeList:
     """Class to exclude certain paths (given as a list of regexes) from tracing requests"""
@@ -98,3 +105,23 @@ def remove_url_credentials(url: str) -> str:
     except ValueError:  # an unparseable url was passed
         pass
     return url
+
+
+def normalise_request_header_name(header: str) -> str:
+    key = header.lower().replace("-", "_")
+    return f"http.request.header.{key}"
+
+
+def normalise_response_header_name(header: str) -> str:
+    key = header.lower().replace("-", "_")
+    return f"http.response.header.{key}"
+
+
+def get_custom_headers(env_var: str) -> List[str]:
+    custom_headers = environ.get(env_var, [])
+    if custom_headers:
+        custom_headers = [
+            custom_headers.strip()
+            for custom_headers in custom_headers.split(",")
+        ]
+    return custom_headers
diff --git a/util/opentelemetry-util-http/tests/test_capture_custom_headers.py b/util/opentelemetry-util-http/tests/test_capture_custom_headers.py
@@ -0,0 +1,67 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import patch
+
+from opentelemetry.test.test_base import TestBase
+from opentelemetry.util.http import (
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST,
+    OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE,
+    get_custom_headers,
+    normalise_request_header_name,
+    normalise_response_header_name,
+)
+
+
+class TestCaptureCustomHeaders(TestBase):
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST: "User-Agent,Test-Header"
+        },
+    )
+    def test_get_custom_request_header(self):
+        custom_headers_to_capture = get_custom_headers(
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_REQUEST
+        )
+        self.assertEqual(
+            custom_headers_to_capture, ["User-Agent", "Test-Header"]
+        )
+
+    @patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "content-type,content-length,test-header"
+        },
+    )
+    def test_get_custom_response_header(self):
+        custom_headers_to_capture = get_custom_headers(
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE
+        )
+        self.assertEqual(
+            custom_headers_to_capture,
+            [
+                "content-type",
+                "content-length",
+                "test-header",
+            ],
+        )
+
+    def test_normalise_request_header_name(self):
+        key = normalise_request_header_name("Test-Header")
+        self.assertEqual(key, "http.request.header.test_header")
+
+    def test_normalise_response_header_name(self):
+        key = normalise_response_header_name("Test-Header")
+        self.assertEqual(key, "http.response.header.test_header")

Original file line number	Diff line number	Diff line change
`@@ -275,7 +275,7 @@ def process_response(self, request, response):`
`275`	`275`	`add_response_attributes(`
`276`	`276`	`span,`
`277`	`277`	`f"{response.status_code} {response.reason_phrase}",`
`278`		`- response,`
	`278`	`+ response.items(),`
`279`	`279`	`)`
`280`	`280`
`281`	`281`	`propagator = get_global_response_propagator()`
Original file line number	Diff line number	Diff line change
`@@ -161,7 +161,7 @@ def trace_tween(request):`
`161`	`161`	`otel_wsgi.add_response_attributes(`
`162`	`162`	`span,`
`163`	`163`	`response_or_exception.status,`
`164`		`- response_or_exception.headers,`
	`164`	`+ response_or_exception.headerlist,`
`165`	`165`	`)`
`166`	`166`
`167`	`167`	`propagator = get_global_response_propagator()`