Repository: tornadoweb/tornado
Pull request #1419 (2015-04-14 15:08:30)

Add max_body_size parameter to SimpleHTTPClient
This change allows the SimpleHTTPClient to receive responses larger then `max_buffer_size`. This should allow downloading of very large files without increasing max_buffer_size to unreasonable values.

Intended usage:

```python
#!/usr/bin/env python
import tornado.gen
import tornado.ioloop
import tornado.httpclient

@tornado.gen.coroutine
def main():

    # download 1.1 GB file
    url = 'http://releases.ubuntu.com/14.10/ubuntu-14.10-desktop-amd64.iso'
    http = tornado.httpclient.AsyncHTTPClient(max_body_size=1024*1024*1024*5)

    with open('/dev/null', 'wb') as f:

        response = yield http.fetch(url,
                                    streaming_callback=f.write,
                                    request_timeout=3600)

tornado.ioloop.IOLoop.current().run_sync(main)
```


Patch:
diff --git a/tornado/simple_httpclient.py b/tornado/simple_httpclient.py
index 6321a81d..0319ff44 100644
--- a/tornado/simple_httpclient.py
+++ b/tornado/simple_httpclient.py
@@ -60,7 +60,8 @@ class SimpleAsyncHTTPClient(AsyncHTTPClient):
     """
     def initialize(self, io_loop, max_clients=10,
                    hostname_mapping=None, max_buffer_size=104857600,
-                   resolver=None, defaults=None, max_header_size=None):
+                   resolver=None, defaults=None, max_header_size=None,
+                   max_body_size=None):
         """Creates a AsyncHTTPClient.
 
         Only a single AsyncHTTPClient instance exists per IOLoop
@@ -88,6 +89,7 @@ class SimpleAsyncHTTPClient(AsyncHTTPClient):
         self.waiting = {}
         self.max_buffer_size = max_buffer_size
         self.max_header_size = max_header_size
+        self.max_body_size = max_body_size
         # TCPClient could create a Resolver for us, but we have to do it
         # ourselves to support hostname_mapping.
         if resolver:
@@ -142,7 +144,7 @@ class SimpleAsyncHTTPClient(AsyncHTTPClient):
         self._connection_class()(
             self.io_loop, self, request, release_callback,
             final_callback, self.max_buffer_size, self.tcp_client,
-            self.max_header_size)
+            self.max_header_size, self.max_body_size)
 
     def _release_fetch(self, key):
         del self.active[key]
@@ -170,7 +172,7 @@ class _HTTPConnection(httputil.HTTPMessageDelegate):
 
     def __init__(self, io_loop, client, request, release_callback,
                  final_callback, max_buffer_size, tcp_client,
-                 max_header_size):
+                 max_header_size, max_body_size):
         self.start_time = io_loop.time()
         self.io_loop = io_loop
         self.client = client
@@ -180,6 +182,7 @@ class _HTTPConnection(httputil.HTTPMessageDelegate):
         self.max_buffer_size = max_buffer_size
         self.tcp_client = tcp_client
         self.max_header_size = max_header_size
+        self.max_body_size = max_body_size
         self.code = None
         self.headers = None
         self.chunks = []
@@ -368,6 +371,7 @@ class _HTTPConnection(httputil.HTTPMessageDelegate):
             HTTP1ConnectionParameters(
                 no_keep_alive=True,
                 max_header_size=self.max_header_size,
+                max_body_size=self.max_body_size,
                 decompress=self.request.decompress_response),
             self._sockaddr)
         return connection
diff --git a/tornado/websocket.py b/tornado/websocket.py
index 581ff7b7..adf238be 100644
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -908,7 +908,7 @@ class WebSocketClientConnection(simple_httpclient._HTTPConnection):
         self.tcp_client = TCPClient(io_loop=io_loop)
         super(WebSocketClientConnection, self).__init__(
             io_loop, None, request, lambda: None, self._on_http_response,
-            104857600, self.tcp_client, 65536)
+            104857600, self.tcp_client, 65536, 104857600)
 
     def close(self, code=None, reason=None):
         """Closes the websocket connection.


Test patch:
diff --git a/tornado/test/simple_httpclient_test.py b/tornado/test/simple_httpclient_test.py
index fdc3cadf..c0de22b7 100644
--- a/tornado/test/simple_httpclient_test.py
+++ b/tornado/test/simple_httpclient_test.py
@@ -631,3 +631,49 @@ class MaxHeaderSizeTest(AsyncHTTPTestCase):
         with ExpectLog(gen_log, "Unsatisfiable read"):
             response = self.fetch('/large')
         self.assertEqual(response.code, 599)
+
+
+class MaxBodySizeTest(AsyncHTTPTestCase):
+    def get_app(self):
+        class SmallBody(RequestHandler):
+            def get(self):
+                self.write("a"*1024*64)
+
+        class LargeBody(RequestHandler):
+            def get(self):
+                self.write("a"*1024*100)
+
+        return Application([('/small', SmallBody),
+                            ('/large', LargeBody)])
+
+    def get_http_client(self):
+        return SimpleAsyncHTTPClient(io_loop=self.io_loop, max_body_size=1024*64)
+
+    def test_small_body(self):
+        response = self.fetch('/small')
+        response.rethrow()
+        self.assertEqual(response.body, b'a'*1024*64)
+
+    def test_large_body(self):
+        with ExpectLog(gen_log, "Malformed HTTP message from None: Content-Length too long"):
+            response = self.fetch('/large')
+        self.assertEqual(response.code, 599)
+
+
+class MaxBufferSizeTest(AsyncHTTPTestCase):
+    def get_app(self):
+
+        class LargeBody(RequestHandler):
+            def get(self):
+                self.write("a"*1024*100)
+
+        return Application([('/large', LargeBody)])
+
+    def get_http_client(self):
+        # 100KB body with 64KB buffer
+        return SimpleAsyncHTTPClient(io_loop=self.io_loop, max_body_size=1024*100, max_buffer_size=1024*64)
+
+    def test_large_body(self):
+        response = self.fetch('/large')
+        response.rethrow()
+        self.assertEqual(response.body, b'a'*1024*100)