Merge branch 'main' into avara1986/PPSEC-55288-split-listens

ddtrace/_monkey.py

@@ -5,6 +5,7 @@
 
 from wrapt.importer import when_imported
 
+from .appsec._iast._utils import _is_iast_enabled
 from .internal import telemetry
 from .internal.logger import get_logger
 from .internal.utils import formats
@@ -225,7 +226,7 @@ def patch_all(**patch_modules):
  modules.update(patch_modules)
 
  patch(raise_errors=False, **modules)
- if asm_config._iast_enabled:
+ if _is_iast_enabled():
  from ddtrace.appsec._iast._patch_modules import patch_iast
  from ddtrace.appsec.iast import enable_iast_propagation
 

ddtrace/appsec/_iast/__init__.py

@@ -28,14 +28,11 @@ def wrapped_function(wrapped, instance, args, kwargs):
  return wrapped(*args, **kwargs)
 """ # noqa: RST201, RST213, RST210
 import inspect
-import os
 import sys
 
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.module import ModuleWatchdog
-from ddtrace.internal.utils.formats import asbool
 
-from .._constants import IAST
 from ._overhead_control_engine import OverheadControl
 from ._utils import _is_iast_enabled
 
@@ -73,19 +70,19 @@ def ddtrace_iast_flask_patch():
 
 def enable_iast_propagation():
  """Add IAST AST patching in the ModuleWatchdog"""
- if asbool(os.getenv(IAST.ENV, "false")):
- from ddtrace.appsec._iast._utils import _is_python_version_supported
-
- if _is_python_version_supported():
- from ddtrace.appsec._iast._ast.ast_patching import _should_iast_patch
- from ddtrace.appsec._iast._loader import _exec_iast_patched_module
+ # DEV: These imports are here to avoid _ast.ast_patching import in the top level
+ # because they are slow and affect serverless startup time
+ from ddtrace.appsec._iast._ast.ast_patching import _should_iast_patch
+ from ddtrace.appsec._iast._loader import _exec_iast_patched_module
 
-  log.debug("IAST enabled")
-  ModuleWatchdog.register_pre_exec_module_hook(_should_iast_patch, _exec_iast_patched_module)
+ log.debug("IAST enabled")
+ ModuleWatchdog.register_pre_exec_module_hook(_should_iast_patch, _exec_iast_patched_module)
 
 
 def disable_iast_propagation():
  """Remove IAST AST patching from the ModuleWatchdog. Only for testing proposes"""
+ # DEV: These imports are here to avoid _ast.ast_patching import in the top level
+ # because they are slow and affect serverless startup time
  from ddtrace.appsec._iast._ast.ast_patching import _should_iast_patch
  from ddtrace.appsec._iast._loader import _exec_iast_patched_module
 

ddtrace/appsec/_iast/_taint_tracking/__init__.py

@@ -1,6 +1,8 @@
 from io import BytesIO
 from io import StringIO
+import itertools
 from typing import Any
+from typing import Sequence
 from typing import Tuple
 
 from ddtrace.internal._unpatched import _threading as threading
@@ -63,7 +65,6 @@
  new_pyobject_id = ops.new_pyobject_id
  set_ranges_from_values = ops.set_ranges_from_values
 
-
 __all__ = [
  "OriginType",
  "Source",
@@ -140,7 +141,7 @@ def is_pyobject_tainted(pyobject: Any) -> bool:
  return False
 
 
-def taint_pyobject(pyobject: Any, source_name: Any, source_value: Any, source_origin=None) -> Any:
+def _taint_pyobject_base(pyobject: Any, source_name: Any, source_value: Any, source_origin=None) -> Any:
  if not is_iast_request_enabled():
  return pyobject
 
@@ -166,13 +167,25 @@ def taint_pyobject(pyobject: Any, source_name: Any, source_value: Any, source_or
 
  try:
  pyobject_newid = set_ranges_from_values(pyobject, pyobject_len, source_name, source_value, source_origin)
- _set_metric_iast_executed_source(source_origin)
  return pyobject_newid
  except ValueError as e:
  log.debug("Tainting object error (pyobject type %s): %s", type(pyobject), e, exc_info=True)
  return pyobject
 
 
+def taint_pyobject(pyobject: Any, source_name: Any, source_value: Any, source_origin=None) -> Any:
+ try:
+ if source_origin is None:
+ source_origin = OriginType.PARAMETER
+
+ res = _taint_pyobject_base(pyobject, source_name, source_value, source_origin)
+ _set_metric_iast_executed_source(source_origin)
+ return res
+ except ValueError as e:
+ log.debug("Tainting object error (pyobject type %s): %s", type(pyobject), e)
+ return pyobject
+
+
 def taint_pyobject_with_ranges(pyobject: Any, ranges: Tuple) -> bool:
  if not is_iast_request_enabled():
  return False
@@ -244,3 +257,36 @@ def trace_calls_and_returns(frame, event, arg):
  return
 
  threading.settrace(trace_calls_and_returns)
+
+
+def copy_ranges_to_string(s: str, ranges: Sequence[TaintRange]) -> str:
+ for r in ranges:
+ if s in r.source.value:
+ s = _taint_pyobject_base(
+ pyobject=s, source_name=r.source.name, source_value=r.source.value, source_origin=r.source.origin
+ )
+ break
+ else:
+ # no total match found, maybe partial match, just take the first one
+ s = _taint_pyobject_base(
+ pyobject=s,
+ source_name=ranges[0].source.name,
+ source_value=ranges[0].source.value,
+ source_origin=ranges[0].source.origin,
+ )
+ return s
+
+
+# Given a list of ranges, try to match them with the iterable and return a new iterable with a new range applied that
+# matched the original one Source. If no range matches, take the Source from the first one.
+def copy_ranges_to_iterable_with_strings(iterable: Sequence[str], ranges: Sequence[TaintRange]) -> Sequence[str]:
+ iterable_type = type(iterable)
+
+ new_result = []
+ # do this so it doesn't consume a potential generator
+ items, items_backup = itertools.tee(iterable)
+ for i in items_backup:
+ i = copy_ranges_to_string(i, ranges)
+ new_result.append(i)
+
+ return iterable_type(new_result) # type: ignore[call-arg]

ddtrace/appsec/_iast/_taint_tracking/aspects.py

@@ -40,6 +40,8 @@
 from .._taint_tracking import common_replace
 from .._taint_tracking import copy_and_shift_ranges_from_strings
 from .._taint_tracking import copy_ranges_from_strings
+from .._taint_tracking import copy_ranges_to_iterable_with_strings
+from .._taint_tracking import copy_ranges_to_string
 from .._taint_tracking import get_ranges
 from .._taint_tracking import get_tainted_ranges
 from .._taint_tracking import iast_taint_log_error
@@ -943,25 +945,26 @@ def re_findall_aspect(
  args = args[(flag_added_args or 1) :]
  result = orig_function(*args, **kwargs)
 
- if not isinstance(self, (Pattern, ModuleType)):
- # This is not the sub we're looking for
- return result
- elif isinstance(self, ModuleType):
- if self.__name__ != "re" or self.__package__ not in ("", "re"):
+ try:
+ if not isinstance(self, (Pattern, ModuleType)):
+ # This is not the sub we're looking for
+ return result
+ elif isinstance(self, ModuleType):
+ if self.__name__ != "re" or self.__package__ not in ("", "re"):
+ return result
+ # In this case, the first argument is the pattern
+ # which we don't need to check for tainted ranges
+ args = args[1:]
+ elif not isinstance(result, list) or not len(result):
  return result
- # In this case, the first argument is the pattern
- # which we don't need to check for tainted ranges
- args = args[1:]
- elif not isinstance(result, list) or not len(result):
- return result
 
- if len(args) >= 1:
- string = args[0]
- if is_pyobject_tainted(string):
- for i in result:
- if len(i):
-  # Taint results
-  copy_and_shift_ranges_from_strings(string, i, 0, len(i))
+  if len(args) >= 1:
+  string = args[0]
+  ranges = get_tainted_ranges(string)
+ if ranges:
+ result = copy_ranges_to_iterable_with_strings(result, ranges)
+ except Exception as e:
+ iast_taint_log_error("re_findall_aspect. {}".format(e))
 
  return result
 
@@ -1171,11 +1174,11 @@ def re_groups_aspect(orig_function: Optional[Callable], flag_added_args: int, *a
  if not result or not isinstance(self, Match) or not is_pyobject_tainted(self):
  return result
 
- for group in result:
- if group is not None:
-  copy_and_shift_ranges_from_strings(self, group, 0, len(group))
-
- return result
+ try:
+ return copy_ranges_to_iterable_with_strings(result, get_ranges(self))
+ except Exception as e:
+ iast_taint_log_error("re_groups_aspect. {}".format(e))
+  return result
 
 
 def re_group_aspect(orig_function: Optional[Callable], flag_added_args: int, *args: Any, **kwargs: Any) -> Any:
@@ -1193,12 +1196,13 @@ def re_group_aspect(orig_function: Optional[Callable], flag_added_args: int, *ar
  if not result or not isinstance(self, Match) or not is_pyobject_tainted(self):
  return result
 
- if isinstance(result, tuple):
- for group in result:
- if group is not None:
- copy_and_shift_ranges_from_strings(self, group, 0, len(group))
- else:
- copy_and_shift_ranges_from_strings(self, result, 0, len(result))
+ try:
+ if isinstance(result, tuple):
+ result = copy_ranges_to_iterable_with_strings(result, get_ranges(self))
+ else:
+ result = copy_ranges_to_string(result, get_ranges(self))
+ except Exception as e:
+ iast_taint_log_error("re_group_aspect. {}".format(e))
 
  return result
 
@@ -1219,13 +1223,16 @@ def re_expand_aspect(orig_function: Optional[Callable], flag_added_args: int, *a
  # No need to taint the result
  return result
 
- if not is_pyobject_tainted(self) and len(args) and not is_pyobject_tainted(args[0]):
- # Nothing tainted, no need to taint the result either
- return result
+ try:
+ if not is_pyobject_tainted(self) and len(args) and not is_pyobject_tainted(args[0]):
+ # Nothing tainted, no need to taint the result either
+ return result
 
- if is_pyobject_tainted(self):
- copy_and_shift_ranges_from_strings(self, result, 0, len(result))
- elif is_pyobject_tainted(args[0]):
- copy_and_shift_ranges_from_strings(args[0], result, 0, len(result))
+ if is_pyobject_tainted(self):
+ result = copy_ranges_to_string(result, get_ranges(self))
+ elif is_pyobject_tainted(args[0]):
+ result = copy_ranges_to_string(result, get_ranges(args[0]))
+ except Exception as e:
+ iast_taint_log_error("re_expand_aspect. {}".format(e))
 
  return result

ddtrace/appsec/_iast/_utils.py

@@ -1,3 +1,4 @@
+from functools import lru_cache
 import os
 import sys
 from typing import List
@@ -9,6 +10,7 @@
 from ddtrace.settings.asm import config as asm_config
 
 
+@lru_cache(maxsize=1)
 def _is_python_version_supported() -> bool:
  # IAST supports Python versions 3.6 to 3.12
  return (3, 6, 0) <= sys.version_info < (3, 13, 0)

ddtrace/bootstrap/preload.py

@@ -6,6 +6,7 @@
 import os # noqa:I001
 
 from ddtrace import config # noqa:F401
+from ddtrace.appsec._iast._utils import _is_iast_enabled
 from ddtrace.settings.profiling import config as profiling_config # noqa:F401
 from ddtrace.internal.logger import get_logger # noqa:F401
 from ddtrace.internal.module import ModuleWatchdog # noqa:F401
@@ -71,7 +72,7 @@ def register_post_preload(func: t.Callable) -> None:
 if config._runtime_metrics_enabled:
  RuntimeWorker.enable()
 
-if asbool(os.getenv("DD_IAST_ENABLED", False)):
+if _is_iast_enabled():
  """
  This is the entry point for the IAST instrumentation. `enable_iast_propagation` is called on patch_all function
  too but patch_all depends of DD_TRACE_ENABLED environment variable. This is the reason why we need to call it

ddtrace/contrib/internal/elasticsearch/patch.py

@@ -142,11 +142,6 @@ def _perform_request(func, instance, args, kwargs):
 
  span.set_tag(SPAN_MEASURED_KEY)
 
- # Only instrument if trace is sampled or if we haven't tried to sample yet
- if span.context.sampling_priority is not None and span.context.sampling_priority <= 0:
- yield func(*args, **kwargs)
- return
-
  method, target = args
  params = kwargs.get("params")
  body = kwargs.get("body")

pyproject.toml

@@ -107,6 +107,7 @@ exclude = '''
  | ddtrace/internal/datadog/profiling/ddup/_ddup.pyx$
  | ddtrace/vendor/
  | ddtrace/appsec/_iast/_taint_tracking/_vendor/
+ | ddtrace/appsec/_iast/_taint_tracking/cmake-build-debug/
  | ddtrace/_version.py
  | \.eggs
  | \.git

releasenotes/notes/fix-elasticsearch-always-populate-span-tags-c2a212ec706b3f4b.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+ - |
+ elasticsearch: this fix resolves an issue where span tags were not fully populated on "sampled" spans, causing metric dimensions to be incorrect when spans were prematurely marked as sampled, including resource_name.