fix: Capture user info at right moment

untitaker · untitaker · commit 98399ca5f03f · 2018-08-28T20:48:34.000+02:00
diff --git a/sentry_sdk/integrations/_wsgi.py b/sentry_sdk/integrations/_wsgi.py
@@ -208,6 +208,10 @@ def event_processor(event):
             # if the code below fails halfway through we at least have some data
             request_info = event.setdefault("request", {})
 
+            if _should_send_default_pii():
+                user_info = event.setdefault("user", {})
+                user_info["ip_address"] = get_client_ip(environ)
+
             if "query_string" not in request_info:
                 request_info["query_string"] = environ.get("QUERY_STRING")
 
diff --git a/sentry_sdk/integrations/django.py b/sentry_sdk/integrations/django.py
@@ -12,7 +12,7 @@
 
 from sentry_sdk import get_current_hub, capture_exception
 from sentry_sdk.hub import _internal_exceptions, _should_send_default_pii
-from ._wsgi import RequestExtractor, get_client_ip, run_wsgi_app
+from ._wsgi import RequestExtractor, run_wsgi_app
 from . import Integration
 
 
@@ -139,10 +139,7 @@ def size_of_file(self, file):
 
 
 def _set_user_info(request, event):
-    if "user" in event:
-        return
-
-    event["user"] = user_info = {"ip_address": get_client_ip(request.META)}
+    user_info = event.setdefault("user", {})
 
     user = getattr(request, "user", None)
 
diff --git a/sentry_sdk/integrations/flask.py b/sentry_sdk/integrations/flask.py
@@ -8,15 +8,16 @@
 from . import Integration
 
 try:
-    from flask_login import current_user as current_user_proxy
+    import flask_login
 except ImportError:
-    current_user_proxy = None
+    flask_login = None
 
 from flask import Flask, _request_ctx_stack, _app_ctx_stack
 from flask.signals import (
     appcontext_pushed,
     appcontext_tearing_down,
     got_request_exception,
+    request_started,
 )
 
 
@@ -29,6 +30,7 @@ def __init__(self):
     def install(self, client):
         appcontext_pushed.connect(_push_appctx)
         appcontext_tearing_down.connect(_pop_appctx)
+        request_started.connect(_request_started)
         got_request_exception.connect(_capture_exception)
 
         old_app = Flask.__call__
@@ -45,61 +47,25 @@ def _push_appctx(*args, **kwargs):
     # always want to push scope regardless of whether WSGI app might already
     # have (not the case for CLI for example)
     get_current_hub().push_scope()
+    get_current_hub().add_event_processor(_make_user_event_processor)
 
-    app = getattr(_app_ctx_stack.top, "app", None)
 
-    try:
-        weak_request = weakref.ref(_request_ctx_stack.top.request)
-    except AttributeError:
-        weak_request = lambda: None
+def _pop_appctx(*args, **kwargs):
+    get_current_hub().pop_scope_unsafe()
 
-    try:
-        weak_user = weakref.ref(current_user_proxy._get_current_object())
-    except (AttributeError, RuntimeError, TypeError):
-        weak_user = lambda: None
 
+def _request_started(sender, **kwargs):
+    weak_request = weakref.ref(_request_ctx_stack.top.request)
+    app = _app_ctx_stack.top.app
     get_current_hub().add_event_processor(
-        lambda: _make_event_processor(app, weak_request, weak_user)
+        lambda: _make_request_event_processor(app, weak_request)
     )
 
 
-def _pop_appctx(*args, **kwargs):
-    get_current_hub().pop_scope_unsafe()
-
-
 def _capture_exception(sender, exception, **kwargs):
     capture_exception(exception)
 
 
-def _make_event_processor(app, weak_request, weak_user):
-    client_options = get_current_hub().client.options
-
-    def event_processor(event):
-        request = weak_request()
-
-        # if the request is gone we are fine not logging the data from
-        # it.  This might happen if the processor is pushed away to
-        # another thread.
-        if request is not None:
-            if "transaction" not in event:
-                try:
-                    event["transaction"] = request.url_rule.endpoint
-                except Exception:
-                    pass
-
-            with _internal_exceptions():
-                FlaskRequestExtractor(request).extract_into_event(event, client_options)
-
-        if _should_send_default_pii():
-            with _internal_exceptions():
-                _set_user_info(request, weak_user(), event)
-
-        with _internal_exceptions():
-            _process_frames(app, event)
-
-    return event_processor
-
-
 def _process_frames(app, event):
     for frame in event.iter_frames():
         if "in_app" in frame:
@@ -139,28 +105,57 @@ def size_of_file(self, file):
         return file.content_length
 
 
-def _set_user_info(request, user, event):
-    if "user" in event:
-        return
-
-    if request:
-        try:
-            ip_address = request.access_route[0]
-        except IndexError:
-            ip_address = request.remote_addr
-    else:
-        ip_address = None
-
-    user_info = {"id": None, "ip_address": ip_address}
-
-    try:
-        user_info["id"] = user.get_id()
-        # TODO: more configurable user attrs here
-    except AttributeError:
-        # might happen if:
-        # - flask_login could not be imported
-        # - flask_login is not configured
-        # - no user is logged in
-        pass
+def _make_request_event_processor(app, weak_request):
+    client_options = get_current_hub().client.options
+
+    def inner(event):
+        with _internal_exceptions():
+            _process_frames(app, event)
+
+        request = weak_request()
+
+        # if the request is gone we are fine not logging the data from
+        # it.  This might happen if the processor is pushed away to
+        # another thread.
+        if request is None:
+            return
+
+        if "transaction" not in event:
+            try:
+                event["transaction"] = request.url_rule.endpoint
+            except Exception:
+                pass
+
+        with _internal_exceptions():
+            FlaskRequestExtractor(request).extract_into_event(event, client_options)
+
+    return inner
+
+
+def _make_user_event_processor():
+    def inner(event):
+        if flask_login is None or not _should_send_default_pii():
+            return
+
+        user = flask_login.current_user
+        if user is None:
+            return
+
+        with _internal_exceptions():
+            # Access this object as late as possible as accessing the user
+            # is relatively costly
+
+            user_info = event.setdefault("user", {})
+
+            if user_info.get("id", None) is None:
+                try:
+                    user_info["id"] = user.get_id()
+                    # TODO: more configurable user attrs here
+                except AttributeError:
+                    # might happen if:
+                    # - flask_login could not be imported
+                    # - flask_login is not configured
+                    # - no user is logged in
+                    pass
 
-    event["user"] = user_info
+    return inner
diff --git a/tests/integrations/flask/test_flask.py b/tests/integrations/flask/test_flask.py
@@ -74,7 +74,7 @@ def index():
 def test_flask_login_not_installed(sentry_init, app, capture_events, monkeypatch):
     sentry_init(integrations=[flask_sentry.FlaskIntegration()])
 
-    monkeypatch.setattr(flask_sentry, "current_user", None)
+    monkeypatch.setattr(flask_sentry, "flask_login", None)
 
     events = capture_events()
 
@@ -88,7 +88,7 @@ def test_flask_login_not_installed(sentry_init, app, capture_events, monkeypatch
 def test_flask_login_not_configured(sentry_init, app, capture_events, monkeypatch):
     sentry_init(integrations=[flask_sentry.FlaskIntegration()])
 
-    assert flask_sentry.current_user is not None
+    assert flask_sentry.flask_login
 
     events = capture_events()
     client = app.test_client()
