Merge pull request #602 from mozilla-services/feat/501

feat: Limit the size of allowed HTTP bodies & headers

autopush/main.py

@@ -32,6 +32,7 @@
 )
 from autopush.web.simplepush import SimplePushHandler
 from autopush.web.webpush import WebPushHandler
+from autopush.web.limitedhttpconnection import LimitedHTTPConnection
 
 
 shared_config_files = [
@@ -553,6 +554,8 @@ def endpoint_main(sysargs=None, use_files=True):
         default_host=settings.hostname, debug=args.debug,
         log_function=skip_request_logging
     )
+    site.protocol = LimitedHTTPConnection
+    site.protocol.maxData = settings.max_data
     mount_health_handlers(site, settings)
 
     settings.metrics.start()

autopush/tests/test_limitedhttpconnection.py

@@ -0,0 +1,51 @@
+from io import BytesIO
+
+from mock import Mock
+from twisted.trial import unittest
+from nose.tools import eq_
+
+from autopush.web.limitedhttpconnection import (
+    LimitedHTTPConnection,
+)
+
+
+class TestLimitedHttpConnection(unittest.TestCase):
+    def test_lineRecieved(self):
+        mock_transport = Mock()
+        conn = LimitedHTTPConnection()
+        conn.factory = Mock()
+        conn.factory.settings = {}
+        conn.makeConnection(mock_transport)
+        conn._on_headers = Mock()
+
+        conn.maxHeaders = 2
+        conn.lineReceived("line 1")
+        eq_(conn._headersbuffer, ["line 1\r\n"])
+        conn.lineReceived("line 2")
+        conn.lineReceived("line 3")
+        mock_transport.loseConnection.assert_called()
+        conn.lineReceived("")
+        eq_(conn._headersbuffer, [])
+        conn._on_headers.assert_called()
+        eq_(conn._on_headers.call_args[0][0],
+            "line 1\r\nline 2\r\n")
+
+    def test_rawDataReceived(self):
+        mock_transport = Mock()
+        conn = LimitedHTTPConnection()
+        conn.factory = Mock()
+        conn.factory.settings = {}
+        conn.makeConnection(mock_transport)
+        conn._on_headers = Mock()
+        conn._on_request_body = Mock()
+        conn._contentbuffer = BytesIO()
+
+        conn.maxData = 10
+        conn.rawDataReceived("12345")
+        conn._contentbuffer = BytesIO()
+        conn.content_length = 3
+        conn.rawDataReceived("12345")
+        eq_(False, mock_transport.loseConnection.called)
+        conn._on_request_body.assert_called()
+        conn.rawDataReceived("12345678901")
+        mock_transport.loseConnection.assert_called()

autopush/web/limitedhttpconnection.py

@@ -0,0 +1,54 @@
+from cyclone import httpserver
+from twisted.logger import Logger
+
+
+class LimitedHTTPConnection(httpserver.HTTPConnection):
+    """
+    Limit the amount of data being sent to a reasonable amount.
+
+    twisted already limits TCP streamed chunk reads to 65K, with
+    ~16k per header line. By default, we'll limit the number of
+    header lines to 100, and the maximum amount of data for the body
+    to be 4K.
+
+    """
+    maxHeaders = 100
+    maxData = 1024*4
+
+    def lineReceived(self, line):
+        """Process a header line of data, ensuring we have not exceeded the
+        max number of allowable headers.
+
+        :param line: raw header line
+        """
+        if line:
+            if len(self._headersbuffer) == self.maxHeaders:
+                Logger().warn("Too many headers sent, terminating connection")
+                return self.lineLengthExceeded(line)
+            self._headersbuffer.append(line + self.delimiter)
+        else:
+            buff = "".join(self._headersbuffer)
+            self._headersbuffer = []
+            self._on_headers(buff)
+
+    def rawDataReceived(self, data):
+        """Process a raw chunk of data, ensuring we have not exceeded the
+        max size of a data block
+
+        :param data: raw data block
+        """
+        if len(data) > self.maxData:
+            Logger().warn("Too much data sent, terminating connection")
+            return self.lineLengthExceeded(data)
+        if self.content_length is not None:
+            data, rest = data[:self.content_length], data[self.content_length:]
+            self.content_length -= len(data)
+        else:
+            rest = ''
+
+        self._contentbuffer.write(data)
+        if self.content_length <= 0:
+            self._contentbuffer.seek(0, 0)
+            self._on_request_body(self._contentbuffer.read())
+            self._content_length = self._contentbuffer = None
+            self.setLineMode(rest)