gh-135056: Add a --header CLI option to http.server (#135057)

Support custom headers in `python -m http.server` and `http.server.SimpleHTTPRequestHandler`.

Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com>
Co-authored-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com>

Lib/test/test_httpservers.py

@@ -540,8 +540,16 @@ def test_err(self):
         self.assertIn(f"{t.status_client_error}404", lines[1])
 
 
+class CustomHeaderSimpleHTTPRequestHandler(SimpleHTTPRequestHandler):
+    extra_response_headers = None
+
+    def __init__(self, *args, **kwargs):
+        kwargs.setdefault('extra_response_headers', self.extra_response_headers)
+        super().__init__(*args, **kwargs)
+
+
 class SimpleHTTPServerTestCase(BaseTestCase):
-    class request_handler(NoLogRequestHandler, SimpleHTTPRequestHandler):
+    class request_handler(NoLogRequestHandler, CustomHeaderSimpleHTTPRequestHandler):
         pass
 
     def setUp(self):
@@ -898,6 +906,65 @@ def test_path_without_leading_slash(self):
         self.assertEqual(response.getheader("Location"),
                          self.tempdir_name + "/?hi=1")
 
+    def test_extra_response_headers_list_dir(self):
+        with mock.patch.object(self.request_handler, 'extra_response_headers', [
+            ('X-Test1', 'test1'),
+            ('X-Test2', 'test2'),
+        ]):
+            response = self.request(self.base_url + '/')
+            self.assertEqual(response.status, 200)
+            self.assertEqual(response.getheader("X-Test1"), 'test1')
+            self.assertEqual(response.getheader("X-Test2"), 'test2')
+
+    def test_extra_response_headers_get_file(self):
+        with mock.patch.object(self.request_handler, 'extra_response_headers', [
+            ('Set-Cookie', 'test1=value1'),
+            ('Set-Cookie', 'test2=value2'),
+            ('X-Test1', 'value3'),
+        ]):
+            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.status, 200)
+            self.assertEqual(response.getheader("Set-Cookie"),
+                                                'test1=value1, test2=value2')
+            self.assertEqual(response.getheader("X-Test1"), 'value3')
+
+    def test_extra_response_headers_missing_on_404(self):
+        with mock.patch.object(self.request_handler, 'extra_response_headers', [
+            ('X-Test1', 'value'),
+        ]):
+            response = self.request(self.base_url + '/missing.html')
+            self.assertEqual(response.status, 404)
+            self.assertEqual(response.getheader("X-Test1"), None)
+
+    def test_extra_response_headers_dont_overwrite_default_headers(self):
+        with mock.patch.object(self.request_handler, 'extra_response_headers', [
+            ('Content-Type', 'test/not_allowed'),
+            ('Server', 'not_allowed'),
+            ('Set-Cookie', 'test=allowed'),
+        ]):
+            # The Content-Type header should not be overwritten by the extra_response_headers
+            # But cookies in the extra_allowed_duplicate_headers are allowed,
+            # including Set-Cookie
+            response = self.request(self.base_url + '/')
+            self.assertEqual(response.status, 200)
+            self.assertNotEqual(response.getheader("Content-Type"), 'test/not_allowed')
+            self.assertNotEqual(response.getheader("Server"), 'not_allowed')
+            self.assertEqual(response.getheader("Set-Cookie"), 'test=allowed')
+
+    def test_multiple_requests_dont_duplicate_extra_response_headers(self):
+        with mock.patch.object(self.request_handler, 'extra_response_headers', [
+            ('x-test', 'test-value'),
+        ]):
+            response = self.request(self.base_url + '/')
+            self.assertEqual(response.status, 200)
+            self.assertEqual(response.getheader("x-test"), 'test-value')
+            response = self.request(self.base_url + '/')
+            self.assertEqual(response.status, 200)
+            self.assertEqual(response.getheader("x-test"), 'test-value')
+
 
 class SocketlessRequestHandler(SimpleHTTPRequestHandler):
     def __init__(self, directory=None):
@@ -1458,6 +1525,21 @@ def test_content_type_flag(self, mock_func):
                 mock_func.assert_called_once_with(**call_args)
                 mock_func.reset_mock()
 
+    @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')
+        mock_func.assert_called_once_with(**call_args)
+        mock_func.reset_mock()
+
+    def test_extra_header_flag_too_few_args(self):
+        with self.assertRaises(SystemExit):
+            self.invoke_httpd('--header', 'h1')
+
+    def test_extra_header_flag_too_many_args(self):
+        with self.assertRaises(SystemExit):
+            self.invoke_httpd('--header', 'h1', 'v1', 'h2')
+
     @unittest.skipIf(ssl is None, "requires ssl")
     @mock.patch('http.server.test')
     def test_tls_cert_and_key_flags(self, mock_func):
@@ -1541,6 +1623,30 @@ def test_unknown_flag(self, _):
         self.assertEqual(stdout.getvalue(), '')
         self.assertIn('error', stderr.getvalue())
 
+    @mock.patch('http.server.test')
+    def test_extra_response_headers_arg(self, mock_test):
+        # Call the main function with extra response headers cli args
+        server._main(
+            ['-H', 'Set-Cookie', 'k=v', '-H', 'Set-Cookie', 'k2=v2:v3 v4', '8080']
+        )
+        # Get the ServerClass (DualStackServerMixin subclass) that _main()
+        # passed to test(), and verify its finish_request passes
+        # extra_response_headers to the handler.
+        _, kwargs = mock_test.call_args
+        server_class = kwargs['ServerClass']
+
+        mock_handler_class = mock.MagicMock()
+        mock_server = mock.Mock()
+        mock_server.RequestHandlerClass = mock_handler_class
+        server_class.finish_request(mock_server, mock.Mock(), '127.0.0.1')
+        mock_handler_class.assert_called_once_with(
+            mock.ANY, mock.ANY, mock_server,
+            directory=mock.ANY,
+            extra_response_headers=[
+                ['Set-Cookie', 'k=v'], ['Set-Cookie', 'k2=v2:v3 v4']
+            ]
+        )
+
 
 class CommandLineRunTimeTestCase(unittest.TestCase):
     served_data = os.urandom(32)