[APPSEC-61865] (WIP) Add inferred span and AppSec calls for API Gateway

.rubocop.yml

@@ -1,5 +1,7 @@
 AllCops:
   TargetRubyVersion: 3.2
+  Exclude:
+    - 'test/**/*'
 
 Metrics/MethodLength:
   Max: 20

lib/datadog/lambda.rb

@@ -48,6 +48,9 @@ def self.configure_apm
         c.tracing.instrument :aws if trace_managed_services?
 
         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
 

lib/datadog/lambda/appsec.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Lambda
+    module AppSec
+      class << self
+        def on_start(event, trace:, span:)
+          return unless enabled?
+
+          create_context(trace, span)
+          return unless Datadog::AppSec::Context.active
+
+          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)
+          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
+          context.export_request_telemetry
+        rescue StandardError => e
+          Datadog::Utils.logger.debug "failed to finish AppSec: #{e}"
+        ensure
+          Datadog::AppSec::Context.deactivate if context
+        end
+
+        private
+
+        def enabled?
+          defined?(Datadog::AppSec) &&
+            Datadog::AppSec.respond_to?(:enabled?) &&
+            Datadog::AppSec.enabled?
+        end
+
+        def create_context(trace, span)
+          return if trace.nil? || span.nil?
+
+          security_engine = Datadog::AppSec.security_engine
+          return unless security_engine
+
+          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

lib/datadog/lambda/inferred_span.rb

@@ -0,0 +1,88 @@
+# 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
+      EVENT_SOURCES = [ApiGatewayV1, ApiGatewayV2].freeze
+      ARN_REGION_INDEX = 3
+      ARN_SPLIT_LIMIT = 5
+
+      class << self
+        def create(event, request_context, trace_digest)
+          klass = EVENT_SOURCES.find { |event_source| event_source.match?(event) }
+          return unless klass
+
+          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
+        end
+
+        private
+
+        def start_span(event_source, request_context:, trace_digest:)
+          resource = "#{event_source.method} #{event_source.resource_path}"
+
+          tags = {
+            'http.method' => event_source.method,
+            'http.url' => http_url_for(event_source),
+            'http.route' => event_source.resource_path,
+            'endpoint' => event_source.path,
+            'resource_names' => resource,
+            'span.kind' => 'server',
+            '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'
+          }
+
+          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: 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] = 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
+end

lib/datadog/lambda/inferred_span/api_gateway_v1.rb

@@ -0,0 +1,42 @@
+# 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
+        class << self
+          def match?(payload)
+            api_gateway?(payload) && payload.key?('httpMethod')
+          end
+
+          private
+
+          def api_gateway?(payload)
+            payload.is_a?(Hash) &&
+              payload.key?('requestContext') && payload['requestContext'].key?('stage')
+          end
+        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

lib/datadog/lambda/inferred_span/api_gateway_v2.rb

@@ -0,0 +1,43 @@
+# 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
+        class << self
+          def match?(payload)
+            api_gateway?(payload) && payload.key?('routeKey')
+          end
+
+          private
+
+          def api_gateway?(payload)
+            payload.is_a?(Hash) &&
+              payload.key?('requestContext') && payload['requestContext'].key?('stage')
+          end
+        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[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

lib/datadog/lambda/trace/listener.rb

@@ -11,12 +11,13 @@
 require 'datadog/lambda/trace/context'
 require 'datadog/lambda/trace/patch_http'
 require 'datadog/lambda/trace/ddtrace'
+require 'datadog/lambda/inferred_span'
+require 'datadog/lambda/appsec'
 
 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
@@ -26,8 +27,11 @@ 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:)
+        @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}"
@@ -43,19 +47,26 @@ 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)
+        @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)
 
         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
+        )
       end
-      # rubocop:enable Metrics/AbcSize
+      # rubocop:enable Metrics/AbcSize, Metrics/MethodLength
 
       def on_end(response:, request_context:)
-        Datadog::Utils.send_end_invocation_request(response:, span_id: @trace.id, request_context:)
-        @trace&.finish
+        Datadog::Lambda::AppSec.on_finish(response)
+        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
+        @inferred_span&.finish
       end
 
       private