From a7960bcb3575fd9fcd5180986ddcdba1caa46dd5 Mon Sep 17 00:00:00 2001 From: Matthew Wood Date: Mon, 18 Mar 2013 16:06:24 -0600 Subject: [PATCH] Fixed #18972 -- Refactored bundled wsgi server's chunking algorithm. Thanks to amosonn at yahoo.com for the report, @doda for the initial patch and @datagrok for the revamped logic and test case. --- django/core/servers/basehttp.py | 26 ++++++-------- tests/builtin_server/tests.py | 63 ++++++++++++++++++++++++++++----- 2 files changed, 66 insertions(+), 23 deletions(-) diff --git a/django/core/servers/basehttp.py b/django/core/servers/basehttp.py index 8e390bce27..d329221ce4 100644 --- a/django/core/servers/basehttp.py +++ b/django/core/servers/basehttp.py @@ -9,7 +9,7 @@ been reviewed for security issues. DON'T USE IT FOR PRODUCTION USE! from __future__ import unicode_literals -import os +from io import BytesIO import socket import sys import traceback @@ -26,7 +26,13 @@ from django.core.wsgi import get_wsgi_application from django.utils.module_loading import import_by_path from django.utils import six -__all__ = ['WSGIServer', 'WSGIRequestHandler'] +__all__ = ('WSGIServer', 'WSGIRequestHandler', 'MAX_SOCKET_CHUNK_SIZE') + + +# If data is too large, socket will choke, so write chunks no larger than 32MB +# at a time. The rationale behind the 32MB can be found on Django's Trac: +# https://code.djangoproject.com/ticket/5596#comment:4 +MAX_SOCKET_CHUNK_SIZE = 32 * 1024 * 1024 # 32 MB def get_internal_wsgi_application(): @@ -78,19 +84,9 @@ class ServerHandler(simple_server.ServerHandler, object): self.bytes_sent += len(data) # XXX check Content-Length and truncate if too many bytes written? - - # If data is too large, socket will choke, so write chunks no larger - # than 32MB at a time. - length = len(data) - if length > 33554432: - offset = 0 - while offset < length: - chunk_size = min(33554432, length) - self._write(data[offset:offset+chunk_size]) - self._flush() - offset += chunk_size - else: - self._write(data) + data = BytesIO(data) + for chunk in iter(lambda: data.read(MAX_SOCKET_CHUNK_SIZE), b''): + self._write(chunk) self._flush() def error_output(self, environ, start_response): diff --git a/tests/builtin_server/tests.py b/tests/builtin_server/tests.py index 041bb3c319..662466a110 100644 --- a/tests/builtin_server/tests.py +++ b/tests/builtin_server/tests.py @@ -2,22 +2,18 @@ from __future__ import unicode_literals from io import BytesIO -from django.core.servers.basehttp import ServerHandler +from django.core.servers.basehttp import ServerHandler, MAX_SOCKET_CHUNK_SIZE from django.utils.unittest import TestCase -# -# Tests for #9659: wsgi.file_wrapper in the builtin server. -# We need to mock a couple of handlers and keep track of what -# gets called when using a couple kinds of WSGI apps. -# class DummyHandler(object): - def log_request(*args, **kwargs): + def log_request(self, *args, **kwargs): pass + class FileWrapperHandler(ServerHandler): def __init__(self, *args, **kwargs): - ServerHandler.__init__(self, *args, **kwargs) + super(FileWrapperHandler, self).__init__(*args, **kwargs) self.request_handler = DummyHandler() self._used_sendfile = False @@ -25,17 +21,24 @@ class FileWrapperHandler(ServerHandler): self._used_sendfile = True return True + def wsgi_app(environ, start_response): start_response(str('200 OK'), [(str('Content-Type'), str('text/plain'))]) return [b'Hello World!'] + def wsgi_app_file_wrapper(environ, start_response): start_response(str('200 OK'), [(str('Content-Type'), str('text/plain'))]) return environ['wsgi.file_wrapper'](BytesIO(b'foo')) + class WSGIFileWrapperTests(TestCase): """ Test that the wsgi.file_wrapper works for the builting server. + + Tests for #9659: wsgi.file_wrapper in the builtin server. + We need to mock a couple of handlers and keep track of what + gets called when using a couple kinds of WSGI apps. """ def test_file_wrapper_uses_sendfile(self): @@ -53,3 +56,47 @@ class WSGIFileWrapperTests(TestCase): self.assertFalse(handler._used_sendfile) self.assertEqual(handler.stdout.getvalue().splitlines()[-1], b'Hello World!') self.assertEqual(handler.stderr.getvalue(), b'') + + +class WriteChunkCounterHandler(ServerHandler): + """ + Server handler that counts the number of chunks written after headers were + sent. Used to make sure large response body chunking works properly. + """ + + def __init__(self, *args, **kwargs): + super(WriteChunkCounterHandler, self).__init__(*args, **kwargs) + self.request_handler = DummyHandler() + self.headers_written = False + self.write_chunk_counter = 0 + + def send_headers(self): + super(WriteChunkCounterHandler, self).send_headers() + self.headers_written = True + + def _write(self, data): + if self.headers_written: + self.write_chunk_counter += 1 + self.stdout.write(data) + + +def send_big_data_app(environ, start_response): + start_response(str('200 OK'), [(str('Content-Type'), str('text/plain'))]) + # Return a blob of data that is 1.5 times the maximum chunk size. + return [b'x' * (MAX_SOCKET_CHUNK_SIZE + MAX_SOCKET_CHUNK_SIZE // 2)] + + +class ServerHandlerChunksProperly(TestCase): + """ + Test that the ServerHandler chunks data properly. + + Tests for #18972: The logic that performs the math to break data into + 32MB (MAX_SOCKET_CHUNK_SIZE) chunks was flawed, BUT it didn't actually + cause any problems. + """ + + def test_chunked_data(self): + env = {'SERVER_PROTOCOL': 'HTTP/1.0'} + handler = WriteChunkCounterHandler(None, BytesIO(), BytesIO(), env) + handler.run(send_big_data_app) + self.assertEqual(handler.write_chunk_counter, 2)