Merge branch 'main' into chris.agocs/patch_botocore_stepfunctions

Author: agocs

.github/workflows/stale.yml

@@ -17,11 +17,9 @@ jobs:
           # DEV: GitHub Actions have an API rate limit of 1000 operations per hour per repository
           #      This limit is shared across all actions
           operations-per-run: 200
-          days-before-stale: 180
           days-before-close: 180
           exempt-issue-labels: 'proposal'
           exempt-pr-labels: 'proposal'
-          remove-stale-when-updated: true
           close-issue-message: |
             This issue has been automatically closed after six months of inactivity. If it's a
             feature request, it has been added to the maintainers' internal backlog and will be

.github/workflows/testrunner.yml

@@ -3,8 +3,7 @@ name: Testrunner
 on:
   push:
     branches:
-      - '1.x'
-  pull_request:
+      - 'main'
     paths:
       - 'docker/**'
 

ddtrace/appsec/_asm_request_context.py

@@ -412,14 +412,14 @@ def _on_set_request_tags(request, span, flask_config):
     if _is_iast_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 LazyTaintDict
+        from ddtrace.appsec._iast._taint_utils import taint_structure
 
         _set_metric_iast_instrumented_source(OriginType.COOKIE_NAME)
         _set_metric_iast_instrumented_source(OriginType.COOKIE)
-
-        request.cookies = LazyTaintDict(
+        request.cookies = taint_structure(
             request.cookies,
-            origins=(OriginType.COOKIE_NAME, OriginType.COOKIE),
+            OriginType.COOKIE_NAME,
+            OriginType.COOKIE,
             override_pyobject_tainted=True,
         )
 

ddtrace/appsec/_constants.py

@@ -75,6 +75,7 @@ class IAST(metaclass=Constant_Class):
     ENV = "DD_IAST_ENABLED"
     ENV_DEBUG = "_DD_IAST_DEBUG"
     TELEMETRY_REPORT_LVL = "DD_IAST_TELEMETRY_VERBOSITY"
+    LAZY_TAINT = "_DD_IAST_LAZY_TAINT"
     JSON = "_dd.iast.json"
     ENABLED = "_dd.iast.enabled"
     CONTEXT_KEY = "_iast_data"

ddtrace/appsec/_handlers.py

@@ -210,16 +210,13 @@ def _on_django_func_wrapped(fn_args, fn_kwargs, first_arg_expected_type, *_):
         from ddtrace.appsec._iast._taint_tracking import OriginType  # noqa: F401
         from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
         from ddtrace.appsec._iast._taint_tracking import taint_pyobject
-        from ddtrace.appsec._iast._taint_utils import LazyTaintDict
+        from ddtrace.appsec._iast._taint_utils import taint_structure
 
         http_req = fn_args[0]
 
-        if not isinstance(http_req.COOKIES, LazyTaintDict):
-            http_req.COOKIES = LazyTaintDict(http_req.COOKIES, origins=(OriginType.COOKIE_NAME, OriginType.COOKIE))
-        if not isinstance(http_req.GET, LazyTaintDict):
-            http_req.GET = LazyTaintDict(http_req.GET, origins=(OriginType.PARAMETER_NAME, OriginType.PARAMETER))
-        if not isinstance(http_req.POST, LazyTaintDict):
-            http_req.POST = LazyTaintDict(http_req.POST, origins=(OriginType.BODY, OriginType.BODY))
+        http_req.COOKIES = taint_structure(http_req.COOKIES, OriginType.COOKIE_NAME, OriginType.COOKIE)
+        http_req.GET = taint_structure(http_req.GET, OriginType.PARAMETER_NAME, OriginType.PARAMETER)
+        http_req.POST = taint_structure(http_req.POST, OriginType.BODY, OriginType.BODY)
         if not is_pyobject_tainted(getattr(http_req, "_body", None)):
             http_req._body = taint_pyobject(
                 http_req.body,
@@ -228,10 +225,7 @@ def _on_django_func_wrapped(fn_args, fn_kwargs, first_arg_expected_type, *_):
                 source_origin=OriginType.BODY,
             )
 
-        if not isinstance(http_req.META, LazyTaintDict):
-            http_req.META = LazyTaintDict(http_req.META, origins=(OriginType.HEADER_NAME, OriginType.HEADER))
-        if not isinstance(http_req.headers, LazyTaintDict):
-            http_req.headers = LazyTaintDict(http_req.headers, origins=(OriginType.HEADER_NAME, OriginType.HEADER))
+        http_req.headers = taint_structure(http_req.headers, OriginType.HEADER_NAME, OriginType.HEADER)
         http_req.path = taint_pyobject(
             http_req.path, source_name="path", source_value=http_req.path, source_origin=OriginType.PATH
         )
@@ -247,6 +241,7 @@ def _on_django_func_wrapped(fn_args, fn_kwargs, first_arg_expected_type, *_):
             source_value=http_req.path,
             source_origin=OriginType.PATH,
         )
+        http_req.META = taint_structure(http_req.META, OriginType.HEADER_NAME, OriginType.HEADER)
         if fn_kwargs:
             try:
                 for k, v in fn_kwargs.items():
@@ -264,7 +259,7 @@ def _on_wsgi_environ(wrapped, _instance, args, kwargs):
 
         from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_source
         from ddtrace.appsec._iast._taint_tracking import OriginType  # noqa: F401
-        from ddtrace.appsec._iast._taint_utils import LazyTaintDict
+        from ddtrace.appsec._iast._taint_utils import taint_structure
 
         _set_metric_iast_instrumented_source(OriginType.HEADER_NAME)
         _set_metric_iast_instrumented_source(OriginType.HEADER)
@@ -277,9 +272,7 @@ def _on_wsgi_environ(wrapped, _instance, args, kwargs):
         _set_metric_iast_instrumented_source(OriginType.PARAMETER_NAME)
         _set_metric_iast_instrumented_source(OriginType.BODY)
 
-        return wrapped(
-            *((LazyTaintDict(args[0], origins=(OriginType.HEADER_NAME, OriginType.HEADER)),) + args[1:]), **kwargs
-        )
+        return wrapped(*((taint_structure(args[0], OriginType.HEADER_NAME, OriginType.HEADER),) + args[1:]), **kwargs)
 
     return wrapped(*args, **kwargs)
 

ddtrace/appsec/_iast/_patches/json_tainting.py

@@ -7,6 +7,7 @@
 from .._patch import try_wrap_function_wrapper
 from .._taint_utils import LazyTaintDict
 from .._taint_utils import LazyTaintList
+from .._taint_utils import taint_structure
 
 
 log = get_logger(__name__)
@@ -24,8 +25,9 @@ def unpatch_iast():
     # type: () -> None
     set_module_unpatched("json", default_attr=_DEFAULT_ATTR)
     try_unwrap("json", "loads")
-    try_unwrap("json.encoder", "JSONEncoder.default")
-    try_unwrap("simplejson.encoder", "JSONEncoder.default")
+    if asm_config._iast_lazy_taint:
+        try_unwrap("json.encoder", "JSONEncoder.default")
+        try_unwrap("simplejson.encoder", "JSONEncoder.default")
 
 
 def patch():
@@ -34,8 +36,9 @@ def patch():
     if not set_and_check_module_is_patched("json", default_attr=_DEFAULT_ATTR):
         return
     try_wrap_function_wrapper("json", "loads", wrapped_loads)
-    try_wrap_function_wrapper("json.encoder", "JSONEncoder.default", patched_json_encoder_default)
-    try_wrap_function_wrapper("simplejson.encoder", "JSONEncoder.default", patched_json_encoder_default)
+    if asm_config._iast_lazy_taint:
+        try_wrap_function_wrapper("json.encoder", "JSONEncoder.default", patched_json_encoder_default)
+        try_wrap_function_wrapper("simplejson.encoder", "JSONEncoder.default", patched_json_encoder_default)
 
 
 def wrapped_loads(wrapped, instance, args, kwargs):
@@ -54,9 +57,9 @@ def wrapped_loads(wrapped, instance, args, kwargs):
                 # take the first source as main source
                 source = ranges[0].source
                 if isinstance(obj, dict):
-                    obj = LazyTaintDict(obj, origins=(source.origin, source.origin))
+                    obj = taint_structure(obj, source.origin, source.origin)
                 elif isinstance(obj, list):
-                    obj = LazyTaintList(obj, origins=(source.origin, source.origin))
+                    obj = taint_structure(obj, source.origin, source.origin)
                 elif isinstance(obj, (str, bytes, bytearray)):
                     obj = taint_pyobject(obj, source.name, source.value, source.origin)
                 pass

ddtrace/appsec/_iast/_taint_utils.py

@@ -1,7 +1,13 @@
 #!/usr/bin/env python3
 from collections import abc
+import dataclasses
+from typing import Any
+from typing import List
+from typing import Optional
+from typing import Union
 
 from ddtrace.internal.logger import get_logger
+from ddtrace.settings.asm import config as asm_config
 
 
 DBAPI_INTEGRATIONS = ("sqlite", "psycopg", "mysql", "mariadb")
@@ -10,6 +16,124 @@
 log = get_logger(__name__)
 
 
+# Non Lazy Tainting
+
+
+@dataclasses.dataclass
+class _DeepTaintCommand:
+    pre: bool
+    source_key: str
+    obj: Any
+    store_struct: Union[list, dict]
+    key: Optional[List[str]] = None
+    struct: Optional[Union[list, dict]] = None
+    is_key: bool = False
+
+    def store(self, value):
+        if isinstance(self.store_struct, list):
+            self.store_struct.append(value)
+        elif isinstance(self.store_struct, dict):
+            key = self.key[0] if self.key else None
+            self.store_struct[key] = value
+        else:
+            raise ValueError(f"store_struct of type {type(self.store_struct)}")
+
+    def post(self, struct):
+        return self.__class__(False, self.source_key, self.obj, self.store_struct, self.key, struct)
+
+
+def build_new_tainted_object_from_generic_object(initial_object, wanted_object):
+    if initial_object.__class__ is wanted_object.__class__:
+        return wanted_object
+    #### custom tailor actions
+    wanted_type = initial_object.__class__.__module__, initial_object.__class__.__name__
+    if wanted_type == ("builtins", "tuple"):
+        return tuple(wanted_object)
+    # Django
+    if wanted_type == ("django.http.request", "HttpHeaders"):
+        res = initial_object.__class__({})
+        res._store = {k.lower(): (k, v) for k, v in wanted_object.items()}
+        return res
+    if wanted_type == ("django.http.request", "QueryDict"):
+        res = initial_object.__class__()
+        for k, v in wanted_object.items():
+            dict.__setitem__(res, k, v)
+        return res
+    # Flask 2+
+    if wanted_type == ("werkzeug.datastructures.structures", "ImmutableMultiDict"):
+        return initial_object.__class__(wanted_object)
+    # Flask 1
+    if wanted_type == ("werkzeug.datastructures", "ImmutableMultiDict"):
+        return initial_object.__class__(wanted_object)
+
+    # if the class is unknown, return the initial object
+    # this may prevent interned string to be tainted but ensure
+    # that normal behavior of the code is not changed.
+    return initial_object
+
+
+def taint_structure(main_obj, source_key, source_value, override_pyobject_tainted=False):
+    """taint any structured object
+    use a queue like mechanism to avoid recursion
+    Best effort: mutate mutable structures and rebuild immutable ones if possible
+    """
+    from ._taint_tracking import is_pyobject_tainted
+    from ._taint_tracking import taint_pyobject
+
+    if not main_obj:
+        return main_obj
+
+    main_res = []
+    try:
+        # fifo contains tuple (pre/post:bool, source key, object to taint,
+        # key to use, struct to store result, struct to )
+        stack = [_DeepTaintCommand(True, source_key, main_obj, main_res)]
+        while stack:
+            command = stack.pop()
+            if command.pre:  # first processing of the object
+                if not command.obj:
+                    command.store(command.obj)
+                elif isinstance(command.obj, (str, bytes, bytearray)):
+                    if override_pyobject_tainted or not is_pyobject_tainted(command.obj):
+                        new_obj = taint_pyobject(
+                            pyobject=command.obj,
+                            source_name=command.source_key,
+                            source_value=command.obj,
+                            source_origin=source_key if command.is_key else source_value,
+                        )
+                        command.store(new_obj)
+                    else:
+                        command.store(command.obj)
+                elif isinstance(command.obj, abc.Mapping):
+                    res = {}
+                    stack.append(command.post(res))
+                    # use dict fondamental enumeration if possible to bypass any override of custom classes
+                    iterable = dict.items(command.obj) if isinstance(command.obj, dict) else command.obj.items()
+                    todo = []
+                    for k, v in list(iterable):
+                        key_store = []
+                        todo.append(_DeepTaintCommand(True, k, k, key_store, is_key=True))
+                        todo.append(_DeepTaintCommand(True, k, v, res, key_store))
+                    stack.extend(reversed(todo))
+                elif isinstance(command.obj, abc.Sequence):
+                    res = []
+                    stack.append(command.post(res))
+                    todo = [_DeepTaintCommand(True, command.source_key, v, res) for v in command.obj]
+                    stack.extend(reversed(todo))
+                else:
+                    command.store(command.obj)
+            else:
+                command.store(build_new_tainted_object_from_generic_object(command.obj, command.struct))
+    except BaseException:
+        log.debug("taint_structure error", exc_info=True)
+        pass
+    finally:
+        return main_res[0] if main_res else main_obj
+
+
+# Lazy Tainting
+
+
 def _is_tainted_struct(obj):
     return hasattr(obj, "_origins")
 
@@ -402,3 +526,13 @@ def check_tainted_args(args, kwargs, tracer, integration_name, method):
         return len(args) and args[0] and is_pyobject_tainted(args[0])
 
     return False
+
+
+if asm_config._iast_lazy_taint:
+    # redefining taint_structure to use lazy object if required
+
+    def taint_structure(main_obj, source_key, source_value, override_pyobject_tainted=False):  # noqa: F811
+        if isinstance(main_obj, abc.Mapping):
+            return LazyTaintDict(main_obj, source_key, source_value, override_pyobject_tainted)
+        elif isinstance(main_obj, abc.Sequence):
+            return LazyTaintList(main_obj, source_key, source_value, override_pyobject_tainted)

ddtrace/contrib/langchain/patch.py

@@ -555,7 +555,7 @@ def traced_embedding(langchain, pin, func, instance, args, kwargs):
         # langchain currently does not support token tracking for OpenAI embeddings:
         #  https://github.com/hwchase17/langchain/issues/945
         embeddings = func(*args, **kwargs)
-        if isinstance(embeddings, list) and isinstance(embeddings[0], list):
+        if isinstance(embeddings, list) and embeddings and isinstance(embeddings[0], list):
             for idx, embedding in enumerate(embeddings):
                 span.set_metric("langchain.response.outputs.%d.embedding_length" % idx, len(embedding))
         else: