diff --git a/ddtrace/appsec/_asm_request_context.py b/ddtrace/appsec/_asm_request_context.py index 8a1d99679c6..83343d33068 100644 --- a/ddtrace/appsec/_asm_request_context.py +++ b/ddtrace/appsec/_asm_request_context.py @@ -16,6 +16,7 @@ from ddtrace.appsec._constants import APPSEC from ddtrace.appsec._constants import EXPLOIT_PREVENTION from ddtrace.appsec._constants import SPAN_DATA_NAMES +from ddtrace.appsec._iast_request_state import iast_suppress_context from ddtrace.appsec._metrics import UNKNOWN_VERSION from ddtrace.appsec._metrics import report_waf_run_error from ddtrace.appsec._metrics import report_waf_truncation @@ -788,7 +789,5 @@ def asm_listen() -> None: def iast_disabled_taint_sources() -> "contextlib.AbstractContextManager[None]": if asm_config._iast_enabled: - from ddtrace.appsec._iast._iast_request_context_base import iast_suppress_context - return iast_suppress_context() return contextlib.nullcontext() diff --git a/ddtrace/appsec/_exploit_prevention/stack_traces.py b/ddtrace/appsec/_exploit_prevention/stack_traces.py index 903e85a01b6..3b08a162ecd 100644 --- a/ddtrace/appsec/_exploit_prevention/stack_traces.py +++ b/ddtrace/appsec/_exploit_prevention/stack_traces.py @@ -6,7 +6,6 @@ from typing import Optional from ddtrace._trace.span import Span -from ddtrace.appsec import _asm_request_context from ddtrace.appsec._constants import STACK_TRACE from ddtrace.internal import span_bus from ddtrace.internal.settings.asm import config as asm_config @@ -17,6 +16,13 @@ ) + (os.path.sep + "wrapt" + os.path.sep,) +def _get_stack_span() -> Optional[Span]: + span = span_bus.get_span() + if span is not None: + return span._service_entry_span + return span_bus.get_root_span() + + def report_stack( message: Optional[str] = None, span: Optional[Span] = None, @@ -24,26 +30,18 @@ def report_stack( 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. - """ + """Report an AppSec stack trace to a span.""" 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 + if namespace == STACK_TRACE.IAST and not asm_config._iast_enabled: return False if namespace == STACK_TRACE.IAST and asm_config._iast_use_root_span: span = span_bus.get_root_span() - - if span is None: - span = _asm_request_context.get_entry_span() + elif span is None: + span = _get_stack_span() if span is None or stack_id is None: return False @@ -51,7 +49,6 @@ def report_stack( 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 total_length >= asm_config._ep_max_stack_traces: return False @@ -62,7 +59,6 @@ def report_stack( if frame.frame.f_code.co_name == crop_stack: crop_index = i + 1 break - # Strip any remaining ddtrace/wrapt frames from the top of the stack while crop_index < len(stack) and any(d in stack[crop_index].filename for d in _INTERNAL_FRAMES): crop_index += 1 if crop_index: diff --git a/ddtrace/appsec/_iast/_iast_request_context_base.py b/ddtrace/appsec/_iast/_iast_request_context_base.py index 9eb47dde35f..f687e304ac8 100644 --- a/ddtrace/appsec/_iast/_iast_request_context_base.py +++ b/ddtrace/appsec/_iast/_iast_request_context_base.py @@ -1,5 +1,3 @@ -import contextlib -import contextvars from typing import Optional from ddtrace.appsec._constants import IAST @@ -7,10 +5,14 @@ from ddtrace.appsec._iast._iast_env import IASTEnvironment from ddtrace.appsec._iast._iast_env import _get_iast_env from ddtrace.appsec._iast._overhead_control_engine import oce -from ddtrace.appsec._iast._taint_tracking._context import debug_num_tainted_objects +from ddtrace.appsec._iast._request_taint_state import _num_objects_tainted_in_request from ddtrace.appsec._iast._taint_tracking._context import finish_request_context from ddtrace.appsec._iast._taint_tracking._context import start_request_context from ddtrace.appsec._iast.sampling.vulnerability_detection import update_global_vulnerability_limit +from ddtrace.appsec._iast_request_state import IAST_CONTEXT +from ddtrace.appsec._iast_request_state import _get_iast_context_id +from ddtrace.appsec._iast_request_state import iast_suppress_context as iast_suppress_context # noqa: F401 +from ddtrace.appsec._iast_request_state import is_iast_request_enabled as is_iast_request_enabled from ddtrace.internal import core from ddtrace.internal.logger import get_logger from ddtrace.internal.settings.asm import config as asm_config @@ -18,31 +20,6 @@ log = get_logger(__name__) -# Stopgap module for providing ASM context for the blocking features wrapping some contextvars. - -IAST_CONTEXT: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar("iast_var", default=None) - -# Keep source suppression separate from IAST_CONTEXT. Clearing the -# request context id disables request-scoped taint queries and propagation and -# can send no-context queries through unsafe native fallback paths. -_IAST_TAINT_SOURCES_SUPPRESSED: contextvars.ContextVar[bool] = contextvars.ContextVar( - "iast_taint_sources_suppressed", default=False -) - - -@contextlib.contextmanager -def iast_suppress_context(): - """Temporarily disable IAST taint *source* generation for the current context.""" - token = _IAST_TAINT_SOURCES_SUPPRESSED.set(True) - try: - yield - finally: - _IAST_TAINT_SOURCES_SUPPRESSED.reset(token) - - -def _is_iast_taint_source_enabled() -> bool: - return not _IAST_TAINT_SOURCES_SUPPRESSED.get() - def _set_span_tag_iast_request_tainted(span): total_objects_tainted = _num_objects_tainted_in_request() @@ -102,11 +79,6 @@ def _iast_start_request(span=None) -> Optional[int]: return context_id -def _get_iast_context_id() -> Optional[int]: - """Retrieve the current IAST context identifier from the ContextVar.""" - return IAST_CONTEXT.get() - - def _iast_finish_request(span=None, shoud_update_global_vulnerability_limit: bool = True) -> bool: """Finalize the IAST request context and optionally update global limits. @@ -129,22 +101,6 @@ def _iast_finish_request(span=None, shoud_update_global_vulnerability_limit: boo return False -def is_iast_request_enabled() -> bool: - """Check whether IAST is currently operating within an active request context.""" - return _get_iast_context_id() is not None - - -def _num_objects_tainted_in_request() -> int: - """Get the count of tainted objects tracked in the active IAST request context. - - Useful for span metrics and internal telemetry. - """ - context_id = _get_iast_context_id() - if context_id is not None: - return debug_num_tainted_objects(context_id) - return 0 - - def get_hash_object_tracking_len(): env = _get_iast_env() if env: diff --git a/ddtrace/appsec/_iast/_metrics.py b/ddtrace/appsec/_iast/_metrics.py index 86ffe4c6145..0596e74d714 100644 --- a/ddtrace/appsec/_iast/_metrics.py +++ b/ddtrace/appsec/_iast/_metrics.py @@ -5,7 +5,7 @@ from ddtrace.appsec._constants import TELEMETRY_INFORMATION_VERBOSITY from ddtrace.appsec._constants import TELEMETRY_MANDATORY_VERBOSITY from ddtrace.appsec._deduplications import deduplication -from ddtrace.appsec._iast._iast_request_context_base import _num_objects_tainted_in_request +from ddtrace.appsec._iast._request_taint_state import _num_objects_tainted_in_request from ddtrace.appsec._iast._taint_tracking import OriginType from ddtrace.appsec._iast._taint_tracking import origin_to_str from ddtrace.appsec._iast._utils import _is_iast_debug_enabled diff --git a/ddtrace/appsec/_iast/_request_taint_state.py b/ddtrace/appsec/_iast/_request_taint_state.py new file mode 100644 index 00000000000..9ea288d9224 --- /dev/null +++ b/ddtrace/appsec/_iast/_request_taint_state.py @@ -0,0 +1,11 @@ +from ddtrace.appsec._iast._taint_tracking._context import debug_num_tainted_objects +from ddtrace.appsec._iast_request_state import _get_iast_context_id + + +def _num_objects_tainted_in_request() -> int: + """Return the number of objects tainted in the current request.""" + context_id = _get_iast_context_id() + if context_id is not None: + num_tainted: int = debug_num_tainted_objects(context_id) + return num_tainted + return 0 diff --git a/ddtrace/appsec/_iast/_span_metrics.py b/ddtrace/appsec/_iast/_span_metrics.py index e2048d63209..e4a3990582f 100644 --- a/ddtrace/appsec/_iast/_span_metrics.py +++ b/ddtrace/appsec/_iast/_span_metrics.py @@ -1,7 +1,7 @@ from ddtrace.appsec._constants import IAST_SPAN_TAGS from ddtrace.appsec._iast._iast_env import _get_iast_env -from ddtrace.appsec._iast._iast_request_context_base import _num_objects_tainted_in_request from ddtrace.appsec._iast._metrics import _metric_key_as_snake_case +from ddtrace.appsec._iast._request_taint_state import _num_objects_tainted_in_request def _set_span_tag_iast_request_tainted(span): diff --git a/ddtrace/appsec/_iast/_taint_tracking/__init__.py b/ddtrace/appsec/_iast/_taint_tracking/__init__.py index e40809c5a68..ded0837c206 100644 --- a/ddtrace/appsec/_iast/_taint_tracking/__init__.py +++ b/ddtrace/appsec/_iast/_taint_tracking/__init__.py @@ -43,35 +43,24 @@ from ddtrace.appsec._iast._taint_tracking._native.taint_tracking import shift_taint_ranges # noqa: F401 from ddtrace.appsec._iast._taint_tracking._native.taint_tracking import str_to_origin # noqa: F401 from ddtrace.appsec._iast._taint_tracking._native.taint_tracking import taint_range as TaintRange # noqa: F401 +from ddtrace.appsec._iast_request_state import _get_iast_context_id from ddtrace.internal.logger import get_logger log = get_logger(__name__) -# AIDEV-NOTE: _get_iast_context_id is imported lazily — a top-level import here -# circularly bootstraps via _iast_request_context_base -> _taint_tracking._context -# -> _taint_tracking/__init__.py. The cached module-global avoids the per-call -# import dance on this hot path. -_CACHE_GET_IAST_CONTEXT_ID = None +# AIDEV-NOTE: Resolve the active native context through the lightweight request-state +# module. Keeping lifecycle imports out of taint tracking avoids circular initialization +# and makes this hot path a direct ContextVar lookup. def _current_iast_context_id() -> Optional[int]: - global _CACHE_GET_IAST_CONTEXT_ID - if _CACHE_GET_IAST_CONTEXT_ID is None: - from ddtrace.appsec._iast._iast_request_context_base import _get_iast_context_id - - _CACHE_GET_IAST_CONTEXT_ID = _get_iast_context_id - return _CACHE_GET_IAST_CONTEXT_ID() + return _get_iast_context_id() def get_ranges(string_input: Any, context_id: Optional[int] = None) -> Any: if context_id is None: - global _CACHE_GET_IAST_CONTEXT_ID - if _CACHE_GET_IAST_CONTEXT_ID is None: - from ddtrace.appsec._iast._iast_request_context_base import _get_iast_context_id - - _CACHE_GET_IAST_CONTEXT_ID = _get_iast_context_id - context_id = _CACHE_GET_IAST_CONTEXT_ID() + context_id = _get_iast_context_id() if context_id is None: return [] return _native_get_ranges(string_input, context_id) diff --git a/ddtrace/appsec/_iast/_taint_tracking/_taint_objects.py b/ddtrace/appsec/_iast/_taint_tracking/_taint_objects.py index bf746ed6c0a..756dcc7d71c 100644 --- a/ddtrace/appsec/_iast/_taint_tracking/_taint_objects.py +++ b/ddtrace/appsec/_iast/_taint_tracking/_taint_objects.py @@ -4,8 +4,6 @@ from ddtrace.appsec._constants import IAST from ddtrace.appsec._constants import IAST_SPAN_TAGS -from ddtrace.appsec._iast._iast_request_context_base import _get_iast_context_id -from ddtrace.appsec._iast._iast_request_context_base import _is_iast_taint_source_enabled from ddtrace.appsec._iast._logs import iast_propagation_debug_log from ddtrace.appsec._iast._metrics import _set_metric_iast_executed_source from ddtrace.appsec._iast._span_metrics import increment_iast_span_metric @@ -13,6 +11,8 @@ from ddtrace.appsec._iast._taint_tracking import TaintRange from ddtrace.appsec._iast._taint_tracking import set_ranges from ddtrace.appsec._iast._taint_tracking._taint_objects_base import _taint_pyobject_base +from ddtrace.appsec._iast_request_state import _get_iast_context_id +from ddtrace.appsec._iast_request_state import _is_iast_taint_source_enabled from ddtrace.internal.logger import get_logger diff --git a/ddtrace/appsec/_iast/_taint_tracking/_taint_objects_base.py b/ddtrace/appsec/_iast/_taint_tracking/_taint_objects_base.py index 775d644a769..687d5ba58f3 100644 --- a/ddtrace/appsec/_iast/_taint_tracking/_taint_objects_base.py +++ b/ddtrace/appsec/_iast/_taint_tracking/_taint_objects_base.py @@ -1,7 +1,6 @@ from typing import Any from ddtrace.appsec._constants import IAST -from ddtrace.appsec._iast._iast_request_context_base import _get_iast_context_id from ddtrace.appsec._iast._logs import iast_propagation_debug_log from ddtrace.appsec._iast._logs import iast_propagation_error_log from ddtrace.appsec._iast._taint_tracking import OriginType @@ -9,6 +8,7 @@ from ddtrace.appsec._iast._taint_tracking import origin_to_str from ddtrace.appsec._iast._taint_tracking import taint_pyobject from ddtrace.appsec._iast._taint_tracking._context import is_in_taint_map +from ddtrace.appsec._iast_request_state import _get_iast_context_id def _taint_pyobject_base(pyobject: Any, source_name: Any, source_value: Any, source_origin=None, contextid=None) -> Any: diff --git a/ddtrace/appsec/_iast/_taint_tracking/aspects.py b/ddtrace/appsec/_iast/_taint_tracking/aspects.py index bfb960af069..99282a79051 100644 --- a/ddtrace/appsec/_iast/_taint_tracking/aspects.py +++ b/ddtrace/appsec/_iast/_taint_tracking/aspects.py @@ -16,7 +16,6 @@ from typing import Union from ddtrace.appsec._constants import IAST -from ddtrace.appsec._iast._iast_request_context_base import is_iast_request_enabled from ddtrace.appsec._iast._logs import iast_propagation_error_log from ddtrace.appsec._iast._taint_tracking import TagMappingMode from ddtrace.appsec._iast._taint_tracking import TaintRange @@ -50,6 +49,7 @@ from ddtrace.appsec._iast._taint_tracking._taint_objects_base import get_tainted_ranges from ddtrace.appsec._iast._taint_tracking._taint_objects_base import is_pyobject_tainted from ddtrace.appsec._iast._taint_utils import taint_structure +from ddtrace.appsec._iast_request_state import is_iast_request_enabled TEXT_TYPES = Union[str, bytes, bytearray] diff --git a/ddtrace/appsec/_iast/reporter.py b/ddtrace/appsec/_iast/reporter.py index 50c84e3e5e2..64ab27f02cd 100644 --- a/ddtrace/appsec/_iast/reporter.py +++ b/ddtrace/appsec/_iast/reporter.py @@ -9,11 +9,11 @@ from ddtrace.appsec._constants import STACK_TRACE from ddtrace.appsec._exploit_prevention.stack_traces import report_stack from ddtrace.appsec._iast._evidence_redaction._sensitive_handler import sensitive_handler -from ddtrace.appsec._iast._iast_request_context_base import is_iast_request_enabled from ddtrace.appsec._iast._utils import _get_source_index from ddtrace.appsec._iast.constants import VULN_INSECURE_HASHING_TYPE from ddtrace.appsec._iast.constants import VULN_WEAK_CIPHER_TYPE from ddtrace.appsec._iast.constants import VULN_WEAK_RANDOMNESS +from ddtrace.appsec._iast_request_state import is_iast_request_enabled from ddtrace.internal.logger import get_logger from ddtrace.internal.settings import env from ddtrace.internal.settings.asm import config as asm_config diff --git a/ddtrace/appsec/_iast_request_state.py b/ddtrace/appsec/_iast_request_state.py new file mode 100644 index 00000000000..6fc65605c73 --- /dev/null +++ b/ddtrace/appsec/_iast_request_state.py @@ -0,0 +1,38 @@ +import contextlib +import contextvars +from typing import Iterator +from typing import Optional + + +IAST_CONTEXT: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar("iast_var", default=None) + +# Keep source suppression separate from IAST_CONTEXT. Clearing the request context id disables +# request-scoped taint queries and propagation and can send no-context queries through unsafe native +# fallback paths. +_IAST_TAINT_SOURCES_SUPPRESSED: contextvars.ContextVar[bool] = contextvars.ContextVar( + "iast_taint_sources_suppressed", default=False +) + + +@contextlib.contextmanager +def iast_suppress_context() -> Iterator[None]: + """Temporarily disable IAST taint source generation for the current context.""" + token = _IAST_TAINT_SOURCES_SUPPRESSED.set(True) + try: + yield + finally: + _IAST_TAINT_SOURCES_SUPPRESSED.reset(token) + + +def _is_iast_taint_source_enabled() -> bool: + return not _IAST_TAINT_SOURCES_SUPPRESSED.get() + + +def _get_iast_context_id() -> Optional[int]: + """Return the native taint context identifier for the current request.""" + return IAST_CONTEXT.get() + + +def is_iast_request_enabled() -> bool: + """Check whether IAST is operating within an active request context.""" + return _get_iast_context_id() is not None diff --git a/tests/appsec/appsec/test_stack_traces.py b/tests/appsec/appsec/test_stack_traces.py new file mode 100644 index 00000000000..c58aad8e6d3 --- /dev/null +++ b/tests/appsec/appsec/test_stack_traces.py @@ -0,0 +1,35 @@ +from ddtrace.appsec._constants import STACK_TRACE +from ddtrace.appsec._exploit_prevention.stack_traces import report_stack +from tests.utils import override_global_config + + +def test_report_stack_uses_service_entry_span(tracer): + config = { + "_asm_enabled": True, + "_ep_enabled": True, + "_ep_stack_trace_enabled": True, + } + + with override_global_config(config): + with tracer.trace("request", service="service") as request_span: + with tracer.trace("child", service="service") as child_span: + assert report_stack(stack_id="1") + + assert child_span._get_struct_tag(STACK_TRACE.TAG) is None + assert request_span._get_struct_tag(STACK_TRACE.TAG)[STACK_TRACE.RASP][0]["id"] == "1" + + +def test_report_iast_stack_uses_root_span(tracer): + config = { + "_ep_stack_trace_enabled": True, + "_iast_enabled": True, + "_iast_use_root_span": True, + } + + with override_global_config(config): + with tracer.trace("root") as root_span: + with tracer.trace("child", service="child") as child_span: + assert report_stack(stack_id="1", namespace=STACK_TRACE.IAST) + + assert child_span._get_struct_tag(STACK_TRACE.TAG) is None + assert root_span._get_struct_tag(STACK_TRACE.TAG)[STACK_TRACE.IAST][0]["id"] == "1"