gh-135056: Use response_headers only in SimpleHTTPRequestHandler

Author: aisipos

Doc/library/http.server.rst

@@ -375,8 +375,9 @@ instantiation, of which this module provides three different variants:
       The *directory* parameter accepts a :term:`path-like object`.
 
    .. versionchanged:: next
-      The *response_headers* parameter accepts an optional dictionary of
-      additional HTTP headers to add to each response.
+      Added *response_headers*, which accepts an optional dictionary of
+      additional HTTP headers to add to each successful HTTP status 200
+      response. All other status code responses will not include these headers.
 
    A lot of the work, such as parsing the request, is done by the base class
    :class:`BaseHTTPRequestHandler`.  This class implements the :func:`do_GET`

Lib/http/server.py

@@ -117,24 +117,13 @@ class HTTPServer(socketserver.TCPServer):
     allow_reuse_address = True    # Seems to make sense in testing environment
     allow_reuse_port = True
 
-    def __init__(self, *args, response_headers=None, **kwargs):
-        self.response_headers = response_headers
-        super().__init__(*args, **kwargs)
-
     def server_bind(self):
         """Override server_bind to store the server name."""
         socketserver.TCPServer.server_bind(self)
         host, port = self.server_address[:2]
         self.server_name = socket.getfqdn(host)
         self.server_port = port
 
-    def finish_request(self, request, client_address):
-        """Finish one request by instantiating RequestHandlerClass."""
-        args = (request, client_address, self)
-        kwargs = {}
-        if hasattr(self, 'response_headers'):
-            kwargs['response_headers'] = self.response_headers
-        self.RequestHandlerClass(request, client_address, self, **kwargs)
 
 class ThreadingHTTPServer(socketserver.ThreadingMixIn, HTTPServer):
     daemon_threads = True
@@ -143,7 +132,7 @@ class ThreadingHTTPServer(socketserver.ThreadingMixIn, HTTPServer):
 class HTTPSServer(HTTPServer):
     def __init__(self, server_address, RequestHandlerClass,
                  bind_and_activate=True, *, certfile, keyfile=None,
-                 password=None, alpn_protocols=None, **http_server_kwargs):
+                 password=None, alpn_protocols=None):
         try:
             import ssl
         except ImportError:
@@ -161,8 +150,7 @@ def __init__(self, server_address, RequestHandlerClass,
 
         super().__init__(server_address,
                          RequestHandlerClass,
-                         bind_and_activate,
-                         **http_server_kwargs)
+                         bind_and_activate)
 
     def server_activate(self):
         """Wrap the socket in SSLSocket."""
@@ -726,6 +714,13 @@ def do_HEAD(self):
         if f:
             f.close()
 
+    def send_custom_response_headers(self):
+        """Send the headers stored in self.response_headers"""
+        # User specified response_headers
+        if self.response_headers is not None:
+            for header, value in self.response_headers.items():
+                self.send_header(header, value)
+
     def send_head(self):
         """Common code for GET and HEAD commands.
 
@@ -749,10 +744,6 @@ def send_head(self):
                 new_url = urllib.parse.urlunsplit(new_parts)
                 self.send_header("Location", new_url)
                 self.send_header("Content-Length", "0")
-                # User specified response_headers
-                if self.response_headers is not None:
-                    for header, value in self.response_headers.items():
-                        self.send_header(header, value)
                 self.end_headers()
                 return None
             for index in self.index_pages:
@@ -812,9 +803,7 @@ def send_head(self):
             self.send_header("Content-Length", str(fs[6]))
             self.send_header("Last-Modified",
                 self.date_time_string(fs.st_mtime))
-            if self.response_headers is not None:
-                for header, value in self.response_headers.items():
-                    self.send_header(header, value)
+            self.send_custom_response_headers()
             self.end_headers()
             return f
         except:
@@ -879,6 +868,7 @@ def list_directory(self, path):
         self.send_response(HTTPStatus.OK)
         self.send_header("Content-type", "text/html; charset=%s" % enc)
         self.send_header("Content-Length", str(len(encoded)))
+        self.send_custom_response_headers()
         self.end_headers()
         return f
 
@@ -990,8 +980,7 @@ def _get_best_family(*address):
 def test(HandlerClass=BaseHTTPRequestHandler,
          ServerClass=ThreadingHTTPServer,
          protocol="HTTP/1.0", port=8000, bind=None,
-         tls_cert=None, tls_key=None, tls_password=None,
-         response_headers=None):
+         tls_cert=None, tls_key=None, tls_password=None):
     """Test the HTTP request handler class.
 
     This runs an HTTP server on port 8000 (or the port argument).
@@ -1002,10 +991,9 @@ def test(HandlerClass=BaseHTTPRequestHandler,
 
     if tls_cert:
         server = ServerClass(addr, HandlerClass, certfile=tls_cert,
-                             keyfile=tls_key, password=tls_password,
-                             response_headers=response_headers)
+                             keyfile=tls_key, password=tls_password)
     else:
-        server = ServerClass(addr, HandlerClass, response_headers=response_headers)
+        server = ServerClass(addr, HandlerClass)
 
     with server as httpd:
         host, port = httpd.socket.getsockname()[:2]
@@ -1067,6 +1055,10 @@ def _main(args=None):
         except OSError as e:
             parser.error(f"Failed to read TLS password file: {e}")
 
+    response_headers = {}
+    for header, value in args.header or []:
+        response_headers[header] = value
+
     # ensure dual-stack is not disabled; ref #38907
     class DualStackServerMixin:
 
@@ -1080,18 +1072,14 @@ def server_bind(self):
         def finish_request(self, request, client_address):
             self.RequestHandlerClass(request, client_address, self,
                                      directory=args.directory,
-                                     response_headers=self.response_headers)
+                                     response_headers=response_headers)
 
     class HTTPDualStackServer(DualStackServerMixin, ThreadingHTTPServer):
         pass
     class HTTPSDualStackServer(DualStackServerMixin, ThreadingHTTPSServer):
         pass
 
     ServerClass = HTTPSDualStackServer if args.tls_cert else HTTPDualStackServer
-    response_headers = {}
-    for header, value in args.header or []:
-        response_headers[header] = value
-
 
     test(
         HandlerClass=SimpleHTTPRequestHandler,
@@ -1102,7 +1090,6 @@ class HTTPSDualStackServer(DualStackServerMixin, ThreadingHTTPSServer):
         tls_cert=args.tls_cert,
         tls_key=args.tls_key,
         tls_password=tls_key_password,
-        response_headers=response_headers or None
     )
 
 

Lib/socketserver.py

@@ -757,7 +757,7 @@ class BaseRequestHandler:
 
     """
 
-    def __init__(self, request, client_address, server, **kwargs):
+    def __init__(self, request, client_address, server):
         self.request = request
         self.client_address = client_address
         self.server = server

Lib/test/test_httpservers.py

@@ -81,24 +81,21 @@ def test_https_server_raises_runtime_error(self):
 
 
 class TestServerThread(threading.Thread):
-    def __init__(self, test_object, request_handler, tls=None, server_kwargs=None):
+    def __init__(self, test_object, request_handler, tls=None):
         threading.Thread.__init__(self)
         self.request_handler = request_handler
         self.test_object = test_object
         self.tls = tls
-        self.server_kwargs = server_kwargs or {}
 
     def run(self):
         if self.tls:
             certfile, keyfile, password = self.tls
             self.server = create_https_server(
                 certfile, keyfile, password,
                 request_handler=self.request_handler,
-                **self.server_kwargs
             )
         else:
-            self.server = HTTPServer(('localhost', 0), self.request_handler,
-                                     **self.server_kwargs)
+            self.server = HTTPServer(('localhost', 0), self.request_handler)
         self.test_object.HOST, self.test_object.PORT = self.server.socket.getsockname()
         self.test_object.server_started.set()
         self.test_object = None
@@ -116,14 +113,12 @@ class BaseTestCase(unittest.TestCase):
 
     # Optional tuple (certfile, keyfile, password) to use for HTTPS servers.
     tls = None
-    server_kwargs = None
 
     def setUp(self):
         self._threads = threading_helper.threading_setup()
         os.environ = os_helper.EnvironmentVarGuard()
         self.server_started = threading.Event()
-        self.thread = TestServerThread(self, self.request_handler, self.tls,
-                                       self.server_kwargs)
+        self.thread = TestServerThread(self, self.request_handler, self.tls)
         self.thread.start()
         self.server_started.wait()
 
@@ -470,8 +465,14 @@ def test_err(self):
         self.assertEndsWith(lines[1], '"ERROR / HTTP/1.1" 404 -')
 
 
+class CustomHeaderSimpleHTTPRequestHandler(SimpleHTTPRequestHandler):
+    custom_headers = None
+    def __init__(self, *args, directory=None, response_headers=None, **kwargs):
+        super().__init__(*args, response_headers=self.custom_headers, **kwargs)
+
+
 class SimpleHTTPServerTestCase(BaseTestCase):
-    class request_handler(NoLogRequestHandler, SimpleHTTPRequestHandler):
+    class request_handler(NoLogRequestHandler, CustomHeaderSimpleHTTPRequestHandler):
         pass
 
     def setUp(self):
@@ -828,6 +829,26 @@ def test_path_without_leading_slash(self):
         self.assertEqual(response.getheader("Location"),
                          self.tempdir_name + "/?hi=1")
 
+    def test_custom_headers_list_dir(self):
+        with mock.patch.object(self.request_handler, 'custom_headers', new={
+            'X-Test1': 'test1',
+            'X-Test2': 'test2',
+        }):
+            response = self.request(self.base_url + '/')
+            self.assertEqual(response.getheader("X-Test1"), 'test1')
+            self.assertEqual(response.getheader("X-Test2"), 'test2')
+
+    def test_custom_headers_get_file(self):
+        with mock.patch.object(self.request_handler, 'custom_headers', new={
+            'X-Test1': 'test1',
+            'X-Test2': 'test2',
+        }):
+            data = b"Dummy index file\r\n"
+            with open(os.path.join(self.tempdir_name, 'index.html'), 'wb') as f:
+                f.write(data)
+            response = self.request(self.base_url + '/')
+            self.assertEqual(response.getheader("X-Test1"), 'test1')
+            self.assertEqual(response.getheader("X-Test2"), 'test2')
 
 class SocketlessRequestHandler(SimpleHTTPRequestHandler):
     def __init__(self, directory=None):
@@ -1311,7 +1332,6 @@ class CommandLineTestCase(unittest.TestCase):
         'tls_cert': None,
         'tls_key': None,
         'tls_password': None,
-        'response_headers': None,
     }
 
     def setUp(self):
@@ -1379,13 +1399,8 @@ def test_protocol_flag(self, mock_func):
 
     @mock.patch('http.server.test')
     def test_header_flag(self, mock_func):
+        call_args = self.args
         self.invoke_httpd('--header', 'h1', 'v1', '-H', 'h2', 'v2')
-        call_args = self.args | dict(
-            response_headers={
-                'h1': 'v1',
-                'h2': 'v2'
-            }
-        )
         mock_func.assert_called_once_with(**call_args)
         mock_func.reset_mock()