elasticapm/contrib/asgi.py

#  BSD 3-Clause License
#
#  Copyright (c) 2022, Elasticsearch BV
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#  * Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
#
#  * Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
#  * Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import functools
import urllib.parse
from typing import TYPE_CHECKING, Optional, Tuple, Union

if TYPE_CHECKING:
    from asgiref.typing import ASGIApplication, ASGIReceiveCallable, ASGISendCallable, Scope, ASGISendEvent

import elasticapm
from elasticapm import Client, get_client, instrument
from elasticapm.conf import constants
from elasticapm.contrib.asyncio.traces import set_context
from elasticapm.utils import default_ports, encoding
from elasticapm.utils.disttracing import TraceParent


def wrap_send(send, middleware):
    @functools.wraps(send)
    async def wrapped_send(message) -> None:
        if message.get("type") == "http.response.start":
            await set_context(lambda: middleware.get_data_from_response(message, constants.TRANSACTION), "response")
            result = "HTTP {}xx".format(message["status"] // 100)
            elasticapm.set_transaction_result(result, override=False)
            elasticapm.set_transaction_outcome(http_status_code=message["status"], override=False)
        await send(message)

    return wrapped_send


class ASGITracingMiddleware:
    __slots__ = ("_app", "client")

    def __init__(self, app: "ASGIApplication", **options) -> None:
        self._app = app
        client = get_client()
        if not client:
            client = Client(**options)
        self.client = client
        if self.client.config.instrument and self.client.config.enabled:
            instrument()

    async def __call__(self, scope: "Scope", receive: "ASGIReceiveCallable", send: "ASGISendCallable") -> None:
        if scope["type"] != "http":
            await self._app(scope, receive, send)
            return
        send = wrap_send(send, self)
        wrapped_receive = receive
        url, url_dict = self.get_url(scope)
        body = None
        if not self.client.should_ignore_url(url):
            self.client.begin_transaction(
                transaction_type="request", trace_parent=TraceParent.from_headers(scope["headers"])
            )
            self.set_transaction_name(scope["method"], url)
            if scope["method"] in constants.HTTP_WITH_BODY and self.client.config.capture_body != "off":
                messages = []
                more_body = True
                while more_body:
                    message = await receive()
                    messages.append(message)
                    more_body = message.get("more_body", False)

                body_raw = b"".join([message.get("body", b"") for message in messages])
                body = str(body_raw, errors="ignore")

                # Dispatch to the ASGI callable
                async def new_wrapped_receive():
                    if messages:
                        return messages.pop(0)

                    # Once that's done we can just await any other messages.
                    return await receive()

                wrapped_receive = new_wrapped_receive

            await set_context(lambda: self.get_data_from_request(scope, constants.TRANSACTION, body), "request")

        try:
            await self._app(scope, wrapped_receive, send)
            elasticapm.set_transaction_outcome(constants.OUTCOME.SUCCESS, override=False)
            return
        except Exception as exc:
            self.client.capture_exception()
            elasticapm.set_transaction_result("HTTP 5xx", override=False)
            elasticapm.set_transaction_outcome(constants.OUTCOME.FAILURE, override=True)
            elasticapm.set_context({"status_code": 500}, "response")
            raise exc from None
        finally:
            self.client.end_transaction()

    def get_headers(self, scope_or_message: Union["Scope", "ASGISendEvent"]) -> dict[str, str]:
        headers = {}
        for k, v in scope_or_message.get("headers", {}):
            key = k.decode("latin1")
            val = v.decode("latin1")
            if key in headers:
                headers[key] = f"{headers[key]}, {val}"
            else:
                headers[key] = val
        return headers

    def get_url(self, scope: "Scope", host: Optional[str] = None) -> Tuple[str, dict[str, str]]:
        url_dict = {}
        scheme = scope.get("scheme", "http")
        server = scope.get("server", None)
        path = scope.get("root_path", "") + scope.get("path", "")

        url_dict["protocol"] = scheme + ":"

        if host:
            url = f"{scheme}://{host}{path}"
            url_dict["hostname"] = host
        elif server is not None:
            host, port = server
            url_dict["hostname"] = host
            if port:
                url_dict["port"] = port
            default_port = default_ports.get(scheme, None)
            if port != default_port:
                url = f"{scheme}://{host}:{port}{path}"
            else:
                url = f"{scheme}://{host}{path}"
        else:
            url = path
        qs = scope.get("query_string")
        if qs:
            query = "?" + urllib.parse.unquote(qs.decode("latin-1"))
            url += query
            url_dict["search"] = encoding.keyword_field(query)
        url_dict["full"] = encoding.keyword_field(url)
        return url, url_dict

    def get_ip(self, scope: "Scope", headers: dict) -> Optional[str]:
        x_forwarded_for = headers.get("x-forwarded-for")
        remote_addr = headers.get("remote-addr")
        ip: Optional[str] = None
        if x_forwarded_for:
            ip = x_forwarded_for.split(",")[0]
        elif remote_addr:
            ip = remote_addr
        elif scope.get("client"):
            ip = scope.get("client")[0]
        return ip

    async def get_data_from_request(self, scope: "Scope", event_type: str, body: Optional[str]) -> dict:
        """Loads data from incoming request for APM capturing.

        Args:
            request (Request)
            config (Config)
            event_type (str)
            body (str)

        Returns:
            dict
        """
        headers = self.get_headers(scope)
        result = {
            "method": scope["method"],
            "socket": {"remote_address": self.get_ip(scope, headers)},
            "cookies": headers.pop("cookies", {}),
        }
        if self.client.config.capture_headers:
            result["headers"] = headers
        if body and self.client.config.capture_body in ("all", event_type):
            result["body"] = body
        url, url_dict = self.get_url(scope)
        result["url"] = url_dict

        return result

    async def get_data_from_response(self, message: dict, event_type: str) -> dict:
        """Loads data from response for APM capturing.

        Args:
            message (dict)
            config (Config)
            event_type (str)

        Returns:
            dict
        """
        result = {}

        if "status" in message:
            result["status_code"] = message["status"]

        if self.client.config.capture_headers and "headers" in message:
            headers = self.get_headers(message)
            if headers:
                result["headers"] = headers

        return result

    def set_transaction_name(self, method: str, url: str) -> None:
        """
        Default implementation sets transaction name to "METHOD unknown route".
        Subclasses may add framework specific naming.
        """
        elasticapm.set_transaction_name(f"{method.upper()} unknown route")