From 9ae095700e6785cfcd924d1746e6b2190e415c4e Mon Sep 17 00:00:00 2001 From: Alberto Vara Date: Fri, 20 Sep 2024 13:17:55 +0200 Subject: [PATCH 1/2] chore(iast): fastapi body --- ddtrace/appsec/_iast/_patch.py | 14 ++++++- .../fastapi/test_fastapi_appsec_iast.py | 42 ++++++++++++++++++- 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/ddtrace/appsec/_iast/_patch.py b/ddtrace/appsec/_iast/_patch.py index 6d3226f983..8fba9ee5a1 100644 --- a/ddtrace/appsec/_iast/_patch.py +++ b/ddtrace/appsec/_iast/_patch.py @@ -171,7 +171,7 @@ def _on_iast_fastapi_patch(): _set_metric_iast_instrumented_source(OriginType.PATH) # Body source - try_wrap_function_wrapper("starlette.requests", "Request.__init__", _iast_instrument_starlette_request) + try_wrap_function_wrapper("starlette.requests", "Request.body", _iast_instrument_starlette_request_body) _set_metric_iast_instrumented_source(OriginType.BODY) # Instrumented on _iast_starlette_scope_taint @@ -212,6 +212,18 @@ async def wrapped_property_call(): instance.__class__.receive = property(receive) +async def _iast_instrument_starlette_request_body(wrapped, instance, args, kwargs): + from ddtrace.appsec._iast._taint_tracking import OriginType + from ddtrace.appsec._iast._taint_tracking import origin_to_str + from ddtrace.appsec._iast._taint_tracking import taint_pyobject + + result = await wrapped(*args, **kwargs) + + return taint_pyobject( + result, source_name=origin_to_str(OriginType.PATH), source_value=result, source_origin=OriginType.BODY + ) + + def _iast_instrument_starlette_scope(scope): from ddtrace.appsec._iast._taint_tracking import OriginType from ddtrace.appsec._iast._taint_tracking import taint_pyobject diff --git a/tests/contrib/fastapi/test_fastapi_appsec_iast.py b/tests/contrib/fastapi/test_fastapi_appsec_iast.py index 7538e832b0..ffcb361525 100644 --- a/tests/contrib/fastapi/test_fastapi_appsec_iast.py +++ b/tests/contrib/fastapi/test_fastapi_appsec_iast.py @@ -301,7 +301,7 @@ async def test_route(request: Request): @pytest.mark.usefixtures("setup_core_ok_after_test") -def test_path_body_source(fastapi_application, client, tracer, test_spans): +def test_path_body_receive_source(fastapi_application, client, tracer, test_spans): @fastapi_application.post("/index.html") async def test_route(request: Request): from ddtrace.appsec._iast._taint_tracking import get_tainted_ranges @@ -341,6 +341,46 @@ async def test_route(request: Request): assert result["ranges_origin"] == "http.request.body" +@pytest.mark.usefixtures("setup_core_ok_after_test") +def test_path_body_body_source(fastapi_application, client, tracer, test_spans): + @fastapi_application.post("/index.html") + async def test_route(request: Request): + from ddtrace.appsec._iast._taint_tracking import get_tainted_ranges + from ddtrace.appsec._iast._taint_tracking import origin_to_str + + body = await request.body() + ranges_result = get_tainted_ranges(body) + + return JSONResponse( + { + "result": str(body, encoding="utf-8"), + "is_tainted": len(ranges_result), + "ranges_start": ranges_result[0].start, + "ranges_length": ranges_result[0].length, + "ranges_origin": origin_to_str(ranges_result[0].source.origin), + } + ) + + # test if asgi middleware is ok without any callback registered + # core.reset_listeners(event_id="asgi.request.parse.body") + + with override_global_config(dict(_iast_enabled=True)), override_env(IAST_ENV): + # disable callback + _aux_appsec_prepare_tracer(tracer) + resp = client.post( + "/index.html", + data='{"name": "yqrweytqwreasldhkuqwgervflnmlnli"}', + headers={"Content-Type": "application/json"}, + ) + assert resp.status_code == 200 + result = json.loads(get_response_body(resp)) + assert result["result"] == '{"name": "yqrweytqwreasldhkuqwgervflnmlnli"}' + assert result["is_tainted"] == 1 + assert result["ranges_start"] == 0 + assert result["ranges_length"] == 44 + assert result["ranges_origin"] == "http.request.body" + + @pytest.mark.skip(reason="Pydantic not supported yet APPSEC-52941") def test_path_body_source_pydantic(fastapi_application, client, tracer, test_spans): from pydantic import BaseModel From b5e483d05a29a5eeedc702d9b3370cf83ff6fd19 Mon Sep 17 00:00:00 2001 From: Alberto Vara Date: Fri, 20 Sep 2024 13:44:31 +0200 Subject: [PATCH 2/2] chore(iast): fastapi body --- ddtrace/appsec/_iast/_patch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ddtrace/appsec/_iast/_patch.py b/ddtrace/appsec/_iast/_patch.py index 8fba9ee5a1..f3b5dd6da2 100644 --- a/ddtrace/appsec/_iast/_patch.py +++ b/ddtrace/appsec/_iast/_patch.py @@ -171,6 +171,7 @@ def _on_iast_fastapi_patch(): _set_metric_iast_instrumented_source(OriginType.PATH) # Body source + try_wrap_function_wrapper("starlette.requests", "Request.__init__", _iast_instrument_starlette_request) try_wrap_function_wrapper("starlette.requests", "Request.body", _iast_instrument_starlette_request_body) _set_metric_iast_instrumented_source(OriginType.BODY)