gh-131724: Add a new max_headers param to HTTP/HTTPSConnection

Author: aeurielesn

Lib/http/client.py

@@ -209,22 +209,24 @@ def getallmatchingheaders(self, name):
                 lst.append(line)
         return lst
 
-def _read_headers(fp):
+def _read_headers(fp, max_headers):
     """Reads potential header lines into a list from a file pointer.
 
     Length of line is limited by _MAXLINE, and number of
-    headers is limited by _MAXHEADERS.
+    headers is limited by max_headers.
     """
     headers = []
+    if max_headers is None:
+        max_headers = _MAXHEADERS
     while True:
         line = fp.readline(_MAXLINE + 1)
         if len(line) > _MAXLINE:
             raise LineTooLong("header line")
-        headers.append(line)
-        if len(headers) > _MAXHEADERS:
-            raise HTTPException("got more than %d headers" % _MAXHEADERS)
         if line in (b'\r\n', b'\n', b''):
             break
+        headers.append(line)
+        if len(headers) > max_headers:
+            raise HTTPException("got more than %d headers" % max_headers)
     return headers
 
 def _parse_header_lines(header_lines, _class=HTTPMessage):
@@ -241,10 +243,10 @@ def _parse_header_lines(header_lines, _class=HTTPMessage):
     hstring = b''.join(header_lines).decode('iso-8859-1')
     return email.parser.Parser(_class=_class).parsestr(hstring)
 
-def parse_headers(fp, _class=HTTPMessage):
+def parse_headers(fp, _class=HTTPMessage, _max_headers=None):
     """Parses only RFC2822 headers from a file pointer."""
 
-    headers = _read_headers(fp)
+    headers = _read_headers(fp, _max_headers)
     return _parse_header_lines(headers, _class)
 
 
@@ -320,7 +322,7 @@ def _read_status(self):
             raise BadStatusLine(line)
         return version, status, reason
 
-    def begin(self):
+    def begin(self, _max_headers=None):
         if self.headers is not None:
             # we've already started reading the response
             return
@@ -331,7 +333,7 @@ def begin(self):
             if status != CONTINUE:
                 break
             # skip the header from the 100 response
-            skipped_headers = _read_headers(self.fp)
+            skipped_headers = _read_headers(self.fp, _max_headers)
             if self.debuglevel > 0:
                 print("headers:", skipped_headers)
             del skipped_headers
@@ -346,7 +348,9 @@ def begin(self):
         else:
             raise UnknownProtocol(version)
 
-        self.headers = self.msg = parse_headers(self.fp)
+        self.headers = self.msg = parse_headers(
+            self.fp, _max_headers=_max_headers
+        )
 
         if self.debuglevel > 0:
             for hdr, val in self.headers.items():
@@ -864,7 +868,7 @@ def _get_content_length(body, method):
         return None
 
     def __init__(self, host, port=None, timeout=socket._GLOBAL_DEFAULT_TIMEOUT,
-                 source_address=None, blocksize=8192):
+                 source_address=None, blocksize=8192, max_headers=None):
         self.timeout = timeout
         self.source_address = source_address
         self.blocksize = blocksize
@@ -877,6 +881,9 @@ def __init__(self, host, port=None, timeout=socket._GLOBAL_DEFAULT_TIMEOUT,
         self._tunnel_port = None
         self._tunnel_headers = {}
         self._raw_proxy_headers = None
+        if max_headers is None:
+            max_headers = _MAXHEADERS
+        self.max_headers = max_headers
 
         (self.host, self.port) = self._get_hostport(host, port)
 
@@ -969,7 +976,7 @@ def _tunnel(self):
         try:
             (version, code, message) = response._read_status()
 
-            self._raw_proxy_headers = _read_headers(response.fp)
+            self._raw_proxy_headers = _read_headers(response.fp, max_headers=self.max_headers)
 
             if self.debuglevel > 0:
                 for header in self._raw_proxy_headers:
@@ -1426,7 +1433,7 @@ def getresponse(self):
 
         try:
             try:
-                response.begin()
+                response.begin(_max_headers=self.max_headers)
             except ConnectionError:
                 self.close()
                 raise
@@ -1457,10 +1464,12 @@ class HTTPSConnection(HTTPConnection):
 
         def __init__(self, host, port=None,
                      *, timeout=socket._GLOBAL_DEFAULT_TIMEOUT,
-                     source_address=None, context=None, blocksize=8192):
+                     source_address=None, context=None, blocksize=8192,
+                     max_headers=None):
             super(HTTPSConnection, self).__init__(host, port, timeout,
                                                   source_address,
-                                                  blocksize=blocksize)
+                                                  blocksize=blocksize,
+                                                  max_headers=max_headers)
             if context is None:
                 context = _create_https_context(self._http_vsn)
             self._context = context

Lib/test/test_httplib.py

@@ -386,6 +386,50 @@ def test_headers_debuglevel(self):
         self.assertEqual(lines[2], "header: Second: val1")
         self.assertEqual(lines[3], "header: Second: val2")
 
+    def test_max_response_headers(self):
+        max_headers = client._MAXHEADERS + 20
+        headers = [f"Name{i}: Value{i}".encode() for i in range(max_headers)]
+        body = b"HTTP/1.1 200 OK\r\n" + b"\r\n".join(headers)
+
+        with self.subTest(max_headers=None):
+            sock = FakeSocket(body)
+            resp = client.HTTPResponse(sock)
+            with self.assertRaisesRegex(
+                client.HTTPException, f"got more than 100 headers"
+            ):
+                resp.begin()
+
+        with self.subTest(max_headers=max_headers):
+            sock = FakeSocket(body)
+            resp = client.HTTPResponse(sock)
+            resp.begin(_max_headers=max_headers)
+
+    def test_max_connection_headers(self):
+        max_headers = client._MAXHEADERS + 20
+        headers = (
+            f"Name{i}: Value{i}".encode() for i in range(max_headers - 1)
+        )
+        body = (
+            b"HTTP/1.1 200 OK\r\n"
+            + b"\r\n".join(headers)
+            + b"\r\nContent-Length: 12\r\n\r\nDummy body\r\n"
+        )
+
+        with self.subTest(max_headers=None):
+            conn = client.HTTPConnection("example.com")
+            conn.sock = FakeSocket(body)
+            conn.request("GET", "/")
+            with self.assertRaisesRegex(
+                client.HTTPException, f"got more than {client._MAXHEADERS} headers"
+            ):
+                response = conn.getresponse()
+
+        with self.subTest(max_headers=None):
+            conn = client.HTTPConnection("example.com", max_headers=max_headers)
+            conn.sock = FakeSocket(body)
+            conn.request("GET", "/")
+            response = conn.getresponse()
+            response.read()
 
 class HttpMethodTests(TestCase):
     def test_invalid_method_names(self):