fix(asm): fix receive wrapper for asgi [backport 2.6] (#8701)

- Fix receive wrapper for asynchronous requests with more than one
message received.
- Add a test_streaming (only for FastAPI) to cover the fix
- Add one other possible large request body for threat blocking tests on
request body

## Checklist

- [x] Change(s) are motivated and described in the PR description
- [x] Testing strategy is described if automated tests are not included
in the PR
- [x] Risks are described (performance impact, potential for breakage,
maintainability)
- [x] Change is maintainable (easy to change, telemetry, documentation)
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed or label `changelog/no-changelog` is set
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/))
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))
- [x] If this PR changes the public interface, I've notified
`@DataDog/apm-tees`.
- [x] If change touches code that signs or publishes builds or packages,
or handles credentials of any kind, I've requested a review from
`@DataDog/security-design-and-guidance`.

## Reviewer Checklist

- [x] Title is accurate
- [x] All changes are related to the pull request's stated goal
- [x] Description motivates each change
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- [x] Testing strategy adequately addresses listed risks
- [x] Change is maintainable (easy to change, telemetry, documentation)
- [x] Release note makes sense to a user of the library
- [x] Author has acknowledged and discussed the performance implications
of this PR as reported in the benchmarks PR comment
- [x] Backport labels are set in a manner that is consistent with the
[release branch maintenance

policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

(cherry picked from commit 28a03cc)

Author: christophe-papazian

ddtrace/appsec/_handlers.py

@@ -95,8 +95,11 @@ async def _on_asgi_request_parse_body(receive, headers):
         data_received = await receive()
         body = data_received.get("body", b"")
 
-        async def receive():
-            return data_received
+        async def receive_wrapped(once=[True]):
+            if once[0]:
+                once[0] = False
+                return data_received
+            return await receive()
 
         content_type = headers.get("content-type") or headers.get("Content-Type")
         try:
@@ -111,9 +114,9 @@ async def receive():
                 req_body = None
             else:
                 req_body = parse_form_multipart(body.decode(), headers) or None
-            return receive, req_body
+            return receive_wrapped, req_body
         except BaseException:
-            return receive, None
+            return receive_wrapped, None
 
     return receive, None
 

releasenotes/notes/fix_asgi_receive_wrapper-685522a1af75d59d.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    ASM: This fix resolves an issue where the asgi middleware could crash with a RuntimeError "Unexpected message received".

tests/appsec/contrib_appsec/fastapi_app/app.py

@@ -1,9 +1,11 @@
+import asyncio
 from typing import Optional
 
 from fastapi import FastAPI
 from fastapi import Request
 from fastapi.responses import HTMLResponse
 from fastapi.responses import JSONResponse
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 
 
@@ -83,4 +85,13 @@ async def new_service(service_name: str, request: Request):  # noqa: B008
         ddtrace.Pin.override(app, service=service_name, tracer=ddtrace.tracer)
         return HTMLResponse(service_name, 200)
 
+    async def slow_numbers(minimum, maximum):
+        for number in range(minimum, maximum):
+            yield "%d" % number
+            await asyncio.sleep(0.25)
+
+    @app.get("/stream/")
+    async def stream():
+        return StreamingResponse(slow_numbers(0, 10), media_type="text/html")
+
     return app

tests/appsec/contrib_appsec/utils.py

@@ -660,13 +660,20 @@ def test_request_suspicious_request_block_match_response_headers(
                 assert get_tag(http.STATUS_CODE) == "200"
                 assert get_tag(APPSEC.JSON) is None
 
+    LARGE_BODY = {
+        f"key_{i}": {f"key_{i}_{j}": {f"key_{i}_{j}_{k}": f"value_{i}_{j}_{k}" for k in range(4)} for j in range(4)}
+        for i in range(254)
+    }
+    LARGE_BODY["attack"] = "yqrweytqwreasldhkuqwgervflnmlnli"
+
     @pytest.mark.parametrize("asm_enabled", [True, False])
     @pytest.mark.parametrize(
         ("body", "content_type", "blocked"),
         [
             # json body must be blocked
             ('{"attack": "yqrweytqwreasldhkuqwgervflnmlnli"}', "application/json", "tst-037-003"),
             ('{"attack": "yqrweytqwreasldhkuqwgervflnmlnli"}', "text/json", "tst-037-003"),
+            (json.dumps(LARGE_BODY), "text/json", "tst-037-003"),
             # xml body must be blocked
             (
                 '<?xml version="1.0" encoding="UTF-8"?><attack>yqrweytqwreasldhkuqwgervflnmlnli</attack>',
@@ -687,6 +694,7 @@ def test_request_suspicious_request_block_match_response_headers(
             # other values must not be blocked
             ('{"attack": "zqrweytqwreasldhkuqxgervflnmlnli"}', "application/json", False),
         ],
+        ids=["json", "text_json", "json_large", "xml", "form", "form_multipart", "text", "no_attack"],
     )
     def test_request_suspicious_request_block_match_request_body(
         self, interface: Interface, get_tag, asm_enabled, root_span, body, content_type, blocked
@@ -1001,6 +1009,25 @@ def test_global_callback_list_length(self, interface):
             # only two global callbacks are expected for API Security and Nested Events
             assert len(_asm_request_context.GLOBAL_CALLBACKS.get(_asm_request_context._CONTEXT_CALL, [])) == 2
 
+    @pytest.mark.parametrize("asm_enabled", [True, False])
+    @pytest.mark.parametrize("metastruct", [True, False])
+    def test_stream_response(
+        self,
+        interface: Interface,
+        get_tag,
+        asm_enabled,
+        metastruct,
+        root_span,
+    ):
+        if interface.name != "fastapi":
+            raise pytest.skip("only fastapi tests have support for stream response")
+        with override_global_config(
+            dict(_asm_enabled=asm_enabled, _use_metastruct_for_triggers=metastruct)
+        ), override_env(dict(DD_APPSEC_RULES=rules.RULES_SRB)):
+            self.update_tracer(interface)
+            response = interface.client.get("/stream/")
+            assert self.body(response) == "0123456789"
+
 
 @contextmanager
 def test_tracer():