diff --git a/autopush/main.py b/autopush/main.py index 04907306..eee864d1 100644 --- a/autopush/main.py +++ b/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() diff --git a/autopush/tests/test_limitedhttpconnection.py b/autopush/tests/test_limitedhttpconnection.py new file mode 100644 index 00000000..60c50074 --- /dev/null +++ b/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() diff --git a/autopush/web/limitedhttpconnection.py b/autopush/web/limitedhttpconnection.py new file mode 100644 index 00000000..323fdbf9 --- /dev/null +++ b/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)