getsentry · untitaker · Sep 25, 2018 · Sep 24, 2018 · Sep 25, 2018 · Sep 25, 2018
@@ -228,15 +228,19 @@ def event_processor(event, hint):
                 request_info["env"] = dict(get_environ(environ))
 
             if "headers" not in request_info:
-                request_info["headers"] = dict(get_headers(environ))
-                if not _should_send_default_pii():
-                    request_info["headers"] = {
-                        k: v
-                        for k, v in request_info["headers"].items()
-                        if k.lower().replace("_", "-")
-                        not in ("set-cookie", "cookie", "authorization")
-                    }
+                request_info["headers"] = _filter_headers(dict(get_headers(environ)))
 
         return event
 
     return event_processor
+
+
+def _filter_headers(headers):
+    if _should_send_default_pii():
+        return headers
+
+    return {
+        k: v
+        for k, v in headers.items()
+        if k.lower().replace("_", "-") not in ("set-cookie", "cookie", "authorization")
+    }
@@ -0,0 +1,109 @@
+import sys
+
+from sentry_sdk import configure_scope
+from sentry_sdk.hub import Hub, _should_send_default_pii
+from sentry_sdk._compat import reraise
+from sentry_sdk.utils import (
+    AnnotatedValue,
+    capture_internal_exceptions,
+    event_from_exception,
+)
+from sentry_sdk.integrations import Integration
+from sentry_sdk.integrations._wsgi import _filter_headers
+
+import __main__ as lambda_bootstrap
+
+
+class AwsLambdaIntegration(Integration):
+    identifier = "aws_lambda"
+
+    def install(self):
+        old_make_final_handler = lambda_bootstrap.make_final_handler
+
+        def sentry_make_final_handler(*args, **kwargs):
+            handler = old_make_final_handler(*args, **kwargs)
+
+            def sentry_handler(event, context, *args, **kwargs):
+                hub = Hub.current
+
+                with hub.push_scope():
+                    with capture_internal_exceptions():
+                        with configure_scope() as scope:
+                            scope.transaction = context.function_name
+                            scope.add_event_processor(
+                                _make_request_event_processor(event, context)
+                            )
+
+                    try:
+                        return handler(event, context, *args, **kwargs)
+                    except Exception:
+                        exc_info = sys.exc_info()
+                        event, hint = event_from_exception(
+                            exc_info,
+                            with_locals=hub.client.options["with_locals"],
+                            mechanism={"type": "aws_lambda", "handled": False},
+                        )
+
+                        hub.capture_event(event, hint=hint)
+                        reraise(*exc_info)
+                    finally:
+                        client = hub.client
+                        # Flush out the event queue before AWS kills the
+                        # process. This is not threadsafe.
+                        if client is not None:
+                            # make new transport with empty queue
+                            new_transport = client.transport.copy()
+                            client.close()
+                            client.transport = new_transport
+
+            return sentry_handler
+
+        lambda_bootstrap.make_final_handler = sentry_make_final_handler
+
+
+def _make_request_event_processor(aws_event, aws_context):
+    def event_processor(event, hint):
+        extra = event.setdefault("extra", {})
+        extra["lambda"] = {
+            "remaining_time_in_millis": aws_context.get_remaining_time_in_millis(),
+            "function_name": aws_context.function_name,
+            "function_version": aws_context.function_version,
+            "invoked_function_arn": aws_context.invoked_function_arn,
+            "aws_request_id": aws_context.aws_request_id,
+        }
+
+        request = event.setdefault("request", {})
+
+        if "httpMethod" in aws_event and "method" not in request:
+            request["method"] = aws_event["httpMethod"]
+        if "url" not in request:
+            request["url"] = _get_url(aws_event, aws_context)
+        if "queryStringParameters" in aws_event and "query_string" not in request:
+            request["query_string"] = aws_event["queryStringParameters"]
+        if "headers" in aws_event and "headers" not in request:
+            request["headers"] = _filter_headers(aws_event["headers"])
+        if aws_event.get("body", None):
+            # Unfortunately couldn't find a way to get structured body from AWS
+            # event. Meaning every body is unstructured to us.
+            request["data"] = AnnotatedValue("", {"rem": [["!raw", "x", 0, 0]]})
+
+        if _should_send_default_pii():
+            user_info = event.setdefault("user", {})
+            if "id" not in user_info:
+                user_info["id"] = aws_event.get("identity", {}).get("userArn")
+            if "ip_address" not in user_info:
+                user_info["ip_address"] = aws_event.get("identity", {}).get("sourceIp")
+
+        return event
+
+    return event_processor
+
+
+def _get_url(event, context):
+    path = event.get("path", None)
+    headers = event.get("headers", {})
+    host = headers.get("Host", None)
+    proto = headers.get("X-Forwarded-Proto", None)
+    if proto and host and path:
+        return "{}://{}{}".format(proto, host, path)
+    return "awslambda:///{}".format(context.function_name)
@@ -62,6 +62,15 @@ def kill(self):
         """Forcefully kills the transport."""
         pass
 
+    def copy(self):
+        """Copy the transport.
+
+        The returned transport should behave completely independent from the
+        previous one.  It still may share HTTP connection pools, but not share
+        any state such as internal queues.
+        """
+        return self
+
     def __del__(self):
         try:
             self.kill()
@@ -83,6 +92,7 @@ def __init__(self, options):
         )
         self._disabled_until = None
         self._retry = urllib3.util.Retry()
+        self.options = options
 
     def _send_event(self, event):
         if self._disabled_until is not None:
@@ -143,6 +153,11 @@ def kill(self):
         logger.debug("Killing HTTP transport")
         self._worker.kill()
 
+    def copy(self):
+        transport = type(self)(self.options)
+        transport._pool = self._pool
+        return transport
+
 
 class _FunctionTransport(Transport):
     def __init__(self, func):

@@ -0,0 +1,183 @@
+import base64
+import json
+import os
+import shutil
+import subprocess
+import sys
+import uuid
+
+import pytest
+
+boto3 = pytest.importorskip("boto3")
+
+LAMBDA_TEMPLATE = """
+from __future__ import print_function
+
+from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration
+import sentry_sdk
+import json
+from sentry_sdk.transport import Transport
+
+class TestTransport(Transport):
+    def __init__(self):
+        Transport.__init__(self)
+        self._queue = []
+
+    def capture_event(self, event):
+        self._queue.append(event)
+
+    def shutdown(self, timeout, callback=None):
+        # Delay event output like this to test proper shutdown
+        for event in self._queue:
+            print("EVENT:", json.dumps(event))
+
+sentry_sdk.init(
+    "http://bogus@example.com/2",
+    transport=TestTransport(),
+    integrations=[AwsLambdaIntegration()],
+    **{extra_init_args}
+)
+
+
+def test_handler(event, context):
+{code}
+"""
+
+
+@pytest.fixture
+def lambda_client():
+    if "AWS_ACCESS_KEY_ID" not in os.environ:
+        pytest.skip("AWS environ vars not set")
+
+    return boto3.client(
+        "lambda",
+        aws_access_key_id=os.environ["AWS_ACCESS_KEY_ID"],
+        aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
+        region_name="us-east-1",
+    )
+
+
+@pytest.fixture(params=["python3.6", "python2.7"])
+def run_lambda_function(tmpdir, lambda_client, request, assert_semaphore_acceptance):
+    def inner(lambda_body, payload, extra_init_args=None):
+        tmpdir.ensure_dir("lambda_tmp").remove()
+        tmp = tmpdir.ensure_dir("lambda_tmp")
+
+        # https://docs.aws.amazon.com/lambda/latest/dg/lambda-python-how-to-create-deployment-package.html
+        tmp.join("test_lambda.py").write(
+            LAMBDA_TEMPLATE.format(
+                code="\n".join("    " + x.strip() for x in lambda_body.splitlines()),
+                extra_init_args=repr(extra_init_args or {}),
+            )
+        )
+        tmp.join("setup.cfg").write("[install]\nprefix=")
+        subprocess.check_call([sys.executable, "setup.py", "sdist", "-d", str(tmpdir)])
+        subprocess.check_call("pip install ../*.tar.gz -t .", cwd=str(tmp), shell=True)
+        shutil.make_archive(tmpdir.join("ball"), "zip", str(tmp))
+
+        fn_name = "test_function_{}".format(uuid.uuid4())
+
+        lambda_client.create_function(
+            FunctionName=fn_name,
+            Runtime=request.param,
+            Role=os.environ["AWS_IAM_ROLE"],
+            Handler="test_lambda.test_handler",
+            Code={"ZipFile": tmpdir.join("ball.zip").read(mode="rb")},
+            Description="Created as part of testsuite for getsentry/sentry-python",
+        )
+
+        @request.addfinalizer
+        def delete_function():
+            lambda_client.delete_function(FunctionName=fn_name)
+
+        response = lambda_client.invoke(
+            FunctionName=fn_name,
+            InvocationType="RequestResponse",
+            LogType="Tail",
+            Payload=payload,
+        )
+
+        assert 200 <= response["StatusCode"] < 300, response
+
+        events = []
+
+        for line in base64.b64decode(response["LogResult"]).splitlines():
+            print("AWS:", line)
+            if not line.startswith(b"EVENT: "):
+                continue
+            line = line[len(b"EVENT: ") :]
+            events.append(json.loads(line.decode("utf-8")))
+            assert_semaphore_acceptance(events[-1])
+
+        return events, response
+
+    return inner
+
+
+def test_basic(run_lambda_function):
+    events, response = run_lambda_function(
+        'raise Exception("something went wrong")\n', b'{"foo": "bar"}'
+    )
+
+    assert response["FunctionError"] == "Unhandled"
+
+    event, = events
+    assert event["level"] == "error"
+    exception, = event["exception"]["values"]
+    assert exception["type"] == "Exception"
+    assert exception["value"] == "something went wrong"
+
+    frame1, = exception["stacktrace"]["frames"]
+    assert frame1["filename"] == "test_lambda.py"
+    assert frame1["abs_path"] == "/var/task/test_lambda.py"
+    assert frame1["function"] == "test_handler"
+
+    assert frame1["in_app"] is True
+
+    assert exception["mechanism"] == {"type": "aws_lambda", "handled": False}
+
+    assert event["extra"]["lambda"]["function_name"].startswith("test_function_")
+
+
+def test_request_data(run_lambda_function):
+    events, _response = run_lambda_function(
+        'sentry_sdk.capture_message("hi")\nreturn "ok"',
+        payload=b"""
+        {
+          "resource": "/asd",
+          "path": "/asd",
+          "httpMethod": "GET",
+          "headers": {
+            "Host": "iwsz2c7uwi.execute-api.us-east-1.amazonaws.com",
+            "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.13; rv:62.0) Gecko/20100101 Firefox/62.0",
+            "X-Forwarded-Proto": "https"
+          },
+          "queryStringParameters": {
+            "bonkers": "true"
+          },
+          "pathParameters": null,
+          "stageVariables": null,
+          "requestContext": {
+            "identity": {
+              "sourceIp": "213.47.147.207",
+              "userArn": "42"
+            }
+          },
+          "body": null,
+          "isBase64Encoded": false
+        }
+        """,
+    )
+
+    event, = events
+
+    assert event["request"] == {
+        "headers": {
+            "Host": "iwsz2c7uwi.execute-api.us-east-1.amazonaws.com",
+            "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.13; rv:62.0) Gecko/20100101 Firefox/62.0",
+            "X-Forwarded-Proto": "https",
+        },
+        "method": "GET",
+        "query_string": {"bonkers": "true"},
+        "url": "https://iwsz2c7uwi.execute-api.us-east-1.amazonaws.com/asd",
+    }
@@ -21,6 +21,7 @@ envlist =
 
     {py2.7,py3.7}-requests
     {pypy,py2.7,py3.5,py3.6,py3.7,py3.8}-celery-4
+    py3.7-aws_lambda
     {pypy,py2.7}-celery-3
 
 
@@ -44,7 +45,8 @@ deps =
     flask-dev: git+https://github.com/pallets/flask.git#egg=flask
     celery-3: Celery>=3.1,<4.0
     celery-4: Celery>=4.0,<5.0
-    requests: requests>=2.0<3.0
+    requests: requests>=2.0
+    aws_lambda: boto3
     flask: flask-login
     linters: black
     linters: flake8
@@ -55,6 +57,11 @@ setenv =
     flask: TESTPATH=tests/integrations/flask
     celery: TESTPATH=tests/integrations/celery
     requests: TESTPATH=tests/integrations/requests
+    aws_lambda: TESTPATH=tests/integrations/aws_lambda
+passenv =
+    AWS_ACCESS_KEY_ID
+    AWS_SECRET_ACCESS_KEY
+    AWS_IAM_ROLE
 usedevelop = True
 extras =
     flask: flask