fix(asm): blocking response fixes post RC1 (#5264)

## Description

- Ensure a subset of tags are always set on blocked requests.
- Fix Flask sending headers before the actual blocking response on
Suspicious Request matches.
- Ensure the content-type is correctly set on blocking response.

Thanks @christophe-papazian for the help.

## 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] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [X] Change is maintainable (easy to change, telemetry, documentation).
- [X] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/contributing.html#Release-Note-Guidelines)
are followed.
- [X] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [X] Author is aware of the performance implications of this PR as
reported in the benchmarks PR comment.

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer is aware of, and discussed the performance implications
of this PR as reported in the benchmarks PR comment.

---------

Signed-off-by: Juanjo Alvarez <[email protected]>
Co-authored-by: Federico Mon <[email protected]>
Co-authored-by: Alberto Vara <[email protected]>

Author: 3 people

ddtrace/contrib/asgi/middleware.py

@@ -224,7 +224,11 @@ async def wrapped_send(message):
                         span.finish()
 
             async def wrapped_blocked_send(message):
-                accept_header = headers.get("accept")
+                accept_header = (
+                    "text/html"
+                    if "text/html" in _asm_request_context.get_headers().get("accept", "").lower()
+                    else "text/json"
+                )
                 if span and message.get("type") == "http.response.start":
                     message["headers"] = [
                         (

ddtrace/contrib/django/patch.py

@@ -42,6 +42,8 @@
 from . import utils
 from .. import trace_utils
 from ...internal.utils import get_argument_value
+from ..trace_utils import _get_request_header_user_agent
+from ..trace_utils import _set_url_tag
 
 
 log = get_logger(__name__)
@@ -354,8 +356,27 @@ def wrapped_factory(func, instance, args, kwargs):
     return func(*args, **kwargs)
 
 
-def _block_request_callable(span):
+def _set_block_tags(request, request_headers, span):
+    try:
+        span.set_tag_str(http.STATUS_CODE, "403")
+        span.set_tag_str(http.METHOD, request.method)
+        url = utils.get_request_uri(request)
+        query = request.META.get("QUERY_STRING", "")
+        _set_url_tag(config.django, span, url, query)
+        if query and config.django.trace_query_string:
+            span.set_tag_str(http.QUERY_STRING, query)
+        user_agent = _get_request_header_user_agent(request_headers)
+        if user_agent:
+            span.set_tag_str(http.USER_AGENT, user_agent)
+    except Exception as e:
+        log.warning("Could not set some span tags on blocked request: %s", str(e))  # noqa: G200
+
+
+def _block_request_callable(request, request_headers, span):
+    # This is used by user-id blocking to block responses. It could be called
+    # at any point so it's a callable stored in the ASM context.
     _context.set_item("http.request.blocked", True, span=span)
+    _set_block_tags(request, request_headers, span)
     raise PermissionDenied()
 
 
@@ -387,7 +408,9 @@ def traced_get_response(django, pin, func, instance, args, kwargs):
             service=trace_utils.int_service(pin, config.django),
             span_type=SpanTypes.WEB,
         ) as span:
-            _asm_request_context.set_block_request_callable(functools.partial(_block_request_callable, span))
+            _asm_request_context.set_block_request_callable(
+                functools.partial(_block_request_callable, request, request_headers, span)
+            )
             span.set_tag_str(COMPONENT, config.django.integration_name)
 
             # set span.kind to the type of request being performed
@@ -399,10 +422,11 @@ def traced_get_response(django, pin, func, instance, args, kwargs):
             response = None
 
             def blocked_response():
-                response = HttpResponseForbidden(
-                    appsec_utils._get_blocked_template(request_headers.get("Accept")), content_type="text/plain"
-                )
-                response.content = appsec_utils._get_blocked_template(request_headers.get("Accept"))
+                ctype = "text/html" if "text/html" in request_headers.get("Accept", "").lower() else "text/json"
+                content = appsec_utils._get_blocked_template(ctype)
+                response = HttpResponseForbidden(content, content_type=ctype)
+                response.content = content
+                utils._after_request_tags(pin, span, request, response)
                 return response
 
             try:
@@ -460,7 +484,6 @@ def blocked_response():
                     # [Suspicious Request Blocking on response]
                     if _context.get_item("http.request.blocked", span=span):
                         response = blocked_response()
-                        utils._after_request_tags(pin, span, request, response)
                         return response
 
 

ddtrace/contrib/flask/patch.py

@@ -34,10 +34,13 @@
 from ...constants import SPAN_MEASURED_KEY
 from ...contrib.wsgi.wsgi import _DDWSGIMiddlewareBase
 from ...ext import SpanTypes
+from ...ext import http
 from ...internal.compat import maybe_stringify
 from ...internal.logger import get_logger
 from ...internal.utils import get_argument_value
 from ...internal.utils.version import parse_version
+from ..trace_utils import _get_request_header_user_agent
+from ..trace_utils import _set_url_tag
 from ..trace_utils import unwrap as _u
 from .helpers import get_current_app
 from .helpers import simple_tracer
@@ -123,13 +126,23 @@ def _traced_start_response(self, start_response, req_span, app_span, status_code
             req_span, config.flask, status_code=code, response_headers=headers, route=req_span.get_tag(FLASK_URL_RULE)
         )
 
-        result = start_response(status_code, headers)
         if config._appsec_enabled and not _context.get_item("http.request.blocked", span=req_span):
             log.debug("Flask WAF call for Suspicious Request Blocking on response")
             _asm_request_context.call_waf_callback()
             if _context.get_item("http.request.blocked", span=req_span):
-                # response code must be set here or it will be too late
-                result = start_response("403 FORBIDDEN", [])
+                # response code must be set here, or it will be too late
+                ctype = (
+                    "text/html"
+                    if "text/html" in _asm_request_context.get_headers().get("Accept", "").lower()
+                    else "text/json"
+                )
+                response_headers = [("content-type", ctype)]
+                result = start_response("403 FORBIDDEN", response_headers)
+                trace_utils.set_http_meta(req_span, config.flask, status_code="403", response_headers=response_headers)
+            else:
+                result = start_response(status_code, headers)
+        else:
+            result = start_response(status_code, headers)
         return result
 
     def _request_span_modifier(self, span, environ, parsed_headers=None):
@@ -563,9 +576,30 @@ def _wrap(code_or_exception, f):
     return _wrap(*args, **kwargs)
 
 
-def _block_request_callable(headers, span):
+def _set_block_tags(span):
+    span.set_tag_str(http.STATUS_CODE, "403")
+    request = flask.request
+    try:
+        base_url = getattr(request, "base_url", None)
+        query_string = getattr(request, "query_string", None)
+        if base_url and query_string:
+            _set_url_tag(config.flask, span, base_url, query_string)
+        if query_string and config.flask.trace_query_string:
+            span.set_tag_str(http.QUERY_STRING, query_string)
+        if request.method is not None:
+            span.set_tag_str(http.METHOD, request.method)
+        user_agent = _get_request_header_user_agent(request.headers)
+        if user_agent:
+            span.set_tag_str(http.USER_AGENT, user_agent)
+    except Exception as e:
+        log.warning("Could not set some span tags on blocked request: %s", str(e))  # noqa: G200
+
+
+def _block_request_callable(span):
+    request = flask.request
     _context.set_item("http.request.blocked", True, span=span)
-    ctype = headers.get("Accept") or "text/json"
+    _set_block_tags(span)
+    ctype = "text/html" if "text/html" in request.headers.get("Accept", "").lower() else "text/json"
     abort(flask.Response(utils._get_blocked_template(ctype), content_type=ctype, status=403))
 
 
@@ -585,15 +619,12 @@ def _traced_request(pin, wrapped, instance, args, kwargs):
         # This call may be unnecessary since we try to add the tags earlier
         # We just haven't been able to confirm this yet
         _set_request_tags(span)
-        request = flask.request
 
         with pin.tracer.trace(
             ".".join(("flask", name)),
             service=trace_utils.int_service(pin, config.flask, pin),
         ) as request_span:
-            _asm_request_context.set_block_request_callable(
-                functools.partial(_block_request_callable, request.headers, span)
-            )
+            _asm_request_context.set_block_request_callable(functools.partial(_block_request_callable, span))
             request_span.set_tag_str(COMPONENT, config.flask.integration_name)
 
             request_span._ignore_exception(werkzeug.exceptions.NotFound)

ddtrace/contrib/trace_utils.py

@@ -404,6 +404,18 @@ def ext_service(pin, int_config, default=None):
     return default
 
 
+def _set_url_tag(integration_config, span, url, query):
+    # type: (IntegrationConfig, Span, str, str) -> None
+
+    if integration_config.http_tag_query_string:  # Tagging query string in http.url
+        if config.global_query_string_obfuscation_disabled:  # No redacting of query strings
+            span.set_tag_str(http.URL, url)
+        else:  # Redact query strings
+            span.set_tag_str(http.URL, redact_url(url, config._obfuscation_query_string_pattern, query))
+    else:  # Not tagging query string in http.url
+        span.set_tag_str(http.URL, strip_query_string(url))
+
+
 def set_http_meta(
     span,  # type: Span
     integration_config,  # type: IntegrationConfig
@@ -446,14 +458,7 @@ def set_http_meta(
 
     if url is not None:
         url = _sanitized_url(url)
-
-        if integration_config.http_tag_query_string:  # Tagging query string in http.url
-            if config.global_query_string_obfuscation_disabled:  # No redacting of query strings
-                span.set_tag_str(http.URL, url)
-            else:  # Redact query strings
-                span.set_tag_str(http.URL, redact_url(url, config._obfuscation_query_string_pattern, query))
-        else:  # Not tagging query string in http.url
-            span.set_tag_str(http.URL, strip_query_string(url))
+        _set_url_tag(integration_config, span, url, query)
 
     if status_code is not None:
         try:

ddtrace/contrib/wsgi/wsgi.py

@@ -4,6 +4,10 @@
 
 from ddtrace.appsec import _asm_request_context
 
+from ...appsec._constants import SPAN_DATA_NAMES
+from ..trace_utils import _get_request_header_user_agent
+from ..trace_utils import _set_url_tag
+
 
 if TYPE_CHECKING:  # pragma: no cover
     from typing import Any
@@ -119,6 +123,29 @@ def _response_span_name(self):
         "Returns the name of a response span. Example: `flask.response`"
         raise NotImplementedError
 
+    def _make_block_content(self, environ, headers, span):
+        ctype = "text/html" if "text/html" in headers.get("Accept", "").lower() else "text/json"
+        content = utils._get_blocked_template(ctype).encode("UTF-8")
+        try:
+            span.set_tag_str(SPAN_DATA_NAMES.RESPONSE_HEADERS_NO_COOKIES + ".content-length", str(len(content)))
+            span.set_tag_str(SPAN_DATA_NAMES.RESPONSE_HEADERS_NO_COOKIES + ".content-type", ctype)
+            span.set_tag_str(http.STATUS_CODE, "403")
+            url = construct_url(environ)
+            query_string = environ.get("QUERY_STRING")
+            _set_url_tag(self._config, span, url, query_string)
+            if query_string and self._config.trace_query_string:
+                span.set_tag_str(http.QUERY_STRING, query_string)
+            method = environ.get("REQUEST_METHOD")
+            if method:
+                span.set_tag_str(http.METHOD, method)
+            user_agent = _get_request_header_user_agent(headers, headers_are_case_sensitive=True)
+            if user_agent:
+                span.set_tag_str(http.USER_AGENT, user_agent)
+        except Exception as e:
+            log.warning("Could not set some span tags on blocked request: %s", str(e))  # noqa: G200
+
+        return ctype, content
+
     def __call__(self, environ, start_response):
         # type: (Iterable, Callable) -> _TracedIterable
         trace_utils.activate_distributed_headers(self.tracer, int_config=self._config, request_headers=environ)
@@ -138,8 +165,9 @@ def __call__(self, environ, start_response):
             if self.tracer._appsec_enabled:
                 # [IP Blocking]
                 if _context.get_item("http.request.blocked", span=req_span):
-                    start_response("403 FORBIDDEN", [("content-type", headers.get("Accept"))])
-                    closing_iterator = [utils._get_blocked_template(headers.get("Accept")).encode("UTF-8")]
+                    ctype, content = self._make_block_content(environ, headers, req_span)
+                    start_response("403 FORBIDDEN", [("content-type", ctype)])
+                    closing_iterator = [content]
                     not_blocked = False
 
             if not_blocked:
@@ -150,8 +178,9 @@ def __call__(self, environ, start_response):
                 if self.tracer._appsec_enabled:
                     # [Suspicious Request Blocking on request]
                     if _context.get_item("http.request.blocked", span=req_span):
-                        start_response("403 FORBIDDEN", [("content-type", headers.get("Accept"))])
-                        closing_iterator = [utils._get_blocked_template(headers.get("Accept")).encode("UTF-8")]
+                        ctype, content = self._make_block_content(environ, headers, req_span)
+                        start_response("403 FORBIDDEN", [("content-type", ctype)])
+                        closing_iterator = [content]
                         not_blocked = False
 
             if not_blocked:
@@ -174,7 +203,8 @@ def __call__(self, environ, start_response):
                     raise
                 if self.tracer._appsec_enabled and _context.get_item("http.request.blocked", span=req_span):
                     # [Suspicious Request Blocking on response]
-                    closing_iterator = [utils._get_blocked_template(headers.get("Accept")).encode("UTF-8")]
+                    _, content = self._make_block_content(environ, headers, req_span)
+                    closing_iterator = [content]
 
             # start flask.response span. This span will be finished after iter(result) is closed.
             # start_span(child_of=...) is used to ensure correct parenting.

releasenotes/notes/asm-blocking-fixes-fce3bb6a81a2d186.yaml

@@ -0,0 +1,11 @@
+---
+fixes:
+  - |
+    ASM: Solve some corner cases where a Flask blocking request would fail because 
+    headers would be already sent.
+  - |
+    ASM: Solve the content-type not always being correct in blocking responses.
+  - |
+    ASM: Ensure the blocking responses have the following tags: `http.url`, `http.query_string`,
+    `http.useragent`, `http.method`, `http.response.headers.content-type` and 
+    `http.response.headers.content-length`.