getsentry · untitaker · Jul 20, 2018 · Jul 23, 2018 · Jul 23, 2018 · Jul 23, 2018
@@ -1,7 +1,49 @@
+import io
+
 def get_environ(environ):
     """
     Returns our whitelisted environment variables.
     """
     for key in ("REMOTE_ADDR", "SERVER_NAME", "SERVER_PORT"):
         if key in environ:
             yield key, environ[key]
+
+
+def peek_io_stream(stream, input_peek_len=1000):
+    if not hasattr(stream, 'read'):
+        return None
+
+    peek = stream.read(input_peek_len)
+    while len(peek) < input_peek_len:
+        x = stream.read(input_peek_len)
+        if not x:
+            break
+        peek += x
+
+    if isinstance(peek, bytes):
+        peek_io = io.BytesIO(peek)
+    else:
+        peek_io = io.StringIO(peek)
+
+    return peek, ChainedIO(peek_io, stream)
+
+
+class ChainedIO(object):
+    def __init__(self, *ios):
+        self.ios = list(ios)
+        self._last_type = bytes
+
+    def read(self, n=None):
+        if not self.ios:
+            return self._last_type()
+
+        rv = self.ios[0].read(n)
+
+        if n is None or n < 0 or len(rv) < n:  # stream exhausted
+            self.ios = self.ios[1:]  # jump to next stream
+
+        self._last_type = type(rv)
+
+        if n is None or n < 0:
+            rv += self.read(n)
+        return rv
@@ -1,7 +1,12 @@
 from __future__ import absolute_import
 
+import base64
+
 from sentry_sdk import capture_exception, configure_scope, get_current_hub
-from ._wsgi import get_environ
+from sentry_sdk.stripping import strip_string
+from sentry_sdk._compat import text_type
+
+from ._wsgi import get_environ, peek_io_stream
 
 try:
     from flask_login import current_user
@@ -55,30 +60,78 @@ def _before_request(*args, **kwargs):
                 scope.transaction = request.url_rule.endpoint
 
             try:
-                scope.request = _get_request_info()
+                _set_request_info(scope)
             except Exception:
                 get_current_hub().capture_internal_exception()
 
             try:
-                scope.user = _get_user_info()
+                _set_user_info(scope)
             except Exception:
                 get_current_hub().capture_internal_exception()
     except Exception:
         get_current_hub().capture_internal_exception()
 
-
-def _get_request_info():
-    return {
-        "url": "%s://%s%s" % (request.scheme, request.host, request.path),
-        "query_string": request.query_string,
-        "method": request.method,
-        "data": request.get_data(cache=True, as_text=True, parse_form_data=True),
-        "headers": dict(request.headers),
-        "env": get_environ(request.environ),
+def _set_request_info(scope):
+    request_info = {
+        'url': '%s://%s%s' % (
+            request.scheme,
+            request.host,
+            request.path
+        ),
+        'query_string': request.query_string,
+        'method': request.method,
+        'headers': dict(request.headers),
+        'env': dict(get_environ(request.environ)),
+        'cookies': dict(request.cookies)
     }
 
-
-def _get_user_info():
+    scope.request = request_info
+    # if this crashes we at least have the rest of the request already set
+    _set_request_body(request_info, scope)
+
+
+def _set_request_body(request_info, scope):
+    # TODO: peek for length of configured max length
+    peek, request.stream = peek_io_stream(request.stream)
+
+    data = None
+    meta = None
+
+    if request.form:
+        data = request.form
+        if request.mimetype == 'multipart/form-data':
+            ct = 'multipart'
+        else:
+            ct = 'urlencoded'
+        repr = 'structured'
+    elif request.json is not None:
+        data = request.json
+        ct = 'json'
+        repr = 'structured'
+    else:
+        try:
+            if isinstance(peek, text_type):
+                data = peek
+            else:
+                data = peek.decode('utf-8')
+        except UnicodeDecodeError:
+            ct = 'bytes'
+            repr = 'base64'
+            data = base64.b64encode(peek).decode('ascii')
+        else:
+            ct = 'plain'
+            repr = 'other'
+
+        data = strip_string(
+            data,
+            assume_length=request.content_length
+        )
+
+    request_info['data'] = data
+    request_info['data_info'] = {'ct': ct, 'repr': repr}
+
+
+def _set_user_info(scope):
     try:
         ip_address = request.access_route[0]
     except IndexError:
@@ -96,4 +149,4 @@ def _get_user_info():
         # - no user is logged in
         pass
 
-    return user_info
+    scope.user = user_info
@@ -1,3 +1,4 @@
+import json
 import pytest
 
 flask = pytest.importorskip("flask")
@@ -127,3 +128,117 @@ def login():
         assert event.get("user", {}).get("id") is None
     else:
         assert event["user"]["id"] == str(user_id)
+
+
+def test_flask_large_json_request(capture_events, app):
+    data = {'foo': {'bar': 'a' * 2000}}
+    @app.route('/', methods=['POST'])
+    def index():
+        assert request.json == data
+        assert request.data == json.dumps(data).encode('ascii')
+        assert not request.form
+        capture_message('hi')
+        return 'ok'
+
+    client = app.test_client()
+    response = client.post(
+        '/',
+        content_type='application/json',
+        data=json.dumps(data)
+    )
+    assert response.status_code == 200
+
+    event, = capture_events
+    assert event['']['request']['data']['foo']['bar'] == {
+        '': {'len': 2000, 'rem': [['!len', 'x', 509, 512]]}
+    }
+    assert len(event['request']['data']['foo']['bar']) == 512
+    assert event['request']['data_info'] == {
+        'ct': 'json', 'repr': 'structured'
+    }
+
+
+def test_flask_large_formdata_request(capture_events, app):
+    data = {'foo': 'a' * 2000}
+    @app.route('/', methods=['POST'])
+    def index():
+        assert request.form['foo'] == data['foo']
+        assert not request.data
+        assert not request.json
+        capture_message('hi')
+        return 'ok'
+
+    client = app.test_client()
+    response = client.post(
+        '/',
+        data=data
+    )
+    assert response.status_code == 200
+
+    event, = capture_events
+    assert event['']['request']['data']['foo'] == {
+        '': {'len': 2000, 'rem': [['!len', 'x', 509, 512]]}
+    }
+    assert len(event['request']['data']['foo']) == 512
+    assert event['request']['data_info'] == {
+        'ct': 'urlencoded', 'repr': 'structured'
+    }
+
+
+@pytest.mark.parametrize('input_char', [u'a', b'a'])
+def test_flask_large_text_request(input_char, capture_events, app):
+    data = input_char * 2000
+    @app.route('/', methods=['POST'])
+    def index():
+        assert not request.form
+        if isinstance(data, bytes):
+            assert request.data == data
+        else:
+            assert request.data == data.encode('ascii')
+        assert not request.json
+        capture_message('hi')
+        return 'ok'
+
+    client = app.test_client()
+    response = client.post(
+        '/',
+        data=data
+    )
+    assert response.status_code == 200
+
+    event, = capture_events
+    assert event['']['request']['data'] == {
+        '': {'len': 2000, 'rem': [['!len', 'x', 509, 512]]}
+    }
+    assert len(event['request']['data']) == 512
+    assert event['request']['data_info'] == {
+        'ct': 'plain', 'repr': 'other'
+    }
+
+
+
+def test_flask_large_bytes_request(capture_events, app):
+    data = b'\xc3' * 2000
+    @app.route('/', methods=['POST'])
+    def index():
+        assert not request.form
+        assert request.data == data
+        assert not request.json
+        capture_message('hi')
+        return 'ok'
+
+    client = app.test_client()
+    response = client.post(
+        '/',
+        data=data
+    )
+    assert response.status_code == 200
+
+    event, = capture_events
+    assert event['']['request']['data'] == {
+        '': {'len': 2000, 'rem': [['!len', 'x', 509, 512]]}
+    }
+    assert len(event['request']['data']) == 512
+    assert event['request']['data_info'] == {
+        'ct': 'bytes', 'repr': 'base64'
+    }