diff --git a/ddtrace/contrib/_events/web_framework.py b/ddtrace/contrib/_events/web_framework.py index 97794868c6b..1e05ef54ce1 100644 --- a/ddtrace/contrib/_events/web_framework.py +++ b/ddtrace/contrib/_events/web_framework.py @@ -13,6 +13,7 @@ class WebFrameworkEvents(str, Enum): WEB_REQUEST = "web.request" + WEB_REQUEST_STARTING = "web.request.starting" @dataclass diff --git a/ddtrace/contrib/internal/flask/patch.py b/ddtrace/contrib/internal/flask/patch.py index 12ceeb8ec9f..a8658e554ad 100644 --- a/ddtrace/contrib/internal/flask/patch.py +++ b/ddtrace/contrib/internal/flask/patch.py @@ -7,6 +7,7 @@ from werkzeug.exceptions import NotFound from ddtrace.contrib import trace_utils +from ddtrace.contrib._events.web_framework import WebFrameworkEvents from ddtrace.ext import SpanTypes from ddtrace.internal import core from ddtrace.internal.constants import COMPONENT @@ -15,6 +16,7 @@ from ddtrace.internal.schema import schematize_service_name from ddtrace.internal.schema import schematize_url_operation from ddtrace.internal.schema.span_attribute_schema import SpanDirection +from ddtrace.internal.serverless import in_aws_lambda_microvm from ddtrace.internal.settings.appsec_telemetry import config as appsec_telemetry_config from ddtrace.internal.span_bus import span_from_context from ddtrace.internal.utils import get_blocked @@ -391,8 +393,14 @@ def unpatch(): def patched_wsgi_app(wrapped, instance, args, kwargs): environ, start_response = args + script_name = (environ.get("SCRIPT_NAME") or "").rstrip("/") + if in_aws_lambda_microvm(): + path_info = environ.get("PATH_INFO") or "" + core.dispatch( + WebFrameworkEvents.WEB_REQUEST_STARTING.value, (environ.get("REQUEST_METHOD"), script_name + path_info) + ) # Registration is gated on asm_config, not tracing — keep this above the tracing short-circuit. - _collect_routes_once(instance, environ.get("SCRIPT_NAME") or "") + _collect_routes_once(instance, script_name) if not is_tracing_enabled(): return wrapped(*args, **kwargs) middleware = _FlaskWSGIMiddleware(wrapped, None, config.flask) diff --git a/ddtrace/internal/serverless/__init__.py b/ddtrace/internal/serverless/__init__.py index 86fb136b538..a92d576d15a 100644 --- a/ddtrace/internal/serverless/__init__.py +++ b/ddtrace/internal/serverless/__init__.py @@ -20,6 +20,12 @@ def has_aws_lambda_agent_extension(): return path.exists("/opt/extensions/datadog-agent") +def in_aws_lambda_microvm(): + # type: () -> bool + """Returns whether the environment is an AWS Lambda MicroVM.""" + return bool(env.get("AWS_LAMBDA_MICROVM_IMAGE_ARN", "").strip()) + + def in_gcp_function(): # type: () -> bool """Returns whether the environment is a GCP Function. diff --git a/tests/contrib/flask/test_microvm_identity_refresh.py b/tests/contrib/flask/test_microvm_identity_refresh.py new file mode 100644 index 00000000000..d9643480a0a --- /dev/null +++ b/tests/contrib/flask/test_microvm_identity_refresh.py @@ -0,0 +1,93 @@ +import mock + +from ddtrace.contrib._events.web_framework import WebFrameworkEvents +from ddtrace.contrib.internal.flask.patch import patched_wsgi_app +from ddtrace.internal import core + +from . import BaseFlaskTestCase + + +REQUEST_STARTING_PATH = "/web-request-starting" +MICROVM_RUNTIME_PREFIX = "/aws/lambda-microvms/runtime/v1" +MICROVM_RUN_HOOK_PATH = "/run" + + +class FlaskMicrovmIdentityRefreshTestCase(BaseFlaskTestCase): + """patched_wsgi_app() must dispatch every request's method/path before tracing starts. + + The matching logic itself is tested in tests/tracer/runtime/test_runtime_id.py. + """ + + def test_microvm_run_hook_request(self): + """No route is registered at the hook path: wsgi_app() runs before routing, so it + must fire even on a 404 -- the stronger, more general form of this check (whether the + route matches doesn't change what gets dispatched). + """ + with ( + mock.patch("ddtrace.contrib.internal.flask.patch.in_aws_lambda_microvm", return_value=True), + mock.patch("ddtrace.contrib.internal.flask.patch.core.dispatch", wraps=core.dispatch) as m, + ): + res = self.client.post(REQUEST_STARTING_PATH) + + self.assertEqual(res.status_code, 404) + m.assert_any_call(WebFrameworkEvents.WEB_REQUEST_STARTING.value, ("POST", REQUEST_STARTING_PATH)) + + def test_other_request_does_not_dispatch_outside_microvm(self): + @self.app.route("/") + def index(): + return "ok", 200 + + with ( + mock.patch("ddtrace.contrib.internal.flask.patch.in_aws_lambda_microvm", return_value=False), + mock.patch("ddtrace.contrib.internal.flask.patch.core.dispatch", wraps=core.dispatch) as m, + ): + res = self.client.get("/") + + self.assertEqual(res.status_code, 200) + assert all(call.args[0] != WebFrameworkEvents.WEB_REQUEST_STARTING.value for call in m.call_args_list) + + def test_dispatches_script_name_prefixed_request_path(self): + with ( + mock.patch("ddtrace.contrib.internal.flask.patch.in_aws_lambda_microvm", return_value=True), + mock.patch("ddtrace.contrib.internal.flask.patch.core.dispatch", wraps=core.dispatch) as m, + ): + res = self.client.post(MICROVM_RUN_HOOK_PATH, environ_overrides={"SCRIPT_NAME": MICROVM_RUNTIME_PREFIX}) + + self.assertEqual(res.status_code, 404) + m.assert_any_call( + WebFrameworkEvents.WEB_REQUEST_STARTING.value, + ("POST", MICROVM_RUNTIME_PREFIX + MICROVM_RUN_HOOK_PATH), + ) + + def test_pre_request_event_dispatches_before_wsgi_middleware(self): + events = [] + environ = {"REQUEST_METHOD": "POST", "PATH_INFO": REQUEST_STARTING_PATH, "SCRIPT_NAME": ""} + + def start_response(status, headers, exc_info=None): + pass + + def wrapped(environ, start_response): + return [] + + def dispatch(name, args): + if name == WebFrameworkEvents.WEB_REQUEST_STARTING.value: + events.append("starting") + + class WSGIMiddleware: + def __init__(self, app, tracer, integration_config): + pass + + def __call__(self, environ, start_response): + events.append("middleware") + return [] + + with ( + mock.patch("ddtrace.contrib.internal.flask.patch.in_aws_lambda_microvm", return_value=True), + mock.patch("ddtrace.contrib.internal.flask.patch.core.dispatch", side_effect=dispatch), + mock.patch("ddtrace.contrib.internal.flask.patch._collect_routes_once"), + mock.patch("ddtrace.contrib.internal.flask.patch.is_tracing_enabled", return_value=True), + mock.patch("ddtrace.contrib.internal.flask.patch._FlaskWSGIMiddleware", WSGIMiddleware), + ): + patched_wsgi_app(wrapped, self.app, (environ, start_response), {}) + + assert events == ["starting", "middleware"] diff --git a/tests/internal/test_serverless.py b/tests/internal/test_serverless.py index ff1da9ade19..6d79bc2c357 100644 --- a/tests/internal/test_serverless.py +++ b/tests/internal/test_serverless.py @@ -1,5 +1,6 @@ import pytest +from ddtrace.internal.serverless import in_aws_lambda_microvm from ddtrace.internal.serverless import in_azure_function from ddtrace.internal.serverless import in_gcp_function from tests.utils import override_env @@ -19,6 +20,21 @@ def test_not_gcp_function(): assert in_gcp_function() is False +def test_is_aws_lambda_microvm(): + with override_env(dict(AWS_LAMBDA_MICROVM_IMAGE_ARN="arn:aws:lambda:us-east-1::runtime:python3.12")): + assert in_aws_lambda_microvm() is True + + +def test_not_aws_lambda_microvm(): + assert in_aws_lambda_microvm() is False + + +@pytest.mark.parametrize("image_arn", ["", " "]) +def test_blank_aws_lambda_microvm_image_arn(image_arn): + with override_env(dict(AWS_LAMBDA_MICROVM_IMAGE_ARN=image_arn)): + assert in_aws_lambda_microvm() is False + + def test_is_azure_function(): with override_env(dict(FUNCTIONS_WORKER_RUNTIME="python", FUNCTIONS_EXTENSION_VERSION="2")): assert in_azure_function() is True