Merge branch 'main' into michael.zhao/dsm-ckpt-all-records

Author: michael-zhao459

Dockerfile

@@ -23,8 +23,6 @@ RUN rm -rf ./python/lib/$runtime/site-packages/setuptools
 RUN rm -rf ./python/lib/$runtime/site-packages/jsonschema/tests
 RUN rm -f ./python/lib/$runtime/site-packages/ddtrace/appsec/_iast/_taint_tracking/*.so
 RUN rm -f ./python/lib/$runtime/site-packages/ddtrace/appsec/_iast/_stacktrace*.so
-RUN rm -f ./python/lib/$runtime/site-packages/ddtrace/internal/datadog/profiling/libdd_wrapper*.so
-RUN rm -f ./python/lib/$runtime/site-packages/ddtrace/internal/datadog/profiling/ddup/_ddup.*.so
 # _stack_v2 may not exist for some versions of ddtrace (e.g. under python 3.13)
 RUN rm -f ./python/lib/$runtime/site-packages/ddtrace/internal/datadog/profiling/stack_v2/_stack_v2.*.so
 # remove *.dist-info directories except any entry_points.txt files

datadog_lambda/asm.py

@@ -1,9 +1,12 @@
-from copy import deepcopy
 import logging
+import urllib.parse
+from copy import deepcopy
 from typing import Any, Dict, List, Optional, Union
 
 from ddtrace.contrib.internal.trace_utils import _get_request_header_client_ip
 from ddtrace.internal import core
+from ddtrace.internal.utils import get_blocked
+from ddtrace.internal.utils import http as http_utils
 from ddtrace.trace import Span
 
 from datadog_lambda.trigger import (
@@ -50,6 +53,7 @@ def asm_set_context(event_source: _EventSource):
     This allows the AppSecSpanProcessor to know information about the event
     at the moment the span is created and skip it when not relevant.
     """
+
     if event_source.event_type not in _http_event_types:
         core.set_item("appsec_skip_next_lambda_event", True)
 
@@ -126,6 +130,14 @@ def asm_start_request(
         span.set_tag_str("http.client_ip", request_ip)
         span.set_tag_str("network.client.ip", request_ip)
 
+    # Encode the parsed query and append it to reconstruct the original raw URI expected by AppSec.
+    if parsed_query:
+        try:
+            encoded_query = urllib.parse.urlencode(parsed_query, doseq=True)
+            raw_uri += "?" + encoded_query  # type: ignore
+        except Exception:
+            pass
+
     core.dispatch(
         # The matching listener is registered in ddtrace.appsec._handlers
         "aws_lambda.start_request",
@@ -182,3 +194,47 @@ def asm_start_response(
             response_headers,
         ),
     )
+
+    if isinstance(response, dict) and "statusCode" in response:
+        body = response.get("body")
+    else:
+        body = response
+
+    core.dispatch(
+        # The matching listener is registered in ddtrace.appsec._handlers
+        "aws_lambda.parse_body",
+        (body,),
+    )
+
+
+def get_asm_blocked_response(
+    event_source: _EventSource,
+) -> Optional[Dict[str, Any]]:
+    """Get the blocked response for the given event source."""
+    if event_source.event_type not in _http_event_types:
+        return None
+
+    blocked = get_blocked()
+    if not blocked:
+        return None
+
+    desired_type = blocked.get("type", "auto")
+    if desired_type == "none":
+        content_type = "text/plain; charset=utf-8"
+        content = ""
+    else:
+        content_type = blocked.get("content-type", "application/json")
+        content = http_utils._get_blocked_template(content_type)
+
+    response_headers = {
+        "content-type": content_type,
+    }
+    if "location" in blocked:
+        response_headers["location"] = blocked["location"]
+
+    return {
+        "statusCode": blocked.get("status_code", 403),
+        "headers": response_headers,
+        "body": content,
+        "isBase64Encoded": False,
+    }

datadog_lambda/tracing.py

@@ -4,6 +4,7 @@
 # Copyright 2019 Datadog, Inc.
 import logging
 import os
+import re
 import traceback
 import ujson as json
 from datetime import datetime, timezone
@@ -828,15 +829,31 @@ def create_service_mapping(val):
     return new_service_mapping
 
 
-def determine_service_name(service_mapping, specific_key, generic_key, default_value):
-    service_name = service_mapping.get(specific_key)
-    if service_name is None:
-        service_name = service_mapping.get(generic_key, default_value)
-    return service_name
+def determine_service_name(
+    service_mapping, specific_key, generic_key, extracted_key, fallback=None
+):
+    # Check for mapped service (specific key first, then generic key)
+    mapped_service = service_mapping.get(specific_key) or service_mapping.get(
+        generic_key
+    )
+    if mapped_service:
+        return mapped_service
+
+    # Check if AWS service representation is disabled
+    aws_service_representation = os.environ.get(
+        "DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED", ""
+    ).lower()
+    if aws_service_representation in ("false", "0"):
+        return fallback
+
+    # Use extracted_key if it exists and is not empty, otherwise use fallback
+    return (
+        extracted_key.strip() if extracted_key and extracted_key.strip() else fallback
+    )
 
 
 # Initialization code
-service_mapping_str = os.getenv("DD_SERVICE_MAPPING", "")
+service_mapping_str = os.environ.get("DD_SERVICE_MAPPING", "")
 service_mapping = create_service_mapping(service_mapping_str)
 
 _dd_origin = {"_dd.origin": "lambda"}
@@ -960,6 +977,7 @@ def create_inferred_span_from_api_gateway_websocket_event(
         "http.url": http_url,
         "endpoint": endpoint,
         "resource_names": endpoint,
+        "span.kind": "server",
         "apiid": api_id,
         "apiname": api_id,
         "stage": request_context.get("stage"),
@@ -1018,6 +1036,7 @@ def create_inferred_span_from_api_gateway_event(
         "endpoint": path,
         "http.method": method,
         "resource_names": resource,
+        "span.kind": "server",
         "apiid": api_id,
         "apiname": api_id,
         "stage": request_context.get("stage"),
@@ -1122,12 +1141,13 @@ def create_inferred_span_from_sqs_event(event, context):
     event_source_arn = event_record.get("eventSourceARN")
     queue_name = event_source_arn.split(":")[-1]
     service_name = determine_service_name(
-        service_mapping, queue_name, "lambda_sqs", "sqs"
+        service_mapping, queue_name, "lambda_sqs", queue_name, "sqs"
     )
     attrs = event_record.get("attributes") or {}
     tags = {
         "operation_name": "aws.sqs",
         "resource_names": queue_name,
+        "span.kind": "server",
         "queuename": queue_name,
         "event_source_arn": event_source_arn,
         "receipt_handle": event_record.get("receiptHandle"),
@@ -1189,11 +1209,12 @@ def create_inferred_span_from_sns_event(event, context):
     topic_arn = sns_message.get("TopicArn")
     topic_name = topic_arn.split(":")[-1]
     service_name = determine_service_name(
-        service_mapping, topic_name, "lambda_sns", "sns"
+        service_mapping, topic_name, "lambda_sns", topic_name, "sns"
     )
     tags = {
         "operation_name": "aws.sns",
         "resource_names": topic_name,
+        "span.kind": "server",
         "topicname": topic_name,
         "topic_arn": topic_arn,
         "message_id": sns_message.get("MessageId"),
@@ -1224,15 +1245,16 @@ def create_inferred_span_from_kinesis_event(event, context):
     event_record = get_first_record(event)
     event_source_arn = event_record.get("eventSourceARN")
     event_id = event_record.get("eventID")
-    stream_name = event_source_arn.split(":")[-1]
+    stream_name = re.sub(r"^stream/", "", (event_source_arn or "").split(":")[-1])
     shard_id = event_id.split(":")[0]
     service_name = determine_service_name(
-        service_mapping, stream_name, "lambda_kinesis", "kinesis"
+        service_mapping, stream_name, "lambda_kinesis", stream_name, "kinesis"
     )
     kinesis = event_record.get("kinesis") or {}
     tags = {
         "operation_name": "aws.kinesis",
         "resource_names": stream_name,
+        "span.kind": "server",
         "streamname": stream_name,
         "shardid": shard_id,
         "event_source_arn": event_source_arn,
@@ -1259,12 +1281,13 @@ def create_inferred_span_from_dynamodb_event(event, context):
     event_source_arn = event_record.get("eventSourceARN")
     table_name = event_source_arn.split("/")[1]
     service_name = determine_service_name(
-        service_mapping, table_name, "lambda_dynamodb", "dynamodb"
+        service_mapping, table_name, "lambda_dynamodb", table_name, "dynamodb"
     )
     dynamodb_message = event_record.get("dynamodb") or {}
     tags = {
         "operation_name": "aws.dynamodb",
         "resource_names": table_name,
+        "span.kind": "server",
         "tablename": table_name,
         "event_source_arn": event_source_arn,
         "event_id": event_record.get("eventID"),
@@ -1293,11 +1316,12 @@ def create_inferred_span_from_s3_event(event, context):
     obj = s3.get("object") or {}
     bucket_name = bucket.get("name")
     service_name = determine_service_name(
-        service_mapping, bucket_name, "lambda_s3", "s3"
+        service_mapping, bucket_name, "lambda_s3", bucket_name, "s3"
     )
     tags = {
         "operation_name": "aws.s3",
         "resource_names": bucket_name,
+        "span.kind": "server",
         "event_name": event_record.get("eventName"),
         "bucketname": bucket_name,
         "bucket_arn": bucket.get("arn"),
@@ -1323,11 +1347,12 @@ def create_inferred_span_from_s3_event(event, context):
 def create_inferred_span_from_eventbridge_event(event, context):
     source = event.get("source")
     service_name = determine_service_name(
-        service_mapping, source, "lambda_eventbridge", "eventbridge"
+        service_mapping, source, "lambda_eventbridge", source, "eventbridge"
     )
     tags = {
         "operation_name": "aws.eventbridge",
         "resource_names": source,
+        "span.kind": "server",
         "detail_type": event.get("detail-type"),
     }
     InferredSpanInfo.set_tags(
@@ -1401,9 +1426,21 @@ def create_function_execution_span(
         tags["_dd.parent_source"] = trace_context_source
     tags.update(trigger_tags)
     tracer.set_tags(_dd_origin)
+    # Determine service name based on config and env var
+    if config.service:
+        service_name = config.service
+    else:
+        aws_service_representation = os.environ.get(
+            "DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED", ""
+        ).lower()
+        if aws_service_representation in ("false", "0"):
+            service_name = "aws.lambda"
+        else:
+            service_name = function_name if function_name else "aws.lambda"
+
     span = tracer.trace(
         "aws.lambda",
-        service="aws.lambda",
+        service=service_name,
         resource=function_name,
         span_type="serverless",
     )

datadog_lambda/version.py

@@ -1 +1 @@
-__version__ = "7.112.0"
+__version__ = "8.114.0.dev0"

datadog_lambda/wrapper.py

@@ -11,6 +11,7 @@
 
 from datadog_lambda.asm import asm_set_context, asm_start_response, asm_start_request
 from datadog_lambda.dsm import set_dsm_context
+from ddtrace.internal._exceptions import BlockingException
 from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import (
     set_cold_start,
@@ -47,6 +48,14 @@
     extract_http_status_code_tag,
 )
 
+if config.appsec_enabled:
+    from datadog_lambda.asm import (
+        asm_set_context,
+        asm_start_response,
+        asm_start_request,
+        get_asm_blocked_response,
+    )
+
 if config.profiling_enabled:
     from ddtrace.profiling import profiler
 
@@ -121,6 +130,7 @@ def __init__(self, func):
             self.span = None
             self.inferred_span = None
             self.response = None
+            self.blocking_response = None
 
             if config.profiling_enabled:
                 self.prof = profiler.Profiler(env=config.env, service=config.service)
@@ -160,8 +170,12 @@ def __call__(self, event, context, **kwargs):
         """Executes when the wrapped function gets called"""
         self._before(event, context)
         try:
+            if self.blocking_response:
+                return self.blocking_response
             self.response = self.func(event, context, **kwargs)
             return self.response
+        except BlockingException:
+            self.blocking_response = get_asm_blocked_response(self.event_source)
         except Exception:
             from datadog_lambda.metric import submit_errors_metric
 
@@ -172,6 +186,8 @@ def __call__(self, event, context, **kwargs):
             raise
         finally:
             self._after(event, context)
+            if self.blocking_response:
+                return self.blocking_response
 
     def _inject_authorizer_span_headers(self, request_id):
         reference_span = self.inferred_span if self.inferred_span else self.span
@@ -204,6 +220,7 @@ def _inject_authorizer_span_headers(self, request_id):
     def _before(self, event, context):
         try:
             self.response = None
+            self.blocking_response = None
             set_cold_start(init_timestamp_ns)
 
             if not should_use_extension:
@@ -257,6 +274,7 @@ def _before(self, event, context):
                 )
                 if config.appsec_enabled:
                     asm_start_request(self.span, event, event_source, self.trigger_tags)
+                    self.blocking_response = get_asm_blocked_response(self.event_source)
             else:
                 set_correlation_ids()
             if config.profiling_enabled and is_new_sandbox():
@@ -290,20 +308,24 @@ def _after(self, event, context):
                 if status_code:
                     self.span.set_tag("http.status_code", status_code)
 
-                if config.appsec_enabled:
+                if config.appsec_enabled and not self.blocking_response:
                     asm_start_response(
                         self.span,
                         status_code,
                         self.event_source,
                         response=self.response,
                     )
+                    self.blocking_response = get_asm_blocked_response(self.event_source)
 
                 self.span.finish()
 
             if self.inferred_span:
                 if status_code:
                     self.inferred_span.set_tag("http.status_code", status_code)
 
+                if self.trigger_tags and (route := self.trigger_tags.get("http.route")):
+                    self.inferred_span.set_tag("http.route", route)
+
                 if config.service:
                     self.inferred_span.set_tag("peer.service", config.service)
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "datadog_lambda"
-version = "7.112.0"
+version = "8.114.0.dev0"
 description = "The Datadog AWS Lambda Library"
 authors = ["Datadog, Inc. <dev@datadoghq.com>"]
 license = "Apache-2.0"
@@ -28,7 +28,7 @@ classifiers = [
 python = ">=3.8.0,<4"
 datadog = ">=0.51.0,<1.0.0"
 wrapt = "^1.11.2"
-ddtrace = ">=3.10.2,<4"
+ddtrace = ">=3.11.0,<4"
 ujson = ">=5.9.0"
 botocore = { version = "^1.34.0", optional = true }
 requests = { version ="^2.22.0", optional = true }

scripts/check_layer_size.sh

@@ -8,8 +8,8 @@
 # Compares layer size to threshold, and fails if below that threshold
 
 set -e
-MAX_LAYER_COMPRESSED_SIZE_KB=$(expr 6 \* 1024)
-MAX_LAYER_UNCOMPRESSED_SIZE_KB=$(expr 15 \* 1024)
+MAX_LAYER_COMPRESSED_SIZE_KB=$(expr 8 \* 1024)
+MAX_LAYER_UNCOMPRESSED_SIZE_KB=$(expr 21 \* 1024)
 
 
 LAYER_FILES_PREFIX="datadog_lambda_py"