Issue #7776: Fix ``Host:'' header and reconnection when using http.client.HTTPConnection.set_tunnel().

Patch by Nikolaus Rath.

1 files changed, 48 insertions, 2 deletions

diff --git a/Lib/test/test_httplib.py b/Lib/test/test_httplib.py
index 30b6c0c..22f7329 100644
--- a/Lib/test/test_httplib.py
+++ b/Lib/test/test_httplib.py
@@ -21,13 +21,15 @@ CACERT_svn_python_org = os.path.join(here, 'https_svn_python_org_root.pem')
 HOST = support.HOST
 
 class FakeSocket:
-    def __init__(self, text, fileclass=io.BytesIO):
+    def __init__(self, text, fileclass=io.BytesIO, host=None, port=None):
         if isinstance(text, str):
             text = text.encode("ascii")
         self.text = text
         self.fileclass = fileclass
         self.data = b''
         self.sendall_calls = 0
+        self.host = host
+        self.port = port
 
     def sendall(self, data):
         self.sendall_calls += 1
@@ -38,6 +40,9 @@ class FakeSocket:
             raise client.UnimplementedFileMode()
         return self.fileclass(self.text)
 
+    def close(self):
+        pass
+
 class EPipeSocket(FakeSocket):
 
     def __init__(self, text, pipe_trigger):
@@ -970,10 +975,51 @@ class HTTPResponseTest(TestCase):
         header = self.resp.getheader('No-Such-Header',default=42)
         self.assertEqual(header, 42)
 
+class TunnelTests(TestCase):
+
+    def test_connect(self):
+        response_text = (
+            'HTTP/1.0 200 OK\r\n\r\n' # Reply to CONNECT
+            'HTTP/1.1 200 OK\r\n' # Reply to HEAD
+            'Content-Length: 42\r\n\r\n'
+        )
+
+        def create_connection(address, timeout=None, source_address=None):
+            return FakeSocket(response_text, host=address[0],
+                              port=address[1])
+
+        conn = client.HTTPConnection('proxy.com')
+        conn._create_connection = create_connection
+
+        # Once connected, we shouldn't be able to tunnel anymore
+        conn.connect()
+        self.assertRaises(RuntimeError, conn.set_tunnel,
+                          'destination.com')
+
+        # But if we close the connection, we're good
+        conn.close()
+        conn.set_tunnel('destination.com')
+        conn.request('HEAD', '/', '')
+
+        self.assertEqual(conn.sock.host, 'proxy.com')
+        self.assertEqual(conn.sock.port, 80)
+        self.assertTrue(b'CONNECT destination.com' in conn.sock.data)
+        self.assertTrue(b'Host: destination.com' in conn.sock.data)
+
+        # This test should be removed when CONNECT gets the HTTP/1.1 blessing
+        self.assertTrue(b'Host: proxy.com' not in conn.sock.data)
+
+        conn.close()
+        conn.request('PUT', '/', '')
+        self.assertEqual(conn.sock.host, 'proxy.com')
+        self.assertEqual(conn.sock.port, 80)
+        self.assertTrue(b'CONNECT destination.com' in conn.sock.data)
+        self.assertTrue(b'Host: destination.com' in conn.sock.data)
+
 def test_main(verbose=None):
     support.run_unittest(HeaderTests, OfflineTest, BasicTest, TimeoutTest,
                          HTTPSTest, RequestBodyTest, SourceAddressTest,
-                         HTTPResponseTest)
+                         HTTPResponseTest, TunnelTests)
 
 if __name__ == '__main__':
     test_main()