Add HTTP span tags to the AWS Lambda invocation span

.github/CODEOWNERS

@@ -49,7 +49,7 @@
 /dd-java-agent/instrumentation/azure-functions/                         @DataDog/apm-serverless
 /dd-java-agent/instrumentation/azure-functions-1.2.2/                   @DataDog/apm-serverless
 /dd-trace-core/src/main/java/datadog/trace/lambda/                      @DataDog/apm-serverless
-/dd-trace-core/src/test/groovy/datadog/trace/lambda/                    @DataDog/apm-serverless
+/dd-trace-core/src/test/*/datadog/trace/lambda/                         @DataDog/apm-serverless
 /utils/container-utils/                                                 @DataDog/apm-serverless
 **/InferredProxy*.java                                                  @DataDog/apm-serverless
 **/InferredProxy*.groovy                                                @DataDog/apm-serverless

dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/LambdaHandlerInstrumentationTest.java

@@ -17,6 +17,7 @@
 import com.amazonaws.services.lambda.runtime.LambdaLogger;
 import datadog.trace.agent.test.AbstractInstrumentationTest;
 import datadog.trace.api.DDSpanTypes;
+import datadog.trace.api.DDTags;
 import datadog.trace.api.function.TriConsumer;
 import datadog.trace.api.function.TriFunction;
 import datadog.trace.api.gateway.Flow;
@@ -26,6 +27,7 @@
 import datadog.trace.api.gateway.SubscriptionService;
 import datadog.trace.bootstrap.ActiveSubsystems;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
 import datadog.trace.test.junit.utils.config.WithConfig;
 import java.io.ByteArrayInputStream;
@@ -416,6 +418,48 @@ void responseAndRequestCallbacksAreBothInvoked() throws IOException {
     assertTraces(trace(span().type(DDSpanTypes.SERVERLESS).error(false)));
   }
 
+  @Test
+  void invocationSpanCarriesHttpTags() throws IOException {
+    String eventJson =
+        "{"
+            + "\"resource\": \"/api/users/{id}\","
+            + "\"path\": \"/api/users/123\","
+            + "\"httpMethod\": \"GET\","
+            + "\"queryStringParameters\": {\"q\": \"hello\"},"
+            + "\"headers\": {\"Host\": \"api.example.com\","
+            + "              \"User-Agent\": \"test-agent\"},"
+            + "\"requestContext\": {"
+            + "  \"httpMethod\": \"GET\","
+            + "  \"requestId\": \"req-tags\","
+            + "  \"domainName\": \"api.example.com\","
+            + "  \"identity\": {\"sourceIp\": \"127.0.0.1\"}"
+            + "}"
+            + "}";
+
+    ByteArrayInputStream input =
+        new ByteArrayInputStream(eventJson.getBytes(StandardCharsets.UTF_8));
+    ByteArrayOutputStream output = new ByteArrayOutputStream();
+    new HandlerStreamingWithApiGwResponse().handleRequest(input, output, newContext());
+
+    assertTraces(
+        trace(
+            span()
+                .type(DDSpanTypes.SERVERLESS)
+                .error(false)
+                .tags(
+                    defaultTags(),
+                    tag("request_id", is(REQUEST_ID)),
+                    tag(Tags.HTTP_METHOD, is("GET")),
+                    // The tracer tags http.url without the query string; QueryObfuscator
+                    // obfuscates http.query.string and re-appends it as the trace is serialised
+                    tag(Tags.HTTP_URL, is("https://api.example.com/api/users/123?q=hello")),
+                    tag(DDTags.HTTP_QUERY, is("q=hello")),
+                    tag(Tags.HTTP_USER_AGENT, is("test-agent")),
+                    tag(Tags.HTTP_ROUTE, is("/api/users/{id}")),
+                    tag(Tags.HTTP_HOSTNAME, is("api.example.com")),
+                    tag(Tags.HTTP_STATUS, is(200)))));
+  }
+
   @Test
   void responseCallbacksFireBeforeRequestEnded() throws IOException {
     List<String> callOrder = new ArrayList<>();

dd-trace-core/src/main/java/datadog/trace/lambda/LambdaAppSecHandler.java

@@ -3,12 +3,16 @@
 import static datadog.trace.api.gateway.Events.EVENTS;
 import static datadog.trace.lambda.LambdaEventParser.MAX_EVENT_SIZE;
 import static datadog.trace.lambda.LambdaEventParser.buildFullPath;
+import static datadog.trace.lambda.LambdaEventParser.findHeader;
 import static datadog.trace.lambda.LambdaEventParser.parseEvent;
 import static datadog.trace.lambda.LambdaEventParser.parseJsonValue;
 import static datadog.trace.lambda.LambdaEventParser.parseResponse;
 
 import datadog.logging.RatelimitedLogger;
+import datadog.trace.api.Config;
+import datadog.trace.api.DDTags;
 import datadog.trace.api.ProductTraceSource;
+import datadog.trace.api.TagMap;
 import datadog.trace.api.appsec.AppSecContext;
 import datadog.trace.api.function.TriConsumer;
 import datadog.trace.api.gateway.BlockResponseFunction;
@@ -23,9 +27,11 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.ClientIpAddressData;
+import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter;
+import datadog.trace.bootstrap.instrumentation.api.URIUtils;
 import datadog.trace.lambda.LambdaEventParser.LambdaRequestData;
 import datadog.trace.lambda.LambdaEventParser.LambdaResponseData;
 import datadog.trace.lambda.LambdaEventParser.LambdaTriggerType;
@@ -42,8 +48,9 @@
 import org.slf4j.LoggerFactory;
 
 /**
- * Handles AppSec processing for AWS Lambda invocations. Extracts Lambda event data and invokes
- * AppSec gateway callbacks.
+ * Handles AppSec processing for AWS Lambda invocations: invokes the AppSec gateway callbacks for
+ * the event and the handler response, and derives the HTTP span tags from the event. Payload
+ * parsing is delegated to {@link LambdaEventParser}.
  */
 public class LambdaAppSecHandler {
 
@@ -55,12 +62,13 @@ public class LambdaAppSecHandler {
   private static final ThreadLocal<LambdaTriggerType> CURRENT_TRIGGER_TYPE = new ThreadLocal<>();
 
   /**
-   * Process AppSec request data at the start of a Lambda invocation. Extract event data and invokes
-   * all relevant AppSec gateway callbacks.
+   * Processes AppSec request data at the start of a Lambda invocation: invokes all relevant AppSec
+   * gateway callbacks on the parsed event, and, for recognised HTTP triggers, applies the HTTP tags
+   * to the returned context so they land on the invocation span at creation.
    *
    * @param event the Lambda event object
-   * @return AgentSpanContext containing AppSec data, or null if AppSec is disabled or processing
-   *     fails
+   * @return a {@link TagContext} carrying the AppSec request context and the HTTP tags, or null if
+   *     AppSec is disabled, the event is not a parseable payload, or processing fails
    */
   public static AgentSpanContext processRequestStart(Object event) {
     if (!ActiveSubsystems.APPSEC_ACTIVE) {
@@ -83,15 +91,28 @@ public static AgentSpanContext processRequestStart(Object event) {
         return null;
       }
       CURRENT_TRIGGER_TYPE.set(eventData.triggerType);
-      return processAppSecRequestData(eventData);
+      // v2 payloads carry the request line verbatim; the others expose the path and a decoded
+      // parameter map only, so the query string has to be rebuilt from them
+      String fullPath = eventData.rawUri;
+      if (fullPath == null && eventData.path != null) {
+        fullPath = buildFullPath(eventData.path, eventData.queryParameters);
+      }
+      LambdaURIDataAdapter uriAdapter =
+          new LambdaURIDataAdapter(fullPath, eventData.headers, eventData.host);
+      AgentSpanContext context = processAppSecRequestData(eventData, uriAdapter);
+      if (context instanceof TagContext && eventData.triggerType.isHttp()) {
+        applyHttpTags((TagContext) context, eventData, uriAdapter);
+      }
+      return context;
     } catch (Exception e) {
       log.debug("Failed to process AppSec request data", e);
       return null;
     }
   }
 
   /**
-   * Invokes the requestEnded gateway callback to add AppSec data to the span.
+   * Invokes the requestEnded gateway callback to add AppSec data to the span, propagates the
+   * sampling decision of trace-tagging rules, and clears the per-invocation state.
    *
    * @param span the current span
    */
@@ -129,8 +150,9 @@ public static void processRequestEnd(AgentSpan span) {
   }
 
   /**
-   * Process response data through WAF before the request context is closed. Extracts status code,
-   * headers, and body from the Lambda response and fires the corresponding gateway events.
+   * Processes response data through the WAF before the request context is closed: fires the
+   * response gateway events with the status code, headers and body parsed from the Lambda response,
+   * and sets {@code http.status_code} on the span. Only applies to recognised HTTP triggers.
    *
    * @param span the current span
    * @param result the Lambda handler result (expected to be a ByteArrayOutputStream)
@@ -184,6 +206,13 @@ public static void processResponseData(AgentSpan span, Object result) {
         // (statusCode remains 0, so the responseStarted guard below will not fire).
       }
 
+      // The only HTTP tag set on the exit path: the status does not exist at span creation.
+      if (responseData.statusCode > 0) {
+        span.setHttpStatusCode(responseData.statusCode);
+        boolean isError = Config.get().getHttpServerErrorStatuses().get(responseData.statusCode);
+        span.setError(isError, ErrorPriorities.HTTP_SERVER_DECORATOR);
+      }
+
       RequestContext requestContext = span.getRequestContext();
       if (requestContext == null) {
         log.debug("Span has no RequestContext, skipping response processing");
@@ -237,11 +266,13 @@ public static void processResponseData(AgentSpan span, Object result) {
   }
 
   /**
-   * Merge AppSec context data into extension context.
+   * Merges the AppSec request context data and the HTTP tags into the context returned by the
+   * Lambda Extension, which is the one that survives and seeds the invocation span.
    *
-   * @param extensionContext context from extension
-   * @param appSecContext context containing AppSec data
-   * @return merged context
+   * @param extensionContext context from the extension, may be null when no extension is in the
+   *     path
+   * @param appSecContext context returned by {@link #processRequestStart(Object)}, may be null
+   * @return the surviving context: the extension one when both are present
    */
   public static AgentSpanContext mergeContexts(
       AgentSpanContext extensionContext, AgentSpanContext appSecContext) {
@@ -261,6 +292,15 @@ public static AgentSpanContext mergeContexts(
         if (appSecData != null) {
           merged.withRequestContextDataAppSec(appSecData);
         }
+        // The extension context is the one that survives, so the HTTP tags applied to the AppSec
+        // context have to be carried over: CoreTracer copies them onto the span at creation.
+        // The AppSec-derived values win on a key collision. No collision is reachable today: the
+        // extension context only carries tags for headers mapped through
+        // DD_TRACE_REQUEST_HEADER_TAGS
+        // (ContextInterpreter.handleTags), and those would have to be mapped onto an http.* key.
+        for (TagMap.EntryReader tag : extracted.getTags()) {
+          merged.putTag(tag.tag(), tag.stringValue());
+        }
         return merged;
       }
 
@@ -271,7 +311,54 @@ public static AgentSpanContext mergeContexts(
     return extensionContext;
   }
 
-  private static AgentSpanContext processAppSecRequestData(LambdaRequestData eventData) {
+  /**
+   * Writes the HTTP tags derived from the Lambda event onto the context that will seed the
+   * invocation span. Transcribed from {@code HttpServerDecorator.doOnRequest}, minus the client IP
+   * tags, {@code span.kind} and {@code http.fragment}.
+   */
+  static void applyHttpTags(TagContext ctx, LambdaRequestData req, LambdaURIDataAdapter url) {
+    // The synthetic "WEBSOCKET" method stays inside the AppSec path; none is fabricated here.
+    if (req.method != null && req.triggerType != LambdaTriggerType.API_GATEWAY_V2_WEBSOCKET) {
+      ctx.putTag(Tags.HTTP_METHOD, req.method);
+    }
+
+    if (req.host != null) {
+      // No query string: QueryObfuscator obfuscates DDTags.HTTP_QUERY and re-appends it here.
+      ctx.putTag(
+          Tags.HTTP_URL, URIUtils.buildURL(url.scheme(), url.host(), url.port(), url.path()));
+    }
+
+    String query = url.rawQuery();
+    if (query != null && !query.isEmpty() && Config.get().isHttpServerTagQueryString()) {
+      ctx.putTag(DDTags.HTTP_QUERY, query);
+    }
+
+    String userAgent = findHeader(req.headers, "user-agent");
+    if (userAgent != null) {
+      ctx.putTag(Tags.HTTP_USER_AGENT, userAgent);
+    }
+
+    if (req.route != null) {
+      ctx.putTag(Tags.HTTP_ROUTE, req.route);
+    }
+
+    // Deliberately a different host from the one in http.url, as in the decorator
+    String forwardedHost = findHeader(req.headers, "x-forwarded-host");
+    String hostname = forwardedHost != null ? forwardedHost : req.host;
+    if (hostname != null) {
+      ctx.putTag(Tags.HTTP_HOSTNAME, hostname);
+    }
+  }
+
+  /**
+   * Fires the request-phase gateway callbacks against a {@link TemporaryRequestContext}, since the
+   * span does not exist yet, and returns the context carrying the resulting AppSec request context.
+   *
+   * @return the context to hand back to the tracer, or null if AppSec registered no {@code
+   *     requestStarted} callback
+   */
+  private static AgentSpanContext processAppSecRequestData(
+      LambdaRequestData eventData, LambdaURIDataAdapter uriAdapter) {
     AgentTracer.TracerAPI tracer = AgentTracer.get();
     Supplier<Flow<Object>> requestStartedCallback =
         tracer.getCallbackProvider(RequestContextSlot.APPSEC).getCallback(EVENTS.requestStarted());
@@ -298,9 +385,6 @@ private static AgentSpanContext processAppSecRequestData(LambdaRequestData event
                     .getCallbackProvider(RequestContextSlot.APPSEC)
                     .getCallback(EVENTS.requestMethodUriRaw());
         if (methodUriCallback != null) {
-          // Reconstruct full path with query string for AppSec analysis
-          String fullPath = buildFullPath(eventData.path, eventData.queryParameters);
-          LambdaURIDataAdapter uriAdapter = new LambdaURIDataAdapter(fullPath, eventData.headers);
           methodUriCallback.apply(requestContext, eventData.method, uriAdapter);
         } else {
           log.debug("requestMethodUriRaw callback is null");