Test workflow

Author: florentinl

.github/workflows/build_layer.yml

@@ -0,0 +1,32 @@
+name: build_layers
+
+on:
+  push:
+    branches:
+      - "main"
+      - "florentinl/APPSEC-58224/block-requests"
+  pull_request:
+
+jobs:
+  build:
+    runs-on: ${{ matrix.arch == 'arm64' && 'ubuntu-24.04-arm' || 'ubuntu-latest' }}
+
+    strategy:
+      matrix:
+        arch: [arm64, amd64]
+        python_version: ["3.8", "3.9", "3.10", "3.11", "3.12", "3.13"]
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Build layer for Python ${{ matrix.python_version }} on ${{ matrix.arch }}
+        run: |
+          echo "Building layer for Python ${{ matrix.python_version }} on ${{ matrix.arch }}"
+          ARCH=${{ matrix.arch }} PYTHON_VERSION=${{ matrix.python_version }} ./scripts/build_layers.sh
+
+      - name: Upload layer artifact
+        uses: actions/upload-artifact@v4
+        with:
+          path: .layers/datadog_lambda_py-${{ matrix.arch }}-${{ matrix.python_version }}.zip
+          name: datadog-lambda-python-${{ matrix.python_version }}-${{ matrix.arch }}

datadog_lambda/asm.py

@@ -1,10 +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, set_blocked
+from ddtrace.internal.utils import http as http_utils
 from ddtrace.trace import Span
 
 from datadog_lambda.trigger import (
@@ -51,8 +53,20 @@ 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.
     """
+
+    def block_request_callable():
+        """Callable to block the request."""
+        blocked_config = get_blocked()
+        print("Blocking request before", blocked_config)
+        set_blocked(blocked_config)
+        print("Blocking request after", blocked_config)
+
     if event_source.event_type not in _http_event_types:
         core.set_item("appsec_skip_next_lambda_event", True)
+        core.set_item(
+            "block_request_callable",
+            block_request_callable,
+        )
 
 
 def asm_start_request(
@@ -127,6 +141,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)
+        except Exception:
+            pass
+        raw_uri += "?" + encoded_query  # type: ignore
+
     core.dispatch(
         # The matching listener is registered in ddtrace.appsec._handlers
         "aws_lambda.start_request",
@@ -193,7 +215,33 @@ def get_asm_blocked_response(
         return None
 
     blocked = get_blocked()
-    if blocked:
-        set_blocked(blocked)
-        return blocked.get("type", "auto")
-    return None
+    if not blocked:
+        return None
+
+    desired_type = blocked.get("type", "auto")
+    print("desired_type", desired_type)
+    if desired_type == "none":
+        content = "".encode()
+        response_headers = {
+            "content-type": "text/plain; charset=utf-8",
+        }
+    else:
+        ctype = blocked.get("content_type", "application/json")
+        print("ctype", ctype)
+        content = http_utils._get_blocked_template(ctype).encode()
+        response_headers = {
+            "content-type": ctype,
+        }
+
+    status = blocked.get("status", 403)
+
+    location = blocked.get("location")
+    if location:
+        response_headers["location"] = location
+
+    return {
+        "statusCode": status,
+        "headers": response_headers,
+        "body": content,
+        "isBase64Encoded": False,
+    }

datadog_lambda/wrapper.py

@@ -9,6 +9,7 @@
 from importlib import import_module
 from time import time_ns
 
+from ddtrace.internal._exceptions import BlockingException
 from datadog_lambda.asm import (
     asm_set_context,
     asm_start_response,
@@ -126,6 +127,7 @@ def __init__(self, func):
             self.span = None
             self.inferred_span = None
             self.response = None
+            self.asm_blocked = False
 
             if config.profiling_enabled:
                 self.prof = profiler.Profiler(env=config.env, service=config.service)
@@ -164,17 +166,19 @@ def __init__(self, func):
     def __call__(self, event, context, **kwargs):
         """Executes when the wrapped function gets called"""
         self._before(event, context)
-        if config.appsec_enabled:
-            blocking_response = get_asm_blocked_response(self.event_source)
-            if blocking_response:
-                return blocking_response
         try:
-            self.response = self.func(event, context, **kwargs)
-            if config.appsec_enabled:
+            if config.appsec_enabled and (
+                blocking_response := get_asm_blocked_response(self.event_source)
+            ):
+                return blocking_response
+
+            try:
+                self.response = self.func(event, context, **kwargs)
+                return self.response
+            except BlockingException:
                 blocking_response = get_asm_blocked_response(self.event_source)
-                if blocking_response:
-                    return blocking_response
-            return self.response
+                return blocking_response
+
         except Exception:
             from datadog_lambda.metric import submit_errors_metric
 
@@ -185,6 +189,9 @@ def __call__(self, event, context, **kwargs):
             raise
         finally:
             self._after(event, context)
+            if self.asm_blocked:
+                self.asm_blocked = False
+                return self.blocking_response
 
     def _inject_authorizer_span_headers(self, request_id):
         reference_span = self.inferred_span if self.inferred_span else self.span
@@ -310,6 +317,10 @@ def _after(self, event, context):
                         response=self.response,
                     )
 
+                    if blocking_response := get_asm_blocked_response(self.event_source):
+                        self.asm_blocked = True
+                        self.blocking_response = blocking_response
+
                 self.span.finish()
 
             if self.inferred_span: