Correct Serverless Distributed Tracing Logic (#870)

* Fix serverless logic for distributed tracing

* Test stubs

* Collapse testing changes

* Add negative testing to regular DT test suite

* Apply linter fixes

* [Mega-Linter] Apply linters fixes

---------

Co-authored-by: TimPansino <[email protected]>

Author: TimPansino

newrelic/api/transaction.py

@@ -1039,7 +1039,9 @@ def _create_distributed_trace_data(self):
 
         settings = self._settings
         account_id = settings.account_id
-        trusted_account_key = settings.trusted_account_key
+        trusted_account_key = settings.trusted_account_key or (
+            self._settings.serverless_mode.enabled and self._settings.account_id
+        )
         application_id = settings.primary_application_id
 
         if not (account_id and application_id and trusted_account_key and settings.distributed_tracing.enabled):
@@ -1130,7 +1132,10 @@ def _can_accept_distributed_trace_headers(self):
             return False
 
         settings = self._settings
-        if not (settings.distributed_tracing.enabled and settings.trusted_account_key):
+        trusted_account_key = settings.trusted_account_key or (
+            self._settings.serverless_mode.enabled and self._settings.account_id
+        )
+        if not (settings.distributed_tracing.enabled and trusted_account_key):
             return False
 
         if self._distributed_trace_state:
@@ -1176,10 +1181,13 @@ def _accept_distributed_trace_payload(self, payload, transport_type="HTTP"):
 
             settings = self._settings
             account_id = data.get("ac")
+            trusted_account_key = settings.trusted_account_key or (
+                self._settings.serverless_mode.enabled and self._settings.account_id
+            )
 
             # If trust key doesn't exist in the payload, use account_id
             received_trust_key = data.get("tk", account_id)
-            if settings.trusted_account_key != received_trust_key:
+            if trusted_account_key != received_trust_key:
                 self._record_supportability("Supportability/DistributedTrace/AcceptPayload/Ignored/UntrustedAccount")
                 if settings.debug.log_untrusted_distributed_trace_keys:
                     _logger.debug(
@@ -1288,8 +1296,10 @@ def accept_distributed_trace_headers(self, headers, transport_type="HTTP"):
                 tracestate = ensure_str(tracestate)
                 try:
                     vendors = W3CTraceState.decode(tracestate)
-                    tk = self._settings.trusted_account_key
-                    payload = vendors.pop(tk + "@nr", "")
+                    trusted_account_key = self._settings.trusted_account_key or (
+                        self._settings.serverless_mode.enabled and self._settings.account_id
+                    )
+                    payload = vendors.pop(trusted_account_key + "@nr", "")
                     self.tracing_vendors = ",".join(vendors.keys())
                     self.tracestate = vendors.text(limit=31)
                 except:
@@ -1298,7 +1308,7 @@ def accept_distributed_trace_headers(self, headers, transport_type="HTTP"):
                     # Remove trusted new relic header if available and parse
                     if payload:
                         try:
-                            tracestate_data = NrTraceState.decode(payload, tk)
+                            tracestate_data = NrTraceState.decode(payload, trusted_account_key)
                         except:
                             tracestate_data = None
                         if tracestate_data:

tests/agent_features/test_distributed_tracing.py

@@ -30,6 +30,7 @@
 
 from newrelic.api.application import application_instance
 from newrelic.api.background_task import BackgroundTask, background_task
+from newrelic.api.external_trace import ExternalTrace
 from newrelic.api.time_trace import current_trace
 from newrelic.api.transaction import (
     accept_distributed_trace_headers,
@@ -361,3 +362,65 @@ def test_current_span_id_inside_transaction():
 def test_current_span_id_outside_transaction():
     span_id = current_span_id()
     assert span_id is None
+
+
+@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
+def test_outbound_dt_payload_generation(trusted_account_key):
+    @override_application_settings(
+        {
+            "distributed_tracing.enabled": True,
+            "account_id": "1",
+            "trusted_account_key": trusted_account_key,
+            "primary_application_id": "1",
+        }
+    )
+    @background_task(name="test_outbound_dt_payload_generation")
+    def _test_outbound_dt_payload_generation():
+        transaction = current_transaction()
+        payload = ExternalTrace.generate_request_headers(transaction)
+        if trusted_account_key:
+            assert payload
+            # Ensure trusted account key present as vendor
+            assert dict(payload)["tracestate"].startswith("1@nr=")
+        else:
+            assert not payload
+
+    _test_outbound_dt_payload_generation()
+
+
+@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
+def test_inbound_dt_payload_acceptance(trusted_account_key):
+    @override_application_settings(
+        {
+            "distributed_tracing.enabled": True,
+            "account_id": "1",
+            "trusted_account_key": trusted_account_key,
+            "primary_application_id": "1",
+        }
+    )
+    @background_task(name="_test_inbound_dt_payload_acceptance")
+    def _test_inbound_dt_payload_acceptance():
+        transaction = current_transaction()
+
+        payload = {
+            "v": [0, 1],
+            "d": {
+                "ty": "Mobile",
+                "ac": "1",
+                "tk": "1",
+                "ap": "2827902",
+                "pa": "5e5733a911cfbc73",
+                "id": "7d3efb1b173fecfa",
+                "tr": "d6b4ba0c3a712ca",
+                "ti": 1518469636035,
+                "tx": "8703ff3d88eefe9d",
+            },
+        }
+
+        result = transaction.accept_distributed_trace_payload(payload)
+        if trusted_account_key:
+            assert result
+        else:
+            assert not result
+
+    _test_inbound_dt_payload_acceptance()

tests/agent_features/test_serverless_mode.py

@@ -13,7 +13,16 @@
 # limitations under the License.
 
 import json
+
 import pytest
+from testing_support.fixtures import override_generic_settings
+from testing_support.validators.validate_serverless_data import validate_serverless_data
+from testing_support.validators.validate_serverless_metadata import (
+    validate_serverless_metadata,
+)
+from testing_support.validators.validate_serverless_payload import (
+    validate_serverless_payload,
+)
 
 from newrelic.api.application import application_instance
 from newrelic.api.background_task import background_task
@@ -22,23 +31,14 @@
 from newrelic.api.transaction import current_transaction
 from newrelic.core.config import global_settings
 
-from testing_support.fixtures import override_generic_settings
-from testing_support.validators.validate_serverless_data import (
-        validate_serverless_data)
-from testing_support.validators.validate_serverless_payload import (
-        validate_serverless_payload)
-from testing_support.validators.validate_serverless_metadata import (
-        validate_serverless_metadata)
 
-
-@pytest.fixture(scope='function')
+@pytest.fixture(scope="function")
 def serverless_application(request):
     settings = global_settings()
     orig = settings.serverless_mode.enabled
     settings.serverless_mode.enabled = True
 
-    application_name = 'Python Agent Test (test_serverless_mode:%s)' % (
-            request.node.name)
+    application_name = "Python Agent Test (test_serverless_mode:%s)" % (request.node.name)
     application = application_instance(application_name)
     application.activate()
 
@@ -48,17 +48,18 @@ def serverless_application(request):
 
 
 def test_serverless_payload(capsys, serverless_application):
-
-    @override_generic_settings(serverless_application.settings, {
-        'distributed_tracing.enabled': True,
-    })
+    @override_generic_settings(
+        serverless_application.settings,
+        {
+            "distributed_tracing.enabled": True,
+        },
+    )
     @validate_serverless_data(
-            expected_methods=('metric_data', 'analytic_event_data'),
-            forgone_methods=('preconnect', 'connect', 'get_agent_commands'))
+        expected_methods=("metric_data", "analytic_event_data"),
+        forgone_methods=("preconnect", "connect", "get_agent_commands"),
+    )
     @validate_serverless_payload()
-    @background_task(
-            application=serverless_application,
-            name='test_serverless_payload')
+    @background_task(application=serverless_application, name="test_serverless_payload")
     def _test():
         transaction = current_transaction()
         assert transaction.settings.serverless_mode.enabled
@@ -75,17 +76,15 @@ def _test():
 
 
 def test_no_cat_headers(serverless_application):
-    @background_task(
-            application=serverless_application,
-            name='test_cat_headers')
+    @background_task(application=serverless_application, name="test_cat_headers")
     def _test_cat_headers():
         transaction = current_transaction()
 
         payload = ExternalTrace.generate_request_headers(transaction)
         assert not payload
 
-        trace = ExternalTrace('testlib', 'http://example.com')
-        response_headers = [('X-NewRelic-App-Data', 'Cookies')]
+        trace = ExternalTrace("testlib", "http://example.com")
+        response_headers = [("X-NewRelic-App-Data", "Cookies")]
         with trace:
             trace.process_response_headers(response_headers)
 
@@ -94,61 +93,66 @@ def _test_cat_headers():
     _test_cat_headers()
 
 
-def test_dt_outbound(serverless_application):
-    @override_generic_settings(serverless_application.settings, {
-        'distributed_tracing.enabled': True,
-        'account_id': '1',
-        'trusted_account_key': '1',
-        'primary_application_id': '1',
-    })
-    @background_task(
-            application=serverless_application,
-            name='test_dt_outbound')
-    def _test_dt_outbound():
+@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
+def test_outbound_dt_payload_generation(serverless_application, trusted_account_key):
+    @override_generic_settings(
+        serverless_application.settings,
+        {
+            "distributed_tracing.enabled": True,
+            "account_id": "1",
+            "trusted_account_key": trusted_account_key,
+            "primary_application_id": "1",
+        },
+    )
+    @background_task(application=serverless_application, name="test_outbound_dt_payload_generation")
+    def _test_outbound_dt_payload_generation():
         transaction = current_transaction()
         payload = ExternalTrace.generate_request_headers(transaction)
         assert payload
-
-    _test_dt_outbound()
-
-
-def test_dt_inbound(serverless_application):
-    @override_generic_settings(serverless_application.settings, {
-        'distributed_tracing.enabled': True,
-        'account_id': '1',
-        'trusted_account_key': '1',
-        'primary_application_id': '1',
-    })
-    @background_task(
-            application=serverless_application,
-            name='test_dt_inbound')
-    def _test_dt_inbound():
+        # Ensure trusted account key or account ID present as vendor
+        assert dict(payload)["tracestate"].startswith("1@nr=")
+
+    _test_outbound_dt_payload_generation()
+
+
+@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
+def test_inbound_dt_payload_acceptance(serverless_application, trusted_account_key):
+    @override_generic_settings(
+        serverless_application.settings,
+        {
+            "distributed_tracing.enabled": True,
+            "account_id": "1",
+            "trusted_account_key": trusted_account_key,
+            "primary_application_id": "1",
+        },
+    )
+    @background_task(application=serverless_application, name="test_inbound_dt_payload_acceptance")
+    def _test_inbound_dt_payload_acceptance():
         transaction = current_transaction()
 
         payload = {
-            'v': [0, 1],
-            'd': {
-                'ty': 'Mobile',
-                'ac': '1',
-                'tk': '1',
-                'ap': '2827902',
-                'pa': '5e5733a911cfbc73',
-                'id': '7d3efb1b173fecfa',
-                'tr': 'd6b4ba0c3a712ca',
-                'ti': 1518469636035,
-                'tx': '8703ff3d88eefe9d',
-            }
+            "v": [0, 1],
+            "d": {
+                "ty": "Mobile",
+                "ac": "1",
+                "tk": "1",
+                "ap": "2827902",
+                "pa": "5e5733a911cfbc73",
+                "id": "7d3efb1b173fecfa",
+                "tr": "d6b4ba0c3a712ca",
+                "ti": 1518469636035,
+                "tx": "8703ff3d88eefe9d",
+            },
         }
 
         result = transaction.accept_distributed_trace_payload(payload)
         assert result
 
-    _test_dt_inbound()
+    _test_inbound_dt_payload_acceptance()
 
 
-@pytest.mark.parametrize('arn_set', (True, False))
+@pytest.mark.parametrize("arn_set", (True, False))
 def test_payload_metadata_arn(serverless_application, arn_set):
-
     # If the session object gathers the arn from the settings object before the
     # lambda handler records it there, then this test will fail.
 
@@ -157,17 +161,17 @@ def test_payload_metadata_arn(serverless_application, arn_set):
 
     arn = None
     if arn_set:
-        arn = 'arrrrrrrrrrRrrrrrrrn'
+        arn = "arrrrrrrrrrRrrrrrrrn"
 
-    settings.aws_lambda_metadata.update({'arn': arn, 'function_version': '$LATEST'})
+    settings.aws_lambda_metadata.update({"arn": arn, "function_version": "$LATEST"})
 
     class Context(object):
         invoked_function_arn = arn
 
-    @validate_serverless_metadata(exact_metadata={'arn': arn})
+    @validate_serverless_metadata(exact_metadata={"arn": arn})
     @lambda_handler(application=serverless_application)
     def handler(event, context):
-        assert settings.aws_lambda_metadata['arn'] == arn
+        assert settings.aws_lambda_metadata["arn"] == arn
         return {}
 
     try: