Enhance boto3 instrumentation to capture HTTP headers in all four flavours

instana/instrumentation/boto3_inst.py

@@ -8,14 +8,25 @@
 import inspect
 
 from ..log import logger
-from ..singletons import tracer
+from ..singletons import tracer, agent
 from ..util.traceutils import get_active_tracer
 
 try:
     import opentracing as ot
     import boto3
     from boto3.s3 import inject
 
+    def extract_custom_headers(span, headers):
+        if agent.options.extra_http_headers is None or headers is None:
+            return
+        try:
+            for custom_header in agent.options.extra_http_headers:
+                if custom_header in headers:
+                    span.set_tag("http.header.%s" % custom_header, headers[custom_header])
+
+        except Exception:
+            logger.debug("extract_custom_headers: ", exc_info=True)
+
 
     def lambda_inject_context(payload, scope):
         """
@@ -35,6 +46,20 @@ def lambda_inject_context(payload, scope):
             logger.debug("non-fatal lambda_inject_context: ", exc_info=True)
 
 
+    @wrapt.patch_function_wrapper("botocore.auth", "SigV4Auth.add_auth")
+    def emit_add_auth_with_instana(wrapped, instance, args, kwargs):
+        active_tracer = get_active_tracer()
+
+        # If we're not tracing, just return;
+        if active_tracer is None:
+            return wrapped(*args, **kwargs)
+
+        span = active_tracer.active_span
+        extract_custom_headers(span, args[0].headers)
+
+        return wrapped(*args, **kwargs)
+
+
     @wrapt.patch_function_wrapper('botocore.client', 'BaseClient._make_api_call')
     def make_api_call_with_instana(wrapped, instance, arg_list, kwargs):
         # pylint: disable=protected-access
@@ -76,6 +101,8 @@ def make_api_call_with_instana(wrapped, instance, arg_list, kwargs):
                         status = http_dict.get('HTTPStatusCode')
                         if status is not None:
                             scope.span.set_tag('http.status_code', status)
+                        headers = http_dict.get('HTTPHeaders')
+                        extract_custom_headers(scope.span, headers)
 
                 return result
             except Exception as exc:

tests/clients/boto3/test_boto3_lambda.py

@@ -2,74 +2,254 @@
 # (c) Copyright Instana Inc. 2020
 
 from __future__ import absolute_import
+import unittest
+import json
 
-import os
 import boto3
-import pytest
-
 # TODO: Remove branching when we drop support for Python 3.7
-import sys
-if sys.version_info >= (3, 8):
+from sys import version_info
+if version_info >= (3, 8):
   from moto import mock_aws
-else:
-  from moto import mock_sqs as mock_aws
 
-from instana.singletons import tracer
+from instana.singletons import tracer, agent
 from ...helpers import get_first_span_by_filter
 
+@unittest.skipIf(version_info < (3, 8), "Test skipped on Python < 3.8")
+class TestLambda(unittest.TestCase):
+    def setUp(self):
+        """ Clear all spans before a test run """
+        self.recorder = tracer.recorder
+        self.recorder.clear_spans()
+        self.mock = mock_aws(config={"lambda": {"use_docker": False}})
+        self.mock.start()
+        self.lambda_region = "us-east-1"
+        self.aws_lambda = boto3.client('lambda', region_name=self.lambda_region)
+        self.function_name = "myfunc"
+
+    def tearDown(self):
+        # Stop Moto after each test
+        self.mock.stop()
+
+
+    def test_lambda_invoke(self):
+        with tracer.start_active_span('test'):
+            result = self.aws_lambda.invoke(FunctionName=self.function_name, Payload=json.dumps({"message": "success"}))
+
+        self.assertEqual(result["StatusCode"], 200)
+        result_payload = json.loads(result["Payload"].read().decode("utf-8"))
+        self.assertIn("message", result_payload)
+        self.assertEqual("success", result_payload["message"])
+
+        spans = tracer.recorder.queued_spans()
+        self.assertEqual(2, len(spans))
+
+        filter = lambda span: span.n == "sdk"
+        test_span = get_first_span_by_filter(spans, filter)
+        self.assertTrue(test_span)
+
+        filter = lambda span: span.n == "boto3"
+        boto_span = get_first_span_by_filter(spans, filter)
+        self.assertTrue(boto_span)
+
+        self.assertEqual(boto_span.t, test_span.t)
+        self.assertEqual(boto_span.p, test_span.s)
+
+        self.assertIsNone(test_span.ec)
+        self.assertIsNone(boto_span.ec)
+
+        self.assertEqual(boto_span.data['boto3']['op'], 'Invoke')
+        endpoint = f'https://lambda.{self.lambda_region}.amazonaws.com'
+        self.assertEqual(boto_span.data['boto3']['ep'], endpoint)
+        self.assertEqual(boto_span.data['boto3']['reg'], self.lambda_region)
+        self.assertIn('FunctionName', boto_span.data['boto3']['payload'])
+        self.assertEqual(boto_span.data['boto3']['payload']['FunctionName'], self.function_name)
+        self.assertEqual(boto_span.data['http']['status'], 200)
+        self.assertEqual(boto_span.data['http']['method'], 'POST')
+        self.assertEqual(boto_span.data['http']['url'], f'{endpoint}:443/Invoke')
+
+
+    def test_request_header_capture_before_call(self):
+        original_extra_http_headers = agent.options.extra_http_headers
+        agent.options.extra_http_headers = ['X-Capture-This', 'X-Capture-That']
 
-@pytest.fixture(scope='function')
-def aws_credentials():
-    """Mocked AWS Credentials for moto."""
-    os.environ['AWS_ACCESS_KEY_ID'] = 'testing'
-    os.environ['AWS_SECRET_ACCESS_KEY'] = 'testing'
-    os.environ['AWS_SECURITY_TOKEN'] = 'testing'
-    os.environ['AWS_SESSION_TOKEN'] = 'testing'
-
-
-@pytest.fixture(scope='function')
-def aws_lambda(aws_credentials):
-    with mock_aws():
-        yield boto3.client('lambda', region_name='us-east-1')
-
-
-def setup_method():
-    """ Clear all spans before a test run """
-    tracer.recorder.clear_spans()
-
-
-@pytest.mark.skip("Lambda mocking requires docker")
-def test_lambda_invoke(aws_lambda):
-    result = None
-
-    with tracer.start_active_span('test'):
-        result = aws_lambda.invoke(FunctionName='arn:aws:lambda:us-west-1:410797082306:function:CanaryInACoalMine')
-
-    assert result
-    assert len(result['Buckets']) == 1
-    assert result['Buckets'][0]['Name'] == 'aws_bucket_name'
-
-    spans = tracer.recorder.queued_spans()
-    assert len(spans) == 2
-
-    filter = lambda span: span.n == "sdk"
-    test_span = get_first_span_by_filter(spans, filter)
-    assert (test_span)
-
-    filter = lambda span: span.n == "boto3"
-    boto_span = get_first_span_by_filter(spans, filter)
-    assert (boto_span)
-
-    assert (boto_span.t == test_span.t)
-    assert (boto_span.p == test_span.s)
-
-    assert (test_span.ec is None)
-    assert (boto_span.ec is None)
+        # Access the event system on the S3 client
+        event_system = self.aws_lambda.meta.events
 
-    assert boto_span.data['boto3']['op'] == 'CreateBucket'
-    assert boto_span.data['boto3']['ep'] == 'https://s3.amazonaws.com'
-    assert boto_span.data['boto3']['reg'] == 'us-east-1'
-    assert boto_span.data['boto3']['payload'] == {'Bucket': 'aws_bucket_name'}
-    assert boto_span.data['http']['status'] == 200
-    assert boto_span.data['http']['method'] == 'POST'
-    assert boto_span.data['http']['url'] == 'https://s3.amazonaws.com:443/CreateBucket'
+        request_headers = {
+                'X-Capture-This': 'this',
+                'X-Capture-That': 'that'
+            }
+
+        # Create a function that adds custom headers
+        def add_custom_header_before_call(params, **kwargs):
+            params['headers'].update(request_headers)
+
+        # Register the function to before-call event.
+        event_system.register('before-call.lambda.Invoke', add_custom_header_before_call)
+
+        with tracer.start_active_span('test'):
+            result = self.aws_lambda.invoke(FunctionName=self.function_name, Payload=json.dumps({"message": "success"}))
+
+        self.assertEqual(result["StatusCode"], 200)
+        result_payload = json.loads(result["Payload"].read().decode("utf-8"))
+        self.assertIn("message", result_payload)
+        self.assertEqual("success", result_payload["message"])
+
+        spans = tracer.recorder.queued_spans()
+        self.assertEqual(2, len(spans))
+
+        filter = lambda span: span.n == "sdk"
+        test_span = get_first_span_by_filter(spans, filter)
+        self.assertTrue(test_span)
+
+        filter = lambda span: span.n == "boto3"
+        boto_span = get_first_span_by_filter(spans, filter)
+        self.assertTrue(boto_span)
+
+        self.assertEqual(boto_span.t, test_span.t)
+        self.assertEqual(boto_span.p, test_span.s)
+
+        self.assertIsNone(test_span.ec)
+        self.assertIsNone(boto_span.ec)
+
+        self.assertEqual(boto_span.data['boto3']['op'], 'Invoke')
+        endpoint = f'https://lambda.{self.lambda_region}.amazonaws.com'
+        self.assertEqual(boto_span.data['boto3']['ep'], endpoint)
+        self.assertEqual(boto_span.data['boto3']['reg'], self.lambda_region)
+        self.assertIn('FunctionName', boto_span.data['boto3']['payload'])
+        self.assertEqual(boto_span.data['boto3']['payload']['FunctionName'], self.function_name)
+        self.assertEqual(boto_span.data['http']['status'], 200)
+        self.assertEqual(boto_span.data['http']['method'], 'POST')
+        self.assertEqual(boto_span.data['http']['url'], f'{endpoint}:443/Invoke')
+
+        self.assertIn("X-Capture-This", boto_span.data["http"]["header"])
+        self.assertEqual("this", boto_span.data["http"]["header"]["X-Capture-This"])
+        self.assertIn("X-Capture-That", boto_span.data["http"]["header"])
+        self.assertEqual("that", boto_span.data["http"]["header"]["X-Capture-That"])
+
+        agent.options.extra_http_headers = original_extra_http_headers
+
+
+    def test_request_header_capture_before_sign(self):
+        original_extra_http_headers = agent.options.extra_http_headers
+        agent.options.extra_http_headers = ['X-Custom-1', 'X-Custom-2']
+
+        # Access the event system on the S3 client
+        event_system = self.aws_lambda.meta.events
+
+        request_headers = {
+                'X-Custom-1': 'Value1',
+                'X-Custom-2': 'Value2'
+            }
+
+        # Create a function that adds custom headers
+        def add_custom_header_before_sign(request, **kwargs):
+            for name, value in request_headers.items():
+                request.headers.add_header(name, value)
+
+        # Register the function to before-sign event.
+        event_system.register_first('before-sign.lambda.Invoke', add_custom_header_before_sign)
+
+        with tracer.start_active_span('test'):
+            result = self.aws_lambda.invoke(FunctionName=self.function_name, Payload=json.dumps({"message": "success"}))
+
+        self.assertEqual(result["StatusCode"], 200)
+        result_payload = json.loads(result["Payload"].read().decode("utf-8"))
+        self.assertIn("message", result_payload)
+        self.assertEqual("success", result_payload["message"])
+
+        spans = tracer.recorder.queued_spans()
+        self.assertEqual(2, len(spans))
+
+        filter = lambda span: span.n == "sdk"
+        test_span = get_first_span_by_filter(spans, filter)
+        self.assertTrue(test_span)
+
+        filter = lambda span: span.n == "boto3"
+        boto_span = get_first_span_by_filter(spans, filter)
+        self.assertTrue(boto_span)
+
+        self.assertEqual(boto_span.t, test_span.t)
+        self.assertEqual(boto_span.p, test_span.s)
+
+        self.assertIsNone(test_span.ec)
+        self.assertIsNone(boto_span.ec)
+
+        self.assertEqual(boto_span.data['boto3']['op'], 'Invoke')
+        endpoint = f'https://lambda.{self.lambda_region}.amazonaws.com'
+        self.assertEqual(boto_span.data['boto3']['ep'], endpoint)
+        self.assertEqual(boto_span.data['boto3']['reg'], self.lambda_region)
+        self.assertIn('FunctionName', boto_span.data['boto3']['payload'])
+        self.assertEqual(boto_span.data['boto3']['payload']['FunctionName'], self.function_name)
+        self.assertEqual(boto_span.data['http']['status'], 200)
+        self.assertEqual(boto_span.data['http']['method'], 'POST')
+        self.assertEqual(boto_span.data['http']['url'], f'{endpoint}:443/Invoke')
+
+        self.assertIn("X-Custom-1", boto_span.data["http"]["header"])
+        self.assertEqual("Value1", boto_span.data["http"]["header"]["X-Custom-1"])
+        self.assertIn("X-Custom-2", boto_span.data["http"]["header"])
+        self.assertEqual("Value2", boto_span.data["http"]["header"]["X-Custom-2"])
+
+        agent.options.extra_http_headers = original_extra_http_headers
+
+
+    def test_response_header_capture(self):
+        original_extra_http_headers = agent.options.extra_http_headers
+        agent.options.extra_http_headers = ['X-Capture-This-Too', 'X-Capture-That-Too']
+
+        # Access the event system on the S3 client
+        event_system = self.aws_lambda.meta.events
+
+        response_headers = {
+            "X-Capture-This-Too": "this too",
+            "X-Capture-That-Too": "that too",
+        }
+
+        # Create a function that sets the custom headers in the after-call event.
+        def modify_after_call_args(parsed, **kwargs):
+            parsed['ResponseMetadata']['HTTPHeaders'].update(response_headers)
+
+        # Register the function to an event
+        event_system.register('after-call.lambda.Invoke', modify_after_call_args)
+
+        with tracer.start_active_span('test'):
+            result = self.aws_lambda.invoke(FunctionName=self.function_name, Payload=json.dumps({"message": "success"}))
+
+        self.assertEqual(result["StatusCode"], 200)
+        result_payload = json.loads(result["Payload"].read().decode("utf-8"))
+        self.assertIn("message", result_payload)
+        self.assertEqual("success", result_payload["message"])
+
+        spans = tracer.recorder.queued_spans()
+        self.assertEqual(2, len(spans))
+
+        filter = lambda span: span.n == "sdk"
+        test_span = get_first_span_by_filter(spans, filter)
+        self.assertTrue(test_span)
+
+        filter = lambda span: span.n == "boto3"
+        boto_span = get_first_span_by_filter(spans, filter)
+        self.assertTrue(boto_span)
+
+        self.assertEqual(boto_span.t, test_span.t)
+        self.assertEqual(boto_span.p, test_span.s)
+
+        self.assertIsNone(test_span.ec)
+        self.assertIsNone(boto_span.ec)
+
+        self.assertEqual(boto_span.data['boto3']['op'], 'Invoke')
+        endpoint = f'https://lambda.{self.lambda_region}.amazonaws.com'
+        self.assertEqual(boto_span.data['boto3']['ep'], endpoint)
+        self.assertEqual(boto_span.data['boto3']['reg'], self.lambda_region)
+        self.assertIn('FunctionName', boto_span.data['boto3']['payload'])
+        self.assertEqual(boto_span.data['boto3']['payload']['FunctionName'], self.function_name)
+        self.assertEqual(boto_span.data['http']['status'], 200)
+        self.assertEqual(boto_span.data['http']['method'], 'POST')
+        self.assertEqual(boto_span.data['http']['url'], f'{endpoint}:443/Invoke')
+
+        self.assertIn("X-Capture-This-Too", boto_span.data["http"]["header"])
+        self.assertEqual("this too", boto_span.data["http"]["header"]["X-Capture-This-Too"])
+        self.assertIn("X-Capture-That-Too", boto_span.data["http"]["header"])
+        self.assertEqual("that too", boto_span.data["http"]["header"]["X-Capture-That-Too"])
+
+        agent.options.extra_http_headers = original_extra_http_headers