tests passing checkpoint - flask "block"

ddtrace/appsec/_asm_request_context.py

@@ -322,6 +322,7 @@ def asm_request_context_manager(
         resources = _DataHandler()
         asm_request_context_set(remote_ip, headers, headers_case_sensitive, block_request_callable)
         core.on("wsgi.block_decided", _on_block_decided)
+        core.on("wsgi._make_block_content", _on_make_block_content)
         try:
             yield resources
         finally:
@@ -330,5 +331,10 @@ def asm_request_context_manager(
         yield None
 
 
+def _on_make_block_content(content, ctype, span):
+    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)
+
+
 def _on_block_decided(callback):
     set_value(_CALLBACKS, "flask_block", callback)

ddtrace/contrib/wsgi/wsgi.py

@@ -5,7 +5,6 @@
 from ddtrace.appsec import _asm_request_context
 from ddtrace.internal.schema.span_attribute_schema import SpanDirection
 
-from ...appsec._constants import SPAN_DATA_NAMES
 from ..trace_utils import _get_request_header_user_agent
 from ..trace_utils import _set_url_tag
 
@@ -138,8 +137,7 @@ 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)
+            core.dispatch("wsgi._make_block_content", [content, ctype, span])
             span.set_tag_str(http.STATUS_CODE, "403")
             url = construct_url(environ)
             query_string = environ.get("QUERY_STRING")
@@ -173,20 +171,17 @@ def __call__(self, environ, start_response):
                 span_type=SpanTypes.WEB,
             )
 
-            if self.tracer._appsec_enabled:
-                # [IP Blocking]
-                if core.get_item(WAF_CONTEXT_NAMES.BLOCKED):
-                    ctype, content = self._make_block_content(environ, headers, req_span)
-                    start_response("403 FORBIDDEN", [("content-type", ctype)])
-                    closing_iterator = [content]
-                    not_blocked = False
+            if core.get_item(WAF_CONTEXT_NAMES.BLOCKED):
+                ctype, content = self._make_block_content(environ, headers, req_span)
+                start_response("403 FORBIDDEN", [("content-type", ctype)])
+                closing_iterator = [content]
+                not_blocked = False
 
-                # [Suspicious Request Blocking on request]
-                def blocked_view():
-                    ctype, content = self._make_block_content(environ, headers, req_span)
-                    return content, 403, [("content-type", ctype)]
+            def blocked_view():
+                ctype, content = self._make_block_content(environ, headers, req_span)
+                return content, 403, [("content-type", ctype)]
 
-                core.dispatch("wsgi.block_decided", [blocked_view])
+            core.dispatch("wsgi.block_decided", [blocked_view])
 
             if not_blocked:
                 req_span.set_tag_str(COMPONENT, self._config.integration_name)
@@ -210,8 +205,7 @@ def blocked_view():
                     app_span.finish()
                     req_span.finish()
                     raise
-                if self.tracer._appsec_enabled and core.get_item(WAF_CONTEXT_NAMES.BLOCKED):
-                    # [Suspicious Request Blocking on request or response]
+                if core.get_item(WAF_CONTEXT_NAMES.BLOCKED):
                     _, content = self._make_block_content(environ, headers, req_span)
                     closing_iterator = [content]