Issue #3709: a flush_headers method to BaseHTTPRequestHandler which manages the

sending of headers to output stream and flushing the internal headers buffer.
Patch contribution by Andrew Schaaf

4 files changed, 92 insertions, 30 deletions

diff --git a/Doc/library/http.server.rst b/Doc/library/http.server.rst
index e3a3a10..d9539f7 100644
--- a/Doc/library/http.server.rst
+++ b/Doc/library/http.server.rst
@@ -179,16 +179,17 @@ of which this module provides three different variants:
 
    .. method:: send_response(code, message=None)
 
-      Sends a response header and logs the accepted request. The HTTP response
-      line is sent, followed by *Server* and *Date* headers. The values for
-      these two headers are picked up from the :meth:`version_string` and
-      :meth:`date_time_string` methods, respectively.
+      Adds a response header to the headers buffer and logs the accepted
+      request. The HTTP response line is sent, followed by *Server* and
+      *Date* headers. The values for these two headers are picked up from
+      the :meth:`version_string` and :meth:`date_time_string` methods,
+      respectively.
 
    .. method:: send_header(keyword, value)
 
-      Stores the HTTP header to an internal buffer which will be written to the
-      output stream when :meth:`end_headers` method is invoked.
-      *keyword* should specify the header keyword, with *value*
+      Adds the HTTP header to an internal buffer which will be written to the
+      output stream when either :meth:`end_headers` or :meth:`flush_headers`
+      is invoked. *keyword* should specify the header keyword, with *value*
       specifying its value.
 
       .. versionchanged:: 3.2 Storing the headers in an internal buffer
@@ -205,11 +206,19 @@ of which this module provides three different variants:
 
    .. method:: end_headers()
 
-      Write the buffered HTTP headers to the output stream and send a blank
-      line, indicating the end of the HTTP headers in the response.
+      Adds a blank line
+      (indicating the end of the HTTP headers in the response)
+      to the headers buffer and calls :meth:`flush_headers()`
 
       .. versionchanged:: 3.2 Writing the buffered headers to the output stream.
 
+   .. method:: flush_headers()
+
+      Finally send the headers to the output stream and flush the internal
+      headers buffer.
+
+      .. versionadded:: 3.3
+
    .. method:: log_request(code='-', size='-')
 
       Logs an accepted (successful) request. *code* should specify the numeric
diff --git a/Lib/http/server.py b/Lib/http/server.py
index 6aacbbd..1d193f8 100644
--- a/Lib/http/server.py
+++ b/Lib/http/server.py
@@ -355,6 +355,7 @@ class BaseHTTPRequestHandler(socketserver.StreamRequestHandler):
 
         """
         self.send_response_only(100)
+        self.flush_headers()
         return True
 
     def handle_one_request(self):
@@ -432,7 +433,8 @@ class BaseHTTPRequestHandler(socketserver.StreamRequestHandler):
             self.wfile.write(content.encode('UTF-8', 'replace'))
 
     def send_response(self, code, message=None):
-        """Send the response header and log the response code.
+        """Add the response header to the headers buffer and log the
+        response code.
 
         Also send two standard headers with the server software
         version and the current date.
@@ -451,11 +453,14 @@ class BaseHTTPRequestHandler(socketserver.StreamRequestHandler):
             else:
                 message = ''
         if self.request_version != 'HTTP/0.9':
-            self.wfile.write(("%s %d %s\r\n" %
-                              (self.protocol_version, code, message)).encode('latin-1', 'strict'))
+            if not hasattr(self, '_headers_buffer'):
+                self._headers_buffer = []
+            self._headers_buffer.append(("%s %d %s\r\n" %
+                    (self.protocol_version, code, message)).encode(
+                        'latin-1', 'strict'))
 
     def send_header(self, keyword, value):
-        """Send a MIME header."""
+        """Send a MIME header to the headers buffer."""
         if self.request_version != 'HTTP/0.9':
             if not hasattr(self, '_headers_buffer'):
                 self._headers_buffer = []
@@ -472,6 +477,10 @@ class BaseHTTPRequestHandler(socketserver.StreamRequestHandler):
         """Send the blank line ending the MIME headers."""
         if self.request_version != 'HTTP/0.9':
             self._headers_buffer.append(b"\r\n")
+            self.flush_headers()
+
+    def flush_headers(self):
+        if hasattr(self, '_headers_buffer'):
             self.wfile.write(b"".join(self._headers_buffer))
             self._headers_buffer = []
 
@@ -1081,6 +1090,7 @@ class CGIHTTPRequestHandler(SimpleHTTPRequestHandler):
             env.setdefault(k, "")
 
         self.send_response(200, "Script output follows")
+        self.flush_headers()
 
         decoded_query = query.replace('+', ' ')
 
diff --git a/Lib/test/test_httpservers.py b/Lib/test/test_httpservers.py
index 39ebc26..1bbaf0e 100644
--- a/Lib/test/test_httpservers.py
+++ b/Lib/test/test_httpservers.py
@@ -461,6 +461,23 @@ class RejectingSocketlessRequestHandler(SocketlessRequestHandler):
         self.send_error(417)
         return False
 
+
+class AuditableBytesIO:
+
+    def __init__(self):
+        self.datas = []
+
+    def write(self, data):
+        self.datas.append(data)
+
+    def getData(self):
+        return b''.join(self.datas)
+
+    @property
+    def numWrites(self):
+        return len(self.datas)
+
+
 class BaseHTTPRequestHandlerTestCase(unittest.TestCase):
     """Test the functionality of the BaseHTTPServer.
 
@@ -527,27 +544,49 @@ class BaseHTTPRequestHandlerTestCase(unittest.TestCase):
         self.verify_get_called()
         self.assertEqual(result[-1], b'<html><body>Data</body></html>\r\n')
 
-    def test_header_buffering(self):
+    def test_header_buffering_of_send_error(self):
 
-        def _readAndReseek(f):
-            pos = f.tell()
-            f.seek(0)
-            data = f.read()
-            f.seek(pos)
-            return data
+        input = BytesIO(b'GET / HTTP/1.1\r\n\r\n')
+        output = AuditableBytesIO()
+        handler = SocketlessRequestHandler()
+        handler.rfile = input
+        handler.wfile = output
+        handler.request_version = 'HTTP/1.1'
+        handler.requestline = ''
+        handler.command = None
+
+        handler.send_error(418)
+        self.assertEqual(output.numWrites, 2)
+
+    def test_header_buffering_of_send_response_only(self):
 
         input = BytesIO(b'GET / HTTP/1.1\r\n\r\n')
-        output = BytesIO()
-        self.handler.rfile = input
-        self.handler.wfile = output
-        self.handler.request_version = 'HTTP/1.1'
+        output = AuditableBytesIO()
+        handler = SocketlessRequestHandler()
+        handler.rfile = input
+        handler.wfile = output
+        handler.request_version = 'HTTP/1.1'
 
-        self.handler.send_header('Foo', 'foo')
-        self.handler.send_header('bar', 'bar')
-        self.assertEqual(_readAndReseek(output), b'')
-        self.handler.end_headers()
-        self.assertEqual(_readAndReseek(output),
-                         b'Foo: foo\r\nbar: bar\r\n\r\n')
+        handler.send_response_only(418)
+        self.assertEqual(output.numWrites, 0)
+        handler.end_headers()
+        self.assertEqual(output.numWrites, 1)
+
+    def test_header_buffering_of_send_header(self):
+
+        input = BytesIO(b'GET / HTTP/1.1\r\n\r\n')
+        output = AuditableBytesIO()
+        handler = SocketlessRequestHandler()
+        handler.rfile = input
+        handler.wfile = output
+        handler.request_version = 'HTTP/1.1'
+
+        handler.send_header('Foo', 'foo')
+        handler.send_header('bar', 'bar')
+        self.assertEqual(output.numWrites, 0)
+        handler.end_headers()
+        self.assertEqual(output.getData(), b'Foo: foo\r\nbar: bar\r\n\r\n')
+        self.assertEqual(output.numWrites, 1)
 
     def test_header_unbuffered_when_continue(self):
 
diff --git a/Misc/NEWS b/Misc/NEWS
index 18a2a55..1c5ce69 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -140,6 +140,10 @@ Core and Builtins
 Library
 -------
 
+- Issue #3709: a flush_headers method to BaseHTTPRequestHandler which manages
+  the sending of headers to output stream and flushing the internal headers
+  buffer. Patch contribution by Andrew Schaaf
+
 - Issue #11743: Rewrite multiprocessing connection classes in pure Python.
 
 - Issue #11164: Stop trying to use _xmlplus in the xml module.