add method get_proxy_response_headers

python · gpshead · May 16, 2023 · May 6, 2023 · May 6, 2023 · May 8, 2023
commit f39feb945454130b30cacf5a1fd087866b85badc
diff --git a/Lib/http/client.py b/Lib/http/client.py
@@ -221,8 +221,9 @@ def _read_headers(fp):
             break
     return headers
 
-def parse_headers(fp, _class=HTTPMessage):
-    """Parses only RFC2822 headers from a file pointer.
+def _parse_header_lines(header_lines, _class=HTTPMessage):
+    """
+    Parses only RFC2822 headers from header lines.
 
     email Parser wants to see strings rather than bytes.
     But a TextIOWrapper around self.rfile would buffer too many bytes
@@ -231,10 +232,15 @@ def parse_headers(fp, _class=HTTPMessage):
     to parse.
 
     """
-    headers = _read_headers(fp)
-    hstring = b''.join(headers).decode('iso-8859-1')
+    hstring = b''.join(header_lines).decode('iso-8859-1')
     return email.parser.Parser(_class=_class).parsestr(hstring)
 
+def parse_headers(fp, _class=HTTPMessage):
+    """Parses only RFC2822 headers from a file pointer."""
+
+    headers = _read_headers(fp)
+    return _parse_header_lines(headers, _class)
+
 
 class HTTPResponse(io.BufferedIOBase):
 
@@ -858,7 +864,7 @@ def __init__(self, host, port=None, timeout=socket._GLOBAL_DEFAULT_TIMEOUT,
         self._tunnel_host = None
         self._tunnel_port = None
         self._tunnel_headers = {}
-        self._proxy_response_headers = None
+        self._raw_proxy_headers = None
 
         (self.host, self.port) = self._get_hostport(host, port)
 
@@ -945,11 +951,11 @@ def _tunnel(self):
         try:
             (version, code, message) = response._read_status()
 
-            self._proxy_response_headers = parse_headers(response.fp)
+            self._raw_proxy_headers = _read_headers(response.fp)
 
             if self.debuglevel > 0:
-                for hdr, val in self._proxy_response_headers.items():
-                    print("header:", hdr + ":", val)
+                for header in self._raw_proxy_headers:
+                    print('header:', header.decode())
 
             if code != http.HTTPStatus.OK:
                 self.close()
@@ -958,6 +964,21 @@ def _tunnel(self):
         finally:
             response.close()
 
+    def get_proxy_response_headers(self):
+        """
+        Returns a dictionary with the headers of the response
+        received from the proxy server to the CONNECT request
+        sent to set the tunnel.
+
+        If the CONNECT request was not sent, the method returns
+        an empty dictionary.
+        """
+        return (
+            _parse_header_lines(self._raw_proxy_headers)
+            if self._raw_proxy_headers is not None
+            else {}
+        )
+
     def connect(self):
         """Connect to the host and port specified in __init__."""
         sys.audit("http.client.connect", self, self.host, self.port)

diff --git a/Lib/test/test_httplib.py b/Lib/test/test_httplib.py
@@ -2401,7 +2401,7 @@ def test_proxy_response_headers(self):
         self.conn.set_tunnel('destination.com')
 
         self.conn.request('PUT', '/', '')
-        headers = self.conn._proxy_response_headers
+        headers = self.conn.get_proxy_response_headers()
         self.assertIn(expected_header, headers.items())
 
     def test_tunnel_leak(self):