chore(iast): refactor iast request context by core.context

benchmarks/appsec_iast_aspects/scenario.py

@@ -6,7 +6,11 @@
 import bm
 from bm.utils import override_env
 
+from ddtrace.appsec._iast import oce
 from ddtrace.appsec._iast._ast.ast_patching import astpatch_module
+from ddtrace.appsec._iast._iast_request_context import end_iast_context
+from ddtrace.appsec._iast._iast_request_context import set_iast_request_enabled
+from ddtrace.appsec._iast._iast_request_context import start_iast_context
 
 
 # Copypasted here from tests.iast.aspects.conftest since the benchmarks can't access tests.*
@@ -19,6 +23,18 @@ def _iast_patched_module(module_name):
  return module_changed
 
 
+def _start_iast_context_and_oce():
+ oce.reconfigure()
+ oce.acquire_request(None)
+ start_iast_context()
+ set_iast_request_enabled(True)
+
+
+def _end_iast_context_and_oce():
+ end_iast_context()
+ oce.release_request()
+
+
 class IAST_Aspects(bm.Scenario):
  iast_enabled: bool
  mod_original_name: str
@@ -27,6 +43,9 @@ class IAST_Aspects(bm.Scenario):
 
  def run(self):
  args = ast.literal_eval(self.args)
+ if self.iast_enabled:
+ with override_env({"DD_IAST_ENABLED": "True"}):
+ _start_iast_context_and_oce()
 
  def _(loops):
  for _ in range(loops):
@@ -40,3 +59,6 @@ def _(loops):
  getattr(module_unpatched, self.function_name)(*args)
 
  yield _
+ if self.iast_enabled:
+ with override_env({"DD_IAST_ENABLED": "True"}):
+ _end_iast_context_and_oce()

benchmarks/appsec_iast_propagation/scenario.py

@@ -1,25 +1,36 @@
 from typing import Any
 
 import bm
-from bm.utils import override_env
 
-
-with override_env({"DD_IAST_ENABLED": "True"}):
- from ddtrace.appsec._iast._taint_tracking import OriginType
- from ddtrace.appsec._iast._taint_tracking import Source
- from ddtrace.appsec._iast._taint_tracking import TaintRange
- from ddtrace.appsec._iast._taint_tracking import create_context
- from ddtrace.appsec._iast._taint_tracking import reset_context
- from ddtrace.appsec._iast._taint_tracking import set_ranges
- from ddtrace.appsec._iast._taint_tracking.aspects import add_aspect
- from ddtrace.appsec._iast._taint_tracking.aspects import join_aspect
+from ddtrace.appsec._iast import oce
+from ddtrace.appsec._iast._iast_request_context import end_iast_context
+from ddtrace.appsec._iast._iast_request_context import set_iast_request_enabled
+from ddtrace.appsec._iast._iast_request_context import start_iast_context
+from ddtrace.appsec._iast._taint_tracking import OriginType
+from ddtrace.appsec._iast._taint_tracking import Source
+from ddtrace.appsec._iast._taint_tracking import TaintRange
+from ddtrace.appsec._iast._taint_tracking import set_ranges
+from ddtrace.appsec._iast._taint_tracking.aspects import add_aspect
+from ddtrace.appsec._iast._taint_tracking.aspects import join_aspect
 
 
 TAINT_ORIGIN = Source(name="sample_name", value="sample_value", origin=OriginType.PARAMETER)
 
 CHECK_RANGES = [TaintRange(0, 3, TAINT_ORIGIN), TaintRange(21, 3, TAINT_ORIGIN), TaintRange(41, 3, TAINT_ORIGIN)]
 
 
+def _start_iast_context_and_oce():
+ oce.reconfigure()
+ oce.acquire_request(None)
+ start_iast_context()
+ set_iast_request_enabled(True)
+
+
+def _end_iast_context_and_oce():
+ end_iast_context()
+ oce.release_request()
+
+
 def taint_pyobject_with_ranges(pyobject: Any, ranges: tuple) -> None:
  set_ranges(pyobject, tuple(ranges))
 
@@ -53,8 +64,6 @@ def aspect_function(internal_loop, tainted):
 
 def new_request(enable_propagation):
  tainted = b"my_string".decode("ascii")
- reset_context()
- create_context()
 
  if enable_propagation:
  taint_pyobject_with_ranges(tainted, (CHECK_RANGES[0],))
@@ -74,6 +83,7 @@ class IastPropagation(bm.Scenario):
  def run(self):
  caller_loop = 10
  if self.iast_enabled:
+ _start_iast_context_and_oce()
  func = aspect_function
  else:
  func = normal_function
@@ -83,3 +93,5 @@ def _(loops):
  launch_function(self.iast_enabled, func, self.internal_loop, caller_loop)
 
  yield _
+ if self.iast_enabled:
+ _end_iast_context_and_oce()

benchmarks/bm/utils.py

@@ -71,10 +71,9 @@ def drop_traces(tracer):
 def drop_telemetry_events():
  # Avoids sending instrumentation telemetry payloads to the agent
  try:
- if telemetry.telemetry_writer.is_periodic:
- telemetry.telemetry_writer.stop()
+ telemetry.telemetry_writer.stop()
  telemetry.telemetry_writer.reset_queues()
- telemetry.telemetry_writer.enable(start_worker_thread=False)
+ telemetry.telemetry_writer.enable()
  except AttributeError:
  # telemetry.telemetry_writer is not defined in this version of dd-trace-py
  # Telemetry events will not be mocked!

ddtrace/appsec/__init__.py

@@ -3,9 +3,11 @@
 
 def load_appsec():
  """Lazily load the appsec module listeners."""
- from ddtrace.appsec._asm_request_context import listen
+ from ddtrace.appsec._asm_request_context import asm_listen
+ from ddtrace.appsec._iast._iast_request_context import iast_listen
 
  global _APPSEC_TO_BE_LOADED
  if _APPSEC_TO_BE_LOADED:
- listen()
+ asm_listen()
+ iast_listen()
  _APPSEC_TO_BE_LOADED = False

ddtrace/appsec/_asm_request_context.py

@@ -2,6 +2,7 @@
 import json
 import re
 import sys
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
 from typing import Dict
@@ -16,8 +17,6 @@
 from ddtrace.appsec._constants import EXPLOIT_PREVENTION
 from ddtrace.appsec._constants import SPAN_DATA_NAMES
 from ddtrace.appsec._constants import WAF_CONTEXT_NAMES
-from ddtrace.appsec._ddwaf import DDWaf_result
-from ddtrace.appsec._iast._utils import _is_iast_enabled
 from ddtrace.appsec._utils import get_triggers
 from ddtrace.internal import core
 from ddtrace.internal._exceptions import BlockingException
@@ -26,6 +25,9 @@
 from ddtrace.settings.asm import config as asm_config
 
 
+if TYPE_CHECKING:
+ from ddtrace.appsec._ddwaf import DDWaf_result
+
 log = get_logger(__name__)
 
 # Stopgap module for providing ASM context for the blocking features wrapping some contextvars.
@@ -56,7 +58,7 @@ class ASM_Environment:
  """
 
  def __init__(self, span: Optional[Span] = None):
- self.root = not in_context()
+ self.root = not in_asm_context()
  if self.root:
  core.add_suppress_exception(BlockingException)
  if span is None:
@@ -92,7 +94,7 @@ def _get_asm_context() -> Optional[ASM_Environment]:
  return core.get_item(_ASM_CONTEXT)
 
 
-def in_context() -> bool:
+def in_asm_context() -> bool:
  return core.get_item(_ASM_CONTEXT) is not None
 
 
@@ -293,7 +295,7 @@ def set_waf_callback(value) -> None:
  set_value(_CALLBACKS, _WAF_CALL, value)
 
 
-def call_waf_callback(custom_data: Optional[Dict[str, Any]] = None, **kwargs) -> Optional[DDWaf_result]:
+def call_waf_callback(custom_data: Optional[Dict[str, Any]] = None, **kwargs) -> Optional["DDWaf_result"]:
  if not asm_config._asm_enabled:
  return None
  callback = get_value(_CALLBACKS, _WAF_CALL)
@@ -451,7 +453,7 @@ def _on_context_ended(ctx):
 def _on_wrapped_view(kwargs):
  return_value = [None, None]
  # if Appsec is enabled, we can try to block as we have the path parameters at that point
- if asm_config._asm_enabled and in_context():
+ if asm_config._asm_enabled and in_asm_context():
  log.debug("Flask WAF call for Suspicious Request Blocking on request")
  if kwargs:
  set_waf_address(REQUEST_PATH_PARAMS, kwargs)
@@ -461,12 +463,14 @@ def _on_wrapped_view(kwargs):
  return_value[0] = callback_block
 
  # If IAST is enabled, taint the Flask function kwargs (path parameters)
+ from ddtrace.appsec._iast._utils import _is_iast_enabled
+
  if _is_iast_enabled() and kwargs:
+ from ddtrace.appsec._iast._iast_request_context import is_iast_request_enabled
  from ddtrace.appsec._iast._taint_tracking import OriginType
  from ddtrace.appsec._iast._taint_tracking import taint_pyobject
- from ddtrace.appsec._iast.processor import AppSecIastSpanProcessor
 
- if not AppSecIastSpanProcessor.is_span_analyzed():
+ if not is_iast_request_enabled():
  return return_value
 
  _kwargs = {}
@@ -479,16 +483,18 @@ def _on_wrapped_view(kwargs):
 
 
 def _on_set_request_tags(request, span, flask_config):
+ from ddtrace.appsec._iast._utils import _is_iast_enabled
+
  if _is_iast_enabled():
+ from ddtrace.appsec._iast._iast_request_context import is_iast_request_enabled
  from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_source
  from ddtrace.appsec._iast._taint_tracking import OriginType
  from ddtrace.appsec._iast._taint_utils import taint_structure
- from ddtrace.appsec._iast.processor import AppSecIastSpanProcessor
 
  _set_metric_iast_instrumented_source(OriginType.COOKIE_NAME)
  _set_metric_iast_instrumented_source(OriginType.COOKIE)
 
- if not AppSecIastSpanProcessor.is_span_analyzed(span._local_root or span):
+ if not is_iast_request_enabled():
  return
 
  request.cookies = taint_structure(
@@ -556,10 +562,11 @@ def _get_headers_if_appsec():
  return get_headers()
 
 
-def listen():
+def asm_listen():
  from ddtrace.appsec._handlers import listen
 
  listen()
+
  core.on("flask.finalize_request.post", _set_headers_and_response)
  core.on("flask.wrapped_view", _on_wrapped_view, "callback_and_args")
  core.on("flask._patched_request", _on_pre_tracedrequest)

ddtrace/appsec/_common_module_patches.py

@@ -81,7 +81,7 @@ def wrapped_open_CFDDB7ABBA9081B6(original_open_callable, instance, args, kwargs
  ):
  try:
  from ddtrace.appsec._asm_request_context import call_waf_callback
- from ddtrace.appsec._asm_request_context import in_context
+ from ddtrace.appsec._asm_request_context import in_asm_context
  from ddtrace.appsec._constants import EXPLOIT_PREVENTION
  except ImportError:
  # open is used during module initialization
@@ -93,7 +93,7 @@ def wrapped_open_CFDDB7ABBA9081B6(original_open_callable, instance, args, kwargs
  filename = os.fspath(filename_arg)
  except Exception:
  filename = ""
- if filename and in_context():
+ if filename and in_asm_context():
  res = call_waf_callback(
  {EXPLOIT_PREVENTION.ADDRESS.LFI: filename},
  crop_trace="wrapped_open_CFDDB7ABBA9081B6",
@@ -126,15 +126,15 @@ def wrapped_open_ED4CF71136E15EBF(original_open_callable, instance, args, kwargs
  ):
  try:
  from ddtrace.appsec._asm_request_context import call_waf_callback
- from ddtrace.appsec._asm_request_context import in_context
+ from ddtrace.appsec._asm_request_context import in_asm_context
  from ddtrace.appsec._constants import EXPLOIT_PREVENTION
  except ImportError:
  # open is used during module initialization
  # and shouldn't be changed at that time
  return original_open_callable(*args, **kwargs)
 
  url = args[0] if args else kwargs.get("fullurl", None)
- if url and in_context():
+ if url and in_asm_context():
  if url.__class__.__name__ == "Request":
  url = url.get_full_url()
  if isinstance(url, str):
@@ -166,15 +166,15 @@ def wrapped_request_D8CB81E472AF98A2(original_request_callable, instance, args,
  ):
  try:
  from ddtrace.appsec._asm_request_context import call_waf_callback
- from ddtrace.appsec._asm_request_context import in_context
+ from ddtrace.appsec._asm_request_context import in_asm_context
  from ddtrace.appsec._constants import EXPLOIT_PREVENTION
  except ImportError:
  # open is used during module initialization
  # and shouldn't be changed at that time
  return original_request_callable(*args, **kwargs)
 
  url = args[1] if len(args) > 1 else kwargs.get("url", None)
- if url and in_context():
+ if url and in_asm_context():
  if isinstance(url, str):
  res = call_waf_callback(
  {EXPLOIT_PREVENTION.ADDRESS.SSRF: url},
@@ -206,12 +206,12 @@ def wrapped_system_5542593D237084A7(original_command_callable, instance, args, k
  ):
  try:
  from ddtrace.appsec._asm_request_context import call_waf_callback
- from ddtrace.appsec._asm_request_context import in_context
+ from ddtrace.appsec._asm_request_context import in_asm_context
  from ddtrace.appsec._constants import EXPLOIT_PREVENTION
  except ImportError:
  return original_command_callable(*args, **kwargs)
 
- if in_context():
+ if in_asm_context():
  res = call_waf_callback(
  {EXPLOIT_PREVENTION.ADDRESS.CMDI: command},
  crop_trace="wrapped_system_5542593D237084A7",
@@ -254,14 +254,14 @@ def execute_4C9BAC8E228EB347(instrument_self, query, args, kwargs) -> None:
  ):
  try:
  from ddtrace.appsec._asm_request_context import call_waf_callback
- from ddtrace.appsec._asm_request_context import in_context
+ from ddtrace.appsec._asm_request_context import in_asm_context
  from ddtrace.appsec._constants import EXPLOIT_PREVENTION
  except ImportError:
  # execute is used during module initialization
  # and shouldn't be changed at that time
  return
 
- if instrument_self and query and in_context():
+ if instrument_self and query and in_asm_context():
  db_type = _DB_DIALECTS.get(
  getattr(instrument_self, "_self_config", {}).get("_dbapi_span_name_prefix", ""), ""
  )

ddtrace/appsec/_constants.py

@@ -120,11 +120,9 @@ class IAST(metaclass=Constant_Class):
  LAZY_TAINT: Literal["_DD_IAST_LAZY_TAINT"] = "_DD_IAST_LAZY_TAINT"
  JSON: Literal["_dd.iast.json"] = "_dd.iast.json"
  ENABLED: Literal["_dd.iast.enabled"] = "_dd.iast.enabled"
- CONTEXT_KEY: Literal["_iast_data"] = "_iast_data"
  PATCH_MODULES: Literal["_DD_IAST_PATCH_MODULES"] = "_DD_IAST_PATCH_MODULES"
  DENY_MODULES: Literal["_DD_IAST_DENY_MODULES"] = "_DD_IAST_DENY_MODULES"
  SEP_MODULES: Literal[","] = ","
- REQUEST_IAST_ENABLED: Literal["_dd.iast.request_enabled"] = "_dd.iast.request_enabled"
  TEXT_TYPES = (str, bytes, bytearray)
  TAINTEABLE_TYPES = (str, bytes, bytearray, Match, BytesIO, StringIO)