Merge branch 'main' into upgrade-latest-flask-version

Author: emmettbutler

.github/workflows/system-tests.yml

@@ -11,33 +11,6 @@ on:
     - cron: '00 04 * * 2-6'
 
 jobs:
-  system-tests-build-agent:
-    runs-on: ubuntu-latest
-    steps:
-
-      - name: Checkout system tests
-        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-        with:
-          persist-credentials: false
-          repository: 'DataDog/system-tests'
-          # Automatically managed, use scripts/update-system-tests-version to update
-          ref: '64afc5f8c1c2aa19adefdf1abf5559b62af8e784'
-
-      - name: Build agent
-        run: ./build.sh -i agent
-
-      - name: Save
-        id: save
-        run: |
-          docker image save system_tests/agent:latest | gzip > agent_${{ github.sha }}.tar.gz
-
-      - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
-        with:
-          name: agent_${{ github.sha }}
-          path: |
-            agent_${{ github.sha }}.tar.gz
-          retention-days: 2
-
   system-tests-build-weblog:
     runs-on: ubuntu-latest
     strategy:
@@ -69,7 +42,7 @@ jobs:
           persist-credentials: false
           repository: 'DataDog/system-tests'
           # Automatically managed, use scripts/update-system-tests-version to update
-          ref: '64afc5f8c1c2aa19adefdf1abf5559b62af8e784'
+          ref: '4dac1cb506ee28bdcac297ce3784f1b217679556'
 
       - name: Checkout dd-trace-py
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
@@ -98,7 +71,7 @@ jobs:
 
   system-tests:
     runs-on: ubuntu-latest
-    needs: [system-tests-build-agent, system-tests-build-weblog]
+    needs: [system-tests-build-weblog]
     strategy:
       matrix:
         weblog-variant: [flask-poc, uwsgi-poc , django-poc, fastapi, python3.12, django-py3.13]
@@ -123,7 +96,7 @@ jobs:
           persist-credentials: false
           repository: 'DataDog/system-tests'
           # Automatically managed, use scripts/update-system-tests-version to update
-          ref: '64afc5f8c1c2aa19adefdf1abf5559b62af8e784'
+          ref: '4dac1cb506ee28bdcac297ce3784f1b217679556'
 
       - name: Build runner
         uses: ./.github/actions/install_runner
@@ -133,16 +106,10 @@ jobs:
           name: ${{ matrix.weblog-variant }}_${{ github.sha }}
           path: images_artifacts/
 
-      - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0
-        with:
-          name: agent_${{ github.sha }}
-          path: images_artifacts/
-
       - name: docker load
         id: docker_load
         run: |
           docker load < images_artifacts/${{ matrix.weblog-variant}}_weblog_${{ github.sha }}.tar.gz
-          docker load < images_artifacts/agent_${{ github.sha }}.tar.gz
 
       - name: Run DEFAULT
         if: always() && steps.docker_load.outcome == 'success' && matrix.scenario == 'other'
@@ -310,7 +277,7 @@ jobs:
           persist-credentials: false
           repository: 'DataDog/system-tests'
           # Automatically managed, use scripts/update-system-tests-version to update
-          ref: '64afc5f8c1c2aa19adefdf1abf5559b62af8e784'
+          ref: '4dac1cb506ee28bdcac297ce3784f1b217679556'
       - name: Checkout dd-trace-py
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:

ddtrace/_monkey.py

@@ -352,7 +352,7 @@ def _patch_all(**patch_modules: bool) -> None:
 
     patch(raise_errors=False, **modules)
     if asm_config._iast_enabled:
-        from ddtrace.appsec._iast._patch_modules import patch_iast
+        from ddtrace.appsec._iast.main import patch_iast
         from ddtrace.appsec.iast import enable_iast_propagation
 
         patch_iast()

ddtrace/appsec/_common_module_patches.py

@@ -42,26 +42,32 @@ def _(module):
         subprocess_patch.patch()
         subprocess_patch.add_str_callback(_RASP_SYSTEM, wrapped_system_5542593D237084A7)
         subprocess_patch.add_lst_callback(_RASP_POPEN, popen_FD233052260D8B4D)
+        log.debug("Patching common modules: subprocess_patch")
 
     if _is_patched:
         return
 
     try_wrap_function_wrapper("builtins", "open", wrapped_open_CFDDB7ABBA9081B6)
     try_wrap_function_wrapper("urllib.request", "OpenerDirector.open", wrapped_open_ED4CF71136E15EBF)
     core.on("asm.block.dbapi.execute", execute_4C9BAC8E228EB347)
+    log.debug("Patching common modules: builtins and urllib.request")
     _is_patched = True
 
 
 def unpatch_common_modules():
     global _is_patched
     if not _is_patched:
         return
+
     try_unwrap("builtins", "open")
     try_unwrap("urllib.request", "OpenerDirector.open")
     try_unwrap("_io", "BytesIO.read")
     try_unwrap("_io", "StringIO.read")
+    subprocess_patch.unpatch()
     subprocess_patch.del_str_callback(_RASP_SYSTEM)
     subprocess_patch.del_lst_callback(_RASP_POPEN)
+
+    log.debug("Unpatching common modules subprocess, builtins and urllib.request")
     _is_patched = False
 
 
@@ -323,7 +329,7 @@ def try_unwrap(module, name):
             apply_patch(parent, attribute, original)
             del _DD_ORIGINAL_ATTRIBUTES[(parent, attribute)]
     except ModuleNotFoundError:
-        pass
+        log.debug("ERROR unwrapping %s.%s ", module, name)
 
 
 def try_wrap_function_wrapper(module_name: str, name: str, wrapper: Callable) -> None:
@@ -332,7 +338,7 @@ def _(module):
         try:
             wrap_object(module, name, FunctionWrapper, (wrapper,))
         except (ImportError, AttributeError):
-            log.debug("ASM patching. Module %s.%s does not exist", module_name, name)
+            log.debug("Module %s.%s does not exist", module_name, name)
 
 
 def wrap_object(module, name, factory, args=(), kwargs=None):

ddtrace/appsec/_handlers.py

@@ -1,14 +1,22 @@
 import io
 import json
+from typing import Any
+from typing import Dict
+from typing import Optional
 
 import xmltodict
 
+from ddtrace._trace.span import Span
+from ddtrace.appsec._asm_request_context import _call_waf
+from ddtrace.appsec._asm_request_context import _call_waf_first
 from ddtrace.appsec._asm_request_context import get_blocked
 from ddtrace.appsec._constants import SPAN_DATA_NAMES
+from ddtrace.appsec._http_utils import extract_cookies_from_headers
+from ddtrace.appsec._http_utils import normalize_headers
+from ddtrace.appsec._http_utils import parse_http_body
 from ddtrace.contrib import trace_utils
 from ddtrace.contrib.internal.trace_utils_base import _get_request_header_user_agent
 from ddtrace.contrib.internal.trace_utils_base import _set_url_tag
-from ddtrace.ext import SpanTypes
 from ddtrace.ext import http
 from ddtrace.internal import core
 from ddtrace.internal.constants import RESPONSE_HEADERS
@@ -53,7 +61,7 @@ def _on_set_http_meta(
     response_headers,
     response_cookies,
 ):
-    if asm_config._asm_enabled and span.span_type == SpanTypes.WEB:
+    if asm_config._asm_enabled and span.span_type in asm_config._asm_http_span_types:
         # avoid circular import
         from ddtrace.appsec._asm_request_context import set_waf_address
 
@@ -77,6 +85,74 @@ def _on_set_http_meta(
                 set_waf_address(k, v)
 
 
+# AWS Lambda
+def _on_lambda_start_request(
+    span: Span,
+    request_headers: Dict[str, str],
+    request_ip: Optional[str],
+    body: Optional[str],
+    is_body_base64: bool,
+    raw_uri: str,
+    route: str,
+    method: str,
+    parsed_query: Dict[str, Any],
+):
+    if not (asm_config._asm_enabled and span.span_type in asm_config._asm_http_span_types):
+        return
+
+    headers = normalize_headers(request_headers)
+    request_body = parse_http_body(headers, body, is_body_base64)
+    request_cookies = extract_cookies_from_headers(headers)
+
+    _on_set_http_meta(
+        span,
+        request_ip,
+        raw_uri,
+        route,
+        method,
+        headers,
+        request_cookies,
+        parsed_query,
+        None,
+        request_body,
+        None,
+        None,
+        None,
+    )
+
+    _call_waf_first(("aws_lambda",))
+
+
+def _on_lambda_start_response(
+    span: Span,
+    status_code: str,
+    response_headers: Dict[str, str],
+):
+    if not (asm_config._asm_enabled and span.span_type in asm_config._asm_http_span_types):
+        return
+
+    waf_headers = normalize_headers(response_headers)
+    response_cookies = extract_cookies_from_headers(waf_headers)
+
+    _on_set_http_meta(
+        span,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        status_code,
+        waf_headers,
+        response_cookies,
+    )
+
+    _call_waf(("aws_lambda",))
+
+
 # ASGI
 
 
@@ -307,6 +383,9 @@ def listen():
 
     core.on("asgi.request.parse.body", _on_asgi_request_parse_body, "await_receive_and_body")
 
+    core.on("aws_lambda.start_request", _on_lambda_start_request)
+    core.on("aws_lambda.start_response", _on_lambda_start_response)
+
     core.on("grpc.server.response.message", _on_grpc_server_response)
     core.on("grpc.server.data", _on_grpc_server_data)
 

ddtrace/appsec/_http_utils.py

@@ -0,0 +1,81 @@
+import base64
+from http.cookies import SimpleCookie
+import json
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import Union
+from urllib.parse import parse_qs
+
+import xmltodict
+
+from ddtrace.internal.utils import http as http_utils
+
+
+def normalize_headers(
+    request_headers: Dict[str, str],
+) -> Dict[str, Optional[str]]:
+    """Normalize headers according to the WAF expectations.
+
+    The WAF expects headers to be lowercased and empty values to be None.
+    """
+    headers: Dict[str, Optional[str]] = {}
+    for key, value in request_headers.items():
+        normalized_key = http_utils.normalize_header_name(key)
+        if value:
+            headers[normalized_key] = str(value).strip()
+        else:
+            headers[normalized_key] = None
+    return headers
+
+
+def parse_http_body(
+    normalized_headers: Dict[str, Optional[str]],
+    body: Optional[str],
+    is_body_base64: bool,
+) -> Union[str, Dict[str, Any], None]:
+    """Parse a request body based on the content-type header."""
+    if body is None:
+        return None
+    if is_body_base64:
+        try:
+            body = base64.b64decode(body).decode()
+        except (ValueError, TypeError):
+            return None
+
+    try:
+        content_type = normalized_headers.get("content-type")
+        if not content_type:
+            return None
+
+        if content_type in ("application/json", "application/vnd.api+json", "text/json"):
+            return json.loads(body)
+        elif content_type in ("application/x-url-encoded", "application/x-www-form-urlencoded"):
+            return parse_qs(body)
+        elif content_type in ("application/xml", "text/xml"):
+            return xmltodict.parse(body)
+        elif content_type.startswith("multipart/form-data"):
+            return http_utils.parse_form_multipart(body, normalized_headers)
+        elif content_type == "text/plain":
+            return body
+        else:
+            return None
+
+    except Exception:
+        return None
+
+
+def extract_cookies_from_headers(
+    normalized_headers: Dict[str, Optional[str]],
+) -> Optional[Dict[str, str]]:
+    """Extract cookies from the WAF headers."""
+    cookie_names = {"cookie", "set-cookie"}
+    for name in cookie_names:
+        if name in normalized_headers:
+            cookie = SimpleCookie()
+            header = normalized_headers[name]
+            del normalized_headers[name]
+            if header:
+                cookie.load(header)
+            return {k: v.value for k, v in cookie.items()}
+    return None