datadog 2.34.0 → 2.36.0

data/ext/datadog_profiling_native_extension/ruby_helpers.c

@@ -21,42 +21,40 @@ void ruby_helpers_init(void) {
   to_s_id = rb_intern("to_s");
 }
 
-// Internal helper for raising pre-formatted syserr exceptions
-static NORETURN(void private_raise_syserr_formatted(int syserr_errno, const char *detailed_message, const char *static_message)) {
-  VALUE exception = rb_syserr_new(syserr_errno, detailed_message);
-  private_raise_exception(exception, static_message);
-}
-
 // Use `raise_syserr` the macro instead, as it provides additional argument checks.
 void private_raise_syserr(int syserr_errno, const char *fmt, ...) {
-  FORMAT_VA_ERROR_MESSAGE(detailed_message, fmt);
-  private_raise_syserr_formatted(syserr_errno, detailed_message, fmt);
+  va_list args;
+  va_start(args, fmt);
+  VALUE detailed_message = rb_vsprintf(fmt, args);
+  va_end(args);
+
+  VALUE exception = rb_syserr_new_str(syserr_errno, detailed_message);
+  private_raise_exception(exception, fmt);
 }
 
 typedef struct {
   VALUE exception_class;
   int syserr_errno;
-  char exception_message[MAX_RAISE_MESSAGE_SIZE];
-  char telemetry_message[MAX_RAISE_MESSAGE_SIZE];
+  const char *format_string;
+  va_list va_args;
 } raise_args;
 
+// Called via rb_thread_call_with_gvl from private_grab_gvl_and_raise.
+// Formats the message with rb_vsprintf (which requires the GVL) and raises.
 static void *trigger_raise(void *raise_arguments) {
   raise_args *args = (raise_args *) raise_arguments;
 
+  VALUE detailed_message = rb_vsprintf(args->format_string, args->va_args);
+
+  VALUE exception;
   if (args->syserr_errno) {
-    private_raise_syserr_formatted(
-      args->syserr_errno,
-      args->exception_message,
-      args->telemetry_message
-    );
+    exception = rb_syserr_new_str(args->syserr_errno, detailed_message);
   } else {
-    private_raise_error_formatted(
-      args->exception_class,
-      args->exception_message,
-      args->telemetry_message
-    );
+    exception = rb_exc_new_str(args->exception_class, detailed_message);
   }
 
+  private_raise_exception(exception, args->format_string);
+
   return NULL;
 }
 
@@ -71,11 +69,17 @@ void private_grab_gvl_and_raise(VALUE exception_class, int syserr_errno, const c
     args.syserr_errno = 0;
   }
 
-  FORMAT_VA_ERROR_MESSAGE(formatted_exception_message, format_string);
-  snprintf(args.exception_message, MAX_RAISE_MESSAGE_SIZE, "%s", formatted_exception_message);
-  snprintf(args.telemetry_message, MAX_RAISE_MESSAGE_SIZE, "%s", format_string);
+  args.format_string = format_string;
+  va_start(args.va_args, format_string);
 
   if (is_current_thread_holding_the_gvl()) {
+    VALUE detailed_message = rb_vsprintf(format_string, args.va_args);
+    va_end(args.va_args);
+
+    VALUE wrapped_message = rb_sprintf(
+      "grab_gvl_and_raise called by thread holding the global VM lock: %"PRIsVALUE,
+      detailed_message
+    );
     char telemetry_message[MAX_RAISE_MESSAGE_SIZE];
     snprintf(
       telemetry_message,
@@ -83,20 +87,14 @@ void private_grab_gvl_and_raise(VALUE exception_class, int syserr_errno, const c
       "grab_gvl_and_raise called by thread holding the global VM lock: %s",
       format_string
     );
-    char exception_message[MAX_RAISE_MESSAGE_SIZE];
-    snprintf(
-      exception_message,
-      MAX_RAISE_MESSAGE_SIZE,
-      "grab_gvl_and_raise called by thread holding the global VM lock: %s",
-      args.exception_message
-    );
-    VALUE exception = rb_exc_new_cstr(rb_eRuntimeError, exception_message);
+    VALUE exception = rb_exc_new_str(rb_eRuntimeError, wrapped_message);
     private_raise_exception(exception, telemetry_message);
   }
 
   rb_thread_call_with_gvl(trigger_raise, &args);
 
-  rb_bug("[ddtrace] Unexpected: Reached the end of grab_gvl_and_raise while raising '%s'\n", args.exception_message);
+  va_end(args.va_args);
+  rb_bug("[ddtrace] Unexpected: Reached the end of grab_gvl_and_raise while raising '%s'\n", format_string);
 }
 
 void private_raise_enforce_syserr(
@@ -107,10 +105,11 @@ void private_raise_enforce_syserr(
   int line,
   const char *function_name
 ) {
+  const char *format = "Failure returned by '%s' at %s:%d:in `%s'";
   if (have_gvl) {
-    rb_exc_raise(rb_syserr_new_str(syserr_errno, rb_sprintf("Failure returned by '%s' at %s:%d:in `%s'", expression, file, line, function_name)));
+    private_raise_exception(rb_syserr_new_str(syserr_errno, rb_sprintf(format, expression, file, line, function_name)), format);
   } else {
-    private_grab_gvl_and_raise(Qnil, syserr_errno, "Failure returned by '%s' at %s:%d:in `%s'", expression, file, line, function_name);
+    private_grab_gvl_and_raise(Qnil, syserr_errno, format, expression, file, line, function_name);
   }
 }
 

data/ext/datadog_profiling_native_extension/stack_recorder.c

@@ -9,6 +9,7 @@
 #include "time_helpers.h"
 #include "heap_recorder.h"
 #include "encoded_profile.h"
+#include "collectors_thread_context.h"
 
 // Used to wrap a ddog_prof_Profile in a Ruby object and expose Ruby-level serialization APIs
 // This file implements the native bits of the Datadog::Profiling::StackRecorder class
@@ -177,6 +178,10 @@ typedef struct {
   heap_recorder *heap_recorder;
   bool heap_clean_after_gc_enabled;
 
+  // When set, _native_serialize will call thread_context_collector_on_serialize on this instance
+  // before serializing, so that threads suspended across the whole profile period still get sampled.
+  VALUE thread_context_collector_instance;
+
   pthread_mutex_t mutex_slot_one;
   profile_slot profile_slot_one;
   pthread_mutex_t mutex_slot_two;
@@ -320,6 +325,7 @@ static VALUE _native_new(VALUE klass) {
   // being leaked.
 
   state->heap_clean_after_gc_enabled = false;
+  state->thread_context_collector_instance = Qnil;
 
   ddog_prof_Slice_SampleType sample_types = {.ptr = all_sample_types, .len = ALL_VALUE_TYPES_COUNT};
 
@@ -391,6 +397,7 @@ static void initialize_profiles(stack_recorder_state *state, ddog_prof_Slice_Sam
 static void stack_recorder_typed_data_mark(void *state_ptr) {
   stack_recorder_state *state = (stack_recorder_state *) state_ptr;
 
+  rb_gc_mark(state->thread_context_collector_instance);
   heap_recorder_mark_pending_recordings(state->heap_recorder);
 }
 
@@ -420,14 +427,12 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
   VALUE heap_samples_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_samples_enabled")));
   VALUE heap_size_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_size_enabled")));
   VALUE heap_sample_every = rb_hash_fetch(options, ID2SYM(rb_intern("heap_sample_every")));
-  VALUE timeline_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("timeline_enabled")));
   VALUE heap_clean_after_gc_enabled = rb_hash_fetch(options, ID2SYM(rb_intern("heap_clean_after_gc_enabled")));
 
   ENFORCE_BOOLEAN(alloc_samples_enabled);
   ENFORCE_BOOLEAN(heap_samples_enabled);
   ENFORCE_BOOLEAN(heap_size_enabled);
   ENFORCE_TYPE(heap_sample_every, T_FIXNUM);
-  ENFORCE_BOOLEAN(timeline_enabled);
   ENFORCE_BOOLEAN(heap_clean_after_gc_enabled);
 
   stack_recorder_state *state;
@@ -440,8 +445,7 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
   uint8_t requested_values_count = ALL_VALUE_TYPES_COUNT -
     (alloc_samples_enabled == Qtrue? 0 : 2) -
     (heap_samples_enabled == Qtrue ? 0 : 1) -
-    (heap_size_enabled == Qtrue ? 0 : 1) -
-    (timeline_enabled == Qtrue ? 0 : 1);
+    (heap_size_enabled == Qtrue ? 0 : 1);
 
   if (requested_values_count == ALL_VALUE_TYPES_COUNT) return Qtrue; // Nothing to do, this is the default
 
@@ -467,6 +471,10 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
   enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_CPU_TIME;
   state->position_for[CPU_TIME_VALUE_ID] = next_enabled_pos++;
 
+  // TIMELINE is always enabled
+  enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_TIMELINE;
+  state->position_for[TIMELINE_VALUE_ID] = next_enabled_pos++;
+
   if (alloc_samples_enabled == Qtrue) {
     enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_ALLOC_SAMPLES;
     state->position_for[ALLOC_SAMPLES_VALUE_ID] = next_enabled_pos++;
@@ -500,13 +508,6 @@ static VALUE _native_initialize(int argc, VALUE *argv, DDTRACE_UNUSED VALUE _sel
     state->heap_recorder = NULL;
   }
 
-  if (timeline_enabled == Qtrue) {
-    enabled_sample_types[next_enabled_pos] = DDOG_PROF_SAMPLE_TYPE_TIMELINE;
-    state->position_for[TIMELINE_VALUE_ID] = next_enabled_pos++;
-  } else {
-    state->position_for[TIMELINE_VALUE_ID] = next_disabled_pos++;
-  }
-
   ddog_prof_Profile_drop(&state->profile_slot_one.profile);
   ddog_prof_Profile_drop(&state->profile_slot_two.profile);
 
@@ -521,12 +522,16 @@ static VALUE _native_serialize(DDTRACE_UNUSED VALUE _self, VALUE recorder_instan
   TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
 
   ddog_Timespec finish_timestamp = system_epoch_now_timespec();
-  // Need to do this while still holding on to the Global VM Lock; see comments on method for why
+  // Need to do this while still holding the Global VM Lock; see comments on method for why
   serializer_set_start_timestamp_for_next_profile(state, finish_timestamp);
 
+  if (state->thread_context_collector_instance != Qnil) {
+    thread_context_collector_on_serialize(state->thread_context_collector_instance);
+  }
+
   long heap_iteration_prep_start_time_ns = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
   // Prepare the iteration on heap recorder we'll be doing outside the GVL. The preparation needs to
-  // happen while holding on to the GVL.
+  // happen while holding the GVL.
   // NOTE: While rare, it's possible for the GVL to be released inside this function (see comments on `heap_recorder_update`)
   // and thus don't assume this is an "atomic" step -- other threads may get some running time in the meanwhile.
   heap_recorder_prepare_iteration(state->heap_recorder);
@@ -554,7 +559,7 @@ static VALUE _native_serialize(DDTRACE_UNUSED VALUE _self, VALUE recorder_instan
     rb_thread_call_without_gvl2(call_serialize_without_gvl, &args, NULL /* No interruption function needed in this case */, NULL /* Not needed */);
   }
 
-  // Cleanup after heap recorder iteration. This needs to happen while holding on to the GVL.
+  // Cleanup after heap recorder iteration. This needs to happen while holding the GVL.
   heap_recorder_finish_iteration(state->heap_recorder);
 
   // NOTE: We are focusing on the serialization time outside of the GVL in this stat here. This doesn't
@@ -639,6 +644,7 @@ void record_sample(VALUE recorder_instance, ddog_prof_Slice_Location locations,
       .values = (ddog_Slice_I64) {.ptr = metric_values, .len = state->enabled_values_count},
       .labels = labels.labels
     },
+    /* To disable end_timestamp_ns to get aggregated profiles in pprof, replace the below with `0` */
     labels.end_timestamp_ns
   );
 
@@ -899,6 +905,9 @@ static ddog_Timespec system_epoch_now_timespec(void) {
 //
 // Assumption: This method gets called BEFORE restarting profiling -- e.g. there are no components attempting to
 // trigger samples at the same time.
+//
+// Note that tests call this method directly in the same process without forking,
+// and in such a case non-current Threads keep running.
 static VALUE _native_reset_after_fork(DDTRACE_UNUSED VALUE self, VALUE recorder_instance) {
   stack_recorder_state *state;
   TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
@@ -1152,3 +1161,10 @@ static VALUE _native_finalize_pending_heap_recordings(DDTRACE_UNUSED VALUE _self
 
   return Qtrue;
 }
+
+void recorder_install_on_serialize(VALUE recorder_instance, VALUE thread_context_collector_instance) {
+  stack_recorder_state *state;
+  TypedData_Get_Struct(recorder_instance, stack_recorder_state, &stack_recorder_typed_data, state);
+
+  state->thread_context_collector_instance = enforce_thread_context_collector_instance(thread_context_collector_instance);
+}

data/ext/datadog_profiling_native_extension/stack_recorder.h

@@ -29,4 +29,5 @@ void record_endpoint(VALUE recorder_instance, uint64_t local_root_span_id, ddog_
 __attribute__((warn_unused_result)) bool track_object(VALUE recorder_instance, VALUE new_object, unsigned int sample_weight, ddog_CharSlice alloc_class);
 void recorder_after_sample(VALUE recorder_instance);
 void recorder_after_gc_step(VALUE recorder_instance);
+void recorder_install_on_serialize(VALUE recorder_instance, VALUE thread_context_collector_instance);
 VALUE enforce_recorder_instance(VALUE object);

data/ext/datadog_profiling_native_extension/unsafe_api_calls_check.h

@@ -4,7 +4,7 @@
 //
 // Specifically, when the profiler is sampling, we're never supposed to call into Ruby code (e.g. methods
 // implemented using Ruby code) or allocate Ruby objects.
-// That's because those events introduce thread switch points, and really we don't the VM switching between threads
+// That's because those events introduce thread switch points, and really we don't want the VM switching between threads
 // in the middle of the profiler sampling. This includes raising exceptions.
 //
 // Raising exceptions as the very last operation, to stop the profiler is ok, but comes a caveat: raising exceptions
@@ -18,10 +18,12 @@
 // in most (all?) thread switch points, Ruby will check for interrupts and run the postponed jobs.
 //
 // Thus, if we set a flag while we're sampling (inside_unsafe_context), trigger the postponed job, and then only unset
-// the flag after sampling, he correct thing to happen is that the postponed job should never see the flag.
+// the flag after sampling, the correct thing to happen is that the postponed job should never see the flag.
 //
 // If, however, we have a bug and there's a thread switch point, our postponed job will see the flag and immediately
 // stop the Ruby VM before further damage happens (and hopefully giving us a stack trace clearly pointing to the culprit).
+//
+// Note that this check currently does not detect Ruby object allocations, as those do not check for interrupts.
 
 void unsafe_api_calls_check_init(void);
 

data/ext/libdatadog_api/datadog_ruby_common.c

@@ -29,16 +29,15 @@ void private_raise_exception(VALUE exception, const char *static_message) {
   rb_exc_raise(exception);
 }
 
-// Helper for raising pre-formatted exceptions
-void private_raise_error_formatted(VALUE exception_class, const char *detailed_message, const char *static_message) {
-  VALUE exception = rb_exc_new_cstr(exception_class, detailed_message);
-  private_raise_exception(exception, static_message);
-}
-
 // Use `raise_error` the macro instead, as it provides additional argument checks.
 void private_raise_error(VALUE exception_class, const char *fmt, ...) {
-  FORMAT_VA_ERROR_MESSAGE(detailed_message, fmt);
-  private_raise_error_formatted(exception_class, detailed_message, fmt);
+  va_list args;
+  va_start(args, fmt);
+  VALUE detailed_message = rb_vsprintf(fmt, args);
+  va_end(args);
+
+  VALUE exception = rb_exc_new_str(exception_class, detailed_message);
+  private_raise_exception(exception, fmt);
 }
 
 VALUE datadog_gem_version(void) {

data/ext/libdatadog_api/datadog_ruby_common.h

@@ -47,11 +47,6 @@ NORETURN(
   __attribute__ ((format (printf, 2, 3)));
 );
 
-// Internal helper for raising pre-formatted exceptions
-NORETURN(
-  void private_raise_error_formatted(VALUE exception_class, const char *detailed_message, const char *static_message)
-);
-
 // Raises an exception with separate telemetry-safe and detailed messages.
 // NOTE: Raising an exception always invokes Ruby code so it requires the GVL and is not compatible with "debug_enter_unsafe_context".
 // @see debug_enter_unsafe_context
@@ -61,13 +56,6 @@ NORETURN(
 
 #define MAX_RAISE_MESSAGE_SIZE 256
 
-#define FORMAT_VA_ERROR_MESSAGE(buf, fmt) \
-  char buf[MAX_RAISE_MESSAGE_SIZE]; \
-  va_list buf##_args; \
-  va_start(buf##_args, fmt); \
-  vsnprintf(buf, MAX_RAISE_MESSAGE_SIZE, fmt, buf##_args); \
-  va_end(buf##_args);
-
 // Helper to retrieve Datadog::VERSION::STRING
 VALUE datadog_gem_version(void);
 

data/ext/libdatadog_extconf_helpers.rb

@@ -10,7 +10,7 @@ module Datadog
   module LibdatadogExtconfHelpers
     # Used to make sure the correct gem version gets loaded, as extconf.rb does not get run with "bundle exec" and thus
     # may see multiple libdatadog versions. See https://github.com/DataDog/dd-trace-rb/pull/2531 for the horror story.
-    LIBDATADOG_VERSION = '~> 33.0.0.1.0'
+    LIBDATADOG_VERSION = '~> 35.0.0.1.0'
 
     # Used as an workaround for a limitation with how dynamic linking works in environments where the datadog gem and
     # libdatadog are moved after the extension gets compiled.

data/lib/datadog/ai_guard/configuration.rb

@@ -71,6 +71,7 @@ module Datadog
               option :app_key do |o|
                 o.type :string, nilable: true
                 o.env Ext::ENV_APP_KEY
+                o.skip_telemetry true
               end
 
               # Request timeout in milliseconds.

data/lib/datadog/ai_guard/contrib/rack/request_middleware.rb

@@ -11,71 +11,85 @@ module Datadog
       module Rack
         # AI Guard Rack middleware.
         #
-        # Inserted into the middleware stack right after
-        # Datadog::Tracing::Contrib::Rack::TraceMiddleware (i.e. nested inside
-        # it). This ordering matters: on the way out of the request, AI Guard's
-        # `ensure` block unwinds *before* Tracing's ensure, while Tracing's
-        # request span is still live. We need that, because Tracing's ensure
-        # calls `request_span.finish`, which builds a frozen `Span` snapshot of
-        # the meta hash — any `set_tag` call after that point mutates the
-        # `SpanOperation` but never reaches the exported `Span`.
+        # At request entry the middleware stores some request attributes for
+        # on the active trace under the `_dd.ai_guard.` prefix.
         #
-        # So while the span is still active, we tag `http.client_ip` and
-        # `network.client.ip` on it — but only when an AI Guard span was
-        # actually recorded during the request.
+        # Later when AI Guard evaluation is performed, those attributes are
+        # mirrored on AI Guard span.
         class RequestMiddleware
-          NETWORK_CLIENT_IP_TAG = "network.client.ip"
-
           def initialize(app, opt = {})
             @app = app
           end
 
           def call(env)
+            trace = Datadog::Tracing.active_trace
+            return @app.call(env) unless trace
+
+            store_anomaly_detection_tags!(trace, env)
+
             @app.call(env)
           ensure
-            tag_client_ip(env) if consume_ai_guard_executed_flag
+            # @type var trace: Datadog::Tracing::TraceSegment?
+            # Steep: https://github.com/soutaro/steep/issues/919
+            if trace
+              tag_client_ip_on_request_span!(trace) if ai_guard_executed?(trace)
+
+              clean_up_ai_guard_temp_tags!(trace)
+            end
           end
 
           private
 
-          # AI Guard's evaluation flow sets `ai_guard.executed` on the trace
-          # whenever an AI Guard span is created during the request. We read
-          # it here to know whether to tag client IP, then clear it so the
-          # internal flag does not propagate to the exported trace.
-          #
-          # `Tracing.active_trace` is publicly typed as `TraceSegment?` but at
-          # runtime returns a `TraceOperation`, which exposes `get_tag` and
-          # `clear_tag`. Pre-existing sig mismatch — hence the steep:ignore.
           # steep:ignore:start
-          def consume_ai_guard_executed_flag
-            trace = Datadog::Tracing.active_trace
-            return false unless trace
-            return false unless trace.get_tag(Datadog::AIGuard::Ext::SERVICE_ENTRY_EXECUTED_TAG) == "1"
+          def store_anomaly_detection_tags!(trace, env)
+            remote_ip = env["REMOTE_ADDR"]
+            trace.set_tag(Datadog::AIGuard::Ext::TRACE_NETWORK_CLIENT_IP_TAG, remote_ip) if remote_ip
+
+            headers = Datadog::Tracing::Contrib::Rack::Header::RequestHeaderCollection.new(env)
+            resolved_client_ip = Datadog::Tracing::ClientIp.extract_client_ip(headers, remote_ip)
+            trace.set_tag(Datadog::AIGuard::Ext::TRACE_HTTP_CLIENT_IP_TAG, resolved_client_ip) if resolved_client_ip
+
+            user_agent = env["HTTP_USER_AGENT"]
+            trace.set_tag(Datadog::AIGuard::Ext::TRACE_HTTP_USERAGENT_TAG, user_agent) if user_agent
+          rescue => e
+            Datadog::AIGuard.telemetry&.report(e, description: "AI Guard: failed to get request attributes")
+          end
+          # steep:ignore:end
+
+          # steep:ignore:start
+          def clean_up_ai_guard_temp_tags!(trace)
+            Ext::TRACE_ANOMALY_DETECTION_TAGS.each do |tag|
+              trace.clear_tag(tag)
+            end
 
-            trace.clear_tag(Datadog::AIGuard::Ext::SERVICE_ENTRY_EXECUTED_TAG)
-            true
+            trace.clear_tag(Datadog::AIGuard::Ext::TRACE_EXECUTED_TAG)
           end
           # steep:ignore:end
 
-          def tag_client_ip(env)
+          # AI Guard's evaluation flow sets `_dd.ai_guard.executed` on the
+          # trace whenever an AI Guard span is created during the request.
+          # steep:ignore:start
+          def ai_guard_executed?(trace)
+            trace.get_tag(Datadog::AIGuard::Ext::TRACE_EXECUTED_TAG) == "1"
+          end
+          # steep:ignore:end
+
+          # steep:ignore:start
+          def tag_client_ip_on_request_span!(trace)
             span = Datadog::Tracing.active_span
             return unless span
 
-            if span.get_tag(Datadog::Tracing::Metadata::Ext::HTTP::TAG_CLIENT_IP).nil?
-              headers = Datadog::Tracing::Contrib::Rack::Header::RequestHeaderCollection.new(env)
-              Datadog::Tracing::ClientIp.set_client_ip_tag!(
-                span,
-                headers: headers,
-                remote_ip: env["REMOTE_ADDR"]
-              )
+            client_ip = trace.get_tag(Datadog::AIGuard::Ext::TRACE_HTTP_CLIENT_IP_TAG)
+            if client_ip && span.get_tag(Datadog::Tracing::Metadata::Ext::HTTP::TAG_CLIENT_IP).nil?
+              span.set_tag(Datadog::Tracing::Metadata::Ext::HTTP::TAG_CLIENT_IP, client_ip)
             end
 
-            if env["REMOTE_ADDR"] && span.get_tag(NETWORK_CLIENT_IP_TAG).nil?
-              span.set_tag(NETWORK_CLIENT_IP_TAG, env["REMOTE_ADDR"])
-            end
+            network_client_ip = trace.get_tag(Datadog::AIGuard::Ext::TRACE_NETWORK_CLIENT_IP_TAG)
+            span["network.client.ip"] = network_client_ip if network_client_ip
           rescue => e
             Datadog::AIGuard.telemetry&.report(e, description: "AI Guard: failed to tag client IP on root span")
           end
+          # steep:ignore:end
         end
       end
     end

data/lib/datadog/ai_guard/evaluation.rb

@@ -16,7 +16,12 @@ module Datadog
               Tracing::Sampling::Ext::Decision::AI_GUARD
             )
             trace.set_tag(Ext::EVENT_TAG, true)
-            trace.set_tag(Ext::SERVICE_ENTRY_EXECUTED_TAG, "1")
+            trace.set_tag(Ext::TRACE_EXECUTED_TAG, "1")
+
+            Ext::TRACE_ANOMALY_DETECTION_TAGS.each do |tag|
+              value = trace.get_tag(tag)
+              span.set_tag(tag, value) if value
+            end
 
             if (last_message = messages.last)
               if last_message.tool_call

data/lib/datadog/ai_guard/ext.rb

@@ -11,8 +11,19 @@ module Datadog
       REASON_TAG = "ai_guard.reason"
       BLOCKED_TAG = "ai_guard.blocked"
       EVENT_TAG = "ai_guard.event"
-      SERVICE_ENTRY_EXECUTED_TAG = "_dd.ai_guard.executed"
+
       METASTRUCT_TAG = "ai_guard"
+
+      TRACE_EXECUTED_TAG = "_dd.ai_guard.executed"
+      TRACE_HTTP_USERAGENT_TAG = "_dd.ai_guard.http.useragent"
+      TRACE_HTTP_CLIENT_IP_TAG = "_dd.ai_guard.http.client_ip"
+      TRACE_NETWORK_CLIENT_IP_TAG = "_dd.ai_guard.network.client.ip"
+
+      TRACE_ANOMALY_DETECTION_TAGS = [
+        TRACE_HTTP_USERAGENT_TAG,
+        TRACE_HTTP_CLIENT_IP_TAG,
+        TRACE_NETWORK_CLIENT_IP_TAG
+      ].freeze
     end
   end
 end

data/lib/datadog/appsec/api_security/route_extractor.rb

@@ -15,6 +15,7 @@ module Datadog
         RAILS_ROUTES_KEY = 'action_dispatch.routes'
         RAILS_PATH_PARAMS_KEY = 'action_dispatch.request.path_parameters'
         RAILS_FORMAT_SUFFIX = '(.:format)'
+        DATADOG_RAILS_ROUTE_KEY = 'datadog.action_dispatch.route'
 
         # HACK: We rely on the fact that each contrib will modify `request.env`
         #       and store information sufficient to compute the canonical
@@ -49,12 +50,20 @@ module Datadog
         #          In Rails < 7.1 it also will not be set even if a route was found,
         #          but in this case `action_dispatch.request.path_parameters` won't be empty.
         def self.route_pattern(request)
+          # NOTE: Requests from contribs like AWS Lambda don't provide a usable
+          #       `::Rack::Request#env`, so infer the route from the path instead
+          unless request.respond_to?(:env)
+            return Tracing::Contrib::Rack::RouteInference.infer(request.path.to_s)
+          end
+
           if request.env.key?(GRAPE_ROUTE_KEY)
             pattern = request.env[GRAPE_ROUTE_KEY][:route_info]&.pattern&.origin
             "#{request.script_name}#{pattern}"
           elsif request.env.key?(SINATRA_ROUTE_KEY)
             pattern = request.env[SINATRA_ROUTE_KEY].split(SINATRA_ROUTE_SEPARATOR, 2)[1]
             "#{request.script_name}#{pattern}"
+          elsif request.env.key?(DATADOG_RAILS_ROUTE_KEY)
+            request.env[DATADOG_RAILS_ROUTE_KEY].path.spec.to_s.delete_suffix(RAILS_FORMAT_SUFFIX)
           elsif request.env.key?(RAILS_ROUTE_KEY)
             request.env[RAILS_ROUTE_KEY].path.spec.to_s.delete_suffix(RAILS_FORMAT_SUFFIX)
           elsif request.env.key?(RAILS_ROUTE_URI_PATTERN_KEY)

data/lib/datadog/appsec/component.rb

@@ -22,7 +22,7 @@ module Datadog
             return
           end
 
-          if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('3.3') && ffi_version < Gem::Version.new('1.16.0')
+          if RubyVersion.is?('>= 3.3') && ffi_version < Gem::Version.new('1.16.0')
             Datadog.logger.warn(
               'AppSec is not supported in Ruby versions above 3.3.0 when using `ffi` versions older than 1.16.0, ' \
               'and will be forcibly disabled due to a memory leak in `ffi`. ' \

data/lib/datadog/appsec/configuration.rb

@@ -134,6 +134,13 @@ module Datadog
                 end
               end
 
+              # NOTE: A value of 0 (or less) disables request body collection
+              option :body_parsing_size_limit do |o|
+                o.type :int
+                o.env 'DD_APPSEC_BODY_PARSING_SIZE_LIMIT' # bytes
+                o.default 10_485_760
+              end
+
               option :waf_debug do |o|
                 o.env 'DD_APPSEC_WAF_DEBUG'
                 o.default false

data/lib/datadog/appsec/contrib/aws_lambda/waf_addresses.rb

@@ -15,6 +15,10 @@ module Datadog
         # Extracts WAF input addresses from normalized AWS Lambda API Gateway event payloads.
         # @api private
         module WAFAddresses
+          BASE64_CHARS_PER_GROUP = 4
+          BASE64_BYTES_PER_GROUP = 3
+          BASE64_PADDING_BYTE = "=".ord
+
           module_function
 
           def from_request(payload)
@@ -27,10 +31,11 @@ module Datadog
               'server.request.uri.raw' => build_fullpath(payload),
               'server.request.headers' => headers,
               'server.request.headers.no_cookies' => headers.dup.tap { |h| h.delete('cookie') },
-              'http.client_ip' => extract_client_ip(payload['source_ip'], headers),
               'server.request.method' => payload['method'],
               'server.request.body' => parse_body(payload, headers),
-              'server.request.path_params' => payload['path_params']
+              'server.request.body.byte_length' => body_byte_length(payload),
+              'server.request.path_params' => payload['path_params'],
+              'http.client_ip' => extract_client_ip(payload['source_ip'], headers)
             }
 
             data.compact!
@@ -44,7 +49,9 @@ module Datadog
             data = {
               'server.response.status' => payload['status_code']&.to_s,
               'server.response.headers' => headers,
-              'server.response.headers.no_cookies' => headers.dup.tap { |h| h.delete('set-cookie') }
+              'server.response.headers.no_cookies' => headers.dup.tap { |h| h.delete('set-cookie') },
+              'server.response.body' => parse_body(payload, headers),
+              'server.response.body.byte_length' => body_byte_length(payload)
             }
 
             data.compact!
@@ -94,7 +101,9 @@ module Datadog
             body = payload['body']
             return unless body
 
-            body = Core::Utils::Base64Codec.strict_decode64(body) if payload['base64_encoded']
+            if (byte_length = body_byte_length(payload))
+              return if byte_length > Datadog.configuration.appsec.body_parsing_size_limit
+            end
 
             content_type = headers['content-type']
             return unless content_type
@@ -102,7 +111,31 @@ module Datadog
             media_type = AppSec::Utils::HTTP::MediaType.parse(content_type)
             return unless media_type
 
+            body = Core::Utils::Base64Codec.strict_decode64(body) if payload['base64_encoded']
             AppSec::Utils::HTTP::Body.parse(body, media_type: media_type)
+          rescue ArgumentError => e
+            AppSec.telemetry.report(e, description: 'AppSec: Failed to decode base64 body')
+
+            nil
+          end
+
+          def body_byte_length(payload)
+            body = payload['body']
+
+            return unless body
+            return body.bytesize unless payload['base64_encoded']
+
+            # NOTE: Base64 packs every 3 bytes into 4 characters and pads the last
+            #       group with up to two "=" bytes. The decoded length is therefore
+            #       derivable from the encoded length, letting us measure the raw
+            #       body size without allocating the decoded string.
+            padding = 0
+            if body.getbyte(-1) == BASE64_PADDING_BYTE
+              padding = 1
+              padding = 2 if body.getbyte(-2) == BASE64_PADDING_BYTE
+            end
+
+            body.bytesize / BASE64_CHARS_PER_GROUP * BASE64_BYTES_PER_GROUP - padding
           end
         end
       end