Merge branch 'main' into christophe-papazian/APPSEC-52972-exploit-pre…

…vention-blocking-support
DataDog · May 2, 2024 · 05b8435 · 05b8435
2 parents 4ca38e7 + a2b1dbb
commit 05b8435
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 2 deletions.
diff --git a/ddtrace/propagation/_database_monitoring.py b/ddtrace/propagation/_database_monitoring.py
@@ -23,6 +23,7 @@
 DBM_DATABASE_SERVICE_NAME_KEY = "dddbs"
 DBM_PEER_HOSTNAME_KEY = "ddh"
 DBM_PEER_DB_NAME_KEY = "dddb"
+DBM_PEER_SERVICE_KEY = "ddprs"
 DBM_ENVIRONMENT_KEY = "dde"
 DBM_VERSION_KEY = "ddpv"
 DBM_TRACE_PARENT_KEY = "traceparent"
@@ -56,12 +57,14 @@ def __init__(
         sql_injector=default_sql_injector,
         peer_hostname_tag="out.host",
         peer_db_name_tag="db.name",
+        peer_service_tag="peer.service",
     ):
         self.sql_pos = sql_pos
         self.sql_kw = sql_kw
         self.sql_injector = sql_injector
         self.peer_hostname_tag = peer_hostname_tag
         self.peer_db_name_tag = peer_db_name_tag
+        self.peer_service_tag = peer_service_tag
 
     def inject(self, dbspan, args, kwargs):
         # run sampling before injection to propagate correct sampling priority
@@ -114,6 +117,10 @@ def _get_dbm_comment(self, db_span):
         if peer_hostname:
             dbm_tags[DBM_PEER_HOSTNAME_KEY] = peer_hostname
 
+        peer_service = db_span.get_tag(self.peer_service_tag)
+        if peer_service:
+            dbm_tags[DBM_PEER_SERVICE_KEY] = peer_service
+
         if dbm_config.propagation_mode == "full":
             db_span.set_tag_str(DBM_TRACE_INJECTED_TAG, "true")
             dbm_tags[DBM_TRACE_PARENT_KEY] = db_span.context._traceparent

diff --git a/tests/contrib/aiomysql/test_aiomysql.py b/tests/contrib/aiomysql/test_aiomysql.py
@@ -230,7 +230,9 @@ class AioMySQLTestCase(AsyncioTestCase):
     TEST_SERVICE = "mysql"
     conn = None
 
-    async def _get_conn_tracer(self):
+    async def _get_conn_tracer(self, tags=None):
+        tags = tags if tags is not None else {}
+
         if not self.conn:
             self.conn = await aiomysql.connect(**AIOMYSQL_CONFIG)
             assert not self.conn.closed
@@ -239,7 +241,7 @@ async def _get_conn_tracer(self):
             assert pin
             # Customize the service
             # we have to apply it on the existing one since new one won't inherit `app`
-            pin.clone(tracer=self.tracer).onto(self.conn)
+            pin.clone(tracer=self.tracer, tags={**tags, **pin.tags}).onto(self.conn)
 
             return self.conn, self.tracer
 
@@ -429,3 +431,26 @@ async def test_aiomysql_dbm_propagation_comment_peer_service_enabled(self):
         await shared_tests._test_dbm_propagation_comment_peer_service_enabled(
             config=AIOMYSQL_CONFIG, cursor=cursor, wrapped_instance=cursor.__wrapped__
         )
+
+    @mark_asyncio
+    @AsyncioTestCase.run_in_subprocess(
+        env_overrides=dict(
+            DD_DBM_PROPAGATION_MODE="service",
+            DD_SERVICE="orders-app",
+            DD_ENV="staging",
+            DD_VERSION="v7343437-d7ac743",
+            DD_TRACE_SPAN_ATTRIBUTE_SCHEMA="v1",
+        )
+    )
+    async def test_aiomysql_dbm_propagation_comment_with_peer_service_tag(self):
+        """tests if dbm comment is set in mysql"""
+        conn, tracer = await self._get_conn_tracer({"peer.service": "peer_service_name"})
+        cursor = await conn.cursor()
+        cursor.__wrapped__ = mock.AsyncMock()
+
+        await shared_tests._test_dbm_propagation_comment_with_peer_service_tag(
+            config=AIOMYSQL_CONFIG,
+            cursor=cursor,
+            wrapped_instance=cursor.__wrapped__,
+            peer_service_name="peer_service_name",
+        )
diff --git a/tests/contrib/shared_tests.py b/tests/contrib/shared_tests.py
@@ -94,3 +94,18 @@ async def _test_dbm_propagation_comment_peer_service_enabled(config, cursor, wra
     await _test_execute(dbm_comment, cursor, wrapped_instance)
     if execute_many:
         await _test_execute_many(dbm_comment, cursor, wrapped_instance)
+
+
+async def _test_dbm_propagation_comment_with_peer_service_tag(
+    config, cursor, wrapped_instance, peer_service_name, execute_many=True
+):
+    """tests if dbm comment is set in mysql"""
+    db_name = config["db"]
+
+    dbm_comment = (
+        f"/*dddb='{db_name}',dddbs='test',dde='staging',ddh='127.0.0.1',ddprs='{peer_service_name}',ddps='orders-app',"
+        "ddpv='v7343437-d7ac743'*/ "
+    )
+    await _test_execute(dbm_comment, cursor, wrapped_instance)
+    if execute_many:
+        await _test_execute_many(dbm_comment, cursor, wrapped_instance)