Merge branch 'main' into patch-1

Author: emmettbutler

ddtrace/_trace/trace_handlers.py

@@ -823,7 +823,7 @@ def _set_span_pointer(span: "Span", span_pointer_description: _SpanPointerDescri
     )
 
 
-def _set_azure_function_tags(span, azure_functions_config, function_name, trigger, span_kind=SpanKind.INTERNAL):
+def _set_azure_function_tags(span, azure_functions_config, function_name, trigger, span_kind):
     span.set_tag_str(COMPONENT, azure_functions_config.integration_name)
     span.set_tag_str(SPAN_KIND, span_kind)
     span.set_tag_str("aas.function.name", function_name)  # codespell:ignore
@@ -857,9 +857,9 @@ def _on_azure_functions_start_response(ctx, azure_functions_config, res, functio
     )
 
 
-def _on_azure_functions_trigger_span_modifier(ctx, azure_functions_config, function_name, trigger):
+def _on_azure_functions_trigger_span_modifier(ctx, azure_functions_config, function_name, trigger, span_kind):
     span = ctx.get_item("trigger_span")
-    _set_azure_function_tags(span, azure_functions_config, function_name, trigger)
+    _set_azure_function_tags(span, azure_functions_config, function_name, trigger, span_kind)
 
 
 def listen():
@@ -968,6 +968,7 @@ def listen():
         "rq.job.perform",
         "rq.job.fetch_many",
         "azure.functions.patched_route_request",
+        "azure.functions.patched_service_bus",
         "azure.functions.patched_timer",
     ):
         core.on(f"context.started.start_span.{context_name}", _start_span)

ddtrace/_trace/tracer.py

@@ -1,5 +1,6 @@
 from contextlib import contextmanager
 import functools
+import inspect
 from inspect import iscoroutinefunction
 from itertools import chain
 import logging
@@ -775,6 +776,56 @@ def flush(self):
         """Flush the buffer of the trace writer. This does nothing if an unbuffered trace writer is used."""
         self._span_aggregator.writer.flush_queue()
 
+    def _wrap_generator(
+        self,
+        f: AnyCallable,
+        span_name: str,
+        service: Optional[str] = None,
+        resource: Optional[str] = None,
+        span_type: Optional[str] = None,
+    ) -> AnyCallable:
+        """Wrap a generator function with tracing."""
+
+        @functools.wraps(f)
+        def func_wrapper(*args, **kwargs):
+            if getattr(self, "_wrap_executor", None):
+                return self._wrap_executor(
+                    self,
+                    f,
+                    args,
+                    kwargs,
+                    span_name,
+                    service=service,
+                    resource=resource,
+                    span_type=span_type,
+                )
+
+            with self.trace(span_name, service=service, resource=resource, span_type=span_type):
+                gen = f(*args, **kwargs)
+                for value in gen:
+                    yield value
+
+        return func_wrapper
+
+    def _wrap_generator_async(
+        self,
+        f: AnyCallable,
+        span_name: str,
+        service: Optional[str] = None,
+        resource: Optional[str] = None,
+        span_type: Optional[str] = None,
+    ) -> AnyCallable:
+        """Wrap a generator function with tracing."""
+
+        @functools.wraps(f)
+        async def func_wrapper(*args, **kwargs):
+            with self.trace(span_name, service=service, resource=resource, span_type=span_type):
+                agen = f(*args, **kwargs)
+                async for value in agen:
+                    yield value
+
+        return func_wrapper
+
     def wrap(
         self,
         name: Optional[str] = None,
@@ -812,6 +863,15 @@ async def coroutine():
             def coroutine():
                 return 'executed'
 
+        >>> # or use it on generators
+            @tracer.wrap()
+            def gen():
+                yield 'executed'
+
+        >>> @tracer.wrap()
+            async def gen():
+                yield 'executed'
+
         You can access the current span using `tracer.current_span()` to set
         tags:
 
@@ -825,10 +885,26 @@ def wrap_decorator(f: AnyCallable) -> AnyCallable:
             # FIXME[matt] include the class name for methods.
             span_name = name if name else "%s.%s" % (f.__module__, f.__name__)
 
-            # detect if the the given function is a coroutine to use the
-            # right decorator; this initial check ensures that the
+            # detect if the the given function is a coroutine and/or a generator
+            # to use the right decorator; this initial check ensures that the
             # evaluation is done only once for each @tracer.wrap
-            if iscoroutinefunction(f):
+            if inspect.isgeneratorfunction(f):
+                func_wrapper = self._wrap_generator(
+                    f,
+                    span_name,
+                    service=service,
+                    resource=resource,
+                    span_type=span_type,
+                )
+            elif inspect.isasyncgenfunction(f):
+                func_wrapper = self._wrap_generator_async(
+                    f,
+                    span_name,
+                    service=service,
+                    resource=resource,
+                    span_type=span_type,
+                )
+            elif iscoroutinefunction(f):
                 # create an async wrapper that awaits the coroutine and traces it
                 @functools.wraps(f)
                 async def func_wrapper(*args, **kwargs):

ddtrace/appsec/_iast/_patch_modules.py

@@ -6,6 +6,7 @@
 from ddtrace.appsec._iast.secure_marks.sanitizers import path_traversal_sanitizer
 from ddtrace.appsec._iast.secure_marks.sanitizers import sqli_sanitizer
 from ddtrace.appsec._iast.secure_marks.sanitizers import xss_sanitizer
+from ddtrace.appsec._iast.secure_marks.validators import ssrf_validator
 from ddtrace.appsec._iast.secure_marks.validators import unvalidated_redirect_validator
 
 
@@ -32,6 +33,7 @@ def patch_iast(patch_modules=IAST_PATCH):
     for module in (m for m, e in patch_modules.items() if e):
         when_imported("hashlib")(_on_import_factory(module, "ddtrace.appsec._iast.taint_sinks.%s", raise_errors=False))
 
+    # CMDI sanitizers
     when_imported("shlex")(
         lambda _: try_wrap_function_wrapper(
             "shlex",
@@ -40,6 +42,11 @@ def patch_iast(patch_modules=IAST_PATCH):
         )
     )
 
+    # SSRF
+    when_imported("django.utils.http")(
+        lambda _: try_wrap_function_wrapper("django.utils.http", "url_has_allowed_host_and_scheme", ssrf_validator)
+    )
+
     # SQL sanitizers
     when_imported("mysql.connector.conversion")(
         lambda _: try_wrap_function_wrapper(

ddtrace/appsec/_iast/secure_marks/validators.py

@@ -93,3 +93,18 @@ def unvalidated_redirect_validator(wrapped: Callable, instance: Any, args: Seque
         True if validation passed, False otherwise
     """
     return create_validator(VulnerabilityType.UNVALIDATED_REDIRECT, wrapped, instance, args, kwargs)
+
+
+def ssrf_validator(wrapped: Callable, instance: Any, args: Sequence, kwargs: dict) -> bool:
+    """Validator for ssrf functions.
+
+    Args:
+        wrapped: The original validator function
+        instance: The instance the function is bound to (if any)
+        args: Positional arguments
+        kwargs: Keyword arguments
+
+    Returns:
+        True if validation passed, False otherwise
+    """
+    return create_validator(VulnerabilityType.SSRF, wrapped, instance, args, kwargs)

ddtrace/appsec/_iast/taint_sinks/ssrf.py

@@ -2,21 +2,19 @@
 
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
 from ddtrace.appsec._iast._logs import iast_error
+from ddtrace.appsec._iast._logs import iast_propagation_sink_point_debug_log
 from ddtrace.appsec._iast._metrics import _set_metric_iast_executed_sink
 from ddtrace.appsec._iast._span_metrics import increment_iast_span_metric
 from ddtrace.appsec._iast._taint_tracking import VulnerabilityType
+from ddtrace.appsec._iast._taint_tracking import get_ranges
 from ddtrace.appsec._iast.constants import VULN_SSRF
 from ddtrace.appsec._iast.taint_sinks._base import VulnerabilityBase
-from ddtrace.internal.logger import get_logger
 from ddtrace.internal.utils import ArgumentError
 from ddtrace.internal.utils import get_argument_value
 from ddtrace.internal.utils.importlib import func_name
 from ddtrace.settings.asm import config as asm_config
 
 
-log = get_logger(__name__)
-
-
 class SSRF(VulnerabilityBase):
     vulnerability_type = VULN_SSRF
     secure_mark = VulnerabilityType.SSRF
@@ -33,24 +31,41 @@ class SSRF(VulnerabilityBase):
 
 
 def _iast_report_ssrf(func: Callable, *args, **kwargs):
+    """
+    Check and report potential SSRF (Server-Side Request Forgery) vulnerabilities in function calls.
+
+    This function analyzes calls to URL-handling functions to detect potential SSRF vulnerabilities.
+    It checks if the URL argument is tainted (user-controlled) and reports it if conditions are met.
+    URL fragments (parts after #) are handled specially - if all tainted parts are in the fragment,
+    no vulnerability is reported.
+    """
     func_key = func_name(func)
     arg_pos, kwarg_name = _FUNC_TO_URL_ARGUMENT.get(func_key, (None, None))
     if arg_pos is None:
-        log.debug("%s not found in list of functions supported for SSRF", func_key)
+        iast_propagation_sink_point_debug_log("%s not found in list of functions supported for SSRF", func_key)
         return
 
     try:
         kw = kwarg_name if kwarg_name else ""
         report_ssrf = get_argument_value(list(args), kwargs, arg_pos, kw)
     except ArgumentError:
-        log.debug("Failed to get URL argument from _FUNC_TO_URL_ARGUMENT dict for function %s", func_key)
+        iast_propagation_sink_point_debug_log(
+            "Failed to get URL argument from _FUNC_TO_URL_ARGUMENT dict for function %s", func_key
+        )
         return
-
     if report_ssrf:
         if asm_config.is_iast_request_enabled:
             try:
                 if SSRF.has_quota() and SSRF.is_tainted_pyobject(report_ssrf):
-                    SSRF.report(evidence_value=report_ssrf)
+                    valid_to_report = True
+                    fragment_start = report_ssrf.find("#")
+                    taint_ranges = get_ranges(report_ssrf)
+                    if fragment_start != -1:
+                        # If all taint ranges are in the fragment, do not report
+                        if all(r.start >= fragment_start for r in taint_ranges):
+                            valid_to_report = False
+                    if valid_to_report:
+                        SSRF.report(evidence_value=report_ssrf)
 
                 # Reports Span Metrics
                 _set_metric_iast_executed_sink(SSRF.vulnerability_type)