fix(asm): fix flask block on path parameters (#5640)

Fix path parameters address handling in Flask framework.

- move the call to the waf from the wsgi layer to the wrapper layer for
views to handle the call just before the request is handled by the
client code but just after the path parameters were computed
- instead of blocking before the request, if a block is occuring, the
call to the view is replaced by a function stored in the asm context
- create a specific wrap_function for views to avoid polluting other
functions with specific instrumentation of IAST and ASM

## 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] PR description includes explicit acknowledgement/acceptance of the
performance implications of this PR as reported in the benchmarks PR
comment.

## Reviewer Checklist

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

Author: christophe-papazian

ddtrace/contrib/flask/patch.py

@@ -49,6 +49,7 @@
 from .helpers import with_instance_pin
 from .wrappers import wrap_function
 from .wrappers import wrap_signal
+from .wrappers import wrap_view
 
 
 try:
@@ -244,9 +245,6 @@ def _request_span_modifier(self, span, environ, parsed_headers=None):
             request_body=req_body,
             peer_ip=request.remote_addr,
         )
-        if config._appsec_enabled:
-            log.debug("Flask WAF call for Suspicious Request Blocking on request")
-            _asm_request_context.call_waf_callback()
 
 
 def patch():
@@ -524,7 +522,7 @@ def _wrap(rule, endpoint=None, view_func=None, **kwargs):
         if view_func:
             # TODO: `if hasattr(view_func, 'view_class')` then this was generated from a `flask.views.View`
             #   should we do something special with these views? Change the name/resource? Add tags?
-            view_func = wrap_function(instance, view_func, name=endpoint, resource=rule)
+            view_func = wrap_view(instance, view_func, name=endpoint, resource=rule)
 
         return wrapped(rule, endpoint=endpoint, view_func=view_func, **kwargs)
 

ddtrace/contrib/flask/wrappers.py

@@ -1,15 +1,21 @@
 from ddtrace import config
+import ddtrace.appsec._asm_request_context as _asmrc
+from ddtrace.appsec._constants import SPAN_DATA_NAMES
 from ddtrace.appsec.iast._util import _is_iast_enabled
 from ddtrace.internal.constants import COMPONENT
 from ddtrace.vendor.wrapt import function_wrapper
 
 from .. import trace_utils
+from ...internal.logger import get_logger
 from ...internal.utils.importlib import func_name
 from ...pin import Pin
 from .helpers import get_current_app
 
 
-def wrap_function(instance, func, name=None, resource=None):
+log = get_logger(__name__)
+
+
+def wrap_view(instance, func, name=None, resource=None):
     """
     Helper function to wrap common flask.app.Flask methods.
 
@@ -26,6 +32,17 @@ def trace_func(wrapped, _instance, args, kwargs):
         with pin.tracer.trace(name, service=trace_utils.int_service(pin, config.flask), resource=resource) as span:
             span.set_tag_str(COMPONENT, config.flask.integration_name)
 
+            # if Appsec is enabled, we can try to block as we have the path parameters at that point
+            if config._appsec_enabled and _asmrc.in_context():
+                log.debug("Flask WAF call for Suspicious Request Blocking on request")
+                if kwargs:
+                    _asmrc.set_waf_address(SPAN_DATA_NAMES.REQUEST_PATH_PARAMS, kwargs)
+                _asmrc.call_waf_callback()
+                if _asmrc.is_blocked():
+                    callback_block = _asmrc.get_value(_asmrc._CALLBACKS, "flask_block")
+                    if callback_block:
+                        return callback_block()
+
             # If IAST is enabled, taint the Flask function kwargs (path parameters)
             if _is_iast_enabled() and kwargs:
                 from ddtrace.appsec.iast._input_info import Input_info
@@ -39,6 +56,27 @@ def trace_func(wrapped, _instance, args, kwargs):
     return trace_func(func)
 
 
+def wrap_function(instance, func, name=None, resource=None):
+    """
+    Helper function to wrap common flask.app.Flask methods.
+
+    This helper will first ensure that a Pin is available and enabled before tracing
+    """
+    if not name:
+        name = func_name(func)
+
+    @function_wrapper
+    def trace_func(wrapped, _instance, args, kwargs):
+        pin = Pin._find(wrapped, _instance, instance, get_current_app())
+        if not pin or not pin.enabled():
+            return wrapped(*args, **kwargs)
+        with pin.tracer.trace(name, service=trace_utils.int_service(pin, config.flask), resource=resource) as span:
+            span.set_tag_str(COMPONENT, config.flask.integration_name)
+            return wrapped(*args, **kwargs)
+
+    return trace_func(func)
+
+
 def wrap_signal(app, signal, func):
     """
     Helper used to wrap signal handlers

ddtrace/contrib/wsgi/wsgi.py

@@ -170,20 +170,18 @@ def __call__(self, environ, start_response):
                     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)]
+
+                _asm_request_context.set_value(_asm_request_context._CALLBACKS, "flask_block", blocked_view)
+
             if not_blocked:
                 req_span.set_tag_str(COMPONENT, self._config.integration_name)
                 # set span.kind to the type of operation being performed
                 req_span.set_tag_str(SPAN_KIND, SpanKind.SERVER)
                 self._request_span_modifier(req_span, environ)
-                if self.tracer._appsec_enabled:
-                    # [Suspicious Request Blocking on request]
-                    if _context.get_item("http.request.blocked", span=req_span):
-                        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:
                 try:
                     app_span = self.tracer.trace(self._application_span_name)
 
@@ -202,7 +200,7 @@ def __call__(self, environ, start_response):
                     req_span.finish()
                     raise
                 if self.tracer._appsec_enabled and _context.get_item("http.request.blocked", span=req_span):
-                    # [Suspicious Request Blocking on response]
+                    # [Suspicious Request Blocking on request or response]
                     _, content = self._make_block_content(environ, headers, req_span)
                     closing_iterator = [content]
 

releasenotes/notes/fix_path_parameters_for_flask_srb-bdac4910826d9c93.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    ASM: This fix resolves an issue where path parameters for the Flask framework were handled at
+    response time instead of at request time for suspicious request blocking. 
+    This close a known issue opened in 1.10.0.