chore(iast): fix stacktrace leak in flask (#12508)

- `werkzeug.DebugTraceback.render_debugger_html` runs outside of
`iast_request_context`. Because of this, when this function is called,
we store the result to report it in the next request when there's
context and the span hasn't been sent yet.
- Support for Werkzeug <= 2.0.3
## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- 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)

Author: avara1986

ddtrace/appsec/_iast/_handlers.py

@@ -438,10 +438,10 @@ def _on_django_finalize_response_pre(ctx, after_request_tags, request, response)
         return
 
     try:
-        from .taint_sinks.stacktrace_leak import asm_check_stacktrace_leak
+        from .taint_sinks.stacktrace_leak import iast_check_stacktrace_leak
 
         content = response.content.decode("utf-8", errors="ignore")
-        asm_check_stacktrace_leak(content)
+        iast_check_stacktrace_leak(content)
     except Exception:
         log.debug("Unexpected exception checking for stacktrace leak", exc_info=True)
 
@@ -470,10 +470,11 @@ def _on_flask_finalize_request_post(response, _):
         return
 
     try:
-        from .taint_sinks.stacktrace_leak import asm_check_stacktrace_leak
+        from .taint_sinks.stacktrace_leak import iast_check_stacktrace_leak
 
         content = response[0].decode("utf-8", errors="ignore")
-        asm_check_stacktrace_leak(content)
+
+        iast_check_stacktrace_leak(content)
     except Exception:
         log.debug("Unexpected exception checking for stacktrace leak", exc_info=True)
 
@@ -483,22 +484,23 @@ def _on_asgi_finalize_response(body, _):
         return
 
     try:
-        from .taint_sinks.stacktrace_leak import asm_check_stacktrace_leak
+        from .taint_sinks.stacktrace_leak import iast_check_stacktrace_leak
 
         content = body.decode("utf-8", errors="ignore")
-        asm_check_stacktrace_leak(content)
+        iast_check_stacktrace_leak(content)
     except Exception:
         log.debug("Unexpected exception checking for stacktrace leak", exc_info=True)
 
 
 def _on_werkzeug_render_debugger_html(html):
-    if not html or not asm_config._iast_enabled or not asm_config.is_iast_request_enabled:
+    # we don't check asm_config.is_iast_request_enabled due to werkzeug.render_debugger_html works outside the request
+    if not html or not asm_config._iast_enabled:
         return
 
     try:
-        from .taint_sinks.stacktrace_leak import asm_check_stacktrace_leak
+        from .taint_sinks.stacktrace_leak import iast_check_stacktrace_leak
 
-        asm_check_stacktrace_leak(html)
+        iast_check_stacktrace_leak(html)
         set_iast_stacktrace_reported(True)
     except Exception:
         log.debug("Unexpected exception checking for stacktrace leak", exc_info=True)

ddtrace/appsec/_iast/_taint_tracking/_errors.py

@@ -9,8 +9,12 @@
 
 
 def iast_taint_log_error(msg):
+    iast_error(msg, default_prefix="[IAST] Propagation error")
+
+
+def iast_error(msg, default_prefix="[IAST] "):
     if _is_iast_debug_enabled():
         stack = inspect.stack()
         frame_info = "\n".join("%s %s" % (frame_info.filename, frame_info.lineno) for frame_info in stack[:7])
-        log.debug("[IAST] Propagation error. %s:\n%s", msg, frame_info)
-    _set_iast_error_metric("[IAST] Propagation error. %s" % msg)
+        log.debug("%s. %s:\n%s", default_prefix, msg, frame_info)
+    _set_iast_error_metric("%s. %s" % (default_prefix, msg))

ddtrace/appsec/_iast/processor.py

@@ -22,8 +22,10 @@ def __post_init__(self) -> None:
     def on_span_start(self, span: Span):
         if span.span_type != SpanTypes.WEB:
             return
-
         _iast_start_request(span)
+        from .taint_sinks.stacktrace_leak import check_and_report_stacktrace_leak
+
+        check_and_report_stacktrace_leak()
 
     def on_span_finish(self, span: Span):
         """Report reported vulnerabilities.

ddtrace/appsec/_iast/taint_sinks/stacktrace_leak.py

@@ -1,12 +1,14 @@
 import os
 import re
 
+from ddtrace.settings.asm import config as asm_config
+
 from ..._constants import IAST_SPAN_TAGS
 from .. import oce
 from .._iast_request_context import set_iast_stacktrace_reported
 from .._metrics import _set_metric_iast_executed_sink
 from .._metrics import increment_iast_span_metric
-from .._taint_tracking._errors import iast_taint_log_error
+from .._taint_tracking._errors import iast_error
 from ..constants import HTML_TAGS_REMOVE
 from ..constants import STACKTRACE_EXCEPTION_REGEX
 from ..constants import STACKTRACE_FILE_LINE
@@ -28,13 +30,39 @@ def asm_report_stacktrace_leak_from_django_debug_page(exc_name, module):
     set_iast_stacktrace_reported(True)
 
 
-def asm_check_stacktrace_leak(content: str) -> None:
+# `werkzeug.DebugTraceback.render_debugger_html` runs outside of `iast_request_context`. Because of this, when
+# this function is called, we store the result to report it in the next request when there's context and the
+# span hasn't been sent yet.
+REPORT_STACKTRACE_LATER = None
+
+
+def check_and_report_stacktrace_leak():
+    report = get_report_stacktrace_later()
+    if report:
+        StacktraceLeak.report(evidence_value=report)
+        set_report_stacktrace_later(None)
+
+
+def set_report_stacktrace_later(evidence):
+    global REPORT_STACKTRACE_LATER
+    REPORT_STACKTRACE_LATER = evidence
+
+
+def get_report_stacktrace_later():
+    return REPORT_STACKTRACE_LATER
+
+
+def iast_check_stacktrace_leak(content: str) -> None:
     if not content:
         return
 
     try:
         # Quick check to avoid the slower operations if on stacktrace
-        if "Traceback (most recent call last):" not in content:
+        if (
+            "Traceback (most recent call last):" not in content
+            and "Traceback <em>(most recent call last)" not in content
+            and '<div class="traceback">' not in content
+        ):
             return
 
         text = HTML_TAGS_REMOVE.sub("", content)
@@ -94,9 +122,11 @@ def asm_check_stacktrace_leak(content: str) -> None:
             if not module_name:
                 module_name = module_path  # fallback: just the path
 
-        increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, StacktraceLeak.vulnerability_type)
         _set_metric_iast_executed_sink(StacktraceLeak.vulnerability_type)
         evidence = "Module: %s\nException: %s" % (module_name.strip(), exception_line.strip())
-        StacktraceLeak.report(evidence_value=evidence)
+        if asm_config.is_iast_request_enabled:
+            StacktraceLeak.report(evidence_value=evidence)
+        else:
+            set_report_stacktrace_later(evidence)
     except Exception as e:
-        iast_taint_log_error("[IAST] error in check stacktrace leak. {}".format(e))
+        iast_error("Taint Sink error. Error in check stacktrace leak. {}".format(e))

ddtrace/contrib/internal/flask/patch.py

@@ -77,6 +77,11 @@ def get_version():
     return get_version_for_package("flask")
 
 
+def get_werkzeug_version():
+    # type: () -> str
+    return get_version_for_package("werkzeug")
+
+
 if _HAS_JSON_MIXIN:
 
     class RequestWithJson(werkzeug.Request, JSONMixin):
@@ -101,6 +106,9 @@ class RequestWithJson(werkzeug.Request, JSONMixin):
 flask_version_str = get_version()
 flask_version = parse_version(flask_version_str)
 
+werkzeug_version_str = get_werkzeug_version()
+werkzeug_version = parse_version(werkzeug_version_str)
+
 
 class _FlaskWSGIMiddleware(_DDWSGIMiddlewareBase):
     _request_call_name = schematize_url_operation("flask.request", protocol="http", direction=SpanDirection.INBOUND)
@@ -224,10 +232,16 @@ def patch():
     _w("flask.templating", "_render", patched_render)
     _w("flask", "render_template", _build_render_template_wrapper("render_template"))
     _w("flask", "render_template_string", _build_render_template_wrapper("render_template_string"))
-    try:
-        _w("werkzeug.debug.tbtools", "DebugTraceback.render_debugger_html", patched_render_debugger_html)
-    except AttributeError:
-        log.debug("Failed to patch DebugTraceback.render_debugger_html, not supported by this werkzeug version")
+    if werkzeug_version >= (2, 1, 0):
+        try:
+            _w("werkzeug.debug.tbtools", "DebugTraceback.render_debugger_html", patched_render_debugger_html)
+        except AttributeError:
+            log.debug("Failed to patch DebugTraceback.render_debugger_html, not supported by this werkzeug version")
+    else:
+        try:
+            _w("werkzeug.debug.tbtools", "Traceback.render_summary", patched_render_debugger_html)
+        except AttributeError:
+            log.debug("Failed to patch Traceback.render_summary, not supported by this werkzeug version")
 
     bp_hooks = [
         "after_app_request",

tests/appsec/app.py

@@ -1,6 +1,5 @@
 """ This Flask application is imported on tests.appsec.appsec_utils.gunicorn_server
 """
-
 import os
 import re
 import subprocess  # nosec
@@ -11,7 +10,9 @@
 
 
 import ddtrace.auto  # noqa: F401  # isort: skip
+from ddtrace import tracer
 from ddtrace.appsec._iast import ddtrace_iast_flask_patch  # noqa: F401
+from ddtrace.internal.utils.formats import asbool
 from tests.appsec.iast_packages.packages.pkg_aiohttp import pkg_aiohttp
 from tests.appsec.iast_packages.packages.pkg_aiosignal import pkg_aiosignal
 from tests.appsec.iast_packages.packages.pkg_annotated_types import pkg_annotated_types
@@ -193,10 +194,21 @@ def iast_cmdi_vulnerability():
     subp.communicate()
     subp.wait()
     resp = Response("OK")
-    resp.set_cookie("insecure", "cookie", secure=True, httponly=True, samesite="None")
     return resp
 
 
+@app.route("/shutdown", methods=["GET"])
+def shutdown_view():
+    tracer._writer.flush_queue()
+    return "OK"
+
+
+@app.route("/iast-stacktrace-leak-vulnerability", methods=["GET"])
+def iast_stacktrace_vulnerability():
+    raise ValueError("Check my stacktrace!")
+    return "OK"
+
+
 @app.route("/iast-weak-hash-vulnerability", methods=["GET"])
 def iast_weak_hash_vulnerability():
     import _md5
@@ -970,5 +982,6 @@ def iast_ast_patching_non_re_search():
 
 if __name__ == "__main__":
     env_port = os.getenv("FLASK_RUN_PORT", 8000)
+    debug = asbool(os.getenv("FLASK_DEBUG", "false"))
     ddtrace_iast_flask_patch()
-    app.run(debug=False, port=env_port)
+    app.run(debug=debug, port=env_port)

tests/appsec/appsec_utils.py

@@ -90,6 +90,7 @@ def appsec_application_server(
     env["DD_REMOTE_CONFIGURATION_ENABLED"] = remote_configuration_enabled
     if token:
         env["_DD_REMOTE_CONFIGURATION_ADDITIONAL_HEADERS"] = "X-Datadog-Test-Session-Token:%s," % (token,)
+        env["_DD_TRACE_WRITER_ADDITIONAL_HEADERS"] = "X-Datadog-Test-Session-Token:{}".format(token)
     if appsec_enabled is not None:
         env["DD_APPSEC_ENABLED"] = appsec_enabled
     if apm_tracing_enabled is not None:

tests/appsec/iast/taint_sinks/test_stacktrace_leak.py

@@ -1,7 +1,12 @@
 import os
 
 from ddtrace.appsec._iast.constants import VULN_STACKTRACE_LEAK
-from ddtrace.appsec._iast.taint_sinks.stacktrace_leak import asm_check_stacktrace_leak
+from ddtrace.appsec._iast.taint_sinks._base import VulnerabilityBase
+from ddtrace.appsec._iast.taint_sinks.stacktrace_leak import check_and_report_stacktrace_leak
+from ddtrace.appsec._iast.taint_sinks.stacktrace_leak import get_report_stacktrace_later
+from ddtrace.appsec._iast.taint_sinks.stacktrace_leak import iast_check_stacktrace_leak
+from tests.appsec.iast.conftest import _end_iast_context_and_oce
+from tests.appsec.iast.conftest import _start_iast_context_and_oce
 from tests.appsec.iast.taint_sinks.conftest import _get_span_report
 
 
@@ -13,8 +18,8 @@ def _load_text_stacktrace():
     return open(os.path.join(os.path.dirname(__file__), "../fixtures/plain_stacktrace.txt")).read()
 
 
-def test_asm_check_stacktrace_leak_html(iast_context_defaults):
-    asm_check_stacktrace_leak(_load_html_django_stacktrace())
+def test_check_stacktrace_leak_html(iast_context_defaults):
+    iast_check_stacktrace_leak(_load_html_django_stacktrace())
     span_report = _get_span_report()
     vulnerabilities = list(span_report.vulnerabilities)
     vulnerabilities_types = [vuln.type for vuln in vulnerabilities]
@@ -26,8 +31,8 @@ def test_asm_check_stacktrace_leak_html(iast_context_defaults):
     )
 
 
-def test_asm_check_stacktrace_leak_text(iast_context_defaults):
-    asm_check_stacktrace_leak(_load_text_stacktrace())
+def test_check_stacktrace_leak_text(iast_context_defaults):
+    iast_check_stacktrace_leak(_load_text_stacktrace())
     span_report = _get_span_report()
     vulnerabilities = list(span_report.vulnerabilities)
     vulnerabilities_types = [vuln.type for vuln in vulnerabilities]
@@ -37,3 +42,47 @@ def test_asm_check_stacktrace_leak_text(iast_context_defaults):
         vulnerabilities[0].evidence.value
         == 'Module: ".usr.local.lib.python3.9.site-packages.constraints.py"\nException: ValueError'
     )
+    VulnerabilityBase._prepare_report._reset_cache()
+
+
+def test_stacktrace_leak_deduplication(iast_context_deduplication_enabled):
+    for num_vuln_expected in [1, 0, 0]:
+        _start_iast_context_and_oce()
+        for _ in range(0, 5):
+            iast_check_stacktrace_leak(_load_text_stacktrace())
+
+        span_report = _get_span_report()
+
+        if num_vuln_expected == 0:
+            assert span_report is None
+        else:
+            assert span_report
+
+            assert len(span_report.vulnerabilities) == num_vuln_expected
+            vulnerability = list(span_report.vulnerabilities)[0]
+            assert vulnerability.type == VULN_STACKTRACE_LEAK
+        _end_iast_context_and_oce()
+    VulnerabilityBase._prepare_report._reset_cache()
+
+
+def test_check_stacktrace_leak_text_outside_context(iast_context_deduplication_enabled):
+    _end_iast_context_and_oce()
+
+    # Report stacktrace outside the context
+    iast_check_stacktrace_leak(_load_text_stacktrace())
+    assert get_report_stacktrace_later() is not None
+    _start_iast_context_and_oce()
+
+    # Check the stacktrace, now with a context, like the beginning of a request
+    check_and_report_stacktrace_leak()
+    span_report = _get_span_report()
+    vulnerabilities = list(span_report.vulnerabilities)
+    vulnerabilities_types = [vuln.type for vuln in vulnerabilities]
+    assert len(vulnerabilities) == 1
+    assert VULN_STACKTRACE_LEAK in vulnerabilities_types
+    assert (
+        vulnerabilities[0].evidence.value
+        == 'Module: ".usr.local.lib.python3.9.site-packages.constraints.py"\nException: ValueError'
+    )
+    assert get_report_stacktrace_later() is None
+    _end_iast_context_and_oce()

tests/appsec/integrations/flask_tests/test_iast_flask_telemetry.py

@@ -6,7 +6,6 @@
 
 
 def test_iast_span_metrics():
-    # TODO: move tests/telemetry/conftest.py::test_agent_session into a common conftest
     with flask_server(iast_enabled="true", token=None, port=8050) as context:
         _, flask_client, pid = context