From 04ef96d18abbe352ac793a0aa3a4e739cdd25278 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Wed, 18 Mar 2026 17:44:28 +0100 Subject: [PATCH 01/24] EXPERIMENTAL --- lib/datadog/lambda/trace/listener.rb | 2 +- lib/datadog/lambda/utils/extension.rb | 21 +++++++++--- test/datadog/lambda/utils/extension.spec.rb | 37 +++++++++++++++++++++ 3 files changed, 55 insertions(+), 5 deletions(-) diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index e4ac564..f76848a 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -54,7 +54,7 @@ def on_start(event:, request_context:, cold_start:) # rubocop:enable Metrics/AbcSize def on_end(response:, request_context:) - Datadog::Utils.send_end_invocation_request(response:, span_id: @trace.id, request_context:) + Datadog::Utils.send_end_invocation_request(response:, span_id: @trace.id, request_context:, span: @trace) @trace&.finish end diff --git a/lib/datadog/lambda/utils/extension.rb b/lib/datadog/lambda/utils/extension.rb index efdf03c..9bfc3b6 100644 --- a/lib/datadog/lambda/utils/extension.rb +++ b/lib/datadog/lambda/utils/extension.rb @@ -57,8 +57,11 @@ def self.send_start_invocation_request(event:, request_context:) Datadog::Utils.logger.debug "failed on start invocation request to extension: #{e}" end + DD_APPSEC_ENABLED_HEADER = 'x-datadog-appsec-enabled' + DD_APPSEC_JSON_HEADER = 'x-datadog-appsec-json' + # rubocop:disable Metrics/AbcSize - def self.send_end_invocation_request(response:, span_id:, request_context:) + def self.send_end_invocation_request(response:, span_id:, request_context:, span: nil) return unless extension_running? request = Net::HTTP::Post.new(END_INVOCATION_URI) @@ -68,14 +71,14 @@ def self.send_end_invocation_request(response:, span_id:, request_context:) trace_digest = Datadog::Tracing.active_trace&.to_digest PROPAGATOR.inject!(trace_digest, request) - # Propagator doesn't inject span_id, so we do it manually - # It is needed for the extension to take this span id request[DD_SPAN_ID_HEADER] = span_id.to_s request[LAMBDA_RUNTIME_AWS_REQUEST_HEADER_ID] = request_context.aws_request_id - # Remove Parent ID if it is the same as the Span ID request.delete(DD_PARENT_ID_HEADER) if request[DD_PARENT_ID_HEADER] == span_id.to_s + + inject_appsec_data(request, span) + Datadog::Utils.logger.debug "End invocation request headers: #{request.to_hash}" Net::HTTP.start(END_INVOCATION_URI.host, END_INVOCATION_URI.port) do |http| @@ -86,6 +89,16 @@ def self.send_end_invocation_request(response:, span_id:, request_context:) end # rubocop:enable Metrics/AbcSize + def self.inject_appsec_data(request, span) + return unless span + + appsec_enabled = span.get_metric('_dd.appsec.enabled') + request[DD_APPSEC_ENABLED_HEADER] = '1' if appsec_enabled + + appsec_json = span.get_tag('_dd.appsec.json') + request[DD_APPSEC_JSON_HEADER] = appsec_json if appsec_json + end + def self.request_headers { # Header used to avoid tracing requests that are internal to diff --git a/test/datadog/lambda/utils/extension.spec.rb b/test/datadog/lambda/utils/extension.spec.rb index eca4a4c..97f0989 100644 --- a/test/datadog/lambda/utils/extension.spec.rb +++ b/test/datadog/lambda/utils/extension.spec.rb @@ -102,6 +102,43 @@ # Call the start request with an empty event Datadog::Utils.send_end_invocation_request(response: nil, span_id: nil, request_context: ctx) end + + it 'forwards AppSec tags from span as headers' do + @trace.set_metric('_dd.appsec.enabled', 1.0) + @trace.set_tag('_dd.appsec.json', '{"triggers":[]}') + + captured_request = nil + http_double = instance_double(Net::HTTP) + allow(Net::HTTP).to receive(:start).and_yield(http_double) + allow(http_double).to receive(:request) do |req| + captured_request = req + Net::HTTPResponse.new('1.1', '200', 'OK') + end + + Datadog::Utils.send_end_invocation_request( + response: nil, span_id: nil, request_context: ctx, span: @trace + ) + + expect(captured_request['x-datadog-appsec-enabled']).to eq('1') + expect(captured_request['x-datadog-appsec-json']).to eq('{"triggers":[]}') + end + + it 'does not send AppSec headers when tags are absent' do + captured_request = nil + http_double = instance_double(Net::HTTP) + allow(Net::HTTP).to receive(:start).and_yield(http_double) + allow(http_double).to receive(:request) do |req| + captured_request = req + Net::HTTPResponse.new('1.1', '200', 'OK') + end + + Datadog::Utils.send_end_invocation_request( + response: nil, span_id: nil, request_context: ctx, span: @trace + ) + + expect(captured_request['x-datadog-appsec-enabled']).to be_nil + expect(captured_request['x-datadog-appsec-json']).to be_nil + end end context 'when extension is not running' do From 1da82aad06afb34550a01fad5229c9781ae7b548 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Thu, 19 Mar 2026 14:14:22 +0100 Subject: [PATCH 02/24] EXPERIMENTAL Add appsec to the request flow --- lib/datadog/lambda/trace/listener.rb | 61 ++++++++++++++++++++++++++- lib/datadog/lambda/utils/extension.rb | 6 +++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index f76848a..695ca19 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -22,6 +22,8 @@ def initialize(handler_name:, function_name:, patch_http:, @handler_name = handler_name @function_name = function_name @merge_xray_traces = merge_xray_traces + @appsec_context = nil + @gateway_request = nil Datadog::Trace.patch_http if patch_http end @@ -43,23 +45,78 @@ def on_start(event:, request_context:, cold_start:) options[:type] = 'serverless' trace_digest = Datadog::Utils.send_start_invocation_request(event:, request_context:) - # Only continue trace from a new one if it exist, or else, - # it will create a new trace, which is not ideal here. options[:continue_from] = trace_digest if trace_digest @trace = Datadog::Tracing.trace('aws.lambda', **options) Datadog::Trace.apply_datadog_trace_context(Datadog::Trace.trace_context) + + start_appsec(event) end # rubocop:enable Metrics/AbcSize def on_end(response:, request_context:) + finish_appsec(response) Datadog::Utils.send_end_invocation_request(response:, span_id: @trace.id, request_context:, span: @trace) @trace&.finish end private + def start_appsec(event) + return unless appsec_enabled? + + ensure_appsec_patched + + security_engine = Datadog::AppSec.security_engine + return unless security_engine + + active_trace = Datadog::Tracing.active_trace + @appsec_context = Datadog::AppSec::Context.activate( + Datadog::AppSec::Context.new(active_trace, @trace, security_engine.new_runner) + ) + + @trace.set_metric(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) + + @gateway_request = Datadog::AppSec::Contrib::AwsLambda::Gateway::Request.new(event) + Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.request.start', @gateway_request) + rescue StandardError => e + Datadog::Utils.logger.debug "failed to start AppSec context: #{e}" + end + + def finish_appsec(response) + return unless @appsec_context + + gateway_response = Datadog::AppSec::Contrib::AwsLambda::Gateway::Response.new( + response, context: @appsec_context + ) + Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.response.start', gateway_response) + + Datadog::AppSec::Event.record(@appsec_context, request: @gateway_request) + + @appsec_context.export_metrics + @appsec_context.export_request_telemetry + rescue StandardError => e + Datadog::Utils.logger.debug "failed to finish AppSec context: #{e}" + ensure + Datadog::AppSec::Context.deactivate + @appsec_context = nil + @gateway_request = nil + end + + def appsec_enabled? + defined?(Datadog::AppSec) && + Datadog::AppSec.respond_to?(:enabled?) && + Datadog::AppSec.enabled? + end + + def ensure_appsec_patched + return if @appsec_patched + + Datadog.configuration.appsec.instrument(:aws_lambda) + @appsec_patched = true + end + def get_option_tags(request_context:, cold_start:) function_arn = request_context.invoked_function_arn.to_s.downcase tk = function_arn.split(':') diff --git a/lib/datadog/lambda/utils/extension.rb b/lib/datadog/lambda/utils/extension.rb index 9bfc3b6..6db5226 100644 --- a/lib/datadog/lambda/utils/extension.rb +++ b/lib/datadog/lambda/utils/extension.rb @@ -92,11 +92,17 @@ def self.send_end_invocation_request(response:, span_id:, request_context:, span def self.inject_appsec_data(request, span) return unless span + Datadog::Utils.logger.debug "[DEBUG:inject_appsec] span=#{span.name} span_id=#{span.id}" + Datadog::Utils.logger.debug "[DEBUG:inject_appsec] all meta keys: #{span.meta.keys rescue 'N/A'}" + Datadog::Utils.logger.debug "[DEBUG:inject_appsec] all metrics keys: #{span.metrics.keys rescue 'N/A'}" + appsec_enabled = span.get_metric('_dd.appsec.enabled') request[DD_APPSEC_ENABLED_HEADER] = '1' if appsec_enabled appsec_json = span.get_tag('_dd.appsec.json') request[DD_APPSEC_JSON_HEADER] = appsec_json if appsec_json + + Datadog::Utils.logger.debug "[DEBUG:inject_appsec] appsec_enabled=#{appsec_enabled.inspect} appsec_json=#{!appsec_json.nil?}" end def self.request_headers From 0b6c3da0e06f38f16e21cdafafe155e2b677fbd5 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Thu, 19 Mar 2026 15:26:14 +0100 Subject: [PATCH 03/24] EXPERIMENTAL Add inferred span --- lib/datadog/lambda/trace/listener.rb | 126 ++++++++++++++++++++++++++- 1 file changed, 125 insertions(+), 1 deletion(-) diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index 695ca19..8210f22 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -24,6 +24,7 @@ def initialize(handler_name:, function_name:, patch_http:, @merge_xray_traces = merge_xray_traces @appsec_context = nil @gateway_request = nil + @inferred_span = nil Datadog::Trace.patch_http if patch_http end @@ -45,7 +46,9 @@ def on_start(event:, request_context:, cold_start:) options[:type] = 'serverless' trace_digest = Datadog::Utils.send_start_invocation_request(event:, request_context:) - options[:continue_from] = trace_digest if trace_digest + + @inferred_span = create_inferred_span(event, request_context, trace_digest) + options[:continue_from] = trace_digest if trace_digest && !@inferred_span @trace = Datadog::Tracing.trace('aws.lambda', **options) @@ -59,10 +62,131 @@ def on_end(response:, request_context:) finish_appsec(response) Datadog::Utils.send_end_invocation_request(response:, span_id: @trace.id, request_context:, span: @trace) @trace&.finish + finish_inferred_span(response) end private + def create_inferred_span(event, request_context, trace_digest) + return unless managed_services_enabled? + + span_name = inferred_span_name(event) + return unless span_name + + rc = event['requestContext'] || {} + domain = rc['domainName'] || '' + api_id = rc['apiId'] || '' + stage = rc['stage'] || '' + + if span_name == 'aws.apigateway' + method = event['httpMethod'] + path = event['path'] || '/' + resource_path = rc['resourcePath'] || path + request_time_ms = rc['requestTimeEpoch'] + user_agent = rc.dig('identity', 'userAgent') + else + http = rc['http'] || {} + method = http['method'] + path = event['rawPath'] || '/' + resource_path = event['routeKey']&.sub(/^[A-Z]+ /, '') || path + request_time_ms = rc['timeEpoch'] + user_agent = http['userAgent'] + end + + resource = "#{method} #{resource_path}" + http_url = domain.empty? ? path : "https://#{domain}#{path}" + + tags = { + 'http.method' => method, + 'http.url' => http_url, + 'http.route' => resource_path, + 'endpoint' => path, + 'resource_names' => resource, + 'span.kind' => 'server', + 'apiid' => api_id, + 'apiname' => api_id, + 'stage' => stage, + 'request_id' => request_context.aws_request_id, + '_inferred_span.synchronicity' => 'sync', + '_inferred_span.tag_source' => 'self', + } + + arn = request_context.invoked_function_arn.to_s + region = arn.split(':')[3] if arn.include?(':') + if region && !api_id.empty? && !stage.empty? + arn_path = span_name == 'aws.apigateway' ? 'restapis' : 'apis' + tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{arn_path}/#{api_id}/stages/#{stage}" + end + + tags['http.useragent'] = user_agent if user_agent + + inferred_options = { + service: domain.empty? ? nil : domain, + resource: resource, + type: 'web', + tags: tags, + } + inferred_options[:continue_from] = trace_digest if trace_digest + inferred_options[:start_time] = Time.at(request_time_ms / 1000.0) if request_time_ms + + span = Datadog::Tracing.trace(span_name, **inferred_options) + span.set_metric('_dd._inferred_span', 1.0) + span + rescue StandardError => e + Datadog::Utils.logger.debug "failed to create inferred span: #{e}" + nil + end + + def finish_inferred_span(response) + return unless @inferred_span + + if @trace + appsec_enabled = @trace.get_metric('_dd.appsec.enabled') + @inferred_span.set_metric('_dd.appsec.enabled', appsec_enabled) if appsec_enabled + + appsec_json = @trace.get_tag('_dd.appsec.json') + @inferred_span.set_tag('_dd.appsec.json', appsec_json) if appsec_json + + appsec_event = @trace.get_tag('appsec.event') + @inferred_span.set_tag('appsec.event', appsec_event) if appsec_event + + origin = @trace.get_tag('_dd.origin') + @inferred_span.set_tag('_dd.origin', origin) if origin + end + + status_code = extract_status_code(response) + @inferred_span.set_tag('http.status_code', status_code.to_s) if status_code + + @inferred_span.finish + rescue StandardError => e + Datadog::Utils.logger.debug "failed to finish inferred span: #{e}" + ensure + @inferred_span = nil + end + + def inferred_span_name(event) + return unless event.is_a?(Hash) + + rc = event['requestContext'] + return unless rc.is_a?(Hash) && rc['stage'] + + if event['httpMethod'] + 'aws.apigateway' + elsif event['routeKey'] + 'aws.httpapi' + end + end + + def managed_services_enabled? + ENV.fetch('DD_TRACE_MANAGED_SERVICES', 'true').downcase != 'false' + end + + def extract_status_code(response) + return unless response.is_a?(Hash) + + response['statusCode'] + end + def start_appsec(event) return unless appsec_enabled? From 0091b674d488511ea4ee23b5cd75e22e8d9aa336 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Tue, 24 Mar 2026 20:00:23 +0100 Subject: [PATCH 04/24] Simplify AppSec to pure gateway push with raw PORO data MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove all AppSec lifecycle management (context creation, Event.record, export_metrics, Context.deactivate) — now owned by dd-trace-rb watcher. Listener becomes a dumb data provider: just pushes raw event/response hashes. Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/datadog/lambda/trace/listener.rb | 35 ++++------------------------ 1 file changed, 5 insertions(+), 30 deletions(-) diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index 8210f22..d008ff5 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -22,8 +22,6 @@ def initialize(handler_name:, function_name:, patch_http:, @handler_name = handler_name @function_name = function_name @merge_xray_traces = merge_xray_traces - @appsec_context = nil - @gateway_request = nil @inferred_span = nil Datadog::Trace.patch_http if patch_http @@ -192,40 +190,17 @@ def start_appsec(event) ensure_appsec_patched - security_engine = Datadog::AppSec.security_engine - return unless security_engine - - active_trace = Datadog::Tracing.active_trace - @appsec_context = Datadog::AppSec::Context.activate( - Datadog::AppSec::Context.new(active_trace, @trace, security_engine.new_runner) - ) - - @trace.set_metric(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) - - @gateway_request = Datadog::AppSec::Contrib::AwsLambda::Gateway::Request.new(event) - Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.request.start', @gateway_request) + Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.request.start', event) rescue StandardError => e - Datadog::Utils.logger.debug "failed to start AppSec context: #{e}" + Datadog::Utils.logger.debug "failed to start AppSec: #{e}" end def finish_appsec(response) - return unless @appsec_context + return unless Datadog::AppSec::Context.active - gateway_response = Datadog::AppSec::Contrib::AwsLambda::Gateway::Response.new( - response, context: @appsec_context - ) - Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.response.start', gateway_response) - - Datadog::AppSec::Event.record(@appsec_context, request: @gateway_request) - - @appsec_context.export_metrics - @appsec_context.export_request_telemetry + Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.response.start', response) rescue StandardError => e - Datadog::Utils.logger.debug "failed to finish AppSec context: #{e}" - ensure - Datadog::AppSec::Context.deactivate - @appsec_context = nil - @gateway_request = nil + Datadog::Utils.logger.debug "failed to finish AppSec: #{e}" end def appsec_enabled? From 40548000541656fe7d578e27dbccae59cf4a4848 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Wed, 25 Mar 2026 14:14:58 +0100 Subject: [PATCH 05/24] Extract AppSec and InferredSpan modules from Listener MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lambda::AppSec owns full context lifecycle (create, activate, record, export, deactivate). Lambda::Trace::InferredSpan owns span creation and finishing. Listener becomes thin orchestrator. AppSec context now targets the inferred span (service-entry), so tags land there directly — tag propagation hack removed. Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/datadog/lambda.rb | 4 +- lib/datadog/lambda/appsec.rb | 60 ++++++++ lib/datadog/lambda/trace/inferred_span.rb | 117 ++++++++++++++++ lib/datadog/lambda/trace/listener.rb | 162 +--------------------- 4 files changed, 186 insertions(+), 157 deletions(-) create mode 100644 lib/datadog/lambda/appsec.rb create mode 100644 lib/datadog/lambda/trace/inferred_span.rb diff --git a/lib/datadog/lambda.rb b/lib/datadog/lambda.rb index 4eee6d4..261175c 100644 --- a/lib/datadog/lambda.rb +++ b/lib/datadog/lambda.rb @@ -169,7 +169,7 @@ def self.do_enhanced_metrics? # Read DD_TRACE_MANAGED_SERVICES environment variable # @return [boolean] true if we should trace AWS services def self.trace_managed_services? - dd_trace_managed_services = ENV[Trace::DD_TRACE_MANAGED_SERVICES] + dd_trace_managed_services = ENV[::Datadog::Trace::DD_TRACE_MANAGED_SERVICES] return true if dd_trace_managed_services.nil? dd_trace_managed_services.downcase == 'true' @@ -191,7 +191,7 @@ def self.initialize_listener return nil end - Trace::Listener.new( + ::Datadog::Trace::Listener.new( handler_name: handler, function_name: function, patch_http: @patch_http, diff --git a/lib/datadog/lambda/appsec.rb b/lib/datadog/lambda/appsec.rb new file mode 100644 index 0000000..96548ef --- /dev/null +++ b/lib/datadog/lambda/appsec.rb @@ -0,0 +1,60 @@ +# frozen_string_literal: true + +module Datadog + module Lambda + module AppSec + class << self + def on_start(event, trace, span) + return unless enabled? + + Datadog.configuration.appsec.instrument(:aws_lambda) + + trace ||= Datadog::Tracing.active_trace + span ||= Datadog::Tracing.active_span + + create_context(trace, span) + + Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.request.start', event) + rescue StandardError => e + Datadog::Utils.logger.debug "failed to start AppSec: #{e}" + end + + def on_finish(response) + context = Datadog::AppSec::Context.active + return unless context + + Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.response.start', response) + + Datadog::AppSec::Event.record(context, request: context.state[:request]) + + context.export_metrics + context.export_request_telemetry + rescue StandardError => e + Datadog::Utils.logger.debug "failed to finish AppSec: #{e}" + ensure + Datadog::AppSec::Context.deactivate + end + + private + + def enabled? + defined?(Datadog::AppSec) && + Datadog::AppSec.respond_to?(:enabled?) && + Datadog::AppSec.enabled? + end + + def create_context(trace, span) + security_engine = Datadog::AppSec.security_engine + return unless security_engine + return if trace.nil? || span.nil? + + Datadog::AppSec::Context.activate( + Datadog::AppSec::Context.new(trace, span, security_engine.new_runner) + ) + + span.set_metric(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) + end + end + end + end +end diff --git a/lib/datadog/lambda/trace/inferred_span.rb b/lib/datadog/lambda/trace/inferred_span.rb new file mode 100644 index 0000000..75fa77c --- /dev/null +++ b/lib/datadog/lambda/trace/inferred_span.rb @@ -0,0 +1,117 @@ +# frozen_string_literal: true + +module Datadog + module Lambda + module Trace + module InferredSpan + class << self + def create(event, request_context, trace_digest) + return unless managed_services_enabled? + + span_name = inferred_span_name(event) + return unless span_name + + rc = event['requestContext'] || {} + domain = rc['domainName'] || '' + api_id = rc['apiId'] || '' + stage = rc['stage'] || '' + + if span_name == 'aws.apigateway' + method = event['httpMethod'] + path = event['path'] || '/' + resource_path = rc['resourcePath'] || path + request_time_ms = rc['requestTimeEpoch'] + user_agent = rc.dig('identity', 'userAgent') + else + http = rc['http'] || {} + method = http['method'] + path = event['rawPath'] || '/' + resource_path = event['routeKey']&.sub(/^[A-Z]+ /, '') || path + request_time_ms = rc['timeEpoch'] + user_agent = http['userAgent'] + end + + resource = "#{method} #{resource_path}" + http_url = domain.empty? ? path : "https://#{domain}#{path}" + + tags = { + 'http.method' => method, + 'http.url' => http_url, + 'http.route' => resource_path, + 'endpoint' => path, + 'resource_names' => resource, + 'span.kind' => 'server', + 'apiid' => api_id, + 'apiname' => api_id, + 'stage' => stage, + 'request_id' => request_context.aws_request_id, + '_inferred_span.synchronicity' => 'sync', + '_inferred_span.tag_source' => 'self', + } + + arn = request_context.invoked_function_arn.to_s + region = arn.split(':')[3] if arn.include?(':') + if region && !api_id.empty? && !stage.empty? + arn_path = span_name == 'aws.apigateway' ? 'restapis' : 'apis' + tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{arn_path}/#{api_id}/stages/#{stage}" + end + + tags['http.useragent'] = user_agent if user_agent + + inferred_options = { + service: domain.empty? ? nil : domain, + resource: resource, + type: 'web', + tags: tags, + } + inferred_options[:continue_from] = trace_digest if trace_digest + inferred_options[:start_time] = Time.at(request_time_ms / 1000.0) if request_time_ms + + span = Datadog::Tracing.trace(span_name, **inferred_options) + span.set_metric('_dd._inferred_span', 1.0) + span + rescue StandardError => e + Datadog::Utils.logger.debug "failed to create inferred span: #{e}" + nil + end + + def finish(inferred_span, response) + return unless inferred_span + + status_code = extract_status_code(response) + inferred_span.set_tag('http.status_code', status_code.to_s) if status_code + + inferred_span.finish + rescue StandardError => e + Datadog::Utils.logger.debug "failed to finish inferred span: #{e}" + end + + private + + def inferred_span_name(event) + return unless event.is_a?(Hash) + + rc = event['requestContext'] + return unless rc.is_a?(Hash) && rc['stage'] + + if event['httpMethod'] + 'aws.apigateway' + elsif event['routeKey'] + 'aws.httpapi' + end + end + + def managed_services_enabled? + ENV.fetch('DD_TRACE_MANAGED_SERVICES', 'true').downcase != 'false' + end + + def extract_status_code(response) + return unless response.is_a?(Hash) + + response['statusCode'] + end + end + end + end + end +end diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index d008ff5..d75d459 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -11,6 +11,8 @@ require 'datadog/lambda/trace/context' require 'datadog/lambda/trace/patch_http' require 'datadog/lambda/trace/ddtrace' +require 'datadog/lambda/trace/inferred_span' +require 'datadog/lambda/appsec' module Datadog module Trace @@ -45,176 +47,26 @@ def on_start(event:, request_context:, cold_start:) trace_digest = Datadog::Utils.send_start_invocation_request(event:, request_context:) - @inferred_span = create_inferred_span(event, request_context, trace_digest) + @inferred_span = Datadog::Lambda::Trace::InferredSpan.create(event, request_context, trace_digest) options[:continue_from] = trace_digest if trace_digest && !@inferred_span @trace = Datadog::Tracing.trace('aws.lambda', **options) Datadog::Trace.apply_datadog_trace_context(Datadog::Trace.trace_context) - start_appsec(event) + Datadog::Lambda::AppSec.on_start(event, @trace, @inferred_span) end # rubocop:enable Metrics/AbcSize def on_end(response:, request_context:) - finish_appsec(response) + Datadog::Lambda::AppSec.on_finish(response) Datadog::Utils.send_end_invocation_request(response:, span_id: @trace.id, request_context:, span: @trace) @trace&.finish - finish_inferred_span(response) - end - - private - - def create_inferred_span(event, request_context, trace_digest) - return unless managed_services_enabled? - - span_name = inferred_span_name(event) - return unless span_name - - rc = event['requestContext'] || {} - domain = rc['domainName'] || '' - api_id = rc['apiId'] || '' - stage = rc['stage'] || '' - - if span_name == 'aws.apigateway' - method = event['httpMethod'] - path = event['path'] || '/' - resource_path = rc['resourcePath'] || path - request_time_ms = rc['requestTimeEpoch'] - user_agent = rc.dig('identity', 'userAgent') - else - http = rc['http'] || {} - method = http['method'] - path = event['rawPath'] || '/' - resource_path = event['routeKey']&.sub(/^[A-Z]+ /, '') || path - request_time_ms = rc['timeEpoch'] - user_agent = http['userAgent'] - end - - resource = "#{method} #{resource_path}" - http_url = domain.empty? ? path : "https://#{domain}#{path}" - - tags = { - 'http.method' => method, - 'http.url' => http_url, - 'http.route' => resource_path, - 'endpoint' => path, - 'resource_names' => resource, - 'span.kind' => 'server', - 'apiid' => api_id, - 'apiname' => api_id, - 'stage' => stage, - 'request_id' => request_context.aws_request_id, - '_inferred_span.synchronicity' => 'sync', - '_inferred_span.tag_source' => 'self', - } - - arn = request_context.invoked_function_arn.to_s - region = arn.split(':')[3] if arn.include?(':') - if region && !api_id.empty? && !stage.empty? - arn_path = span_name == 'aws.apigateway' ? 'restapis' : 'apis' - tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{arn_path}/#{api_id}/stages/#{stage}" - end - - tags['http.useragent'] = user_agent if user_agent - - inferred_options = { - service: domain.empty? ? nil : domain, - resource: resource, - type: 'web', - tags: tags, - } - inferred_options[:continue_from] = trace_digest if trace_digest - inferred_options[:start_time] = Time.at(request_time_ms / 1000.0) if request_time_ms - - span = Datadog::Tracing.trace(span_name, **inferred_options) - span.set_metric('_dd._inferred_span', 1.0) - span - rescue StandardError => e - Datadog::Utils.logger.debug "failed to create inferred span: #{e}" - nil - end - - def finish_inferred_span(response) - return unless @inferred_span - - if @trace - appsec_enabled = @trace.get_metric('_dd.appsec.enabled') - @inferred_span.set_metric('_dd.appsec.enabled', appsec_enabled) if appsec_enabled - - appsec_json = @trace.get_tag('_dd.appsec.json') - @inferred_span.set_tag('_dd.appsec.json', appsec_json) if appsec_json - - appsec_event = @trace.get_tag('appsec.event') - @inferred_span.set_tag('appsec.event', appsec_event) if appsec_event - - origin = @trace.get_tag('_dd.origin') - @inferred_span.set_tag('_dd.origin', origin) if origin - end - - status_code = extract_status_code(response) - @inferred_span.set_tag('http.status_code', status_code.to_s) if status_code - - @inferred_span.finish - rescue StandardError => e - Datadog::Utils.logger.debug "failed to finish inferred span: #{e}" - ensure + Datadog::Lambda::Trace::InferredSpan.finish(@inferred_span, response) @inferred_span = nil end - def inferred_span_name(event) - return unless event.is_a?(Hash) - - rc = event['requestContext'] - return unless rc.is_a?(Hash) && rc['stage'] - - if event['httpMethod'] - 'aws.apigateway' - elsif event['routeKey'] - 'aws.httpapi' - end - end - - def managed_services_enabled? - ENV.fetch('DD_TRACE_MANAGED_SERVICES', 'true').downcase != 'false' - end - - def extract_status_code(response) - return unless response.is_a?(Hash) - - response['statusCode'] - end - - def start_appsec(event) - return unless appsec_enabled? - - ensure_appsec_patched - - Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.request.start', event) - rescue StandardError => e - Datadog::Utils.logger.debug "failed to start AppSec: #{e}" - end - - def finish_appsec(response) - return unless Datadog::AppSec::Context.active - - Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.response.start', response) - rescue StandardError => e - Datadog::Utils.logger.debug "failed to finish AppSec: #{e}" - end - - def appsec_enabled? - defined?(Datadog::AppSec) && - Datadog::AppSec.respond_to?(:enabled?) && - Datadog::AppSec.enabled? - end - - def ensure_appsec_patched - return if @appsec_patched - - Datadog.configuration.appsec.instrument(:aws_lambda) - @appsec_patched = true - end + private def get_option_tags(request_context:, cold_start:) function_arn = request_context.invoked_function_arn.to_s.downcase From 2e62588f7545058f06891c730201dc64de9974c9 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Wed, 25 Mar 2026 15:04:51 +0100 Subject: [PATCH 06/24] Remove debug logging and revert accidental changes - Remove [DEBUG:inject_appsec] debug logging from inject_appsec_data - Restore removed comments in send_end_invocation_request - Revert unnecessary ::Datadog:: namespace prefixes in lambda.rb Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/datadog/lambda.rb | 4 ++-- lib/datadog/lambda/trace/listener.rb | 5 +++-- lib/datadog/lambda/utils/extension.rb | 9 +++------ 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/lib/datadog/lambda.rb b/lib/datadog/lambda.rb index 261175c..4eee6d4 100644 --- a/lib/datadog/lambda.rb +++ b/lib/datadog/lambda.rb @@ -169,7 +169,7 @@ def self.do_enhanced_metrics? # Read DD_TRACE_MANAGED_SERVICES environment variable # @return [boolean] true if we should trace AWS services def self.trace_managed_services? - dd_trace_managed_services = ENV[::Datadog::Trace::DD_TRACE_MANAGED_SERVICES] + dd_trace_managed_services = ENV[Trace::DD_TRACE_MANAGED_SERVICES] return true if dd_trace_managed_services.nil? dd_trace_managed_services.downcase == 'true' @@ -191,7 +191,7 @@ def self.initialize_listener return nil end - ::Datadog::Trace::Listener.new( + Trace::Listener.new( handler_name: handler, function_name: function, patch_http: @patch_http, diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index d75d459..6798ddc 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -48,12 +48,11 @@ def on_start(event:, request_context:, cold_start:) trace_digest = Datadog::Utils.send_start_invocation_request(event:, request_context:) @inferred_span = Datadog::Lambda::Trace::InferredSpan.create(event, request_context, trace_digest) - options[:continue_from] = trace_digest if trace_digest && !@inferred_span + options[:continue_from] = trace_digest if trace_digest && @inferred_span.nil? @trace = Datadog::Tracing.trace('aws.lambda', **options) Datadog::Trace.apply_datadog_trace_context(Datadog::Trace.trace_context) - Datadog::Lambda::AppSec.on_start(event, @trace, @inferred_span) end # rubocop:enable Metrics/AbcSize @@ -61,8 +60,10 @@ def on_start(event:, request_context:, cold_start:) def on_end(response:, request_context:) Datadog::Lambda::AppSec.on_finish(response) Datadog::Utils.send_end_invocation_request(response:, span_id: @trace.id, request_context:, span: @trace) + @trace&.finish Datadog::Lambda::Trace::InferredSpan.finish(@inferred_span, response) + @inferred_span = nil end diff --git a/lib/datadog/lambda/utils/extension.rb b/lib/datadog/lambda/utils/extension.rb index 6db5226..66315ea 100644 --- a/lib/datadog/lambda/utils/extension.rb +++ b/lib/datadog/lambda/utils/extension.rb @@ -71,10 +71,13 @@ def self.send_end_invocation_request(response:, span_id:, request_context:, span trace_digest = Datadog::Tracing.active_trace&.to_digest PROPAGATOR.inject!(trace_digest, request) + # Propagator doesn't inject span_id, so we do it manually + # It is needed for the extension to take this span id request[DD_SPAN_ID_HEADER] = span_id.to_s request[LAMBDA_RUNTIME_AWS_REQUEST_HEADER_ID] = request_context.aws_request_id + # Remove Parent ID if it is the same as the Span ID request.delete(DD_PARENT_ID_HEADER) if request[DD_PARENT_ID_HEADER] == span_id.to_s inject_appsec_data(request, span) @@ -92,17 +95,11 @@ def self.send_end_invocation_request(response:, span_id:, request_context:, span def self.inject_appsec_data(request, span) return unless span - Datadog::Utils.logger.debug "[DEBUG:inject_appsec] span=#{span.name} span_id=#{span.id}" - Datadog::Utils.logger.debug "[DEBUG:inject_appsec] all meta keys: #{span.meta.keys rescue 'N/A'}" - Datadog::Utils.logger.debug "[DEBUG:inject_appsec] all metrics keys: #{span.metrics.keys rescue 'N/A'}" - appsec_enabled = span.get_metric('_dd.appsec.enabled') request[DD_APPSEC_ENABLED_HEADER] = '1' if appsec_enabled appsec_json = span.get_tag('_dd.appsec.json') request[DD_APPSEC_JSON_HEADER] = appsec_json if appsec_json - - Datadog::Utils.logger.debug "[DEBUG:inject_appsec] appsec_enabled=#{appsec_enabled.inspect} appsec_json=#{!appsec_json.nil?}" end def self.request_headers From 416bd754c4738e563f4751fa9034753e36d7fd04 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Wed, 25 Mar 2026 15:35:27 +0100 Subject: [PATCH 07/24] Move AppSec instrumentation to configure_apm, consolidate managed_services check F29: Move `instrument(:aws_lambda)` from per-invocation `AppSec.on_start` to `configure_apm` after yield block. The `instrument` method is a no-op when appsec is disabled (checks `enabled` internally), so no guard needed. Matches Python's pattern of one-time initialization. F30: Reuse `Datadog::Lambda.trace_managed_services?` in InferredSpan instead of duplicating the ENV reading logic. Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/datadog/lambda.rb | 4 +++- lib/datadog/lambda/appsec.rb | 2 -- lib/datadog/lambda/trace/inferred_span.rb | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/lib/datadog/lambda.rb b/lib/datadog/lambda.rb index 4eee6d4..3e6bec8 100644 --- a/lib/datadog/lambda.rb +++ b/lib/datadog/lambda.rb @@ -45,9 +45,11 @@ def self.configure_apm end c.tags = { "_dd.origin": 'lambda' } # Enable AWS SDK instrumentation - c.tracing.instrument :aws if trace_managed_services? + c.tracing.instrument(:aws) if trace_managed_services? yield(c) if block_given? + + c.appsec.instrument(:aws_lambda) end end diff --git a/lib/datadog/lambda/appsec.rb b/lib/datadog/lambda/appsec.rb index 96548ef..cc0d6e1 100644 --- a/lib/datadog/lambda/appsec.rb +++ b/lib/datadog/lambda/appsec.rb @@ -7,8 +7,6 @@ class << self def on_start(event, trace, span) return unless enabled? - Datadog.configuration.appsec.instrument(:aws_lambda) - trace ||= Datadog::Tracing.active_trace span ||= Datadog::Tracing.active_span diff --git a/lib/datadog/lambda/trace/inferred_span.rb b/lib/datadog/lambda/trace/inferred_span.rb index 75fa77c..2195bc8 100644 --- a/lib/datadog/lambda/trace/inferred_span.rb +++ b/lib/datadog/lambda/trace/inferred_span.rb @@ -102,7 +102,7 @@ def inferred_span_name(event) end def managed_services_enabled? - ENV.fetch('DD_TRACE_MANAGED_SERVICES', 'true').downcase != 'false' + Datadog::Lambda.trace_managed_services? end def extract_status_code(response) From 8c09e5f76078443fc36dfcbc4e1cd37ebacdbaec Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Wed, 25 Mar 2026 15:37:57 +0100 Subject: [PATCH 08/24] Add specs for AppSec and InferredSpan modules F31: 44 new examples covering: - AppSec.on_start: enabled/disabled, context creation, gateway push, fallback to active trace/span, missing security engine, error handling - AppSec.on_finish: no active context, gateway push, Event.record, metrics/telemetry export, context deactivation, deactivation on error - InferredSpan.create: managed services disabled, non-hash events, missing requestContext/stage/httpMethod, v1 span name/tags/resource_key, v2 span name/tags/resource_key, trace_digest continuation, empty domain, error handling - InferredSpan.finish: nil span, statusCode tagging, non-hash response, nil response, error handling Co-Authored-By: Claude Opus 4.6 (1M context) --- test/datadog/lambda/appsec.spec.rb | 158 +++++++++++ .../lambda/trace/inferred_span.spec.rb | 247 ++++++++++++++++++ 2 files changed, 405 insertions(+) create mode 100644 test/datadog/lambda/appsec.spec.rb create mode 100644 test/datadog/lambda/trace/inferred_span.spec.rb diff --git a/test/datadog/lambda/appsec.spec.rb b/test/datadog/lambda/appsec.spec.rb new file mode 100644 index 0000000..a1e753f --- /dev/null +++ b/test/datadog/lambda/appsec.spec.rb @@ -0,0 +1,158 @@ +# frozen_string_literal: true + +require 'datadog/lambda' +require 'datadog/lambda/appsec' + +describe Datadog::Lambda::AppSec do + let(:event) { {'httpMethod' => 'GET', 'path' => '/'} } + let(:trace) { double('trace') } + let(:span) { double('span', set_metric: nil) } + let(:gateway) { double('gateway') } + let(:runner) { double('runner') } + let(:security_engine) { double('security_engine', new_runner: runner) } + let(:context) do + double( + 'context', + state: {}, + export_metrics: nil, + export_request_telemetry: nil, + events: [], + ) + end + + before do + allow(Datadog::AppSec::Instrumentation).to receive(:gateway).and_return(gateway) + allow(gateway).to receive(:push) + end + + describe '.on_start' do + context 'when appsec is disabled' do + before { allow(Datadog::AppSec).to receive(:enabled?).and_return(false) } + + it { expect(described_class.on_start(event, trace, span)).to be_nil } + + it 'does not push gateway events' do + described_class.on_start(event, trace, span) + expect(gateway).not_to have_received(:push) + end + end + + context 'when appsec is enabled' do + before do + allow(Datadog::AppSec).to receive(:enabled?).and_return(true) + allow(Datadog::AppSec).to receive(:security_engine).and_return(security_engine) + allow(Datadog::AppSec::Context).to receive(:activate) + allow(Datadog::AppSec::Context).to receive(:new).and_return(context) + end + + it 'creates and activates an AppSec context' do + described_class.on_start(event, trace, span) + expect(Datadog::AppSec::Context).to have_received(:new).with(trace, span, runner) + expect(Datadog::AppSec::Context).to have_received(:activate).with(context) + end + + it 'sets _dd.appsec.enabled metric on span' do + described_class.on_start(event, trace, span) + expect(span).to have_received(:set_metric).with(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) + end + + it 'pushes the request event to the gateway' do + described_class.on_start(event, trace, span) + expect(gateway).to have_received(:push).with('aws_lambda.request.start', event) + end + + context 'when trace and span are nil' do + let(:active_trace) { double('active_trace') } + let(:active_span) { double('active_span', set_metric: nil) } + + before do + allow(Datadog::Tracing).to receive(:active_trace).and_return(active_trace) + allow(Datadog::Tracing).to receive(:active_span).and_return(active_span) + end + + it 'falls back to active trace and span' do + described_class.on_start(event, nil, nil) + expect(Datadog::AppSec::Context).to have_received(:new).with(active_trace, active_span, runner) + end + end + + context 'when security_engine is nil' do + before { allow(Datadog::AppSec).to receive(:security_engine).and_return(nil) } + + it 'does not activate a context' do + described_class.on_start(event, trace, span) + expect(Datadog::AppSec::Context).not_to have_received(:activate) + end + + it 'still pushes the gateway event' do + described_class.on_start(event, trace, span) + expect(gateway).to have_received(:push).with('aws_lambda.request.start', event) + end + end + + context 'when an error occurs' do + before { allow(Datadog::AppSec::Context).to receive(:new).and_raise(StandardError, 'boom') } + + it 'rescues and logs' do + expect { described_class.on_start(event, trace, span) }.not_to raise_error + end + end + end + end + + describe '.on_finish' do + let(:response) { {'statusCode' => 200} } + + context 'when no active context' do + before { allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) } + + it { expect(described_class.on_finish(response)).to be_nil } + + it 'does not push gateway events' do + described_class.on_finish(response) + expect(gateway).not_to have_received(:push) + end + end + + context 'when active context exists' do + before do + allow(Datadog::AppSec::Context).to receive(:active).and_return(context) + allow(Datadog::AppSec::Context).to receive(:deactivate) + allow(Datadog::AppSec::Event).to receive(:record) + end + + it 'pushes the response event to the gateway' do + described_class.on_finish(response) + expect(gateway).to have_received(:push).with('aws_lambda.response.start', response) + end + + it 'records events with the request from context state' do + context.state[:request] = double('request') + described_class.on_finish(response) + expect(Datadog::AppSec::Event).to have_received(:record).with( + context, request: context.state[:request] + ) + end + + it 'exports metrics and telemetry' do + described_class.on_finish(response) + expect(context).to have_received(:export_metrics) + expect(context).to have_received(:export_request_telemetry) + end + + it 'deactivates the context' do + described_class.on_finish(response) + expect(Datadog::AppSec::Context).to have_received(:deactivate) + end + + context 'when an error occurs' do + before { allow(gateway).to receive(:push).and_raise(StandardError, 'boom') } + + it 'still deactivates the context' do + described_class.on_finish(response) + expect(Datadog::AppSec::Context).to have_received(:deactivate) + end + end + end + end +end diff --git a/test/datadog/lambda/trace/inferred_span.spec.rb b/test/datadog/lambda/trace/inferred_span.spec.rb new file mode 100644 index 0000000..680ade7 --- /dev/null +++ b/test/datadog/lambda/trace/inferred_span.spec.rb @@ -0,0 +1,247 @@ +# frozen_string_literal: true + +require 'datadog/lambda' +require 'datadog/lambda/trace/inferred_span' +require_relative '../../lambdacontextversion' + +describe Datadog::Lambda::Trace::InferredSpan do + let(:request_context) { LambdaContextVersion.new } + let(:trace_digest) { nil } + + let(:apigw_v1_event) do + { + 'httpMethod' => 'GET', + 'path' => '/test', + 'requestContext' => { + 'stage' => 'prod', + 'domainName' => 'api.example.com', + 'apiId' => 'abc123', + 'resourcePath' => '/test', + 'requestTimeEpoch' => 1_700_000_000_000, + 'identity' => {'userAgent' => 'TestAgent/1.0', 'sourceIp' => '1.2.3.4'}, + }, + } + end + + let(:apigw_v2_event) do + { + 'rawPath' => '/test', + 'routeKey' => 'GET /test', + 'requestContext' => { + 'stage' => 'prod', + 'domainName' => 'api.example.com', + 'apiId' => 'xyz789', + 'timeEpoch' => 1_700_000_000_000, + 'http' => {'method' => 'GET', 'userAgent' => 'TestAgent/2.0'}, + }, + } + end + + before do + allow(Datadog::Lambda).to receive(:trace_managed_services?).and_return(true) + end + + describe '.create' do + context 'when managed services is disabled' do + before { allow(Datadog::Lambda).to receive(:trace_managed_services?).and_return(false) } + + it { expect(described_class.create(apigw_v1_event, request_context, trace_digest)).to be_nil } + end + + context 'when event is not a Hash' do + it { expect(described_class.create('not a hash', request_context, trace_digest)).to be_nil } + end + + context 'when event has no requestContext' do + it { expect(described_class.create({}, request_context, trace_digest)).to be_nil } + end + + context 'when event has requestContext without stage' do + let(:event) { {'requestContext' => {'apiId' => 'abc'}} } + + it { expect(described_class.create(event, request_context, trace_digest)).to be_nil } + end + + context 'when event has no httpMethod or routeKey' do + let(:event) { {'requestContext' => {'stage' => 'prod'}} } + + it { expect(described_class.create(event, request_context, trace_digest)).to be_nil } + end + + context 'with API Gateway v1 event' do + after { @span&.finish unless @span&.finished? } + + it 'creates an aws.apigateway span' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(@span).not_to be_nil + expect(@span.name).to eq('aws.apigateway') + end + + it 'sets service to domain name' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(@span.service).to eq('api.example.com') + end + + it 'sets resource to method + resource path' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(@span.resource).to eq('GET /test') + end + + it 'sets the inferred span metric' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(@span.get_metric('_dd._inferred_span')).to eq(1.0) + end + + it 'sets http tags' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(@span.get_tag('http.method')).to eq('GET') + expect(@span.get_tag('http.url')).to eq('https://api.example.com/test') + expect(@span.get_tag('http.route')).to eq('/test') + end + + it 'sets api gateway tags' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(@span.get_tag('apiid')).to eq('abc123') + expect(@span.get_tag('stage')).to eq('prod') + end + + it 'sets dd_resource_key with restapis prefix' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(@span.get_tag('dd_resource_key')).to eq( + 'arn:aws:apigateway:us-east-1::/restapis/abc123/stages/prod' + ) + end + + it 'sets user agent tag' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(@span.get_tag('http.useragent')).to eq('TestAgent/1.0') + end + + it 'sets start_time from requestTimeEpoch' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(@span.start_time).to eq(Time.at(1_700_000_000)) + end + + context 'when trace_digest is provided' do + let(:trace_digest) { double('trace_digest') } + let(:captured_kwargs) { {} } + + before do + allow(Datadog::Tracing).to receive(:trace).and_wrap_original do |original, *args, **kwargs| + captured_kwargs.merge!(kwargs) + original.call(*args, **kwargs.except(:continue_from)) + end + end + + it 'passes continue_from to trace' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(captured_kwargs[:continue_from]).to eq(trace_digest) + end + end + + context 'when domain is empty' do + before { apigw_v1_event['requestContext']['domainName'] = '' } + + it 'uses path as http.url' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(@span.get_tag('http.url')).to eq('/test') + end + + it 'sets service to nil' do + @span = described_class.create(apigw_v1_event, request_context, trace_digest) + expect(@span.service).not_to eq('') + end + end + end + + context 'with API Gateway v2 event' do + after { @span&.finish unless @span&.finished? } + + it 'creates an aws.httpapi span' do + @span = described_class.create(apigw_v2_event, request_context, trace_digest) + expect(@span).not_to be_nil + expect(@span.name).to eq('aws.httpapi') + end + + it 'extracts method from http context' do + @span = described_class.create(apigw_v2_event, request_context, trace_digest) + expect(@span.get_tag('http.method')).to eq('GET') + end + + it 'extracts resource path from routeKey' do + @span = described_class.create(apigw_v2_event, request_context, trace_digest) + expect(@span.get_tag('http.route')).to eq('/test') + expect(@span.resource).to eq('GET /test') + end + + it 'sets dd_resource_key with apis prefix' do + @span = described_class.create(apigw_v2_event, request_context, trace_digest) + expect(@span.get_tag('dd_resource_key')).to eq( + 'arn:aws:apigateway:us-east-1::/apis/xyz789/stages/prod' + ) + end + + it 'sets user agent from http context' do + @span = described_class.create(apigw_v2_event, request_context, trace_digest) + expect(@span.get_tag('http.useragent')).to eq('TestAgent/2.0') + end + end + + context 'when an error occurs' do + before do + allow(Datadog::Tracing).to receive(:trace).and_raise(StandardError, 'boom') + end + + it 'returns nil' do + expect(described_class.create(apigw_v1_event, request_context, trace_digest)).to be_nil + end + end + end + + describe '.finish' do + let(:inferred_span) { Datadog::Tracing.trace('test.inferred') } + + after { inferred_span&.finish unless inferred_span&.finished? } + + context 'when inferred_span is nil' do + it { expect(described_class.finish(nil, {})).to be_nil } + end + + context 'with a response containing statusCode' do + it 'sets http.status_code tag and finishes the span' do + described_class.finish(inferred_span, {'statusCode' => 200}) + expect(inferred_span.get_tag('http.status_code')).to eq('200') + expect(inferred_span).to be_finished + end + end + + context 'when response is not a Hash' do + it 'finishes the span without status tag' do + described_class.finish(inferred_span, 'not a hash') + expect(inferred_span.get_tag('http.status_code')).to be_nil + expect(inferred_span).to be_finished + end + end + + context 'when response is nil' do + it 'finishes the span without status tag' do + described_class.finish(inferred_span, nil) + expect(inferred_span.get_tag('http.status_code')).to be_nil + expect(inferred_span).to be_finished + end + end + + context 'when an error occurs' do + let(:inferred_span) { double('span', finished?: true) } + + before do + allow(inferred_span).to receive(:set_tag) + allow(inferred_span).to receive(:finish).and_raise(StandardError, 'boom') + end + + it 'does not raise' do + expect { described_class.finish(inferred_span, {}) }.not_to raise_error + end + end + end +end From 32dfab7571267ae4e5c689c23d5e9d8db9fd755d Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Wed, 25 Mar 2026 15:51:14 +0100 Subject: [PATCH 09/24] Fix trace/span type confusion in AppSec context lifecycle F25: `Datadog::Tracing.trace()` returns a SpanOperation, not a TraceOperation. Rename `@trace` to `@span` in Listener to reflect the actual type. Remove stale class-level `@trace = nil`. AppSec.on_start now always uses `Datadog::Tracing.active_trace` for the TraceOperation (matching Rack's RequestMiddleware pattern) and receives only the target span (inferred span) as an argument, with fallback to `active_span`. F26: Pass `@inferred_span || @span` to `inject_appsec_data` so the extension reads tags from the span where AppSec actually sets them. F32: Variable renamed from `@trace` to `@span`. Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/datadog/lambda/appsec.rb | 4 ++-- lib/datadog/lambda/trace/listener.rb | 13 +++++++---- test/datadog/lambda/appsec.spec.rb | 35 +++++++++++++--------------- 3 files changed, 26 insertions(+), 26 deletions(-) diff --git a/lib/datadog/lambda/appsec.rb b/lib/datadog/lambda/appsec.rb index cc0d6e1..91f303b 100644 --- a/lib/datadog/lambda/appsec.rb +++ b/lib/datadog/lambda/appsec.rb @@ -4,10 +4,10 @@ module Datadog module Lambda module AppSec class << self - def on_start(event, trace, span) + def on_start(event, span = nil) return unless enabled? - trace ||= Datadog::Tracing.active_trace + trace = Datadog::Tracing.active_trace span ||= Datadog::Tracing.active_span create_context(trace, span) diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index 6798ddc..bad846c 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -18,12 +18,12 @@ module Datadog module Trace # TraceListener tracks tracing context information class Listener - @trace = nil def initialize(handler_name:, function_name:, patch_http:, merge_xray_traces:) @handler_name = handler_name @function_name = function_name @merge_xray_traces = merge_xray_traces + @span = nil @inferred_span = nil Datadog::Trace.patch_http if patch_http @@ -50,20 +50,23 @@ def on_start(event:, request_context:, cold_start:) @inferred_span = Datadog::Lambda::Trace::InferredSpan.create(event, request_context, trace_digest) options[:continue_from] = trace_digest if trace_digest && @inferred_span.nil? - @trace = Datadog::Tracing.trace('aws.lambda', **options) + @span = Datadog::Tracing.trace('aws.lambda', **options) Datadog::Trace.apply_datadog_trace_context(Datadog::Trace.trace_context) - Datadog::Lambda::AppSec.on_start(event, @trace, @inferred_span) + Datadog::Lambda::AppSec.on_start(event, @inferred_span) end # rubocop:enable Metrics/AbcSize def on_end(response:, request_context:) Datadog::Lambda::AppSec.on_finish(response) - Datadog::Utils.send_end_invocation_request(response:, span_id: @trace.id, request_context:, span: @trace) + Datadog::Utils.send_end_invocation_request( + response:, span_id: @span.id, request_context:, span: @inferred_span || @span + ) - @trace&.finish + @span&.finish Datadog::Lambda::Trace::InferredSpan.finish(@inferred_span, response) + @span = nil @inferred_span = nil end diff --git a/test/datadog/lambda/appsec.spec.rb b/test/datadog/lambda/appsec.spec.rb index a1e753f..6d42355 100644 --- a/test/datadog/lambda/appsec.spec.rb +++ b/test/datadog/lambda/appsec.spec.rb @@ -5,7 +5,7 @@ describe Datadog::Lambda::AppSec do let(:event) { {'httpMethod' => 'GET', 'path' => '/'} } - let(:trace) { double('trace') } + let(:active_trace) { double('active_trace') } let(:span) { double('span', set_metric: nil) } let(:gateway) { double('gateway') } let(:runner) { double('runner') } @@ -23,16 +23,17 @@ before do allow(Datadog::AppSec::Instrumentation).to receive(:gateway).and_return(gateway) allow(gateway).to receive(:push) + allow(Datadog::Tracing).to receive(:active_trace).and_return(active_trace) end describe '.on_start' do context 'when appsec is disabled' do before { allow(Datadog::AppSec).to receive(:enabled?).and_return(false) } - it { expect(described_class.on_start(event, trace, span)).to be_nil } + it { expect(described_class.on_start(event, span)).to be_nil } it 'does not push gateway events' do - described_class.on_start(event, trace, span) + described_class.on_start(event, span) expect(gateway).not_to have_received(:push) end end @@ -45,33 +46,29 @@ allow(Datadog::AppSec::Context).to receive(:new).and_return(context) end - it 'creates and activates an AppSec context' do - described_class.on_start(event, trace, span) - expect(Datadog::AppSec::Context).to have_received(:new).with(trace, span, runner) + it 'creates context with active_trace and the given span' do + described_class.on_start(event, span) + expect(Datadog::AppSec::Context).to have_received(:new).with(active_trace, span, runner) expect(Datadog::AppSec::Context).to have_received(:activate).with(context) end it 'sets _dd.appsec.enabled metric on span' do - described_class.on_start(event, trace, span) + described_class.on_start(event, span) expect(span).to have_received(:set_metric).with(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) end it 'pushes the request event to the gateway' do - described_class.on_start(event, trace, span) + described_class.on_start(event, span) expect(gateway).to have_received(:push).with('aws_lambda.request.start', event) end - context 'when trace and span are nil' do - let(:active_trace) { double('active_trace') } + context 'when span is not provided' do let(:active_span) { double('active_span', set_metric: nil) } - before do - allow(Datadog::Tracing).to receive(:active_trace).and_return(active_trace) - allow(Datadog::Tracing).to receive(:active_span).and_return(active_span) - end + before { allow(Datadog::Tracing).to receive(:active_span).and_return(active_span) } - it 'falls back to active trace and span' do - described_class.on_start(event, nil, nil) + it 'falls back to active span' do + described_class.on_start(event) expect(Datadog::AppSec::Context).to have_received(:new).with(active_trace, active_span, runner) end end @@ -80,12 +77,12 @@ before { allow(Datadog::AppSec).to receive(:security_engine).and_return(nil) } it 'does not activate a context' do - described_class.on_start(event, trace, span) + described_class.on_start(event, span) expect(Datadog::AppSec::Context).not_to have_received(:activate) end it 'still pushes the gateway event' do - described_class.on_start(event, trace, span) + described_class.on_start(event, span) expect(gateway).to have_received(:push).with('aws_lambda.request.start', event) end end @@ -94,7 +91,7 @@ before { allow(Datadog::AppSec::Context).to receive(:new).and_raise(StandardError, 'boom') } it 'rescues and logs' do - expect { described_class.on_start(event, trace, span) }.not_to raise_error + expect { described_class.on_start(event, span) }.not_to raise_error end end end From 7596aee0c6380fcb074d68f5448120e61d74d5b2 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Wed, 25 Mar 2026 15:57:04 +0100 Subject: [PATCH 10/24] Remove unnecessary change --- lib/datadog/lambda.rb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/datadog/lambda.rb b/lib/datadog/lambda.rb index 3e6bec8..aa64607 100644 --- a/lib/datadog/lambda.rb +++ b/lib/datadog/lambda.rb @@ -45,7 +45,7 @@ def self.configure_apm end c.tags = { "_dd.origin": 'lambda' } # Enable AWS SDK instrumentation - c.tracing.instrument(:aws) if trace_managed_services? + c.tracing.instrument :aws if trace_managed_services? yield(c) if block_given? From 84cb13fb93f3e5cbb600c0b5e1dd47daa6177408 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Wed, 25 Mar 2026 16:52:50 +0100 Subject: [PATCH 11/24] Remove appsec request headers injection --- lib/datadog/lambda/appsec.rb | 1 + lib/datadog/lambda/inferred_span.rb | 129 ++++++++++++++++++++ lib/datadog/lambda/trace/lambda_appsec.rb | 42 +++++++ lib/datadog/lambda/trace/listener.rb | 3 +- lib/datadog/lambda/utils/extension.rb | 17 +-- test/datadog/lambda/appsec.spec.rb | 10 +- test/datadog/lambda/utils/extension.spec.rb | 21 +--- 7 files changed, 184 insertions(+), 39 deletions(-) create mode 100644 lib/datadog/lambda/inferred_span.rb create mode 100644 lib/datadog/lambda/trace/lambda_appsec.rb diff --git a/lib/datadog/lambda/appsec.rb b/lib/datadog/lambda/appsec.rb index 91f303b..bb7b6ca 100644 --- a/lib/datadog/lambda/appsec.rb +++ b/lib/datadog/lambda/appsec.rb @@ -11,6 +11,7 @@ def on_start(event, span = nil) span ||= Datadog::Tracing.active_span create_context(trace, span) + return unless Datadog::AppSec::Context.active Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.request.start', event) rescue StandardError => e diff --git a/lib/datadog/lambda/inferred_span.rb b/lib/datadog/lambda/inferred_span.rb new file mode 100644 index 0000000..7e5e3fc --- /dev/null +++ b/lib/datadog/lambda/inferred_span.rb @@ -0,0 +1,129 @@ +# frozen_string_literal: true + +module Datadog + module Lambda + module InferredSpan + class << self + def create(event, request_context, trace_digest) + return unless managed_services_enabled? + + span_name = inferred_span_name(event) + return unless span_name + + rc = event['requestContext'] || {} + domain = rc['domainName'] || '' + api_id = rc['apiId'] || '' + stage = rc['stage'] || '' + + if span_name == 'aws.apigateway' + method = event['httpMethod'] + path = event['path'] || '/' + resource_path = rc['resourcePath'] || path + request_time_ms = rc['requestTimeEpoch'] + user_agent = rc.dig('identity', 'userAgent') + else + http = rc['http'] || {} + method = http['method'] + path = event['rawPath'] || '/' + resource_path = event['routeKey']&.sub(/^[A-Z]+ /, '') || path + request_time_ms = rc['timeEpoch'] + user_agent = http['userAgent'] + end + + resource = "#{method} #{resource_path}" + http_url = domain.empty? ? path : "https://#{domain}#{path}" + + tags = { + 'http.method' => method, + 'http.url' => http_url, + 'http.route' => resource_path, + 'endpoint' => path, + 'resource_names' => resource, + 'span.kind' => 'server', + 'apiid' => api_id, + 'apiname' => api_id, + 'stage' => stage, + 'request_id' => request_context.aws_request_id, + '_inferred_span.synchronicity' => 'sync', + '_inferred_span.tag_source' => 'self', + } + + arn = request_context.invoked_function_arn.to_s + region = arn.split(':')[3] if arn.include?(':') + if region && !api_id.empty? && !stage.empty? + arn_path = span_name == 'aws.apigateway' ? 'restapis' : 'apis' + tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{arn_path}/#{api_id}/stages/#{stage}" + end + + tags['http.useragent'] = user_agent if user_agent + + inferred_options = { + service: domain.empty? ? nil : domain, + resource: resource, + type: 'web', + tags: tags, + } + inferred_options[:continue_from] = trace_digest if trace_digest + inferred_options[:start_time] = Time.at(request_time_ms / 1000.0) if request_time_ms + + span = Datadog::Tracing.trace(span_name, **inferred_options) + span.set_metric('_dd._inferred_span', 1.0) + span + rescue StandardError => e + Datadog::Utils.logger.debug "failed to create inferred span: #{e}" + nil + end + + def finish(inferred_span, response, trace) + return unless inferred_span + + if trace + appsec_enabled = trace.get_metric('_dd.appsec.enabled') + inferred_span.set_metric('_dd.appsec.enabled', appsec_enabled) if appsec_enabled + + appsec_json = trace.get_tag('_dd.appsec.json') + inferred_span.set_tag('_dd.appsec.json', appsec_json) if appsec_json + + appsec_event = trace.get_tag('appsec.event') + inferred_span.set_tag('appsec.event', appsec_event) if appsec_event + + origin = trace.get_tag('_dd.origin') + inferred_span.set_tag('_dd.origin', origin) if origin + end + + status_code = extract_status_code(response) + inferred_span.set_tag('http.status_code', status_code.to_s) if status_code + + inferred_span.finish + rescue StandardError => e + Datadog::Utils.logger.debug "failed to finish inferred span: #{e}" + end + + private + + def inferred_span_name(event) + return unless event.is_a?(Hash) + + rc = event['requestContext'] + return unless rc.is_a?(Hash) && rc['stage'] + + if event['httpMethod'] + 'aws.apigateway' + elsif event['routeKey'] + 'aws.httpapi' + end + end + + def managed_services_enabled? + ENV.fetch('DD_TRACE_MANAGED_SERVICES', 'true').downcase != 'false' + end + + def extract_status_code(response) + return unless response.is_a?(Hash) + + response['statusCode'] + end + end + end + end +end diff --git a/lib/datadog/lambda/trace/lambda_appsec.rb b/lib/datadog/lambda/trace/lambda_appsec.rb new file mode 100644 index 0000000..aeb2451 --- /dev/null +++ b/lib/datadog/lambda/trace/lambda_appsec.rb @@ -0,0 +1,42 @@ +# frozen_string_literal: true + +module Datadog + module Trace + module LambdaAppSec + class << self + def start(event) + return unless enabled? + + ensure_patched + + Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.request.start', event) + rescue StandardError => e + Datadog::Utils.logger.debug "failed to start AppSec: #{e}" + end + + def finish(response) + return unless Datadog::AppSec::Context.active + + Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.response.start', response) + rescue StandardError => e + Datadog::Utils.logger.debug "failed to finish AppSec: #{e}" + end + + private + + def enabled? + defined?(Datadog::AppSec) && + Datadog::AppSec.respond_to?(:enabled?) && + Datadog::AppSec.enabled? + end + + def ensure_patched + return if @patched + + Datadog.configuration.appsec.instrument(:aws_lambda) + @patched = true + end + end + end + end +end diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index bad846c..f8159dc 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -60,9 +60,10 @@ def on_start(event:, request_context:, cold_start:) def on_end(response:, request_context:) Datadog::Lambda::AppSec.on_finish(response) Datadog::Utils.send_end_invocation_request( - response:, span_id: @span.id, request_context:, span: @inferred_span || @span + response:, span_id: @span.id, request_context: ) + # NOTE: lambda span must finish before inferred span (its parent) @span&.finish Datadog::Lambda::Trace::InferredSpan.finish(@inferred_span, response) diff --git a/lib/datadog/lambda/utils/extension.rb b/lib/datadog/lambda/utils/extension.rb index 66315ea..898ac45 100644 --- a/lib/datadog/lambda/utils/extension.rb +++ b/lib/datadog/lambda/utils/extension.rb @@ -57,11 +57,8 @@ def self.send_start_invocation_request(event:, request_context:) Datadog::Utils.logger.debug "failed on start invocation request to extension: #{e}" end - DD_APPSEC_ENABLED_HEADER = 'x-datadog-appsec-enabled' - DD_APPSEC_JSON_HEADER = 'x-datadog-appsec-json' - # rubocop:disable Metrics/AbcSize - def self.send_end_invocation_request(response:, span_id:, request_context:, span: nil) + def self.send_end_invocation_request(response:, span_id:, request_context:) return unless extension_running? request = Net::HTTP::Post.new(END_INVOCATION_URI) @@ -80,8 +77,6 @@ def self.send_end_invocation_request(response:, span_id:, request_context:, span # Remove Parent ID if it is the same as the Span ID request.delete(DD_PARENT_ID_HEADER) if request[DD_PARENT_ID_HEADER] == span_id.to_s - inject_appsec_data(request, span) - Datadog::Utils.logger.debug "End invocation request headers: #{request.to_hash}" Net::HTTP.start(END_INVOCATION_URI.host, END_INVOCATION_URI.port) do |http| @@ -92,16 +87,6 @@ def self.send_end_invocation_request(response:, span_id:, request_context:, span end # rubocop:enable Metrics/AbcSize - def self.inject_appsec_data(request, span) - return unless span - - appsec_enabled = span.get_metric('_dd.appsec.enabled') - request[DD_APPSEC_ENABLED_HEADER] = '1' if appsec_enabled - - appsec_json = span.get_tag('_dd.appsec.json') - request[DD_APPSEC_JSON_HEADER] = appsec_json if appsec_json - end - def self.request_headers { # Header used to avoid tracing requests that are internal to diff --git a/test/datadog/lambda/appsec.spec.rb b/test/datadog/lambda/appsec.spec.rb index 6d42355..3ff4f4b 100644 --- a/test/datadog/lambda/appsec.spec.rb +++ b/test/datadog/lambda/appsec.spec.rb @@ -43,6 +43,7 @@ allow(Datadog::AppSec).to receive(:enabled?).and_return(true) allow(Datadog::AppSec).to receive(:security_engine).and_return(security_engine) allow(Datadog::AppSec::Context).to receive(:activate) + allow(Datadog::AppSec::Context).to receive(:active).and_return(context) allow(Datadog::AppSec::Context).to receive(:new).and_return(context) end @@ -74,16 +75,19 @@ end context 'when security_engine is nil' do - before { allow(Datadog::AppSec).to receive(:security_engine).and_return(nil) } + before do + allow(Datadog::AppSec).to receive(:security_engine).and_return(nil) + allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) + end it 'does not activate a context' do described_class.on_start(event, span) expect(Datadog::AppSec::Context).not_to have_received(:activate) end - it 'still pushes the gateway event' do + it 'does not push gateway events' do described_class.on_start(event, span) - expect(gateway).to have_received(:push).with('aws_lambda.request.start', event) + expect(gateway).not_to have_received(:push) end end diff --git a/test/datadog/lambda/utils/extension.spec.rb b/test/datadog/lambda/utils/extension.spec.rb index 97f0989..6e0e3a0 100644 --- a/test/datadog/lambda/utils/extension.spec.rb +++ b/test/datadog/lambda/utils/extension.spec.rb @@ -103,7 +103,7 @@ Datadog::Utils.send_end_invocation_request(response: nil, span_id: nil, request_context: ctx) end - it 'forwards AppSec tags from span as headers' do + it 'does not send AppSec headers to the extension' do @trace.set_metric('_dd.appsec.enabled', 1.0) @trace.set_tag('_dd.appsec.json', '{"triggers":[]}') @@ -116,24 +116,7 @@ end Datadog::Utils.send_end_invocation_request( - response: nil, span_id: nil, request_context: ctx, span: @trace - ) - - expect(captured_request['x-datadog-appsec-enabled']).to eq('1') - expect(captured_request['x-datadog-appsec-json']).to eq('{"triggers":[]}') - end - - it 'does not send AppSec headers when tags are absent' do - captured_request = nil - http_double = instance_double(Net::HTTP) - allow(Net::HTTP).to receive(:start).and_yield(http_double) - allow(http_double).to receive(:request) do |req| - captured_request = req - Net::HTTPResponse.new('1.1', '200', 'OK') - end - - Datadog::Utils.send_end_invocation_request( - response: nil, span_id: nil, request_context: ctx, span: @trace + response: nil, span_id: nil, request_context: ctx ) expect(captured_request['x-datadog-appsec-enabled']).to be_nil From 985129434dd9f7dc4924c81132d22236e63c23e9 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Wed, 25 Mar 2026 17:16:52 +0100 Subject: [PATCH 12/24] Ignore test/ directory in RuboCop --- .rubocop.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.rubocop.yml b/.rubocop.yml index 2d9f108..64ecb10 100644 --- a/.rubocop.yml +++ b/.rubocop.yml @@ -1,5 +1,7 @@ AllCops: TargetRubyVersion: 3.2 + Exclude: + - 'test/**/*' Metrics/MethodLength: Max: 20 From 8790c58a6eba47209a2dd1b3f3d0de205a187443 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Wed, 25 Mar 2026 17:22:45 +0100 Subject: [PATCH 13/24] Move InferredSpan outside of Trace folder * Rewrite specs to test behavior instead of internals * Remove obsolete appsec files --- lib/datadog/lambda/inferred_span.rb | 198 +++++------ lib/datadog/lambda/trace/inferred_span.rb | 117 ------- lib/datadog/lambda/trace/lambda_appsec.rb | 42 --- lib/datadog/lambda/trace/listener.rb | 6 +- test/datadog/lambda/appsec.spec.rb | 151 +++++---- test/datadog/lambda/inferred_span.spec.rb | 320 ++++++++++++++++++ .../lambda/trace/inferred_span.spec.rb | 247 -------------- test/datadog/lambda/utils/extension.spec.rb | 106 ++---- 8 files changed, 536 insertions(+), 651 deletions(-) delete mode 100644 lib/datadog/lambda/trace/inferred_span.rb delete mode 100644 lib/datadog/lambda/trace/lambda_appsec.rb create mode 100644 test/datadog/lambda/inferred_span.spec.rb delete mode 100644 test/datadog/lambda/trace/inferred_span.spec.rb diff --git a/lib/datadog/lambda/inferred_span.rb b/lib/datadog/lambda/inferred_span.rb index 7e5e3fc..701bd3c 100644 --- a/lib/datadog/lambda/inferred_span.rb +++ b/lib/datadog/lambda/inferred_span.rb @@ -3,127 +3,113 @@ module Datadog module Lambda module InferredSpan - class << self - def create(event, request_context, trace_digest) - return unless managed_services_enabled? - - span_name = inferred_span_name(event) - return unless span_name - - rc = event['requestContext'] || {} - domain = rc['domainName'] || '' - api_id = rc['apiId'] || '' - stage = rc['stage'] || '' - - if span_name == 'aws.apigateway' - method = event['httpMethod'] - path = event['path'] || '/' - resource_path = rc['resourcePath'] || path - request_time_ms = rc['requestTimeEpoch'] - user_agent = rc.dig('identity', 'userAgent') - else - http = rc['http'] || {} - method = http['method'] - path = event['rawPath'] || '/' - resource_path = event['routeKey']&.sub(/^[A-Z]+ /, '') || path - request_time_ms = rc['timeEpoch'] - user_agent = http['userAgent'] - end - - resource = "#{method} #{resource_path}" - http_url = domain.empty? ? path : "https://#{domain}#{path}" - - tags = { - 'http.method' => method, - 'http.url' => http_url, - 'http.route' => resource_path, - 'endpoint' => path, - 'resource_names' => resource, - 'span.kind' => 'server', - 'apiid' => api_id, - 'apiname' => api_id, - 'stage' => stage, - 'request_id' => request_context.aws_request_id, - '_inferred_span.synchronicity' => 'sync', - '_inferred_span.tag_source' => 'self', - } - - arn = request_context.invoked_function_arn.to_s - region = arn.split(':')[3] if arn.include?(':') - if region && !api_id.empty? && !stage.empty? - arn_path = span_name == 'aws.apigateway' ? 'restapis' : 'apis' - tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{arn_path}/#{api_id}/stages/#{stage}" - end - - tags['http.useragent'] = user_agent if user_agent - - inferred_options = { - service: domain.empty? ? nil : domain, - resource: resource, - type: 'web', - tags: tags, - } - inferred_options[:continue_from] = trace_digest if trace_digest - inferred_options[:start_time] = Time.at(request_time_ms / 1000.0) if request_time_ms - - span = Datadog::Tracing.trace(span_name, **inferred_options) - span.set_metric('_dd._inferred_span', 1.0) - span - rescue StandardError => e - Datadog::Utils.logger.debug "failed to create inferred span: #{e}" - nil + class << self + def create(event, request_context, trace_digest) + return unless managed_services_enabled? + + span_name = inferred_span_name(event) + return unless span_name + + rc = event['requestContext'] || {} + domain = rc['domainName'] || '' + api_id = rc['apiId'] || '' + stage = rc['stage'] || '' + + if span_name == 'aws.apigateway' + method = event['httpMethod'] + path = event['path'] || '/' + resource_path = rc['resourcePath'] || path + request_time_ms = rc['requestTimeEpoch'] + user_agent = rc.dig('identity', 'userAgent') + else + http = rc['http'] || {} + method = http['method'] + path = event['rawPath'] || '/' + resource_path = event['routeKey']&.sub(/^[A-Z]+ /, '') || path + request_time_ms = rc['timeEpoch'] + user_agent = http['userAgent'] end - def finish(inferred_span, response, trace) - return unless inferred_span - - if trace - appsec_enabled = trace.get_metric('_dd.appsec.enabled') - inferred_span.set_metric('_dd.appsec.enabled', appsec_enabled) if appsec_enabled - - appsec_json = trace.get_tag('_dd.appsec.json') - inferred_span.set_tag('_dd.appsec.json', appsec_json) if appsec_json + resource = "#{method} #{resource_path}" + http_url = domain.empty? ? path : "https://#{domain}#{path}" + + tags = { + 'http.method' => method, + 'http.url' => http_url, + 'http.route' => resource_path, + 'endpoint' => path, + 'resource_names' => resource, + 'span.kind' => 'server', + 'apiid' => api_id, + 'apiname' => api_id, + 'stage' => stage, + 'request_id' => request_context.aws_request_id, + '_inferred_span.synchronicity' => 'sync', + '_inferred_span.tag_source' => 'self', + } + + arn = request_context.invoked_function_arn.to_s + region = arn.split(':')[3] if arn.include?(':') + if region && !api_id.empty? && !stage.empty? + arn_path = span_name == 'aws.apigateway' ? 'restapis' : 'apis' + tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{arn_path}/#{api_id}/stages/#{stage}" + end - appsec_event = trace.get_tag('appsec.event') - inferred_span.set_tag('appsec.event', appsec_event) if appsec_event + tags['http.useragent'] = user_agent if user_agent + + inferred_options = { + service: domain.empty? ? nil : domain, + resource: resource, + type: 'web', + tags: tags, + } + inferred_options[:continue_from] = trace_digest if trace_digest + inferred_options[:start_time] = Time.at(request_time_ms / 1000.0) if request_time_ms + + span = Datadog::Tracing.trace(span_name, **inferred_options) + span.set_metric('_dd._inferred_span', 1.0) + span + rescue StandardError => e + Datadog::Utils.logger.debug "failed to create inferred span: #{e}" + nil + end - origin = trace.get_tag('_dd.origin') - inferred_span.set_tag('_dd.origin', origin) if origin - end + def finish(inferred_span, response) + return unless inferred_span - status_code = extract_status_code(response) - inferred_span.set_tag('http.status_code', status_code.to_s) if status_code + status_code = extract_status_code(response) + inferred_span.set_tag('http.status_code', status_code.to_s) if status_code - inferred_span.finish - rescue StandardError => e - Datadog::Utils.logger.debug "failed to finish inferred span: #{e}" - end + inferred_span.finish + rescue StandardError => e + Datadog::Utils.logger.debug "failed to finish inferred span: #{e}" + end - private + private - def inferred_span_name(event) - return unless event.is_a?(Hash) + def inferred_span_name(event) + return unless event.is_a?(Hash) - rc = event['requestContext'] - return unless rc.is_a?(Hash) && rc['stage'] + rc = event['requestContext'] + return unless rc.is_a?(Hash) && rc['stage'] - if event['httpMethod'] - 'aws.apigateway' - elsif event['routeKey'] - 'aws.httpapi' - end + if event['httpMethod'] + 'aws.apigateway' + elsif event['routeKey'] + 'aws.httpapi' end + end - def managed_services_enabled? - ENV.fetch('DD_TRACE_MANAGED_SERVICES', 'true').downcase != 'false' - end + def managed_services_enabled? + Datadog::Lambda.trace_managed_services? + end - def extract_status_code(response) - return unless response.is_a?(Hash) + def extract_status_code(response) + return unless response.is_a?(Hash) - response['statusCode'] - end + response['statusCode'] end end + end end end diff --git a/lib/datadog/lambda/trace/inferred_span.rb b/lib/datadog/lambda/trace/inferred_span.rb deleted file mode 100644 index 2195bc8..0000000 --- a/lib/datadog/lambda/trace/inferred_span.rb +++ /dev/null @@ -1,117 +0,0 @@ -# frozen_string_literal: true - -module Datadog - module Lambda - module Trace - module InferredSpan - class << self - def create(event, request_context, trace_digest) - return unless managed_services_enabled? - - span_name = inferred_span_name(event) - return unless span_name - - rc = event['requestContext'] || {} - domain = rc['domainName'] || '' - api_id = rc['apiId'] || '' - stage = rc['stage'] || '' - - if span_name == 'aws.apigateway' - method = event['httpMethod'] - path = event['path'] || '/' - resource_path = rc['resourcePath'] || path - request_time_ms = rc['requestTimeEpoch'] - user_agent = rc.dig('identity', 'userAgent') - else - http = rc['http'] || {} - method = http['method'] - path = event['rawPath'] || '/' - resource_path = event['routeKey']&.sub(/^[A-Z]+ /, '') || path - request_time_ms = rc['timeEpoch'] - user_agent = http['userAgent'] - end - - resource = "#{method} #{resource_path}" - http_url = domain.empty? ? path : "https://#{domain}#{path}" - - tags = { - 'http.method' => method, - 'http.url' => http_url, - 'http.route' => resource_path, - 'endpoint' => path, - 'resource_names' => resource, - 'span.kind' => 'server', - 'apiid' => api_id, - 'apiname' => api_id, - 'stage' => stage, - 'request_id' => request_context.aws_request_id, - '_inferred_span.synchronicity' => 'sync', - '_inferred_span.tag_source' => 'self', - } - - arn = request_context.invoked_function_arn.to_s - region = arn.split(':')[3] if arn.include?(':') - if region && !api_id.empty? && !stage.empty? - arn_path = span_name == 'aws.apigateway' ? 'restapis' : 'apis' - tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{arn_path}/#{api_id}/stages/#{stage}" - end - - tags['http.useragent'] = user_agent if user_agent - - inferred_options = { - service: domain.empty? ? nil : domain, - resource: resource, - type: 'web', - tags: tags, - } - inferred_options[:continue_from] = trace_digest if trace_digest - inferred_options[:start_time] = Time.at(request_time_ms / 1000.0) if request_time_ms - - span = Datadog::Tracing.trace(span_name, **inferred_options) - span.set_metric('_dd._inferred_span', 1.0) - span - rescue StandardError => e - Datadog::Utils.logger.debug "failed to create inferred span: #{e}" - nil - end - - def finish(inferred_span, response) - return unless inferred_span - - status_code = extract_status_code(response) - inferred_span.set_tag('http.status_code', status_code.to_s) if status_code - - inferred_span.finish - rescue StandardError => e - Datadog::Utils.logger.debug "failed to finish inferred span: #{e}" - end - - private - - def inferred_span_name(event) - return unless event.is_a?(Hash) - - rc = event['requestContext'] - return unless rc.is_a?(Hash) && rc['stage'] - - if event['httpMethod'] - 'aws.apigateway' - elsif event['routeKey'] - 'aws.httpapi' - end - end - - def managed_services_enabled? - Datadog::Lambda.trace_managed_services? - end - - def extract_status_code(response) - return unless response.is_a?(Hash) - - response['statusCode'] - end - end - end - end - end -end diff --git a/lib/datadog/lambda/trace/lambda_appsec.rb b/lib/datadog/lambda/trace/lambda_appsec.rb deleted file mode 100644 index aeb2451..0000000 --- a/lib/datadog/lambda/trace/lambda_appsec.rb +++ /dev/null @@ -1,42 +0,0 @@ -# frozen_string_literal: true - -module Datadog - module Trace - module LambdaAppSec - class << self - def start(event) - return unless enabled? - - ensure_patched - - Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.request.start', event) - rescue StandardError => e - Datadog::Utils.logger.debug "failed to start AppSec: #{e}" - end - - def finish(response) - return unless Datadog::AppSec::Context.active - - Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.response.start', response) - rescue StandardError => e - Datadog::Utils.logger.debug "failed to finish AppSec: #{e}" - end - - private - - def enabled? - defined?(Datadog::AppSec) && - Datadog::AppSec.respond_to?(:enabled?) && - Datadog::AppSec.enabled? - end - - def ensure_patched - return if @patched - - Datadog.configuration.appsec.instrument(:aws_lambda) - @patched = true - end - end - end - end -end diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index f8159dc..34bc114 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -11,7 +11,7 @@ require 'datadog/lambda/trace/context' require 'datadog/lambda/trace/patch_http' require 'datadog/lambda/trace/ddtrace' -require 'datadog/lambda/trace/inferred_span' +require 'datadog/lambda/inferred_span' require 'datadog/lambda/appsec' module Datadog @@ -47,7 +47,7 @@ def on_start(event:, request_context:, cold_start:) trace_digest = Datadog::Utils.send_start_invocation_request(event:, request_context:) - @inferred_span = Datadog::Lambda::Trace::InferredSpan.create(event, request_context, trace_digest) + @inferred_span = Datadog::Lambda::InferredSpan.create(event, request_context, trace_digest) options[:continue_from] = trace_digest if trace_digest && @inferred_span.nil? @span = Datadog::Tracing.trace('aws.lambda', **options) @@ -65,7 +65,7 @@ def on_end(response:, request_context:) # NOTE: lambda span must finish before inferred span (its parent) @span&.finish - Datadog::Lambda::Trace::InferredSpan.finish(@inferred_span, response) + Datadog::Lambda::InferredSpan.finish(@inferred_span, response) @span = nil @inferred_span = nil diff --git a/test/datadog/lambda/appsec.spec.rb b/test/datadog/lambda/appsec.spec.rb index 3ff4f4b..20d0687 100644 --- a/test/datadog/lambda/appsec.spec.rb +++ b/test/datadog/lambda/appsec.spec.rb @@ -4,37 +4,39 @@ require 'datadog/lambda/appsec' describe Datadog::Lambda::AppSec do - let(:event) { {'httpMethod' => 'GET', 'path' => '/'} } + before do + allow(Datadog::AppSec::Instrumentation).to receive(:gateway).and_return(gateway) + allow(gateway).to receive(:push) + allow(Datadog::Tracing).to receive(:active_trace).and_return(active_trace) + end + + let(:gateway) { double('gateway') } let(:active_trace) { double('active_trace') } let(:span) { double('span', set_metric: nil) } - let(:gateway) { double('gateway') } let(:runner) { double('runner') } let(:security_engine) { double('security_engine', new_runner: runner) } - let(:context) do + let(:appsec_context) do double( 'context', state: {}, export_metrics: nil, export_request_telemetry: nil, - events: [], ) end - before do - allow(Datadog::AppSec::Instrumentation).to receive(:gateway).and_return(gateway) - allow(gateway).to receive(:push) - allow(Datadog::Tracing).to receive(:active_trace).and_return(active_trace) - end - describe '.on_start' do + subject(:on_start) { described_class.on_start(event, span) } + + let(:event) { {'httpMethod' => 'GET', 'path' => '/'} } + context 'when appsec is disabled' do before { allow(Datadog::AppSec).to receive(:enabled?).and_return(false) } - it { expect(described_class.on_start(event, span)).to be_nil } - - it 'does not push gateway events' do - described_class.on_start(event, span) - expect(gateway).not_to have_received(:push) + it 'does nothing' do + aggregate_failures('no side effects') do + expect(on_start).to be_nil + expect(gateway).not_to have_received(:push) + end end end @@ -43,33 +45,33 @@ allow(Datadog::AppSec).to receive(:enabled?).and_return(true) allow(Datadog::AppSec).to receive(:security_engine).and_return(security_engine) allow(Datadog::AppSec::Context).to receive(:activate) - allow(Datadog::AppSec::Context).to receive(:active).and_return(context) - allow(Datadog::AppSec::Context).to receive(:new).and_return(context) + allow(Datadog::AppSec::Context).to receive(:active).and_return(appsec_context) + allow(Datadog::AppSec::Context).to receive(:new).and_return(appsec_context) end - it 'creates context with active_trace and the given span' do - described_class.on_start(event, span) - expect(Datadog::AppSec::Context).to have_received(:new).with(active_trace, span, runner) - expect(Datadog::AppSec::Context).to have_received(:activate).with(context) - end + it 'activates context, marks span, and pushes event to gateway' do + on_start - it 'sets _dd.appsec.enabled metric on span' do - described_class.on_start(event, span) - expect(span).to have_received(:set_metric).with(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) - end + aggregate_failures('context lifecycle') do + expect(Datadog::AppSec::Context).to have_received(:new).with(active_trace, span, runner) + expect(Datadog::AppSec::Context).to have_received(:activate).with(appsec_context) + end - it 'pushes the request event to the gateway' do - described_class.on_start(event, span) - expect(gateway).to have_received(:push).with('aws_lambda.request.start', event) + aggregate_failures('span and gateway') do + expect(span).to have_received(:set_metric).with(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) + expect(gateway).to have_received(:push).with('aws_lambda.request.start', event) + end end context 'when span is not provided' do - let(:active_span) { double('active_span', set_metric: nil) } + subject(:on_start) { described_class.on_start(event) } before { allow(Datadog::Tracing).to receive(:active_span).and_return(active_span) } + let(:active_span) { double('active_span', set_metric: nil) } + it 'falls back to active span' do - described_class.on_start(event) + on_start expect(Datadog::AppSec::Context).to have_received(:new).with(active_trace, active_span, runner) end end @@ -80,77 +82,96 @@ allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) end - it 'does not activate a context' do - described_class.on_start(event, span) - expect(Datadog::AppSec::Context).not_to have_received(:activate) + it 'does not activate or push' do + on_start + + aggregate_failures('no side effects') do + expect(Datadog::AppSec::Context).not_to have_received(:activate) + expect(gateway).not_to have_received(:push) + end end + end - it 'does not push gateway events' do - described_class.on_start(event, span) - expect(gateway).not_to have_received(:push) + context 'when trace is nil' do + before do + allow(Datadog::Tracing).to receive(:active_trace).and_return(nil) + allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) + end + + it 'does not activate or push' do + on_start + + aggregate_failures('no side effects') do + expect(Datadog::AppSec::Context).not_to have_received(:activate) + expect(gateway).not_to have_received(:push) + end end end context 'when an error occurs' do before { allow(Datadog::AppSec::Context).to receive(:new).and_raise(StandardError, 'boom') } - it 'rescues and logs' do - expect { described_class.on_start(event, span) }.not_to raise_error - end + it { expect { on_start }.not_to raise_error } end end end describe '.on_finish' do + subject(:on_finish) { described_class.on_finish(response) } + let(:response) { {'statusCode' => 200} } context 'when no active context' do before { allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) } - it { expect(described_class.on_finish(response)).to be_nil } - - it 'does not push gateway events' do - described_class.on_finish(response) - expect(gateway).not_to have_received(:push) + it 'does nothing' do + aggregate_failures('no side effects') do + expect(on_finish).to be_nil + expect(gateway).not_to have_received(:push) + end end end context 'when active context exists' do before do - allow(Datadog::AppSec::Context).to receive(:active).and_return(context) + allow(Datadog::AppSec::Context).to receive(:active).and_return(appsec_context) allow(Datadog::AppSec::Context).to receive(:deactivate) allow(Datadog::AppSec::Event).to receive(:record) end - it 'pushes the response event to the gateway' do - described_class.on_finish(response) - expect(gateway).to have_received(:push).with('aws_lambda.response.start', response) - end + it 'processes response and cleans up' do + on_finish - it 'records events with the request from context state' do - context.state[:request] = double('request') - described_class.on_finish(response) - expect(Datadog::AppSec::Event).to have_received(:record).with( - context, request: context.state[:request] - ) - end + aggregate_failures('gateway and event recording') do + expect(gateway).to have_received(:push).with('aws_lambda.response.start', response) + expect(Datadog::AppSec::Event).to have_received(:record).with(appsec_context, request: nil) + end - it 'exports metrics and telemetry' do - described_class.on_finish(response) - expect(context).to have_received(:export_metrics) - expect(context).to have_received(:export_request_telemetry) + aggregate_failures('telemetry and cleanup') do + expect(appsec_context).to have_received(:export_metrics) + expect(appsec_context).to have_received(:export_request_telemetry) + expect(Datadog::AppSec::Context).to have_received(:deactivate) + end end - it 'deactivates the context' do - described_class.on_finish(response) - expect(Datadog::AppSec::Context).to have_received(:deactivate) + context 'when context has a request in state' do + before { appsec_context.state[:request] = request_data } + + let(:request_data) { double('request') } + + it 'passes request to event recording' do + on_finish + expect(Datadog::AppSec::Event).to have_received(:record).with( + appsec_context, request: request_data + ) + end end context 'when an error occurs' do before { allow(gateway).to receive(:push).and_raise(StandardError, 'boom') } it 'still deactivates the context' do - described_class.on_finish(response) + on_finish expect(Datadog::AppSec::Context).to have_received(:deactivate) end end diff --git a/test/datadog/lambda/inferred_span.spec.rb b/test/datadog/lambda/inferred_span.spec.rb new file mode 100644 index 0000000..6e50639 --- /dev/null +++ b/test/datadog/lambda/inferred_span.spec.rb @@ -0,0 +1,320 @@ +# frozen_string_literal: true + +require 'datadog/lambda' +require 'datadog/lambda/inferred_span' +require_relative '../lambdacontextversion' + +describe Datadog::Lambda::InferredSpan do + before { allow(Datadog::Lambda).to receive(:trace_managed_services?).and_return(true) } + + let(:request_context) { LambdaContextVersion.new } + let(:trace_digest) { nil } + + describe '.create' do + subject(:created_span) { described_class.create(event, request_context, trace_digest) } + + after { created_span&.finish unless created_span&.finished? } + + context 'when managed services is disabled' do + before { allow(Datadog::Lambda).to receive(:trace_managed_services?).and_return(false) } + + let(:event) { {'httpMethod' => 'GET', 'requestContext' => {'stage' => 'prod'}} } + + it { expect(created_span).to be_nil } + end + + context 'when event is not a Hash' do + let(:event) { 'not a hash' } + + it { expect(created_span).to be_nil } + end + + context 'when event has no requestContext' do + let(:event) { {} } + + it { expect(created_span).to be_nil } + end + + context 'when event has requestContext without stage' do + let(:event) { {'requestContext' => {'apiId' => 'abc'}} } + + it { expect(created_span).to be_nil } + end + + context 'when event has no httpMethod or routeKey' do + let(:event) { {'requestContext' => {'stage' => 'prod'}} } + + it { expect(created_span).to be_nil } + end + + context 'with API Gateway v1 event' do + let(:event) do + { + 'httpMethod' => 'GET', + 'path' => '/test', + 'requestContext' => { + 'stage' => 'prod', + 'domainName' => 'api.example.com', + 'apiId' => 'abc123', + 'resourcePath' => '/test', + 'requestTimeEpoch' => 1_700_000_000_000, + 'identity' => {'userAgent' => 'TestAgent/1.0', 'sourceIp' => '1.2.3.4'}, + }, + } + end + + it 'creates a span representing the API Gateway' do + aggregate_failures('span identity') do + expect(created_span.name).to eq('aws.apigateway') + expect(created_span.service).to eq('api.example.com') + expect(created_span.resource).to eq('GET /test') + expect(created_span.type).to eq('web') + expect(created_span.start_time).to eq(Time.at(1_700_000_000)) + end + end + + it 'sets tags for endpoint discovery' do + aggregate_failures('http tags') do + expect(created_span.get_tag('http.method')).to eq('GET') + expect(created_span.get_tag('http.url')).to eq('https://api.example.com/test') + expect(created_span.get_tag('http.route')).to eq('/test') + expect(created_span.get_tag('http.useragent')).to eq('TestAgent/1.0') + expect(created_span.get_tag('span.kind')).to eq('server') + end + end + + it 'sets tags for API Gateway resource correlation' do + aggregate_failures('gateway tags') do + expect(created_span.get_tag('apiid')).to eq('abc123') + expect(created_span.get_tag('stage')).to eq('prod') + expect(created_span.get_tag('request_id')).to eq(request_context.aws_request_id) + expect(created_span.get_tag('dd_resource_key')).to eq( + 'arn:aws:apigateway:us-east-1::/restapis/abc123/stages/prod' + ) + end + end + + it 'marks the span as inferred' do + aggregate_failures('inferred span markers') do + expect(created_span.get_metric('_dd._inferred_span')).to eq(1.0) + expect(created_span.get_tag('_inferred_span.synchronicity')).to eq('sync') + expect(created_span.get_tag('_inferred_span.tag_source')).to eq('self') + end + end + + context 'when trace_digest is provided' do + before do + allow(Datadog::Tracing).to receive(:trace).and_wrap_original do |original, *args, **kwargs| + @captured_kwargs = kwargs + original.call(*args, **kwargs.except(:continue_from)) + end + end + + let(:trace_digest) { double('trace_digest') } + + it 'continues from the existing trace' do + created_span + expect(@captured_kwargs[:continue_from]).to eq(trace_digest) + end + end + + context 'when domain is empty' do + let(:event) do + { + 'httpMethod' => 'GET', + 'path' => '/test', + 'requestContext' => { + 'stage' => 'prod', + 'domainName' => '', + 'apiId' => 'abc123', + 'resourcePath' => '/test', + 'requestTimeEpoch' => 1_700_000_000_000, + }, + } + end + + it { expect(created_span.get_tag('http.url')).to eq('/test') } + it { expect(created_span.service).not_to eq('') } + end + + context 'when requestTimeEpoch is nil' do + let(:event) do + { + 'httpMethod' => 'GET', + 'path' => '/test', + 'requestContext' => { + 'stage' => 'prod', + 'domainName' => 'api.example.com', + 'apiId' => 'abc123', + 'resourcePath' => '/test', + }, + } + end + + it { expect(created_span).not_to be_nil } + end + + context 'when apiId is empty' do + let(:event) do + { + 'httpMethod' => 'GET', + 'path' => '/test', + 'requestContext' => { + 'stage' => 'prod', + 'domainName' => 'api.example.com', + 'apiId' => '', + 'resourcePath' => '/test', + }, + } + end + + it { expect(created_span.get_tag('dd_resource_key')).to be_nil } + end + end + + context 'with API Gateway v2 event' do + let(:event) do + { + 'rawPath' => '/test', + 'routeKey' => 'GET /test', + 'requestContext' => { + 'stage' => 'prod', + 'domainName' => 'api.example.com', + 'apiId' => 'xyz789', + 'timeEpoch' => 1_700_000_000_000, + 'http' => {'method' => 'GET', 'userAgent' => 'TestAgent/2.0'}, + }, + } + end + + it 'creates a span representing the HTTP API' do + aggregate_failures('span identity') do + expect(created_span.name).to eq('aws.httpapi') + expect(created_span.service).to eq('api.example.com') + expect(created_span.resource).to eq('GET /test') + expect(created_span.type).to eq('web') + end + end + + it 'sets tags for endpoint discovery' do + aggregate_failures('http tags') do + expect(created_span.get_tag('http.method')).to eq('GET') + expect(created_span.get_tag('http.url')).to eq('https://api.example.com/test') + expect(created_span.get_tag('http.route')).to eq('/test') + expect(created_span.get_tag('http.useragent')).to eq('TestAgent/2.0') + expect(created_span.get_tag('span.kind')).to eq('server') + end + end + + it 'sets tags for API Gateway resource correlation' do + aggregate_failures('gateway tags') do + expect(created_span.get_tag('apiid')).to eq('xyz789') + expect(created_span.get_tag('stage')).to eq('prod') + expect(created_span.get_tag('dd_resource_key')).to eq( + 'arn:aws:apigateway:us-east-1::/apis/xyz789/stages/prod' + ) + end + end + + context 'when routeKey has no method prefix' do + let(:event) do + { + 'rawPath' => '/test', + 'routeKey' => '$default', + 'requestContext' => { + 'stage' => 'prod', + 'domainName' => 'api.example.com', + 'apiId' => 'xyz789', + 'timeEpoch' => 1_700_000_000_000, + 'http' => {'method' => 'GET'}, + }, + } + end + + it 'uses routeKey as-is for route' do + aggregate_failures('route and resource') do + expect(created_span.get_tag('http.route')).to eq('$default') + expect(created_span.resource).to eq('GET $default') + end + end + end + end + + context 'when an error occurs' do + before { allow(Datadog::Tracing).to receive(:trace).and_raise(StandardError, 'boom') } + + let(:event) do + { + 'httpMethod' => 'GET', + 'path' => '/test', + 'requestContext' => {'stage' => 'prod'}, + } + end + + it { expect(created_span).to be_nil } + end + end + + describe '.finish' do + subject(:finish) { described_class.finish(span, response) } + + after { span&.finish unless span&.finished? } + + let(:span) { Datadog::Tracing.trace('test.inferred') } + let(:response) { {'statusCode' => 200} } + + context 'when span is nil' do + let(:span) { nil } + + it { expect(finish).to be_nil } + end + + context 'when response contains statusCode' do + it 'tags status and finishes the span' do + finish + + aggregate_failures('status and lifecycle') do + expect(span.get_tag('http.status_code')).to eq('200') + expect(span).to be_finished + end + end + end + + context 'when response is not a Hash' do + let(:response) { 'not a hash' } + + it 'finishes without status tag' do + finish + + aggregate_failures('no status, still finished') do + expect(span.get_tag('http.status_code')).to be_nil + expect(span).to be_finished + end + end + end + + context 'when response is nil' do + let(:response) { nil } + + it 'finishes without status tag' do + finish + + aggregate_failures('no status, still finished') do + expect(span.get_tag('http.status_code')).to be_nil + expect(span).to be_finished + end + end + end + + context 'when an error occurs' do + before do + allow(span).to receive(:set_tag) + allow(span).to receive(:finish).and_raise(StandardError, 'boom') + end + + let(:span) { double('span', finished?: true) } + + it { expect { finish }.not_to raise_error } + end + end +end diff --git a/test/datadog/lambda/trace/inferred_span.spec.rb b/test/datadog/lambda/trace/inferred_span.spec.rb deleted file mode 100644 index 680ade7..0000000 --- a/test/datadog/lambda/trace/inferred_span.spec.rb +++ /dev/null @@ -1,247 +0,0 @@ -# frozen_string_literal: true - -require 'datadog/lambda' -require 'datadog/lambda/trace/inferred_span' -require_relative '../../lambdacontextversion' - -describe Datadog::Lambda::Trace::InferredSpan do - let(:request_context) { LambdaContextVersion.new } - let(:trace_digest) { nil } - - let(:apigw_v1_event) do - { - 'httpMethod' => 'GET', - 'path' => '/test', - 'requestContext' => { - 'stage' => 'prod', - 'domainName' => 'api.example.com', - 'apiId' => 'abc123', - 'resourcePath' => '/test', - 'requestTimeEpoch' => 1_700_000_000_000, - 'identity' => {'userAgent' => 'TestAgent/1.0', 'sourceIp' => '1.2.3.4'}, - }, - } - end - - let(:apigw_v2_event) do - { - 'rawPath' => '/test', - 'routeKey' => 'GET /test', - 'requestContext' => { - 'stage' => 'prod', - 'domainName' => 'api.example.com', - 'apiId' => 'xyz789', - 'timeEpoch' => 1_700_000_000_000, - 'http' => {'method' => 'GET', 'userAgent' => 'TestAgent/2.0'}, - }, - } - end - - before do - allow(Datadog::Lambda).to receive(:trace_managed_services?).and_return(true) - end - - describe '.create' do - context 'when managed services is disabled' do - before { allow(Datadog::Lambda).to receive(:trace_managed_services?).and_return(false) } - - it { expect(described_class.create(apigw_v1_event, request_context, trace_digest)).to be_nil } - end - - context 'when event is not a Hash' do - it { expect(described_class.create('not a hash', request_context, trace_digest)).to be_nil } - end - - context 'when event has no requestContext' do - it { expect(described_class.create({}, request_context, trace_digest)).to be_nil } - end - - context 'when event has requestContext without stage' do - let(:event) { {'requestContext' => {'apiId' => 'abc'}} } - - it { expect(described_class.create(event, request_context, trace_digest)).to be_nil } - end - - context 'when event has no httpMethod or routeKey' do - let(:event) { {'requestContext' => {'stage' => 'prod'}} } - - it { expect(described_class.create(event, request_context, trace_digest)).to be_nil } - end - - context 'with API Gateway v1 event' do - after { @span&.finish unless @span&.finished? } - - it 'creates an aws.apigateway span' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(@span).not_to be_nil - expect(@span.name).to eq('aws.apigateway') - end - - it 'sets service to domain name' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(@span.service).to eq('api.example.com') - end - - it 'sets resource to method + resource path' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(@span.resource).to eq('GET /test') - end - - it 'sets the inferred span metric' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(@span.get_metric('_dd._inferred_span')).to eq(1.0) - end - - it 'sets http tags' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(@span.get_tag('http.method')).to eq('GET') - expect(@span.get_tag('http.url')).to eq('https://api.example.com/test') - expect(@span.get_tag('http.route')).to eq('/test') - end - - it 'sets api gateway tags' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(@span.get_tag('apiid')).to eq('abc123') - expect(@span.get_tag('stage')).to eq('prod') - end - - it 'sets dd_resource_key with restapis prefix' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(@span.get_tag('dd_resource_key')).to eq( - 'arn:aws:apigateway:us-east-1::/restapis/abc123/stages/prod' - ) - end - - it 'sets user agent tag' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(@span.get_tag('http.useragent')).to eq('TestAgent/1.0') - end - - it 'sets start_time from requestTimeEpoch' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(@span.start_time).to eq(Time.at(1_700_000_000)) - end - - context 'when trace_digest is provided' do - let(:trace_digest) { double('trace_digest') } - let(:captured_kwargs) { {} } - - before do - allow(Datadog::Tracing).to receive(:trace).and_wrap_original do |original, *args, **kwargs| - captured_kwargs.merge!(kwargs) - original.call(*args, **kwargs.except(:continue_from)) - end - end - - it 'passes continue_from to trace' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(captured_kwargs[:continue_from]).to eq(trace_digest) - end - end - - context 'when domain is empty' do - before { apigw_v1_event['requestContext']['domainName'] = '' } - - it 'uses path as http.url' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(@span.get_tag('http.url')).to eq('/test') - end - - it 'sets service to nil' do - @span = described_class.create(apigw_v1_event, request_context, trace_digest) - expect(@span.service).not_to eq('') - end - end - end - - context 'with API Gateway v2 event' do - after { @span&.finish unless @span&.finished? } - - it 'creates an aws.httpapi span' do - @span = described_class.create(apigw_v2_event, request_context, trace_digest) - expect(@span).not_to be_nil - expect(@span.name).to eq('aws.httpapi') - end - - it 'extracts method from http context' do - @span = described_class.create(apigw_v2_event, request_context, trace_digest) - expect(@span.get_tag('http.method')).to eq('GET') - end - - it 'extracts resource path from routeKey' do - @span = described_class.create(apigw_v2_event, request_context, trace_digest) - expect(@span.get_tag('http.route')).to eq('/test') - expect(@span.resource).to eq('GET /test') - end - - it 'sets dd_resource_key with apis prefix' do - @span = described_class.create(apigw_v2_event, request_context, trace_digest) - expect(@span.get_tag('dd_resource_key')).to eq( - 'arn:aws:apigateway:us-east-1::/apis/xyz789/stages/prod' - ) - end - - it 'sets user agent from http context' do - @span = described_class.create(apigw_v2_event, request_context, trace_digest) - expect(@span.get_tag('http.useragent')).to eq('TestAgent/2.0') - end - end - - context 'when an error occurs' do - before do - allow(Datadog::Tracing).to receive(:trace).and_raise(StandardError, 'boom') - end - - it 'returns nil' do - expect(described_class.create(apigw_v1_event, request_context, trace_digest)).to be_nil - end - end - end - - describe '.finish' do - let(:inferred_span) { Datadog::Tracing.trace('test.inferred') } - - after { inferred_span&.finish unless inferred_span&.finished? } - - context 'when inferred_span is nil' do - it { expect(described_class.finish(nil, {})).to be_nil } - end - - context 'with a response containing statusCode' do - it 'sets http.status_code tag and finishes the span' do - described_class.finish(inferred_span, {'statusCode' => 200}) - expect(inferred_span.get_tag('http.status_code')).to eq('200') - expect(inferred_span).to be_finished - end - end - - context 'when response is not a Hash' do - it 'finishes the span without status tag' do - described_class.finish(inferred_span, 'not a hash') - expect(inferred_span.get_tag('http.status_code')).to be_nil - expect(inferred_span).to be_finished - end - end - - context 'when response is nil' do - it 'finishes the span without status tag' do - described_class.finish(inferred_span, nil) - expect(inferred_span.get_tag('http.status_code')).to be_nil - expect(inferred_span).to be_finished - end - end - - context 'when an error occurs' do - let(:inferred_span) { double('span', finished?: true) } - - before do - allow(inferred_span).to receive(:set_tag) - allow(inferred_span).to receive(:finish).and_raise(StandardError, 'boom') - end - - it 'does not raise' do - expect { described_class.finish(inferred_span, {}) }.not_to raise_error - end - end - end -end diff --git a/test/datadog/lambda/utils/extension.spec.rb b/test/datadog/lambda/utils/extension.spec.rb index 6e0e3a0..5fff63f 100644 --- a/test/datadog/lambda/utils/extension.spec.rb +++ b/test/datadog/lambda/utils/extension.spec.rb @@ -1,7 +1,5 @@ # frozen_string_literal: true -# rubocop:disable Metrics/BlockLength - require 'datadog/lambda' require 'net/http' require_relative '../../lambdacontextversion' @@ -12,42 +10,28 @@ 'x-datadog-parent-id' => '797643193680388254', 'x-datadog-sampling-priority' => '1', 'x-datadog-trace-id' => '4110911582297405557', - 'x-datadog-origin' => 'lambda' + 'x-datadog-origin' => 'lambda', } end - let(:trace_context) do - { - trace_id: '4110911582297405557', - sample_mode: '1', - parent_id: '797643193680388254' - } - end + let(:ctx) { LambdaContextVersion.new } describe '#send_start_invocation_request' do context 'when extension is running' do - ctx = LambdaContextVersion.new - before(:each) do - # Stub the extension_running? method to return true - allow(Datadog::Utils).to receive(:extension_running?).and_return(true) - - # Start tracing for active span + before do + allow(described_class).to receive(:extension_running?).and_return(true) @trace = Datadog::Tracing.trace('aws.lambda') end - after(:each) do - @trace.finish - end + after { @trace.finish } it 'applies trace context from extension' do - # Stub POST request to return a trace context - all_headers = Datadog::Utils.request_headers + all_headers = described_class.request_headers all_headers['lambda-runtime-aws-request-id'] = ctx.aws_request_id expect(Net::HTTP).to receive(:post) - .with(Datadog::Utils::START_INVOCATION_URI, 'null', all_headers) { headers } + .with(described_class::START_INVOCATION_URI, 'null', all_headers) { headers } - # Call the start request with an empty event - digest = Datadog::Utils.send_start_invocation_request(event: nil, request_context: ctx) + digest = described_class.send_start_invocation_request(event: nil, request_context: ctx) expect(digest.trace_id.to_s).to eq('4110911582297405557') expect(digest.span_id.to_s).to eq('797643193680388254') @@ -55,83 +39,63 @@ end it 'skips applying trace context when headers are not present' do - # Stub POST request to return a trace context - all_headers = Datadog::Utils.request_headers + all_headers = described_class.request_headers all_headers['lambda-runtime-aws-request-id'] = ctx.aws_request_id expect(Net::HTTP).to receive(:post) - .with(Datadog::Utils::START_INVOCATION_URI, 'null', all_headers) { {} } + .with(described_class::START_INVOCATION_URI, 'null', all_headers) { {} } - # Call the start request with an empty event - Datadog::Utils.send_start_invocation_request(event: nil, request_context: ctx) + described_class.send_start_invocation_request(event: nil, request_context: ctx) digest = Datadog::Tracing.active_trace.to_digest - expect(digest.trace_id.to_s).not_to eq('4110911582297405557') expect(digest.span_id.to_s).not_to eq('797643193680388254') end end context 'when extension is not running' do - ctx = LambdaContextVersion.new - it 'does nothing' do - result = Datadog::Utils.send_start_invocation_request(event: nil, request_context: ctx) - expect(result).to eq(nil) - end + it { expect(described_class.send_start_invocation_request(event: nil, request_context: ctx)).to be_nil } end end describe '#send_end_invocation_request' do context 'when extension is running' do - ctx = LambdaContextVersion.new - before(:each) do - # Stub the extension_running? method to return true - allow(Datadog::Utils).to receive(:extension_running?).and_return(true) - - # Start tracing for active span + before do + allow(described_class).to receive(:extension_running?).and_return(true) @trace = Datadog::Tracing.trace('aws.lambda') end - after(:each) do - @trace.finish - end + after { @trace.finish } - it 'sends post request as expected' do - # Stub POST request to not do anything + it 'sends post request' do allow(Net::HTTP).to receive(:post) { nil } - - # Call the start request with an empty event - Datadog::Utils.send_end_invocation_request(response: nil, span_id: nil, request_context: ctx) + described_class.send_end_invocation_request(response: nil, span_id: nil, request_context: ctx) end - it 'does not send AppSec headers to the extension' do - @trace.set_metric('_dd.appsec.enabled', 1.0) - @trace.set_tag('_dd.appsec.json', '{"triggers":[]}') - - captured_request = nil - http_double = instance_double(Net::HTTP) - allow(Net::HTTP).to receive(:start).and_yield(http_double) - allow(http_double).to receive(:request) do |req| - captured_request = req - Net::HTTPResponse.new('1.1', '200', 'OK') + context 'when active span has appsec tags' do + before do + @trace.set_metric('_dd.appsec.enabled', 1.0) + @trace.set_tag('_dd.appsec.json', '{"triggers":[]}') + allow(Net::HTTP).to receive(:start).and_yield(http_double) + allow(http_double).to receive(:request) do |req| + @captured_request = req + Net::HTTPResponse.new('1.1', '200', 'OK') + end end - Datadog::Utils.send_end_invocation_request( - response: nil, span_id: nil, request_context: ctx - ) + let(:http_double) { instance_double(Net::HTTP) } - expect(captured_request['x-datadog-appsec-enabled']).to be_nil - expect(captured_request['x-datadog-appsec-json']).to be_nil + it 'does not forward appsec headers to the extension' do + described_class.send_end_invocation_request( + response: nil, span_id: nil, request_context: ctx + ) + expect(@captured_request['x-datadog-appsec-enabled']).to be_nil + expect(@captured_request['x-datadog-appsec-json']).to be_nil + end end end context 'when extension is not running' do - ctx = LambdaContextVersion.new - it 'does nothing' do - result = Datadog::Utils.send_end_invocation_request(response: nil, span_id: nil, request_context: ctx) - expect(result).to eq(nil) - end + it { expect(described_class.send_end_invocation_request(response: nil, span_id: nil, request_context: ctx)).to be_nil } end end end - -# rubocop:enable Metrics/BlockLength From d5940df6ecd9f2c9c0d74e14e2d1699c1c4c4f34 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Thu, 26 Mar 2026 11:51:09 +0100 Subject: [PATCH 14/24] Extract API Gateway parsers and rewrite InferredSpan as dispatcher Split the god method into ApiGatewayV1/V2 parser classes with a uniform interface, and rewire InferredSpan.create as a PARSERS dispatcher + build_span builder. Existing behavior tests pass unchanged. Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/datadog/lambda/inferred_span.rb | 119 ++++++++---------- .../lambda/inferred_span/api_gateway_v1.rb | 38 ++++++ .../lambda/inferred_span/api_gateway_v2.rb | 39 ++++++ .../inferred_span/api_gateway_v1.spec.rb | 65 ++++++++++ .../inferred_span/api_gateway_v2.spec.rb | 80 ++++++++++++ 5 files changed, 274 insertions(+), 67 deletions(-) create mode 100644 lib/datadog/lambda/inferred_span/api_gateway_v1.rb create mode 100644 lib/datadog/lambda/inferred_span/api_gateway_v2.rb create mode 100644 test/datadog/lambda/inferred_span/api_gateway_v1.spec.rb create mode 100644 test/datadog/lambda/inferred_span/api_gateway_v2.spec.rb diff --git a/lib/datadog/lambda/inferred_span.rb b/lib/datadog/lambda/inferred_span.rb index 701bd3c..8c0175c 100644 --- a/lib/datadog/lambda/inferred_span.rb +++ b/lib/datadog/lambda/inferred_span.rb @@ -1,48 +1,63 @@ # frozen_string_literal: true +require_relative 'inferred_span/api_gateway_v1' +require_relative 'inferred_span/api_gateway_v2' + module Datadog module Lambda + # Creates inferred spans representing upstream services + # in the Lambda invocation path (e.g. API Gateway). + # + # @see https://docs.datadoghq.com/tracing/trace_collection/proxy_setup/apigateway/ module InferredSpan + PARSERS = [ApiGatewayV1, ApiGatewayV2].freeze + class << self def create(event, request_context, trace_digest) return unless managed_services_enabled? - span_name = inferred_span_name(event) - return unless span_name - - rc = event['requestContext'] || {} - domain = rc['domainName'] || '' - api_id = rc['apiId'] || '' - stage = rc['stage'] || '' - - if span_name == 'aws.apigateway' - method = event['httpMethod'] - path = event['path'] || '/' - resource_path = rc['resourcePath'] || path - request_time_ms = rc['requestTimeEpoch'] - user_agent = rc.dig('identity', 'userAgent') - else - http = rc['http'] || {} - method = http['method'] - path = event['rawPath'] || '/' - resource_path = event['routeKey']&.sub(/^[A-Z]+ /, '') || path - request_time_ms = rc['timeEpoch'] - user_agent = http['userAgent'] - end + parser = parser_for(event) + return unless parser + + build_span(parser, request_context, trace_digest) + rescue StandardError => e + Datadog::Utils.logger.debug "failed to create inferred span: #{e}" + nil + end + + def finish(inferred_span, response) + return unless inferred_span + + status_code = extract_status_code(response) + inferred_span.set_tag('http.status_code', status_code.to_s) if status_code + + inferred_span.finish + rescue StandardError => e + Datadog::Utils.logger.debug "failed to finish inferred span: #{e}" + end + + private + + def parser_for(event) + klass = PARSERS.find { |parser| parser.match?(event) } + klass&.new(event) + end - resource = "#{method} #{resource_path}" - http_url = domain.empty? ? path : "https://#{domain}#{path}" + def build_span(parser, request_context, trace_digest) + resource = "#{parser.method} #{parser.resource_path}" + domain = parser.domain + http_url = domain.empty? ? parser.path : "https://#{domain}#{parser.path}" tags = { - 'http.method' => method, + 'http.method' => parser.method, 'http.url' => http_url, - 'http.route' => resource_path, - 'endpoint' => path, + 'http.route' => parser.resource_path, + 'endpoint' => parser.path, 'resource_names' => resource, 'span.kind' => 'server', - 'apiid' => api_id, - 'apiname' => api_id, - 'stage' => stage, + 'apiid' => parser.api_id, + 'apiname' => parser.api_id, + 'stage' => parser.stage, 'request_id' => request_context.aws_request_id, '_inferred_span.synchronicity' => 'sync', '_inferred_span.tag_source' => 'self', @@ -50,54 +65,24 @@ def create(event, request_context, trace_digest) arn = request_context.invoked_function_arn.to_s region = arn.split(':')[3] if arn.include?(':') - if region && !api_id.empty? && !stage.empty? - arn_path = span_name == 'aws.apigateway' ? 'restapis' : 'apis' - tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{arn_path}/#{api_id}/stages/#{stage}" + if region && !parser.api_id.empty? && !parser.stage.empty? + tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{parser.arn_path_prefix}/#{parser.api_id}/stages/#{parser.stage}" end - tags['http.useragent'] = user_agent if user_agent + tags['http.useragent'] = parser.user_agent if parser.user_agent - inferred_options = { + options = { service: domain.empty? ? nil : domain, resource: resource, type: 'web', tags: tags, } - inferred_options[:continue_from] = trace_digest if trace_digest - inferred_options[:start_time] = Time.at(request_time_ms / 1000.0) if request_time_ms + options[:continue_from] = trace_digest if trace_digest + options[:start_time] = Time.at(parser.request_time_ms / 1000.0) if parser.request_time_ms - span = Datadog::Tracing.trace(span_name, **inferred_options) + span = Datadog::Tracing.trace(parser.span_name, **options) span.set_metric('_dd._inferred_span', 1.0) span - rescue StandardError => e - Datadog::Utils.logger.debug "failed to create inferred span: #{e}" - nil - end - - def finish(inferred_span, response) - return unless inferred_span - - status_code = extract_status_code(response) - inferred_span.set_tag('http.status_code', status_code.to_s) if status_code - - inferred_span.finish - rescue StandardError => e - Datadog::Utils.logger.debug "failed to finish inferred span: #{e}" - end - - private - - def inferred_span_name(event) - return unless event.is_a?(Hash) - - rc = event['requestContext'] - return unless rc.is_a?(Hash) && rc['stage'] - - if event['httpMethod'] - 'aws.apigateway' - elsif event['routeKey'] - 'aws.httpapi' - end end def managed_services_enabled? diff --git a/lib/datadog/lambda/inferred_span/api_gateway_v1.rb b/lib/datadog/lambda/inferred_span/api_gateway_v1.rb new file mode 100644 index 0000000..472f59b --- /dev/null +++ b/lib/datadog/lambda/inferred_span/api_gateway_v1.rb @@ -0,0 +1,38 @@ +# frozen_string_literal: true + +module Datadog + module Lambda + module InferredSpan + # Parses API Gateway REST API (v1) Lambda proxy integration events + # into a uniform interface. + # + # @see https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-input-format + class ApiGatewayV1 + def self.match?(payload) + api_gateway?(payload) && payload.key?('httpMethod') + end + + private_class_method def self.api_gateway?(payload) + payload.is_a?(Hash) && + payload.key?('requestContext') && payload['requestContext'].key?('stage') + end + + def initialize(payload) + @payload = payload + @request_context = payload.fetch('requestContext', {}) + end + + def span_name = 'aws.apigateway' + def method = @payload['httpMethod'] + def path = @payload.fetch('path', '/') + def resource_path = @request_context.fetch('resourcePath', path) + def domain = @request_context.fetch('domainName', '') + def api_id = @request_context.fetch('apiId', '') + def stage = @request_context.fetch('stage', '') + def request_time_ms = @request_context['requestTimeEpoch'] + def user_agent = @request_context.dig('identity', 'userAgent') + def arn_path_prefix = 'restapis' + end + end + end +end diff --git a/lib/datadog/lambda/inferred_span/api_gateway_v2.rb b/lib/datadog/lambda/inferred_span/api_gateway_v2.rb new file mode 100644 index 0000000..274475c --- /dev/null +++ b/lib/datadog/lambda/inferred_span/api_gateway_v2.rb @@ -0,0 +1,39 @@ +# frozen_string_literal: true + +module Datadog + module Lambda + module InferredSpan + # Parses API Gateway HTTP API (v2) Lambda proxy integration events + # into a uniform interface. + # + # @see https://docs.aws.amazon.com/apigateway/latest/developerguide/http-api-develop-integrations-lambda.html#http-api-develop-integrations-lambda.proxy-format + class ApiGatewayV2 + def self.match?(payload) + api_gateway?(payload) && payload.key?('routeKey') + end + + private_class_method def self.api_gateway?(payload) + payload.is_a?(Hash) && + payload.key?('requestContext') && payload['requestContext'].key?('stage') + end + + def initialize(payload) + @payload = payload + @request_context = payload.fetch('requestContext', {}) + @http = @request_context.fetch('http', {}) + end + + def span_name = 'aws.httpapi' + def method = @http['method'] + def path = @payload.fetch('rawPath', '/') + def resource_path = @payload['routeKey']&.sub(/^[A-Z]+ /, '') || path + def domain = @request_context.fetch('domainName', '') + def api_id = @request_context.fetch('apiId', '') + def stage = @request_context.fetch('stage', '') + def request_time_ms = @request_context['timeEpoch'] + def user_agent = @http['userAgent'] + def arn_path_prefix = 'apis' + end + end + end +end diff --git a/test/datadog/lambda/inferred_span/api_gateway_v1.spec.rb b/test/datadog/lambda/inferred_span/api_gateway_v1.spec.rb new file mode 100644 index 0000000..0f0d132 --- /dev/null +++ b/test/datadog/lambda/inferred_span/api_gateway_v1.spec.rb @@ -0,0 +1,65 @@ +# frozen_string_literal: true + +require 'datadog/lambda/inferred_span/api_gateway_v1' + +RSpec.describe Datadog::Lambda::InferredSpan::ApiGatewayV1 do + subject(:parser) { described_class.new(payload) } + + let(:payload) do + { + 'httpMethod' => 'GET', + 'path' => '/users/42', + 'requestContext' => { + 'stage' => 'prod', + 'domainName' => 'api.example.com', + 'apiId' => 'abc123', + 'resourcePath' => '/users/{id}', + 'requestTimeEpoch' => 1_700_000_000_000, + 'identity' => {'userAgent' => 'TestAgent/1.0'}, + }, + } + end + + describe '.match?' do + it { expect(described_class.match?('not a hash')).to be(false) } + it { expect(described_class.match?({})).to be(false) } + it { expect(described_class.match?({'requestContext' => {'stage' => 'prod'}})).to be(false) } + it { expect(described_class.match?({'httpMethod' => 'GET'})).to be(false) } + + it 'matches a v1 proxy integration event' do + expect( + described_class.match?('httpMethod' => 'GET', 'requestContext' => {'stage' => 'prod'}) + ).to be(true) + end + end + + context 'when all fields are present' do + it { expect(parser.span_name).to eq('aws.apigateway') } + it { expect(parser.method).to eq('GET') } + it { expect(parser.path).to eq('/users/42') } + it { expect(parser.resource_path).to eq('/users/{id}') } + it { expect(parser.domain).to eq('api.example.com') } + it { expect(parser.api_id).to eq('abc123') } + it { expect(parser.stage).to eq('prod') } + it { expect(parser.request_time_ms).to eq(1_700_000_000_000) } + it { expect(parser.user_agent).to eq('TestAgent/1.0') } + it { expect(parser.arn_path_prefix).to eq('restapis') } + end + + context 'when optional fields are missing' do + let(:payload) do + { + 'httpMethod' => 'POST', + 'requestContext' => {'stage' => 'dev'}, + } + end + + it { expect(parser.path).to eq('/') } + it { expect(parser.resource_path).to eq('/') } + it { expect(parser.domain).to eq('') } + it { expect(parser.api_id).to eq('') } + it { expect(parser.stage).to eq('dev') } + it { expect(parser.request_time_ms).to be_nil } + it { expect(parser.user_agent).to be_nil } + end +end diff --git a/test/datadog/lambda/inferred_span/api_gateway_v2.spec.rb b/test/datadog/lambda/inferred_span/api_gateway_v2.spec.rb new file mode 100644 index 0000000..48dd780 --- /dev/null +++ b/test/datadog/lambda/inferred_span/api_gateway_v2.spec.rb @@ -0,0 +1,80 @@ +# frozen_string_literal: true + +require 'datadog/lambda/inferred_span/api_gateway_v2' + +RSpec.describe Datadog::Lambda::InferredSpan::ApiGatewayV2 do + subject(:parser) { described_class.new(payload) } + + let(:payload) do + { + 'rawPath' => '/users/42', + 'routeKey' => 'GET /users/{id}', + 'requestContext' => { + 'stage' => 'prod', + 'domainName' => 'api.example.com', + 'apiId' => 'xyz789', + 'timeEpoch' => 1_700_000_000_000, + 'http' => {'method' => 'GET', 'userAgent' => 'TestAgent/2.0'}, + }, + } + end + + describe '.match?' do + it { expect(described_class.match?('not a hash')).to be(false) } + it { expect(described_class.match?({})).to be(false) } + it { expect(described_class.match?({'requestContext' => {'stage' => 'prod'}})).to be(false) } + it { expect(described_class.match?({'routeKey' => 'GET /test'})).to be(false) } + + it 'matches a v2 proxy integration event' do + expect( + described_class.match?('routeKey' => 'GET /test', 'requestContext' => {'stage' => 'prod'}) + ).to be(true) + end + end + + context 'when all fields are present' do + it { expect(parser.span_name).to eq('aws.httpapi') } + it { expect(parser.method).to eq('GET') } + it { expect(parser.path).to eq('/users/42') } + it { expect(parser.resource_path).to eq('/users/{id}') } + it { expect(parser.domain).to eq('api.example.com') } + it { expect(parser.api_id).to eq('xyz789') } + it { expect(parser.stage).to eq('prod') } + it { expect(parser.request_time_ms).to eq(1_700_000_000_000) } + it { expect(parser.user_agent).to eq('TestAgent/2.0') } + it { expect(parser.arn_path_prefix).to eq('apis') } + end + + context 'when routeKey has no method prefix' do + let(:payload) do + { + 'rawPath' => '/test', + 'routeKey' => '$default', + 'requestContext' => { + 'stage' => 'prod', + 'http' => {'method' => 'GET'}, + }, + } + end + + it { expect(parser.resource_path).to eq('$default') } + end + + context 'when optional fields are missing' do + let(:payload) do + { + 'routeKey' => 'POST /data', + 'requestContext' => {'stage' => 'dev'}, + } + end + + it { expect(parser.path).to eq('/') } + it { expect(parser.resource_path).to eq('/data') } + it { expect(parser.domain).to eq('') } + it { expect(parser.api_id).to eq('') } + it { expect(parser.stage).to eq('dev') } + it { expect(parser.request_time_ms).to be_nil } + it { expect(parser.user_agent).to be_nil } + it { expect(parser.method).to be_nil } + end +end From cf7d919d621282c29d91f8da5652f30c03142ad4 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Thu, 26 Mar 2026 14:31:54 +0100 Subject: [PATCH 15/24] Remove InferredSpan.finish wrapper and http.status_code tag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit http.status_code on the inferred span is not in the RFC and not checked by system tests. Without it, finish() just delegates to span.finish — the listener now calls @inferred_span&.finish directly. Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/datadog/lambda/inferred_span.rb | 17 ------- lib/datadog/lambda/trace/listener.rb | 2 +- test/datadog/lambda/inferred_span.spec.rb | 62 ----------------------- 3 files changed, 1 insertion(+), 80 deletions(-) diff --git a/lib/datadog/lambda/inferred_span.rb b/lib/datadog/lambda/inferred_span.rb index 8c0175c..494563d 100644 --- a/lib/datadog/lambda/inferred_span.rb +++ b/lib/datadog/lambda/inferred_span.rb @@ -25,17 +25,6 @@ def create(event, request_context, trace_digest) nil end - def finish(inferred_span, response) - return unless inferred_span - - status_code = extract_status_code(response) - inferred_span.set_tag('http.status_code', status_code.to_s) if status_code - - inferred_span.finish - rescue StandardError => e - Datadog::Utils.logger.debug "failed to finish inferred span: #{e}" - end - private def parser_for(event) @@ -88,12 +77,6 @@ def build_span(parser, request_context, trace_digest) def managed_services_enabled? Datadog::Lambda.trace_managed_services? end - - def extract_status_code(response) - return unless response.is_a?(Hash) - - response['statusCode'] - end end end end diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index 34bc114..9c09d1b 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -65,7 +65,7 @@ def on_end(response:, request_context:) # NOTE: lambda span must finish before inferred span (its parent) @span&.finish - Datadog::Lambda::InferredSpan.finish(@inferred_span, response) + @inferred_span&.finish @span = nil @inferred_span = nil diff --git a/test/datadog/lambda/inferred_span.spec.rb b/test/datadog/lambda/inferred_span.spec.rb index 6e50639..9b2b377 100644 --- a/test/datadog/lambda/inferred_span.spec.rb +++ b/test/datadog/lambda/inferred_span.spec.rb @@ -255,66 +255,4 @@ end end - describe '.finish' do - subject(:finish) { described_class.finish(span, response) } - - after { span&.finish unless span&.finished? } - - let(:span) { Datadog::Tracing.trace('test.inferred') } - let(:response) { {'statusCode' => 200} } - - context 'when span is nil' do - let(:span) { nil } - - it { expect(finish).to be_nil } - end - - context 'when response contains statusCode' do - it 'tags status and finishes the span' do - finish - - aggregate_failures('status and lifecycle') do - expect(span.get_tag('http.status_code')).to eq('200') - expect(span).to be_finished - end - end - end - - context 'when response is not a Hash' do - let(:response) { 'not a hash' } - - it 'finishes without status tag' do - finish - - aggregate_failures('no status, still finished') do - expect(span.get_tag('http.status_code')).to be_nil - expect(span).to be_finished - end - end - end - - context 'when response is nil' do - let(:response) { nil } - - it 'finishes without status tag' do - finish - - aggregate_failures('no status, still finished') do - expect(span.get_tag('http.status_code')).to be_nil - expect(span).to be_finished - end - end - end - - context 'when an error occurs' do - before do - allow(span).to receive(:set_tag) - allow(span).to receive(:finish).and_raise(StandardError, 'boom') - end - - let(:span) { double('span', finished?: true) } - - it { expect { finish }.not_to raise_error } - end - end end From 4b2b2180c6140d2700c756afd54e99ab35381209 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Thu, 26 Mar 2026 17:30:34 +0100 Subject: [PATCH 16/24] Make AppSec.on_start accept trace and span as explicit keyword arguments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Removes hidden dependency on Datadog::Tracing.active_trace and active_span — the listener now passes both explicitly. Span selection policy (prefer inferred span over lambda span) moves to the caller. Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/datadog/lambda/appsec.rb | 5 +- lib/datadog/lambda/trace/listener.rb | 2 +- test/datadog/lambda/appsec.spec.rb | 96 ++++++++++++++-------------- 3 files changed, 50 insertions(+), 53 deletions(-) diff --git a/lib/datadog/lambda/appsec.rb b/lib/datadog/lambda/appsec.rb index bb7b6ca..ef13cd2 100644 --- a/lib/datadog/lambda/appsec.rb +++ b/lib/datadog/lambda/appsec.rb @@ -4,12 +4,9 @@ module Datadog module Lambda module AppSec class << self - def on_start(event, span = nil) + def on_start(event, trace:, span:) return unless enabled? - trace = Datadog::Tracing.active_trace - span ||= Datadog::Tracing.active_span - create_context(trace, span) return unless Datadog::AppSec::Context.active diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index 9c09d1b..20818fe 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -53,7 +53,7 @@ def on_start(event:, request_context:, cold_start:) @span = Datadog::Tracing.trace('aws.lambda', **options) Datadog::Trace.apply_datadog_trace_context(Datadog::Trace.trace_context) - Datadog::Lambda::AppSec.on_start(event, @inferred_span) + Datadog::Lambda::AppSec.on_start(event, trace: Datadog::Tracing.active_trace, span: @inferred_span || @span) end # rubocop:enable Metrics/AbcSize diff --git a/test/datadog/lambda/appsec.spec.rb b/test/datadog/lambda/appsec.spec.rb index 20d0687..84ec946 100644 --- a/test/datadog/lambda/appsec.spec.rb +++ b/test/datadog/lambda/appsec.spec.rb @@ -3,18 +3,13 @@ require 'datadog/lambda' require 'datadog/lambda/appsec' -describe Datadog::Lambda::AppSec do +RSpec.describe Datadog::Lambda::AppSec do before do allow(Datadog::AppSec::Instrumentation).to receive(:gateway).and_return(gateway) allow(gateway).to receive(:push) - allow(Datadog::Tracing).to receive(:active_trace).and_return(active_trace) end let(:gateway) { double('gateway') } - let(:active_trace) { double('active_trace') } - let(:span) { double('span', set_metric: nil) } - let(:runner) { double('runner') } - let(:security_engine) { double('security_engine', new_runner: runner) } let(:appsec_context) do double( 'context', @@ -25,18 +20,18 @@ end describe '.on_start' do - subject(:on_start) { described_class.on_start(event, span) } + subject(:on_start) { described_class.on_start(event, trace: trace, span: span) } let(:event) { {'httpMethod' => 'GET', 'path' => '/'} } + let(:trace) { double('trace') } + let(:span) { double('span', set_metric: nil) } context 'when appsec is disabled' do before { allow(Datadog::AppSec).to receive(:enabled?).and_return(false) } - it 'does nothing' do - aggregate_failures('no side effects') do - expect(on_start).to be_nil - expect(gateway).not_to have_received(:push) - end + it 'does not push to gateway' do + on_start + expect(gateway).not_to have_received(:push) end end @@ -44,36 +39,27 @@ before do allow(Datadog::AppSec).to receive(:enabled?).and_return(true) allow(Datadog::AppSec).to receive(:security_engine).and_return(security_engine) + allow(Datadog::AppSec::Context).to receive(:new).and_return(appsec_context) allow(Datadog::AppSec::Context).to receive(:activate) allow(Datadog::AppSec::Context).to receive(:active).and_return(appsec_context) - allow(Datadog::AppSec::Context).to receive(:new).and_return(appsec_context) end - it 'activates context, marks span, and pushes event to gateway' do + let(:security_engine) { double('security_engine', new_runner: runner) } + let(:runner) { double('runner') } + + it 'creates and activates context with provided trace and span' do on_start - aggregate_failures('context lifecycle') do - expect(Datadog::AppSec::Context).to have_received(:new).with(active_trace, span, runner) + aggregate_failures do + expect(Datadog::AppSec::Context).to have_received(:new).with(trace, span, runner) expect(Datadog::AppSec::Context).to have_received(:activate).with(appsec_context) - end - - aggregate_failures('span and gateway') do expect(span).to have_received(:set_metric).with(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) - expect(gateway).to have_received(:push).with('aws_lambda.request.start', event) end end - context 'when span is not provided' do - subject(:on_start) { described_class.on_start(event) } - - before { allow(Datadog::Tracing).to receive(:active_span).and_return(active_span) } - - let(:active_span) { double('active_span', set_metric: nil) } - - it 'falls back to active span' do - on_start - expect(Datadog::AppSec::Context).to have_received(:new).with(active_trace, active_span, runner) - end + it 'pushes event to gateway' do + on_start + expect(gateway).to have_received(:push).with('aws_lambda.request.start', event) end context 'when security_engine is nil' do @@ -85,7 +71,7 @@ it 'does not activate or push' do on_start - aggregate_failures('no side effects') do + aggregate_failures do expect(Datadog::AppSec::Context).not_to have_received(:activate) expect(gateway).not_to have_received(:push) end @@ -93,15 +79,29 @@ end context 'when trace is nil' do - before do - allow(Datadog::Tracing).to receive(:active_trace).and_return(nil) - allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) + subject(:on_start) { described_class.on_start(event, trace: nil, span: span) } + + before { allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) } + + it 'does not activate or push' do + on_start + + aggregate_failures do + expect(Datadog::AppSec::Context).not_to have_received(:activate) + expect(gateway).not_to have_received(:push) + end end + end + + context 'when span is nil' do + subject(:on_start) { described_class.on_start(event, trace: trace, span: nil) } + + before { allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) } it 'does not activate or push' do on_start - aggregate_failures('no side effects') do + aggregate_failures do expect(Datadog::AppSec::Context).not_to have_received(:activate) expect(gateway).not_to have_received(:push) end @@ -121,14 +121,12 @@ let(:response) { {'statusCode' => 200} } - context 'when no active context' do + context 'when no active context exists' do before { allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) } - it 'does nothing' do - aggregate_failures('no side effects') do - expect(on_finish).to be_nil - expect(gateway).not_to have_received(:push) - end + it 'does not push to gateway' do + on_finish + expect(gateway).not_to have_received(:push) end end @@ -139,15 +137,19 @@ allow(Datadog::AppSec::Event).to receive(:record) end - it 'processes response and cleans up' do + it 'pushes response and records events' do on_finish - aggregate_failures('gateway and event recording') do + aggregate_failures do expect(gateway).to have_received(:push).with('aws_lambda.response.start', response) expect(Datadog::AppSec::Event).to have_received(:record).with(appsec_context, request: nil) end + end + + it 'exports telemetry and deactivates' do + on_finish - aggregate_failures('telemetry and cleanup') do + aggregate_failures do expect(appsec_context).to have_received(:export_metrics) expect(appsec_context).to have_received(:export_request_telemetry) expect(Datadog::AppSec::Context).to have_received(:deactivate) @@ -161,9 +163,7 @@ it 'passes request to event recording' do on_finish - expect(Datadog::AppSec::Event).to have_received(:record).with( - appsec_context, request: request_data - ) + expect(Datadog::AppSec::Event).to have_received(:record).with(appsec_context, request: request_data) end end From 9dc1bfa69b55b93ac90c5735af04280e224f7c13 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Thu, 26 Mar 2026 17:51:15 +0100 Subject: [PATCH 17/24] Add enabled? guard to AppSec.on_finish and harden create_context Co-Authored-By: Claude Opus 4.6 --- lib/datadog/lambda/appsec.rb | 8 ++++--- lib/datadog/lambda/trace/listener.rb | 8 ++++--- test/datadog/lambda/appsec.spec.rb | 34 ++++++++++++++++++++-------- 3 files changed, 35 insertions(+), 15 deletions(-) diff --git a/lib/datadog/lambda/appsec.rb b/lib/datadog/lambda/appsec.rb index ef13cd2..2922066 100644 --- a/lib/datadog/lambda/appsec.rb +++ b/lib/datadog/lambda/appsec.rb @@ -16,11 +16,12 @@ def on_start(event, trace:, span:) end def on_finish(response) + return unless enabled? + context = Datadog::AppSec::Context.active return unless context Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.response.start', response) - Datadog::AppSec::Event.record(context, request: context.state[:request]) context.export_metrics @@ -28,7 +29,7 @@ def on_finish(response) rescue StandardError => e Datadog::Utils.logger.debug "failed to finish AppSec: #{e}" ensure - Datadog::AppSec::Context.deactivate + Datadog::AppSec::Context.deactivate if context end private @@ -40,9 +41,10 @@ def enabled? end def create_context(trace, span) + return if trace.nil? || span.nil? + security_engine = Datadog::AppSec.security_engine return unless security_engine - return if trace.nil? || span.nil? Datadog::AppSec::Context.activate( Datadog::AppSec::Context.new(trace, span, security_engine.new_runner) diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index 20818fe..9067616 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -29,7 +29,7 @@ def initialize(handler_name:, function_name:, patch_http:, Datadog::Trace.patch_http if patch_http end - # rubocop:disable Metrics/AbcSize + # rubocop:disable Metrics/AbcSize, Metrics/MethodLength def on_start(event:, request_context:, cold_start:) trace_context = Datadog::Trace.extract_trace_context(event, @merge_xray_traces) Datadog::Trace.trace_context = trace_context @@ -53,9 +53,11 @@ def on_start(event:, request_context:, cold_start:) @span = Datadog::Tracing.trace('aws.lambda', **options) Datadog::Trace.apply_datadog_trace_context(Datadog::Trace.trace_context) - Datadog::Lambda::AppSec.on_start(event, trace: Datadog::Tracing.active_trace, span: @inferred_span || @span) + Datadog::Lambda::AppSec.on_start( + event, trace: Datadog::Tracing.active_trace, span: @inferred_span || @span + ) end - # rubocop:enable Metrics/AbcSize + # rubocop:enable Metrics/AbcSize, Metrics/MethodLength def on_end(response:, request_context:) Datadog::Lambda::AppSec.on_finish(response) diff --git a/test/datadog/lambda/appsec.spec.rb b/test/datadog/lambda/appsec.spec.rb index 84ec946..51f6aa1 100644 --- a/test/datadog/lambda/appsec.spec.rb +++ b/test/datadog/lambda/appsec.spec.rb @@ -9,10 +9,10 @@ allow(gateway).to receive(:push) end - let(:gateway) { double('gateway') } + let(:gateway) { instance_double(Datadog::AppSec::Instrumentation::Gateway) } let(:appsec_context) do - double( - 'context', + instance_double( + Datadog::AppSec::Context, state: {}, export_metrics: nil, export_request_telemetry: nil, @@ -23,8 +23,8 @@ subject(:on_start) { described_class.on_start(event, trace: trace, span: span) } let(:event) { {'httpMethod' => 'GET', 'path' => '/'} } - let(:trace) { double('trace') } - let(:span) { double('span', set_metric: nil) } + let(:trace) { instance_double(Datadog::Tracing::TraceOperation) } + let(:span) { instance_double(Datadog::Tracing::SpanOperation, set_metric: nil) } context 'when appsec is disabled' do before { allow(Datadog::AppSec).to receive(:enabled?).and_return(false) } @@ -44,14 +44,14 @@ allow(Datadog::AppSec::Context).to receive(:active).and_return(appsec_context) end - let(:security_engine) { double('security_engine', new_runner: runner) } - let(:runner) { double('runner') } + let(:security_engine) { instance_double(Datadog::AppSec::SecurityEngine::Engine, new_runner: waf_runner) } + let(:waf_runner) { instance_double(Datadog::AppSec::SecurityEngine::Runner) } it 'creates and activates context with provided trace and span' do on_start aggregate_failures do - expect(Datadog::AppSec::Context).to have_received(:new).with(trace, span, runner) + expect(Datadog::AppSec::Context).to have_received(:new).with(trace, span, waf_runner) expect(Datadog::AppSec::Context).to have_received(:activate).with(appsec_context) expect(span).to have_received(:set_metric).with(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) end @@ -121,8 +121,23 @@ let(:response) { {'statusCode' => 200} } + context 'when appsec is disabled' do + before do + allow(Datadog::AppSec).to receive(:enabled?).and_return(false) + allow(Datadog::AppSec::Context).to receive(:active).and_return(appsec_context) + end + + it 'does not push to gateway' do + on_finish + expect(gateway).not_to have_received(:push) + end + end + context 'when no active context exists' do - before { allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) } + before do + allow(Datadog::AppSec).to receive(:enabled?).and_return(true) + allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) + end it 'does not push to gateway' do on_finish @@ -132,6 +147,7 @@ context 'when active context exists' do before do + allow(Datadog::AppSec).to receive(:enabled?).and_return(true) allow(Datadog::AppSec::Context).to receive(:active).and_return(appsec_context) allow(Datadog::AppSec::Context).to receive(:deactivate) allow(Datadog::AppSec::Event).to receive(:record) From c92e23419a6f01693b100802351fdb0240419cfe Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Thu, 26 Mar 2026 18:00:53 +0100 Subject: [PATCH 18/24] Remove managed_services_enabled? gate from InferredSpan DD_TRACE_MANAGED_SERVICES controls outbound AWS SDK call tracing, not inbound inferred spans. The guard was added by mistake during extraction. Co-Authored-By: Claude Opus 4.6 --- lib/datadog/lambda/inferred_span.rb | 5 ----- test/datadog/lambda/inferred_span.spec.rb | 10 ---------- 2 files changed, 15 deletions(-) diff --git a/lib/datadog/lambda/inferred_span.rb b/lib/datadog/lambda/inferred_span.rb index 494563d..a42e79d 100644 --- a/lib/datadog/lambda/inferred_span.rb +++ b/lib/datadog/lambda/inferred_span.rb @@ -14,8 +14,6 @@ module InferredSpan class << self def create(event, request_context, trace_digest) - return unless managed_services_enabled? - parser = parser_for(event) return unless parser @@ -74,9 +72,6 @@ def build_span(parser, request_context, trace_digest) span end - def managed_services_enabled? - Datadog::Lambda.trace_managed_services? - end end end end diff --git a/test/datadog/lambda/inferred_span.spec.rb b/test/datadog/lambda/inferred_span.spec.rb index 9b2b377..0e875de 100644 --- a/test/datadog/lambda/inferred_span.spec.rb +++ b/test/datadog/lambda/inferred_span.spec.rb @@ -5,8 +5,6 @@ require_relative '../lambdacontextversion' describe Datadog::Lambda::InferredSpan do - before { allow(Datadog::Lambda).to receive(:trace_managed_services?).and_return(true) } - let(:request_context) { LambdaContextVersion.new } let(:trace_digest) { nil } @@ -15,14 +13,6 @@ after { created_span&.finish unless created_span&.finished? } - context 'when managed services is disabled' do - before { allow(Datadog::Lambda).to receive(:trace_managed_services?).and_return(false) } - - let(:event) { {'httpMethod' => 'GET', 'requestContext' => {'stage' => 'prod'}} } - - it { expect(created_span).to be_nil } - end - context 'when event is not a Hash' do let(:event) { 'not a hash' } From aa04c33e321def2798ab663fb4b383bcb2ca8e52 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Thu, 26 Mar 2026 18:52:19 +0100 Subject: [PATCH 19/24] Address PR review: trivial fixes across formatting, comments, and specs Co-Authored-By: Claude Opus 4.6 --- lib/datadog/lambda.rb | 1 + lib/datadog/lambda/trace/listener.rb | 4 +--- lib/datadog/lambda/utils/extension.rb | 1 - test/datadog/lambda/appsec.spec.rb | 12 ++++++------ test/datadog/lambda/inferred_span.spec.rb | 7 ++++--- .../lambda/inferred_span/api_gateway_v1.spec.rb | 7 +------ .../lambda/inferred_span/api_gateway_v2.spec.rb | 7 +------ 7 files changed, 14 insertions(+), 25 deletions(-) diff --git a/lib/datadog/lambda.rb b/lib/datadog/lambda.rb index aa64607..8f22232 100644 --- a/lib/datadog/lambda.rb +++ b/lib/datadog/lambda.rb @@ -49,6 +49,7 @@ def self.configure_apm yield(c) if block_given? + # Activation is gated by AppSec.enabled? at runtime — this only registers the integration c.appsec.instrument(:aws_lambda) end end diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index 9067616..f638bb2 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -61,9 +61,7 @@ def on_start(event:, request_context:, cold_start:) def on_end(response:, request_context:) Datadog::Lambda::AppSec.on_finish(response) - Datadog::Utils.send_end_invocation_request( - response:, span_id: @span.id, request_context: - ) + Datadog::Utils.send_end_invocation_request(span_id: @span.id, response:, request_context:) # NOTE: lambda span must finish before inferred span (its parent) @span&.finish diff --git a/lib/datadog/lambda/utils/extension.rb b/lib/datadog/lambda/utils/extension.rb index 898ac45..efdf03c 100644 --- a/lib/datadog/lambda/utils/extension.rb +++ b/lib/datadog/lambda/utils/extension.rb @@ -76,7 +76,6 @@ def self.send_end_invocation_request(response:, span_id:, request_context:) # Remove Parent ID if it is the same as the Span ID request.delete(DD_PARENT_ID_HEADER) if request[DD_PARENT_ID_HEADER] == span_id.to_s - Datadog::Utils.logger.debug "End invocation request headers: #{request.to_hash}" Net::HTTP.start(END_INVOCATION_URI.host, END_INVOCATION_URI.port) do |http| diff --git a/test/datadog/lambda/appsec.spec.rb b/test/datadog/lambda/appsec.spec.rb index 51f6aa1..81c2b4c 100644 --- a/test/datadog/lambda/appsec.spec.rb +++ b/test/datadog/lambda/appsec.spec.rb @@ -50,7 +50,7 @@ it 'creates and activates context with provided trace and span' do on_start - aggregate_failures do + aggregate_failures('context lifecycle') do expect(Datadog::AppSec::Context).to have_received(:new).with(trace, span, waf_runner) expect(Datadog::AppSec::Context).to have_received(:activate).with(appsec_context) expect(span).to have_received(:set_metric).with(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) @@ -71,7 +71,7 @@ it 'does not activate or push' do on_start - aggregate_failures do + aggregate_failures('skipped activation') do expect(Datadog::AppSec::Context).not_to have_received(:activate) expect(gateway).not_to have_received(:push) end @@ -86,7 +86,7 @@ it 'does not activate or push' do on_start - aggregate_failures do + aggregate_failures('skipped activation') do expect(Datadog::AppSec::Context).not_to have_received(:activate) expect(gateway).not_to have_received(:push) end @@ -101,7 +101,7 @@ it 'does not activate or push' do on_start - aggregate_failures do + aggregate_failures('skipped activation') do expect(Datadog::AppSec::Context).not_to have_received(:activate) expect(gateway).not_to have_received(:push) end @@ -156,7 +156,7 @@ it 'pushes response and records events' do on_finish - aggregate_failures do + aggregate_failures('response processing') do expect(gateway).to have_received(:push).with('aws_lambda.response.start', response) expect(Datadog::AppSec::Event).to have_received(:record).with(appsec_context, request: nil) end @@ -165,7 +165,7 @@ it 'exports telemetry and deactivates' do on_finish - aggregate_failures do + aggregate_failures('cleanup') do expect(appsec_context).to have_received(:export_metrics) expect(appsec_context).to have_received(:export_request_telemetry) expect(Datadog::AppSec::Context).to have_received(:deactivate) diff --git a/test/datadog/lambda/inferred_span.spec.rb b/test/datadog/lambda/inferred_span.spec.rb index 0e875de..6ff3ccf 100644 --- a/test/datadog/lambda/inferred_span.spec.rb +++ b/test/datadog/lambda/inferred_span.spec.rb @@ -6,10 +6,9 @@ describe Datadog::Lambda::InferredSpan do let(:request_context) { LambdaContextVersion.new } - let(:trace_digest) { nil } describe '.create' do - subject(:created_span) { described_class.create(event, request_context, trace_digest) } + subject(:created_span) { described_class.create(event, request_context, nil) } after { created_span&.finish unless created_span&.finished? } @@ -93,6 +92,8 @@ end context 'when trace_digest is provided' do + subject(:created_span) { described_class.create(event, request_context, trace_digest) } + before do allow(Datadog::Tracing).to receive(:trace).and_wrap_original do |original, *args, **kwargs| @captured_kwargs = kwargs @@ -100,7 +101,7 @@ end end - let(:trace_digest) { double('trace_digest') } + let(:trace_digest) { instance_double(Datadog::Tracing::TraceDigest) } it 'continues from the existing trace' do created_span diff --git a/test/datadog/lambda/inferred_span/api_gateway_v1.spec.rb b/test/datadog/lambda/inferred_span/api_gateway_v1.spec.rb index 0f0d132..246b0cd 100644 --- a/test/datadog/lambda/inferred_span/api_gateway_v1.spec.rb +++ b/test/datadog/lambda/inferred_span/api_gateway_v1.spec.rb @@ -47,12 +47,7 @@ end context 'when optional fields are missing' do - let(:payload) do - { - 'httpMethod' => 'POST', - 'requestContext' => {'stage' => 'dev'}, - } - end + let(:payload) { {'httpMethod' => 'POST', 'requestContext' => {'stage' => 'dev'}} } it { expect(parser.path).to eq('/') } it { expect(parser.resource_path).to eq('/') } diff --git a/test/datadog/lambda/inferred_span/api_gateway_v2.spec.rb b/test/datadog/lambda/inferred_span/api_gateway_v2.spec.rb index 48dd780..5cd6717 100644 --- a/test/datadog/lambda/inferred_span/api_gateway_v2.spec.rb +++ b/test/datadog/lambda/inferred_span/api_gateway_v2.spec.rb @@ -61,12 +61,7 @@ end context 'when optional fields are missing' do - let(:payload) do - { - 'routeKey' => 'POST /data', - 'requestContext' => {'stage' => 'dev'}, - } - end + let(:payload) { {'routeKey' => 'POST /data', 'requestContext' => {'stage' => 'dev'}} } it { expect(parser.path).to eq('/') } it { expect(parser.resource_path).to eq('/data') } From a5941fa34c8d67fea97713f49eb82a6cc218b75c Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Fri, 27 Mar 2026 11:06:38 +0100 Subject: [PATCH 20/24] Address PR review: small structural fixes across source and specs Co-Authored-By: Claude Opus 4.6 --- lib/datadog/lambda/inferred_span.rb | 11 +-- .../lambda/inferred_span/api_gateway_v1.rb | 16 +-- .../lambda/inferred_span/api_gateway_v2.rb | 18 ++-- lib/datadog/lambda/trace/listener.rb | 8 +- test/datadog/lambda/appsec.spec.rb | 6 +- test/datadog/lambda/inferred_span.spec.rb | 97 ++++++++++--------- 6 files changed, 80 insertions(+), 76 deletions(-) diff --git a/lib/datadog/lambda/inferred_span.rb b/lib/datadog/lambda/inferred_span.rb index a42e79d..2a4b5fd 100644 --- a/lib/datadog/lambda/inferred_span.rb +++ b/lib/datadog/lambda/inferred_span.rb @@ -14,10 +14,10 @@ module InferredSpan class << self def create(event, request_context, trace_digest) - parser = parser_for(event) - return unless parser + klass = PARSERS.find { |parser| parser.match?(event) } + return unless klass - build_span(parser, request_context, trace_digest) + build_span(klass.new(event), request_context, trace_digest) rescue StandardError => e Datadog::Utils.logger.debug "failed to create inferred span: #{e}" nil @@ -25,11 +25,6 @@ def create(event, request_context, trace_digest) private - def parser_for(event) - klass = PARSERS.find { |parser| parser.match?(event) } - klass&.new(event) - end - def build_span(parser, request_context, trace_digest) resource = "#{parser.method} #{parser.resource_path}" domain = parser.domain diff --git a/lib/datadog/lambda/inferred_span/api_gateway_v1.rb b/lib/datadog/lambda/inferred_span/api_gateway_v1.rb index 472f59b..ced9c4c 100644 --- a/lib/datadog/lambda/inferred_span/api_gateway_v1.rb +++ b/lib/datadog/lambda/inferred_span/api_gateway_v1.rb @@ -8,13 +8,17 @@ module InferredSpan # # @see https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-input-format class ApiGatewayV1 - def self.match?(payload) - api_gateway?(payload) && payload.key?('httpMethod') - end + class << self + def match?(payload) + api_gateway?(payload) && payload.key?('httpMethod') + end + + private - private_class_method def self.api_gateway?(payload) - payload.is_a?(Hash) && - payload.key?('requestContext') && payload['requestContext'].key?('stage') + def api_gateway?(payload) + payload.is_a?(Hash) && + payload.key?('requestContext') && payload['requestContext'].key?('stage') + end end def initialize(payload) diff --git a/lib/datadog/lambda/inferred_span/api_gateway_v2.rb b/lib/datadog/lambda/inferred_span/api_gateway_v2.rb index 274475c..cf02fd1 100644 --- a/lib/datadog/lambda/inferred_span/api_gateway_v2.rb +++ b/lib/datadog/lambda/inferred_span/api_gateway_v2.rb @@ -8,13 +8,17 @@ module InferredSpan # # @see https://docs.aws.amazon.com/apigateway/latest/developerguide/http-api-develop-integrations-lambda.html#http-api-develop-integrations-lambda.proxy-format class ApiGatewayV2 - def self.match?(payload) - api_gateway?(payload) && payload.key?('routeKey') - end + class << self + def match?(payload) + api_gateway?(payload) && payload.key?('routeKey') + end + + private - private_class_method def self.api_gateway?(payload) - payload.is_a?(Hash) && - payload.key?('requestContext') && payload['requestContext'].key?('stage') + def api_gateway?(payload) + payload.is_a?(Hash) && + payload.key?('requestContext') && payload['requestContext'].key?('stage') + end end def initialize(payload) @@ -26,7 +30,7 @@ def initialize(payload) def span_name = 'aws.httpapi' def method = @http['method'] def path = @payload.fetch('rawPath', '/') - def resource_path = @payload['routeKey']&.sub(/^[A-Z]+ /, '') || path + def resource_path = @payload['routeKey']&.sub(/\A[A-Z]+ /, '') || path def domain = @request_context.fetch('domainName', '') def api_id = @request_context.fetch('apiId', '') def stage = @request_context.fetch('stage', '') diff --git a/lib/datadog/lambda/trace/listener.rb b/lib/datadog/lambda/trace/listener.rb index f638bb2..226dfd3 100644 --- a/lib/datadog/lambda/trace/listener.rb +++ b/lib/datadog/lambda/trace/listener.rb @@ -23,14 +23,15 @@ def initialize(handler_name:, function_name:, patch_http:, @handler_name = handler_name @function_name = function_name @merge_xray_traces = merge_xray_traces - @span = nil - @inferred_span = nil Datadog::Trace.patch_http if patch_http end # rubocop:disable Metrics/AbcSize, Metrics/MethodLength def on_start(event:, request_context:, cold_start:) + @span = nil + @inferred_span = nil + trace_context = Datadog::Trace.extract_trace_context(event, @merge_xray_traces) Datadog::Trace.trace_context = trace_context Datadog::Utils.logger.debug "extracted trace context #{trace_context}" @@ -66,9 +67,6 @@ def on_end(response:, request_context:) # NOTE: lambda span must finish before inferred span (its parent) @span&.finish @inferred_span&.finish - - @span = nil - @inferred_span = nil end private diff --git a/test/datadog/lambda/appsec.spec.rb b/test/datadog/lambda/appsec.spec.rb index 81c2b4c..4c4dd21 100644 --- a/test/datadog/lambda/appsec.spec.rb +++ b/test/datadog/lambda/appsec.spec.rb @@ -68,7 +68,7 @@ allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) end - it 'does not activate or push' do + it 'skips context activation and gateway push' do on_start aggregate_failures('skipped activation') do @@ -83,7 +83,7 @@ before { allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) } - it 'does not activate or push' do + it 'skips context activation and gateway push' do on_start aggregate_failures('skipped activation') do @@ -98,7 +98,7 @@ before { allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) } - it 'does not activate or push' do + it 'skips context activation and gateway push' do on_start aggregate_failures('skipped activation') do diff --git a/test/datadog/lambda/inferred_span.spec.rb b/test/datadog/lambda/inferred_span.spec.rb index 6ff3ccf..d1ce99d 100644 --- a/test/datadog/lambda/inferred_span.spec.rb +++ b/test/datadog/lambda/inferred_span.spec.rb @@ -5,35 +5,39 @@ require_relative '../lambdacontextversion' describe Datadog::Lambda::InferredSpan do - let(:request_context) { LambdaContextVersion.new } + let(:request_context) do + instance_double( + LambdaContextVersion, + aws_request_id: 'test-request-id', + invoked_function_arn: 'arn:aws:lambda:us-east-1:123456789:function:test-function' + ) + end describe '.create' do - subject(:created_span) { described_class.create(event, request_context, nil) } - - after { created_span&.finish unless created_span&.finished? } + subject(:span) { described_class.create(event, request_context, nil) } context 'when event is not a Hash' do let(:event) { 'not a hash' } - it { expect(created_span).to be_nil } + it { expect(span).to be_nil } end context 'when event has no requestContext' do let(:event) { {} } - it { expect(created_span).to be_nil } + it { expect(span).to be_nil } end context 'when event has requestContext without stage' do let(:event) { {'requestContext' => {'apiId' => 'abc'}} } - it { expect(created_span).to be_nil } + it { expect(span).to be_nil } end context 'when event has no httpMethod or routeKey' do let(:event) { {'requestContext' => {'stage' => 'prod'}} } - it { expect(created_span).to be_nil } + it { expect(span).to be_nil } end context 'with API Gateway v1 event' do @@ -54,30 +58,30 @@ it 'creates a span representing the API Gateway' do aggregate_failures('span identity') do - expect(created_span.name).to eq('aws.apigateway') - expect(created_span.service).to eq('api.example.com') - expect(created_span.resource).to eq('GET /test') - expect(created_span.type).to eq('web') - expect(created_span.start_time).to eq(Time.at(1_700_000_000)) + expect(span.name).to eq('aws.apigateway') + expect(span.service).to eq('api.example.com') + expect(span.resource).to eq('GET /test') + expect(span.type).to eq('web') + expect(span.start_time).to eq(Time.at(1_700_000_000)) end end it 'sets tags for endpoint discovery' do aggregate_failures('http tags') do - expect(created_span.get_tag('http.method')).to eq('GET') - expect(created_span.get_tag('http.url')).to eq('https://api.example.com/test') - expect(created_span.get_tag('http.route')).to eq('/test') - expect(created_span.get_tag('http.useragent')).to eq('TestAgent/1.0') - expect(created_span.get_tag('span.kind')).to eq('server') + expect(span.get_tag('http.method')).to eq('GET') + expect(span.get_tag('http.url')).to eq('https://api.example.com/test') + expect(span.get_tag('http.route')).to eq('/test') + expect(span.get_tag('http.useragent')).to eq('TestAgent/1.0') + expect(span.get_tag('span.kind')).to eq('server') end end it 'sets tags for API Gateway resource correlation' do aggregate_failures('gateway tags') do - expect(created_span.get_tag('apiid')).to eq('abc123') - expect(created_span.get_tag('stage')).to eq('prod') - expect(created_span.get_tag('request_id')).to eq(request_context.aws_request_id) - expect(created_span.get_tag('dd_resource_key')).to eq( + expect(span.get_tag('apiid')).to eq('abc123') + expect(span.get_tag('stage')).to eq('prod') + expect(span.get_tag('request_id')).to eq('test-request-id') + expect(span.get_tag('dd_resource_key')).to eq( 'arn:aws:apigateway:us-east-1::/restapis/abc123/stages/prod' ) end @@ -85,14 +89,14 @@ it 'marks the span as inferred' do aggregate_failures('inferred span markers') do - expect(created_span.get_metric('_dd._inferred_span')).to eq(1.0) - expect(created_span.get_tag('_inferred_span.synchronicity')).to eq('sync') - expect(created_span.get_tag('_inferred_span.tag_source')).to eq('self') + expect(span.get_metric('_dd._inferred_span')).to eq(1.0) + expect(span.get_tag('_inferred_span.synchronicity')).to eq('sync') + expect(span.get_tag('_inferred_span.tag_source')).to eq('self') end end context 'when trace_digest is provided' do - subject(:created_span) { described_class.create(event, request_context, trace_digest) } + subject(:span) { described_class.create(event, request_context, trace_digest) } before do allow(Datadog::Tracing).to receive(:trace).and_wrap_original do |original, *args, **kwargs| @@ -104,7 +108,7 @@ let(:trace_digest) { instance_double(Datadog::Tracing::TraceDigest) } it 'continues from the existing trace' do - created_span + span expect(@captured_kwargs[:continue_from]).to eq(trace_digest) end end @@ -124,8 +128,8 @@ } end - it { expect(created_span.get_tag('http.url')).to eq('/test') } - it { expect(created_span.service).not_to eq('') } + it { expect(span.get_tag('http.url')).to eq('/test') } + it { expect(span.service).not_to eq('') } end context 'when requestTimeEpoch is nil' do @@ -142,7 +146,7 @@ } end - it { expect(created_span).not_to be_nil } + it { expect(span).not_to be_nil } end context 'when apiId is empty' do @@ -159,7 +163,7 @@ } end - it { expect(created_span.get_tag('dd_resource_key')).to be_nil } + it { expect(span.get_tag('dd_resource_key')).to be_nil } end end @@ -180,28 +184,28 @@ it 'creates a span representing the HTTP API' do aggregate_failures('span identity') do - expect(created_span.name).to eq('aws.httpapi') - expect(created_span.service).to eq('api.example.com') - expect(created_span.resource).to eq('GET /test') - expect(created_span.type).to eq('web') + expect(span.name).to eq('aws.httpapi') + expect(span.service).to eq('api.example.com') + expect(span.resource).to eq('GET /test') + expect(span.type).to eq('web') end end it 'sets tags for endpoint discovery' do aggregate_failures('http tags') do - expect(created_span.get_tag('http.method')).to eq('GET') - expect(created_span.get_tag('http.url')).to eq('https://api.example.com/test') - expect(created_span.get_tag('http.route')).to eq('/test') - expect(created_span.get_tag('http.useragent')).to eq('TestAgent/2.0') - expect(created_span.get_tag('span.kind')).to eq('server') + expect(span.get_tag('http.method')).to eq('GET') + expect(span.get_tag('http.url')).to eq('https://api.example.com/test') + expect(span.get_tag('http.route')).to eq('/test') + expect(span.get_tag('http.useragent')).to eq('TestAgent/2.0') + expect(span.get_tag('span.kind')).to eq('server') end end it 'sets tags for API Gateway resource correlation' do aggregate_failures('gateway tags') do - expect(created_span.get_tag('apiid')).to eq('xyz789') - expect(created_span.get_tag('stage')).to eq('prod') - expect(created_span.get_tag('dd_resource_key')).to eq( + expect(span.get_tag('apiid')).to eq('xyz789') + expect(span.get_tag('stage')).to eq('prod') + expect(span.get_tag('dd_resource_key')).to eq( 'arn:aws:apigateway:us-east-1::/apis/xyz789/stages/prod' ) end @@ -224,8 +228,8 @@ it 'uses routeKey as-is for route' do aggregate_failures('route and resource') do - expect(created_span.get_tag('http.route')).to eq('$default') - expect(created_span.resource).to eq('GET $default') + expect(span.get_tag('http.route')).to eq('$default') + expect(span.resource).to eq('GET $default') end end end @@ -242,8 +246,7 @@ } end - it { expect(created_span).to be_nil } + it { expect(span).to be_nil } end end - end From 954a0e754dacb26b95f49de60fba8d69334e2e16 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Fri, 27 Mar 2026 12:36:41 +0100 Subject: [PATCH 21/24] Address PR review: rename PARSERS to EVENT_SOURCES and extract ARN constants Co-Authored-By: Claude Opus 4.6 --- lib/datadog/lambda/inferred_span.rb | 42 +++++++++++++++-------------- 1 file changed, 22 insertions(+), 20 deletions(-) diff --git a/lib/datadog/lambda/inferred_span.rb b/lib/datadog/lambda/inferred_span.rb index 2a4b5fd..10c0fe6 100644 --- a/lib/datadog/lambda/inferred_span.rb +++ b/lib/datadog/lambda/inferred_span.rb @@ -10,11 +10,13 @@ module Lambda # # @see https://docs.datadoghq.com/tracing/trace_collection/proxy_setup/apigateway/ module InferredSpan - PARSERS = [ApiGatewayV1, ApiGatewayV2].freeze + EVENT_SOURCES = [ApiGatewayV1, ApiGatewayV2].freeze + ARN_REGION_INDEX = 3 + ARN_SPLIT_LIMIT = 5 class << self def create(event, request_context, trace_digest) - klass = PARSERS.find { |parser| parser.match?(event) } + klass = EVENT_SOURCES.find { |event_source| event_source.match?(event) } return unless klass build_span(klass.new(event), request_context, trace_digest) @@ -25,44 +27,44 @@ def create(event, request_context, trace_digest) private - def build_span(parser, request_context, trace_digest) - resource = "#{parser.method} #{parser.resource_path}" - domain = parser.domain - http_url = domain.empty? ? parser.path : "https://#{domain}#{parser.path}" + def build_span(event_source, request_context, trace_digest) + resource = "#{event_source.method} #{event_source.resource_path}" + domain = event_source.domain + http_url = domain.empty? ? event_source.path : "https://#{domain}#{event_source.path}" tags = { - 'http.method' => parser.method, + 'http.method' => event_source.method, 'http.url' => http_url, - 'http.route' => parser.resource_path, - 'endpoint' => parser.path, + 'http.route' => event_source.resource_path, + 'endpoint' => event_source.path, 'resource_names' => resource, 'span.kind' => 'server', - 'apiid' => parser.api_id, - 'apiname' => parser.api_id, - 'stage' => parser.stage, + 'apiid' => event_source.api_id, + 'apiname' => event_source.api_id, + 'stage' => event_source.stage, 'request_id' => request_context.aws_request_id, '_inferred_span.synchronicity' => 'sync', - '_inferred_span.tag_source' => 'self', + '_inferred_span.tag_source' => 'self' } arn = request_context.invoked_function_arn.to_s - region = arn.split(':')[3] if arn.include?(':') - if region && !parser.api_id.empty? && !parser.stage.empty? - tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{parser.arn_path_prefix}/#{parser.api_id}/stages/#{parser.stage}" + region = arn.split(':', ARN_SPLIT_LIMIT)[ARN_REGION_INDEX] if arn.include?(':') + if region && !event_source.api_id.empty? && !event_source.stage.empty? + tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{event_source.arn_path_prefix}/#{event_source.api_id}/stages/#{event_source.stage}" end - tags['http.useragent'] = parser.user_agent if parser.user_agent + tags['http.useragent'] = event_source.user_agent if event_source.user_agent options = { service: domain.empty? ? nil : domain, resource: resource, type: 'web', - tags: tags, + tags: tags } options[:continue_from] = trace_digest if trace_digest - options[:start_time] = Time.at(parser.request_time_ms / 1000.0) if parser.request_time_ms + options[:start_time] = Time.at(event_source.request_time_ms / 1000.0) if event_source.request_time_ms - span = Datadog::Tracing.trace(parser.span_name, **options) + span = Datadog::Tracing.trace(event_source.span_name, **options) span.set_metric('_dd._inferred_span', 1.0) span end From 44beb0acb48cb70cc8aadf5fdaa0c1bf646176c5 Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Fri, 27 Mar 2026 14:42:35 +0100 Subject: [PATCH 22/24] Address PR review: test behavior not implementation in appsec spec Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/datadog/lambda/inferred_span.rb | 1 - test/datadog/lambda/appsec.spec.rb | 9 ++------- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/lib/datadog/lambda/inferred_span.rb b/lib/datadog/lambda/inferred_span.rb index 10c0fe6..3c118ce 100644 --- a/lib/datadog/lambda/inferred_span.rb +++ b/lib/datadog/lambda/inferred_span.rb @@ -68,7 +68,6 @@ def build_span(event_source, request_context, trace_digest) span.set_metric('_dd._inferred_span', 1.0) span end - end end end diff --git a/test/datadog/lambda/appsec.spec.rb b/test/datadog/lambda/appsec.spec.rb index 4c4dd21..a078d07 100644 --- a/test/datadog/lambda/appsec.spec.rb +++ b/test/datadog/lambda/appsec.spec.rb @@ -47,14 +47,9 @@ let(:security_engine) { instance_double(Datadog::AppSec::SecurityEngine::Engine, new_runner: waf_runner) } let(:waf_runner) { instance_double(Datadog::AppSec::SecurityEngine::Runner) } - it 'creates and activates context with provided trace and span' do + it 'marks span as appsec-enabled' do on_start - - aggregate_failures('context lifecycle') do - expect(Datadog::AppSec::Context).to have_received(:new).with(trace, span, waf_runner) - expect(Datadog::AppSec::Context).to have_received(:activate).with(appsec_context) - expect(span).to have_received(:set_metric).with(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) - end + expect(span).to have_received(:set_metric).with(Datadog::AppSec::Ext::TAG_APPSEC_ENABLED, 1) end it 'pushes event to gateway' do From eee4562b019d27a99a395cb763c4cafddbe4ab9c Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Fri, 27 Mar 2026 15:50:31 +0100 Subject: [PATCH 23/24] Address PR review: replace captured kwargs with behavioral assertion in continue_from test Co-Authored-By: Claude Opus 4.6 (1M context) --- test/datadog/lambda/inferred_span.spec.rb | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/test/datadog/lambda/inferred_span.spec.rb b/test/datadog/lambda/inferred_span.spec.rb index d1ce99d..4868860 100644 --- a/test/datadog/lambda/inferred_span.spec.rb +++ b/test/datadog/lambda/inferred_span.spec.rb @@ -98,18 +98,16 @@ context 'when trace_digest is provided' do subject(:span) { described_class.create(event, request_context, trace_digest) } - before do - allow(Datadog::Tracing).to receive(:trace).and_wrap_original do |original, *args, **kwargs| - @captured_kwargs = kwargs - original.call(*args, **kwargs.except(:continue_from)) - end - end + before { allow(Datadog::Tracing).to receive(:trace).and_return(span_double) } + let(:span_double) { instance_double(Datadog::Tracing::SpanOperation, set_metric: nil) } let(:trace_digest) { instance_double(Datadog::Tracing::TraceDigest) } - it 'continues from the existing trace' do + it 'passes trace_digest as continue_from' do span - expect(@captured_kwargs[:continue_from]).to eq(trace_digest) + expect(Datadog::Tracing).to have_received(:trace).with( + 'aws.apigateway', hash_including(continue_from: trace_digest) + ) end end From 8862bfc5693114d40c39df2a9b724739d861a4ce Mon Sep 17 00:00:00 2001 From: Sergey Fedorov Date: Fri, 27 Mar 2026 16:11:27 +0100 Subject: [PATCH 24/24] Address PR review: refactor build_span with kwargs and extracted methods Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/datadog/lambda/inferred_span.rb | 40 +++++++++++++++++++---------- 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/lib/datadog/lambda/inferred_span.rb b/lib/datadog/lambda/inferred_span.rb index 3c118ce..bd8fb4c 100644 --- a/lib/datadog/lambda/inferred_span.rb +++ b/lib/datadog/lambda/inferred_span.rb @@ -19,7 +19,7 @@ def create(event, request_context, trace_digest) klass = EVENT_SOURCES.find { |event_source| event_source.match?(event) } return unless klass - build_span(klass.new(event), request_context, trace_digest) + start_span(klass.new(event), request_context: request_context, trace_digest: trace_digest) rescue StandardError => e Datadog::Utils.logger.debug "failed to create inferred span: #{e}" nil @@ -27,14 +27,12 @@ def create(event, request_context, trace_digest) private - def build_span(event_source, request_context, trace_digest) + def start_span(event_source, request_context:, trace_digest:) resource = "#{event_source.method} #{event_source.resource_path}" - domain = event_source.domain - http_url = domain.empty? ? event_source.path : "https://#{domain}#{event_source.path}" tags = { 'http.method' => event_source.method, - 'http.url' => http_url, + 'http.url' => http_url_for(event_source), 'http.route' => event_source.resource_path, 'endpoint' => event_source.path, 'resource_names' => resource, @@ -47,27 +45,43 @@ def build_span(event_source, request_context, trace_digest) '_inferred_span.tag_source' => 'self' } - arn = request_context.invoked_function_arn.to_s - region = arn.split(':', ARN_SPLIT_LIMIT)[ARN_REGION_INDEX] if arn.include?(':') - if region && !event_source.api_id.empty? && !event_source.stage.empty? - tags['dd_resource_key'] = "arn:aws:apigateway:#{region}::/#{event_source.arn_path_prefix}/#{event_source.api_id}/stages/#{event_source.stage}" - end - + resource_key = resource_key_for(event_source, request_context) + tags['dd_resource_key'] = resource_key if resource_key tags['http.useragent'] = event_source.user_agent if event_source.user_agent options = { - service: domain.empty? ? nil : domain, + service: event_source.domain.empty? ? nil : event_source.domain, resource: resource, type: 'web', tags: tags } options[:continue_from] = trace_digest if trace_digest - options[:start_time] = Time.at(event_source.request_time_ms / 1000.0) if event_source.request_time_ms + options[:start_time] = ms_to_time(event_source.request_time_ms) if event_source.request_time_ms span = Datadog::Tracing.trace(event_source.span_name, **options) span.set_metric('_dd._inferred_span', 1.0) span end + + def http_url_for(event_source) + return event_source.path if event_source.domain.empty? + + "https://#{event_source.domain}#{event_source.path}" + end + + def resource_key_for(event_source, request_context) + arn = request_context.invoked_function_arn.to_s + return unless arn.include?(':') + + region = arn.split(':', ARN_SPLIT_LIMIT)[ARN_REGION_INDEX] + return if event_source.api_id.empty? || event_source.stage.empty? + + "arn:aws:apigateway:#{region}::/#{event_source.arn_path_prefix}/#{event_source.api_id}/stages/#{event_source.stage}" + end + + def ms_to_time(ms) + Time.at(ms / 1000.0) + end end end end