feat(asm): add stack trace report for iast (DataDog#11574)

- Add stack trace report to IAST events
- Update test-agent to 1.20.0 to add support for stack traces report
using msgpack (DataDog/images#6234)
- factorize stack trace code already used for exploit prevention, to be
usable for iast too.
- improve code of `ddtrace_iast_flask_patch` so it gets compatible with
the new feature
- improve `test_ddtrace_iast_flask_patch` to check that + operator was
properly replaced.
- add stack trace report unit test on threat tests

System tests will also be enabled with
DataDog/system-tests#3663


APPSEC-55904

## 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)

.circleci/config.templ.yml

@@ -16,7 +16,7 @@ mongo_image: &mongo_image mongo:3.6@sha256:19c11a8f1064fd2bb713ef1270f79a742a184
 httpbin_image: &httpbin_image kennethreitz/httpbin@sha256:2c7abc4803080c22928265744410173b6fea3b898872c01c5fd0f0f9df4a59fb
 vertica_image: &vertica_image vertica/vertica-ce:latest
 rabbitmq_image: &rabbitmq_image rabbitmq:3.7-alpine
-testagent_image: &testagent_image ghcr.io/datadog/dd-apm-test-agent/ddapm-test-agent:v1.17.0
+testagent_image: &testagent_image ghcr.io/datadog/dd-apm-test-agent/ddapm-test-agent:v1.20.0
 
 parameters:
   coverage:

.gitlab/services.yml

@@ -12,7 +12,7 @@
       DD_REMOTE_CONFIGURATION_REFRESH_INTERVAL: 5s
       DD_DOGSTATSD_NON_LOCAL_TRAFFIC: true
   testagent:
-    name: registry.ddbuild.io/images/mirror/dd-apm-test-agent/ddapm-test-agent:v1.17.0
+    name: registry.ddbuild.io/images/mirror/dd-apm-test-agent/ddapm-test-agent:v1.20.0
     alias: testagent
     variables:
       LOG_LEVEL: INFO

ddtrace/appsec/_constants.py

@@ -328,7 +328,6 @@ class DEFAULT(metaclass=Constant_Class):
 
 
 class EXPLOIT_PREVENTION(metaclass=Constant_Class):
-    STACK_TRACES: Literal["_dd.stack"] = "_dd.stack"
     STACK_TRACE_ID: Literal["stack_id"] = "stack_id"
     EP_ENABLED: Literal["DD_APPSEC_RASP_ENABLED"] = "DD_APPSEC_RASP_ENABLED"
     STACK_TRACE_ENABLED: Literal["DD_APPSEC_STACK_TRACE_ENABLED"] = "DD_APPSEC_STACK_TRACE_ENABLED"
@@ -358,3 +357,9 @@ class FINGERPRINTING(metaclass=Constant_Class):
     HEADER = PREFIX + "http.header"
     NETWORK = PREFIX + "http.network"
     SESSION = PREFIX + "session"
+
+
+class STACK_TRACE(metaclass=Constant_Class):
+    RASP = "exploit"
+    IAST = "vulnerability"
+    TAG: Literal["_dd.stack"] = "_dd.stack"

ddtrace/appsec/_exploit_prevention/stack_traces.py

@@ -3,36 +3,48 @@
 from typing import Any
 from typing import Dict
 from typing import Iterable
-from typing import List
 from typing import Optional
 
 from ddtrace._trace.span import Span
-from ddtrace.appsec._constants import EXPLOIT_PREVENTION
+from ddtrace.appsec._constants import STACK_TRACE
 from ddtrace.settings.asm import config as asm_config
 import ddtrace.tracer
 
 
 def report_stack(
-    message: str, span: Optional[Span] = None, crop_stack: Optional[str] = None, stack_id: Optional[str] = None
-):
+    message: Optional[str] = None,
+    span: Optional[Span] = None,
+    crop_stack: Optional[str] = None,
+    stack_id: Optional[str] = None,
+    namespace: str = STACK_TRACE.RASP,
+) -> bool:
     """
     Report a stack trace to the current span.
     This is used to report stack traces for exploit prevention.
     Return the stack id for the reported stack trace to link it in triggers.
     """
-    if not asm_config._ep_enabled or not asm_config._ep_stack_trace_enabled:
-        return None
+    if not asm_config._ep_stack_trace_enabled:
+        # stack trace report disabled
+        return False
+    if namespace == STACK_TRACE.RASP and not (asm_config._asm_enabled and asm_config._ep_enabled):
+        # exploit prevention stack trace with ep disabled
+        return False
+    if namespace == STACK_TRACE.IAST and not (asm_config._iast_enabled):
+        # iast stack trace with iast disabled
+        return False
+
     if span is None:
         span = ddtrace.tracer.current_span()
     if span is None or stack_id is None:
-        return None
+        return False
     root_span = span._local_root or span
-    appsec_traces = root_span.get_struct_tag(EXPLOIT_PREVENTION.STACK_TRACES) or {}
-    exploit: List[Any] = appsec_traces.get("exploit", [])
+    appsec_traces = root_span.get_struct_tag(STACK_TRACE.TAG) or {}
+    current_list = appsec_traces.get(namespace, [])
+    total_length = len(current_list)
 
     # Do not report more than the maximum number of stack traces
-    if asm_config._ep_max_stack_traces and len(exploit) >= asm_config._ep_max_stack_traces:
-        return None
+    if asm_config._ep_max_stack_traces and total_length >= asm_config._ep_max_stack_traces:
+        return False
 
     stack = inspect.stack()
     if crop_stack is not None:
@@ -43,8 +55,9 @@ def report_stack(
     res: Dict[str, Any] = {
         "language": "python",
         "id": stack_id,
-        "message": message,
     }
+    if message is not None:
+        res["message"] = message
     if len(stack) > asm_config._ep_max_stack_trace_depth > 0:
         top_stack = int(asm_config._ep_max_stack_trace_depth * asm_config._ep_stack_top_percent / 100)
         bottom_stack = asm_config._ep_max_stack_trace_depth - top_stack
@@ -61,6 +74,7 @@ def report_stack(
         for i in iterator
     ]
     res["frames"] = frames
-    exploit.append(res)
-    appsec_traces["exploit"] = exploit
-    root_span.set_struct_tag(EXPLOIT_PREVENTION.STACK_TRACES, appsec_traces)
+    current_list.append(res)
+    appsec_traces[namespace] = current_list
+    root_span.set_struct_tag(STACK_TRACE.TAG, appsec_traces)
+    return True

ddtrace/appsec/_iast/__init__.py

@@ -31,6 +31,7 @@ def wrapped_function(wrapped, instance, args, kwargs):
 import inspect
 import os
 import sys
+import types
 
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.module import ModuleWatchdog
@@ -61,7 +62,7 @@ def ddtrace_iast_flask_patch():
     module_name = inspect.currentframe().f_back.f_globals["__name__"]
     module = sys.modules[module_name]
     try:
-        module_path, patched_ast = astpatch_module(module, remove_flask_run=True)
+        module_path, patched_ast = astpatch_module(module)
     except Exception:
         log.debug("Unexpected exception while AST patching", exc_info=True)
         return
@@ -71,8 +72,12 @@ def ddtrace_iast_flask_patch():
         return
 
     compiled_code = compile(patched_ast, module_path, "exec")
+    # creating a new module environment to execute the patched code from scratch
+    new_module = types.ModuleType(module_name)
+    module.__dict__.clear()
+    module.__dict__.update(new_module.__dict__)
+    # executing the compiled code in the new module environment
     exec(compiled_code, module.__dict__)  # nosec B102
-    sys.modules[module_name] = compiled_code
 
 
 _iast_propagation_enabled = False

ddtrace/appsec/_iast/_ast/ast_patching.py

@@ -3,7 +3,6 @@
 import ast
 import codecs
 import os
-import re
 from sys import builtin_module_names
 from sys import version_info
 import textwrap
@@ -388,27 +387,6 @@ def visit_ast(
     return modified_ast
 
 
-_FLASK_INSTANCE_REGEXP = re.compile(r"(\S*)\s*=.*Flask\(.*")
-
-
-def _remove_flask_run(text: Text) -> Text:
-    """
-    Find and remove flask app.run() call. This is used for patching
-    the app.py file and exec'ing to replace the module without creating
-    a new instance.
-    """
-    flask_instance_name = re.search(_FLASK_INSTANCE_REGEXP, text)
-    if not flask_instance_name:
-        return text
-    groups = flask_instance_name.groups()
-    if not groups:
-        return text
-
-    instance_name = groups[-1]
-    new_text = re.sub(instance_name + r"\.run\(.*\)", "pass", text)
-    return new_text
-
-
 _DIR_WRAPPER = textwrap.dedent(
     f"""
 
@@ -442,7 +420,7 @@ def {_PREFIX}set_dir_filter():
 )
 
 
-def astpatch_module(module: ModuleType, remove_flask_run: bool = False) -> Tuple[str, Optional[ast.Module]]:
+def astpatch_module(module: ModuleType) -> Tuple[str, Optional[ast.Module]]:
     module_name = module.__name__
 
     module_origin = origin(module)
@@ -482,9 +460,6 @@ def astpatch_module(module: ModuleType, remove_flask_run: bool = False) -> Tuple
         log.debug("empty file: %s", module_path)
         return "", None
 
-    if remove_flask_run:
-        source_text = _remove_flask_run(source_text)
-
     if not asbool(os.environ.get(IAST.ENV_NO_DIR_PATCH, "false")) and version_info > (3, 7):
         # Add the dir filter so __ddtrace stuff is not returned by dir(module)
         # does not work in 3.7 because it enters into infinite recursion

ddtrace/appsec/_iast/_iast_request_context.py

@@ -50,6 +50,7 @@ def __init__(self, span: Optional[Span] = None):
         self.request_enabled: bool = False
         self.iast_reporter: Optional[IastSpanReporter] = None
         self.iast_span_metrics: Dict[str, int] = {}
+        self.iast_stack_trace_id: int = 0
 
 
 def _get_iast_context() -> Optional[IASTEnvironment]:
@@ -96,6 +97,14 @@ def get_iast_reporter() -> Optional[IastSpanReporter]:
     return None
 
 
+def get_iast_stacktrace_id() -> int:
+    env = _get_iast_context()
+    if env:
+        env.iast_stack_trace_id += 1
+        return env.iast_stack_trace_id
+    return 0
+
+
 def set_iast_request_enabled(request_enabled) -> None:
     env = _get_iast_context()
     if env:

ddtrace/appsec/_iast/reporter.py

@@ -11,6 +11,8 @@
 from typing import Tuple
 import zlib
 
+from ddtrace.appsec._constants import STACK_TRACE
+from ddtrace.appsec._exploit_prevention.stack_traces import report_stack
 from ddtrace.appsec._iast._evidence_redaction import sensitive_handler
 from ddtrace.appsec._iast._utils import _get_source_index
 from ddtrace.appsec._iast.constants import VULN_INSECURE_HASHING_TYPE
@@ -75,9 +77,19 @@ class Vulnerability(NotNoneDictable):
     evidence: Evidence
     location: Location
     hash: int = dataclasses.field(init=False, compare=False, hash=("PYTEST_CURRENT_TEST" in os.environ), repr=False)
+    stackId: Optional[str] = dataclasses.field(init=False, compare=False)
 
     def __post_init__(self):
+        # avoid circular import
+        from ddtrace.appsec._iast._iast_request_context import get_iast_stacktrace_id
+
         self.hash = zlib.crc32(repr(self).encode())
+        stacktrace_id = get_iast_stacktrace_id()
+        self.stackId = None
+        if stacktrace_id:
+            str_id = str(stacktrace_id)
+            if report_stack(stack_id=str_id, namespace=STACK_TRACE.IAST):
+                self.stackId = str_id
 
     def __repr__(self):
         return f"Vulnerability(type='{self.type}', location={self.location})"

ddtrace/appsec/_processor.py

@@ -21,10 +21,12 @@
 from ddtrace.appsec._constants import EXPLOIT_PREVENTION
 from ddtrace.appsec._constants import FINGERPRINTING
 from ddtrace.appsec._constants import SPAN_DATA_NAMES
+from ddtrace.appsec._constants import STACK_TRACE
 from ddtrace.appsec._constants import WAF_ACTIONS
 from ddtrace.appsec._constants import WAF_DATA_NAMES
 from ddtrace.appsec._ddwaf import DDWaf_result
 from ddtrace.appsec._ddwaf.ddwaf_types import ddwaf_context_capsule
+from ddtrace.appsec._exploit_prevention.stack_traces import report_stack
 from ddtrace.appsec._metrics import _set_waf_init_metric
 from ddtrace.appsec._metrics import _set_waf_request_metrics
 from ddtrace.appsec._metrics import _set_waf_updates_metric
@@ -325,10 +327,8 @@ def _waf_action(
                 blocked = parameters
                 blocked[WAF_ACTIONS.TYPE] = "none"
             elif action == WAF_ACTIONS.STACK_ACTION:
-                from ddtrace.appsec._exploit_prevention.stack_traces import report_stack
-
                 stack_trace_id = parameters["stack_id"]
-                report_stack("exploit detected", span, crop_trace, stack_id=stack_trace_id)
+                report_stack("exploit detected", span, crop_trace, stack_id=stack_trace_id, namespace=STACK_TRACE.RASP)
                 for rule in waf_results.data:
                     rule[EXPLOIT_PREVENTION.STACK_TRACE_ID] = stack_trace_id
 

docker-compose.yml

@@ -129,7 +129,7 @@ services:
         volumes:
           - ddagent:/tmp/ddagent:rw
     testagent:
-        image: ghcr.io/datadog/dd-apm-test-agent/ddapm-test-agent:v1.17.0
+        image: ghcr.io/datadog/dd-apm-test-agent/ddapm-test-agent:v1.20.0
         ports:
             - "127.0.0.1:9126:8126"
         volumes:

releasenotes/notes/stack_traces_for_iast-cd2c008168f6181e.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Code Security: This introduces stack trace reports for Code Security.

tests/appsec/contrib_appsec/conftest.py

@@ -45,8 +45,17 @@ def check_waf_timeout(request):
 
 
 @pytest.fixture
-def get_tag(root_span):
-    yield lambda name: root_span().get_tag(name)
+def get_tag(test_spans, root_span):
+    # checking both root spans and web spans for the tag
+    def get(name):
+        for span in test_spans.spans:
+            if span.parent_id is None or span.span_type == "web":
+                res = span.get_tag(name)
+                if res is not None:
+                    return res
+        return root_span().get_tag(name)
+
+    yield get
 
 
 @pytest.fixture

tests/appsec/contrib_appsec/utils.py

@@ -62,9 +62,13 @@ def location(self, response) -> str:
     def body(self, response) -> str:
         raise NotImplementedError
 
+    def get_stack_trace(self, root_span, namespace):
+        appsec_traces = root_span().get_struct_tag(asm_constants.STACK_TRACE.TAG) or {}
+        stacks = appsec_traces.get(namespace, [])
+        return stacks
+
     def check_for_stack_trace(self, root_span):
-        appsec_traces = root_span().get_struct_tag(asm_constants.EXPLOIT_PREVENTION.STACK_TRACES) or {}
-        exploit = appsec_traces.get("exploit", [])
+        exploit = self.get_stack_trace(root_span, "exploit")
         stack_ids = sorted(set(t["id"] for t in exploit))
         triggers = get_triggers(root_span())
         stack_id_in_triggers = sorted(set(t["stack_id"] for t in (triggers or []) if "stack_id" in t))
@@ -1385,9 +1389,9 @@ def validate_top_function(trace):
                 # there may have been multiple evaluations of other rules too
                 assert (("rule_type", endpoint), ("waf_version", DDWAF_VERSION)) in evals
                 if action_level == 2:
-                    assert get_tag("rasp.request.done") is None
+                    assert get_tag("rasp.request.done") is None, get_tag("rasp.request.done")
                 else:
-                    assert get_tag("rasp.request.done") == endpoint
+                    assert get_tag("rasp.request.done") == endpoint, get_tag("rasp.request.done")
                 assert get_metric(APPSEC.RASP_DURATION) is not None
                 assert get_metric(APPSEC.RASP_DURATION_EXT) is not None
                 assert get_metric(APPSEC.RASP_RULE_EVAL) is not None
@@ -1398,7 +1402,7 @@ def validate_top_function(trace):
                     assert "rasp" not in n
                 assert get_triggers(root_span()) is None
                 assert self.check_for_stack_trace(root_span) == []
-                assert get_tag("rasp.request.done") == endpoint
+                assert get_tag("rasp.request.done") == endpoint, get_tag("rasp.request.done")
 
     @pytest.mark.parametrize("asm_enabled", [True, False])
     @pytest.mark.parametrize("auto_events_enabled", [True, False])
@@ -1505,21 +1509,22 @@ def test_fingerprinting(self, interface, root_span, get_tag, asm_enabled, user_a
                 assert get_tag(asm_constants.FINGERPRINTING.SESSION) is None
 
     def test_iast(self, interface, root_span, get_tag):
-        if interface.name == "fastapi" and asm_config._iast_enabled:
-            raise pytest.xfail("fastapi does not fully support IAST for now")
-
         from ddtrace.ext import http
 
-        url = "/rasp/command_injection/?cmd=ls"
+        url = "/rasp/command_injection/?cmd=."
         self.update_tracer(interface)
         response = interface.client.get(url)
         assert self.status(response) == 200
         assert get_tag(http.STATUS_CODE) == "200"
         assert self.body(response).startswith("command_injection endpoint")
+        stack_traces = self.get_stack_trace(root_span, "vulnerability")
         if asm_config._iast_enabled:
             assert get_tag("_dd.iast.json") is not None
+            # checking for iast stack traces
+            assert stack_traces
         else:
             assert get_tag("_dd.iast.json") is None
+            assert stack_traces == []
 
 
 @contextmanager

tests/appsec/integrations/pygoat_tests/test_pygoat.py

@@ -26,6 +26,7 @@
 def client():
     agent_client = requests.session()
     reply = agent_client.get(TESTAGENT_URL + "/start" + TESTAGENT_TOKEN_PARAM, headers=TESTAGENT_HEADERS)
+
     assert reply.status_code == 200
     pygoat_client, token = login_to_pygoat()
 
@@ -65,7 +66,7 @@ def get_traces(agent_client: requests.Session) -> requests.Response:
 def vulnerability_in_traces(vuln_type: str, agent_client: requests.Session) -> bool:
     time.sleep(5)
     traces = get_traces(agent_client)
-    assert traces.status_code == 200
+    assert traces.status_code == 200, traces.text
     traces_list = json.loads(traces.text)
 
     class InnerBreakException(Exception):