fix(appsec): normalize DBAPI queries for SQLi RASP

ddtrace/contrib/_events/dbapi.py

@@ -1,13 +1,17 @@
 from dataclasses import dataclass
+from typing import Union
 
 from ddtrace.internal.core.events import Event
 
 
+DbApiQuery = Union[str, bytes]
+
+
 @dataclass
 class DbApiEvent(Event):
     """A database query shared by instrumentation and product subscribers."""
 
     event_name = "dbapi.query"
 
-    query: str
+    query: DbApiQuery
     span_name_prefix: str

ddtrace/contrib/dbapi.py

@@ -10,6 +10,7 @@
 from ddtrace import config
 from ddtrace._trace.pin import Pin
 from ddtrace.contrib._events.dbapi import DbApiEvent
+from ddtrace.contrib._events.dbapi import DbApiQuery
 from ddtrace.internal import core
 from ddtrace.internal.constants import COMPONENT
 from ddtrace.internal.logger import get_logger
@@ -81,7 +82,7 @@ def __init__(
         )
         self._self_datadog_name = span_name
         self._self_dbapi_span_name_prefix = span_name_prefix
-        self._self_last_execute_operation = None
+        self._self_last_execute_operation: object = None
         self._self_config = cfg or config.dbapi2
         self._self_dbm_propagator = getattr(self._self_config, "_dbm_propagator", None)
         self._self_db_tags = dict(db_tags) if db_tags else {}
@@ -92,6 +93,26 @@ def __iter__(self):
     def __next__(self):
         return self.__wrapped__.__next__()
 
+    def _normalize_dbapi_query(self, query: object) -> Optional[DbApiQuery]:
+        if isinstance(query, (str, bytes)):
+            return query
+        return None
+
+    def _prepare_dbapi_query(self, query: object) -> object:
+        has_listeners = core.has_listeners(DbApiEvent.event_name)
+        normalized_query: Optional[DbApiQuery] = None
+        if isinstance(query, (str, bytes)) or is_tracing_enabled() or has_listeners:
+            try:
+                normalized_query = self._normalize_dbapi_query(query)
+            except Exception:
+                log.debug("Failed to normalize database query", exc_info=True)
+
+        resource = normalized_query if normalized_query is not None else query
+        self._self_last_execute_operation = resource
+        if has_listeners and normalized_query is not None:
+            core.dispatch_event(DbApiEvent(query=normalized_query, span_name_prefix=self._self_dbapi_span_name_prefix))
+        return resource
+
     def _trace_method(self, method, name, resource, extra_tags, dbm_propagator, *args, **kwargs):
         """
         Internal function to trace the call to the underlying cursor method
@@ -142,9 +163,7 @@ def _trace_method(self, method, name, resource, extra_tags, dbm_propagator, *arg
 
     def executemany(self, query, *args, **kwargs):
         """Wraps the cursor.executemany method"""
-        self._self_last_execute_operation = query
-        if isinstance(query, str):
-            core.dispatch_event(DbApiEvent(query=query, span_name_prefix=self._self_dbapi_span_name_prefix))
+        resource = self._prepare_dbapi_query(query)
         # Always return the result as-is
         # DEV: Some libraries return `None`, others `int`, and others the cursor objects
         #      These differences should be overridden at the integration specific layer (e.g. in `sqlite3/patch.py`)
@@ -153,7 +172,7 @@ def executemany(self, query, *args, **kwargs):
         return self._trace_method(
             self.__wrapped__.executemany,
             self._self_datadog_name,
-            query,
+            resource,
             {"sql.executemany": "true"},
             self._self_dbm_propagator,
             query,
@@ -163,17 +182,15 @@ def executemany(self, query, *args, **kwargs):
 
     def execute(self, query, *args, **kwargs):
         """Wraps the cursor.execute method"""
-        self._self_last_execute_operation = query
-        if isinstance(query, str):
-            core.dispatch_event(DbApiEvent(query=query, span_name_prefix=self._self_dbapi_span_name_prefix))
+        resource = self._prepare_dbapi_query(query)
 
         # Always return the result as-is
         # DEV: Some libraries return `None`, others `int`, and others the cursor objects
         #      These differences should be overridden at the integration specific layer (e.g. in `sqlite3/patch.py`)
         return self._trace_method(
             self.__wrapped__.execute,
             self._self_datadog_name,
-            query,
+            resource,
             {},
             self._self_dbm_propagator,
             query,

ddtrace/contrib/dbapi_async.py

@@ -1,7 +1,6 @@
 import inspect
 
 from ddtrace import config
-from ddtrace.contrib._events.dbapi import DbApiEvent
 from ddtrace.internal import core
 from ddtrace.internal.constants import COMPONENT
 from ddtrace.internal.logger import get_logger
@@ -97,9 +96,7 @@ async def _trace_method(self, method, name, resource, extra_tags, dbm_propagator
 
     async def executemany(self, query, *args, **kwargs):
         """Wraps the cursor.executemany method"""
-        self._self_last_execute_operation = query
-        if isinstance(query, str):
-            core.dispatch_event(DbApiEvent(query=query, span_name_prefix=self._self_dbapi_span_name_prefix))
+        resource = self._prepare_dbapi_query(query)
         # Always return the result as-is
         # DEV: Some libraries return `None`, others `int`, and others the cursor objects
         #      These differences should be overridden at the integration specific layer (e.g. in `sqlite3/patch.py`)
@@ -108,7 +105,7 @@ async def executemany(self, query, *args, **kwargs):
         return await self._trace_method(
             self.__wrapped__.executemany,
             self._self_datadog_name,
-            query,
+            resource,
             {"sql.executemany": "true"},
             self._self_dbm_propagator,
             query,
@@ -118,17 +115,15 @@ async def executemany(self, query, *args, **kwargs):
 
     async def execute(self, query, *args, **kwargs):
         """Wraps the cursor.execute method"""
-        self._self_last_execute_operation = query
-        if isinstance(query, str):
-            core.dispatch_event(DbApiEvent(query=query, span_name_prefix=self._self_dbapi_span_name_prefix))
+        resource = self._prepare_dbapi_query(query)
 
         # Always return the result as-is
         # DEV: Some libraries return `None`, others `int`, and others the cursor objects
         #      These differences should be overridden at the integration specific layer (e.g. in `sqlite3/patch.py`)
         return await self._trace_method(
             self.__wrapped__.execute,
             self._self_datadog_name,
-            query,
+            resource,
             {},
             self._self_dbm_propagator,
             query,

ddtrace/contrib/internal/aiomysql/patch.py

@@ -107,15 +107,15 @@ async def _trace_method(self, method, resource, extra_tags, *args, **kwargs):
                     s._set_attribute("db.rownumber", self.rownumber)
 
     async def executemany(self, query, *args, **kwargs):
-        if isinstance(query, str):
+        if isinstance(query, (str, bytes)) and core.has_listeners(DbApiEvent.event_name):
             core.dispatch_event(DbApiEvent(query=query, span_name_prefix="mysql"))
         result = await self._trace_method(
             self.__wrapped__.executemany, query, {"sql.executemany": "true"}, query, *args, **kwargs
         )
         return result
 
     async def execute(self, query, *args, **kwargs):
-        if isinstance(query, str):
+        if isinstance(query, (str, bytes)) and core.has_listeners(DbApiEvent.event_name):
             core.dispatch_event(DbApiEvent(query=query, span_name_prefix="mysql"))
         result = await self._trace_method(self.__wrapped__.execute, query, {}, query, *args, **kwargs)
         return result

ddtrace/contrib/internal/aiopg/connection.py

@@ -1,3 +1,7 @@
+from typing import Optional
+from typing import Protocol
+from typing import runtime_checkable
+
 from aiopg import __version__
 from aiopg.utils import _ContextManager
 import wrapt
@@ -9,19 +13,27 @@
 from ddtrace.contrib import dbapi
 from ddtrace.contrib import trace_utils
 from ddtrace.contrib._events.dbapi import DbApiEvent
+from ddtrace.contrib._events.dbapi import DbApiQuery
 from ddtrace.contrib.internal.trace_utils import set_service_and_source
 from ddtrace.ext import SpanKind
 from ddtrace.ext import SpanTypes
 from ddtrace.ext import db
 from ddtrace.internal import core
 from ddtrace.internal.constants import COMPONENT
+from ddtrace.internal.logger import get_logger
 from ddtrace.internal.schema import schematize_database_operation
 from ddtrace.internal.schema import schematize_service_name
 from ddtrace.internal.utils.version import parse_version
 from ddtrace.trace import tracer
 
 
 AIOPG_VERSION = parse_version(__version__)
+log = get_logger(__name__)
+
+
+@runtime_checkable
+class _StringifiableQuery(Protocol):
+    def as_string(self, context: object) -> str: ...
 
 
 class AIOTracedCursor(wrapt.ObjectProxy):
@@ -32,6 +44,31 @@ def __init__(self, cursor, pin):
         pin.onto(self)
         self._datadog_name = schematize_database_operation("postgres.query", database_provider="postgresql")
 
+    def _normalize_dbapi_query(self, query: object) -> Optional[DbApiQuery]:
+        if isinstance(query, (str, bytes)):
+            return query
+        if isinstance(query, _StringifiableQuery):
+            return query.as_string(self.__wrapped__)
+        return None
+
+    def _prepare_dbapi_query(self, query: object) -> object:
+        has_listeners = core.has_listeners(DbApiEvent.event_name)
+        normalized_query: Optional[DbApiQuery] = None
+        should_normalize = isinstance(query, (str, bytes)) or has_listeners
+        if not should_normalize:
+            pin = Pin.get_from(self)
+            should_normalize = pin is not None and pin.enabled()
+        if should_normalize:
+            try:
+                normalized_query = self._normalize_dbapi_query(query)
+            except Exception:
+                log.debug("Failed to normalize database query", exc_info=True)
+
+        resource = normalized_query if normalized_query is not None else query
+        if has_listeners and normalized_query is not None:
+            core.dispatch_event(DbApiEvent(query=normalized_query, span_name_prefix="postgres"))
+        return resource
+
     async def _trace_method(self, method, resource, extra_tags, *args, **kwargs):
         pin = Pin.get_from(self)
         if not pin or not pin.enabled():
@@ -63,17 +100,15 @@ async def _trace_method(self, method, resource, extra_tags, *args, **kwargs):
     async def executemany(self, query, *args, **kwargs):
         # FIXME[matt] properly handle kwargs here. arg names can be different
         # with different libs.
-        if isinstance(query, str):
-            core.dispatch_event(DbApiEvent(query=query, span_name_prefix="postgres"))
+        resource = self._prepare_dbapi_query(query)
         result = await self._trace_method(
-            self.__wrapped__.executemany, query, {"sql.executemany": "true"}, query, *args, **kwargs
+            self.__wrapped__.executemany, resource, {"sql.executemany": "true"}, query, *args, **kwargs
         )
         return result
 
     async def execute(self, query, *args, **kwargs):
-        if isinstance(query, str):
-            core.dispatch_event(DbApiEvent(query=query, span_name_prefix="postgres"))
-        result = await self._trace_method(self.__wrapped__.execute, query, {}, query, *args, **kwargs)
+        resource = self._prepare_dbapi_query(query)
+        result = await self._trace_method(self.__wrapped__.execute, resource, {}, query, *args, **kwargs)
         return result
 
     async def callproc(self, proc, args):

ddtrace/contrib/internal/asyncpg/patch.py

@@ -193,7 +193,7 @@ def _is_cursor_forward_query(query: str) -> bool:
 
 def _dispatch_dbapi_event(query: Union[str, bytes]) -> None:
     # Cursor.forward() operates on a query already checked when its portal was bound.
-    if isinstance(query, str) and not _is_cursor_forward_query(query):
+    if core.has_listeners(DbApiEvent.event_name) and (isinstance(query, bytes) or not _is_cursor_forward_query(query)):
         core.dispatch_event(DbApiEvent(query=query, span_name_prefix="postgres"))
 
 

ddtrace/contrib/internal/psycopg/cursor.py

@@ -1,28 +1,63 @@
+from typing import Callable
+from typing import Optional
+from typing import Protocol
+from typing import cast
+from typing import runtime_checkable
+
 from ddtrace.contrib import dbapi
+from ddtrace.contrib._events.dbapi import DbApiQuery
+
+
+@runtime_checkable
+class _StringifiableQuery(Protocol):
+    def as_string(self, context: object) -> str: ...
+
+
+@runtime_checkable
+class _TemplateQuery(Protocol):
+    @property
+    def strings(self) -> tuple[str, ...]: ...
 
+    @property
+    def interpolations(self) -> tuple[object, ...]: ...
 
-class Psycopg3TracedCursor(dbapi.TracedCursor):
-    """TracedCursor for psycopg instances"""
+
+class _PsycopgTracedCursor(dbapi.TracedCursor):
+    """Common cursor tracing for psycopg 2 and 3."""
 
     def __init__(self, cursor, cfg, *args, **kwargs):
-        super(Psycopg3TracedCursor, self).__init__(cursor, cfg=cfg, *args, **kwargs)
+        super(_PsycopgTracedCursor, self).__init__(cursor, cfg=cfg, *args, **kwargs)
+
+    def _normalize_dbapi_query(self, query: object) -> Optional[DbApiQuery]:
+        normalized_query = super(_PsycopgTracedCursor, self)._normalize_dbapi_query(query)
+        if normalized_query is not None:
+            return normalized_query
+        if isinstance(query, _StringifiableQuery):
+            return query.as_string(self.__wrapped__)
+        return None
+
+
+class Psycopg3TracedCursor(_PsycopgTracedCursor):
+    """TracedCursor for psycopg 3 instances."""
 
-    def _trace_method(self, method, name, resource, extra_tags, dbm_propagator, *args, **kwargs):
-        # treat Composable resource objects as strings
-        if resource.__class__.__name__ == "SQL" or resource.__class__.__name__ == "Composed":
-            resource = resource.as_string(self.__wrapped__)
-        return super(Psycopg3TracedCursor, self)._trace_method(
-            method, name, resource, extra_tags, dbm_propagator, *args, **kwargs
-        )
+    def _normalize_dbapi_query(self, query: object) -> Optional[DbApiQuery]:
+        normalized_query = super(Psycopg3TracedCursor, self)._normalize_dbapi_query(query)
+        if normalized_query is not None:
+            return normalized_query
+        if isinstance(query, _TemplateQuery):
+            renderer = cast(Optional[Callable[[object, object], str]], self._self_config.get("_query_renderer"))
+            if renderer is not None:
+                return renderer(query, self.__wrapped__)
+        return None
 
 
 class Psycopg3FetchTracedCursor(Psycopg3TracedCursor, dbapi.FetchTracedCursor):
     """Psycopg3FetchTracedCursor for psycopg"""
 
 
-class Psycopg2TracedCursor(Psycopg3TracedCursor):
+class Psycopg2TracedCursor(_PsycopgTracedCursor):
     """TracedCursor for psycopg2"""
 
 
-class Psycopg2FetchTracedCursor(Psycopg3FetchTracedCursor):
+class Psycopg2FetchTracedCursor(Psycopg2TracedCursor, dbapi.FetchTracedCursor):
     """FetchTracedCursor for psycopg2"""

ddtrace/contrib/internal/psycopg/patch.py

@@ -60,6 +60,7 @@ def _psycopg_sql_injector(dbm_comment, sql_statement):
         _dbapi_span_name_prefix="postgres",
         _dbapi_span_operation_name=schematize_database_operation("postgres.query", database_provider="postgresql"),
         _patched_modules=set(),
+        _query_renderer=None,
         trace_fetch_methods=asbool(env.get("DD_PSYCOPG_TRACE_FETCH_METHODS", default=False)),
         trace_connect=asbool(env.get("DD_PSYCOPG_TRACE_CONNECT", default=False)),
         _dbm_propagator=_DBM_Propagator(0, "query", _psycopg_sql_injector),
@@ -127,6 +128,7 @@ def _patch(psycopg_module):
         config.psycopg["_patched_modules"].add(psycopg_module)
     else:
         _get_psycopg3_original_methods()
+        config.psycopg["_query_renderer"] = psycopg_module.sql.as_string
 
         _w(psycopg_module, "connect", patched_connect_factory(psycopg_module))
         _w(psycopg_module, "Cursor", init_cursor_from_connection_factory(psycopg_module))

releasenotes/notes/asm-dbapi-query-normalization-cdc548988f47c744.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    ASM: Fixes an issue where SQL injection attack detection and blocking do not inspect database
+    queries supplied as bytes or adapter-specific objects, including psycopg composables and templates.