[APPSEC-61865] (WIP) Add inferred span and AppSec calls for API Gateway by Strech · Pull Request #136 · DataDog/datadog-lambda-rb

Strech · 2026-03-25T15:23:02Z

What does this PR do?

Compute inferred span for AWS Lambda that suppose to be a service-entry span. In addition - initialize AppSec component (if possible), push events to it and flush captured security events.

Motivation

Part of Endpoint Discovery & Correlation from Inferred Spans RFC

Testing Guidelines

Additional Notes

I would need some assistance to get more context over testing.

Types of changes

Bug fix
New feature
Breaking change
Misc (docs, refactoring, dependency upgrade, etc.)

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) <noreply@anthropic.com>

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) <noreply@anthropic.com>

- 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) <noreply@anthropic.com>

…vices 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) <noreply@anthropic.com>

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) <noreply@anthropic.com>

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) <noreply@anthropic.com>

* Rewrite specs to test behavior instead of internals * Remove obsolete appsec files

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) <noreply@anthropic.com>

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) <noreply@anthropic.com>

Strech · 2026-03-26T14:25:01Z

lib/datadog/lambda/inferred_span/api_gateway_v1.rb

+        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']


No fallback & coersion?

Strech · 2026-03-26T14:25:47Z

lib/datadog/lambda/inferred_span/api_gateway_v1.rb

+        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


Makes sense to do class << self style

Strech · 2026-03-26T14:26:13Z

lib/datadog/lambda/trace/listener.rb

+        @span = nil
+        @inferred_span = nil


Is it needed within instance?

Strech · 2026-03-26T14:27:03Z

lib/datadog/lambda/trace/listener.rb

+        Datadog::Utils.send_end_invocation_request(
+          response:, span_id: @span.id, request_context:
+        )


Single-line is better, and we can reorder empty kwargs to be last

Strech · 2026-03-26T14:28:06Z

lib/datadog/lambda/inferred_span.rb

+        def parser_for(event)
+          klass = PARSERS.find { |parser| parser.match?(event) }
+          klass&.new(event)
+        end


Irresponsible method with no benefits

Strech · 2026-03-26T14:28:39Z

lib/datadog/lambda/inferred_span.rb

+          }
+
+          arn = request_context.invoked_function_arn.to_s
+          region = arn.split(':')[3] if arn.include?(':')


No limit on split and magic number 3

Strech · 2026-03-26T14:30:21Z

lib/datadog/lambda/inferred_span.rb

+        private
+
+        def parser_for(event)
+          klass = PARSERS.find { |parser| parser.match?(event) }


We call it parsers here, but in fact it's value object style, and it's better to word it that way in code and docs as they provide object oriented access to the AWS event payload

Strech · 2026-03-26T16:28:35Z

lib/datadog/lambda/inferred_span.rb

+        def managed_services_enabled?
+          Datadog::Lambda.trace_managed_services?
+        end


Is it needed at all?

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) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

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 <noreply@anthropic.com>

Strech · 2026-03-26T17:04:38Z

lib/datadog/lambda/inferred_span/api_gateway_v2.rb

+        def span_name = 'aws.httpapi'
+        def method = @http['method']
+        def path = @payload.fetch('rawPath', '/')
+        def resource_path = @payload['routeKey']&.sub(/^[A-Z]+ /, '') || path


Does it makes sense to use \A instead of ^ (security)

What this sub here is for?

Strech · 2026-03-26T17:05:27Z

lib/datadog/lambda/utils/extension.rb


      # 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
+


Does it needed?

Strech · 2026-03-26T17:06:34Z

lib/datadog/lambda/appsec.rb

+          return unless context
+
+          Datadog::AppSec::Instrumentation.gateway.push('aws_lambda.response.start', response)
+          Datadog::AppSec::Event.record(context, request: context.state[:request])


Who creates this request in the code?

Seems tests are wrong doing here and over-stubbing

Strech · 2026-03-26T17:08:03Z

lib/datadog/lambda/inferred_span.rb

+          klass&.new(event)
+        end
+
+        def build_span(parser, request_context, trace_digest)


Method is massive now, and signature could use some help of kwargs for secondary arguments

Strech · 2026-03-26T17:08:28Z

lib/datadog/lambda/inferred_span.rb

+            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


Feels like small method and a constant

Strech · 2026-03-26T17:09:22Z

lib/datadog/lambda.rb


        yield(c) if block_given?
+
+        c.appsec.instrument(:aws_lambda)


Makes sense to add comment about instrumentation guard on enabled from AppSec settings class and no auto instrumentation

Strech · 2026-03-26T17:09:51Z

test/datadog/lambda/inferred_span/api_gateway_v1.spec.rb

+      {
+        'httpMethod' => 'POST',
+        'requestContext' => {'stage' => 'dev'},
+      }


Strech · 2026-03-26T17:10:13Z

test/datadog/lambda/inferred_span/api_gateway_v2.spec.rb

+      {
+        'routeKey' => 'POST /data',
+        'requestContext' => {'stage' => 'dev'},
+      }


Strech · 2026-03-26T17:11:06Z

test/datadog/lambda/utils/extension.spec.rb

-      parent_id: '797643193680388254'
-    }
-  end
+  let(:ctx) { LambdaContextVersion.new }


what is needed from this test? Can we write a clean separate test with good standards?

Strech · 2026-03-26T17:12:41Z

test/datadog/lambda/appsec.spec.rb

+        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)


So many stubs, are all of them needed?

Strech · 2026-03-26T17:13:25Z

test/datadog/lambda/appsec.spec.rb

+
+        before { allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) }
+
+        it 'does not activate or push' do


what doesn't activate, trace?

Strech · 2026-03-26T17:13:34Z

test/datadog/lambda/appsec.spec.rb

+        it 'does not activate or push' do
+          on_start
+
+          aggregate_failures do


where is description?

Strech · 2026-03-26T17:14:14Z

test/datadog/lambda/appsec.spec.rb

+      end
+
+      context 'when context has a request in state' do
+        before { appsec_context.state[:request] = request_data }


This stubs stuff, but we miss actual push on the gateway because request data is never given in actual code

Strech · 2026-03-26T17:14:47Z

test/datadog/lambda/inferred_span.spec.rb

+  describe '.create' do
+    subject(:created_span) { described_class.create(event, request_context, trace_digest) }
+
+    after { created_span&.finish unless created_span&.finished? }


Why do we do that here, we are not coding, we controlling environment, we should not have to do safe-navigation in tests at all

Strech · 2026-03-26T17:15:12Z

test/datadog/lambda/inferred_span.spec.rb

+
+describe Datadog::Lambda::InferredSpan do
+  let(:request_context) { LambdaContextVersion.new }
+  let(:trace_digest) { nil }


violates rspec guide

Strech · 2026-03-26T17:15:32Z

test/datadog/lambda/inferred_span.spec.rb

+    end
+
+    context 'when event has no httpMethod or routeKey' do
+      let(:event) { {'requestContext' => {'stage' => 'prod'}} }


should use let/subject instead, hard to read test

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…nstants Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…in continue_from test Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Strech and others added 15 commits March 18, 2026 17:44

EXPERIMENTAL

04ef96d

EXPERIMENTAL Add appsec to the request flow

1da82aa

EXPERIMENTAL Add inferred span

0b6c3da

Remove unnecessary change

7596aee

Remove appsec request headers injection

84cb13f

Ignore test/ directory in RuboCop

9851294

Move InferredSpan outside of Trace folder

8790c58

* Rewrite specs to test behavior instead of internals * Remove obsolete appsec files

Strech commented Mar 26, 2026

View reviewed changes

Strech and others added 3 commits March 26, 2026 17:30

Add enabled? guard to AppSec.on_finish and harden create_context

9dc1bfa

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Remove managed_services_enabled? gate from InferredSpan

c92e234

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 <noreply@anthropic.com>

Strech commented Mar 26, 2026

View reviewed changes

Strech and others added 5 commits March 26, 2026 18:52

Address PR review: trivial fixes across formatting, comments, and specs

aa04c33

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Address PR review: small structural fixes across source and specs

a5941fa

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Address PR review: rename PARSERS to EVENT_SOURCES and extract ARN co…

954a0e7

…nstants Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Address PR review: test behavior not implementation in appsec spec

44beb0a

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Address PR review: replace captured kwargs with behavioral assertion …

eee4562

…in continue_from test Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Strech force-pushed the appsec-61865-add-inferred-span-for-api-gateway branch from 80f3ad5 to eee4562 Compare March 27, 2026 14:52

Address PR review: refactor build_span with kwargs and extracted methods

8862bfc

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>


		# 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


		before { allow(Datadog::AppSec::Context).to receive(:active).and_return(nil) }

		it 'does not activate or push' do

Conversation

Strech commented Mar 25, 2026

What does this PR do?

Motivation

Testing Guidelines

Additional Notes

Types of changes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant