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

Patch by Nikolaus Rath.

1 files changed, 46 insertions, 2 deletions

diff --git a/Lib/test/test_httplib.py b/Lib/test/test_httplib.py
index bb9d75f..72800e5 100644
--- a/Lib/test/test_httplib.py
+++ b/Lib/test/test_httplib.py
@@ -13,10 +13,12 @@ from test import test_support
 HOST = test_support.HOST
 
 class FakeSocket:
-    def __init__(self, text, fileclass=StringIO.StringIO):
+    def __init__(self, text, fileclass=StringIO.StringIO, host=None, port=None):
         self.text = text
         self.fileclass = fileclass
         self.data = ''
+        self.host = host
+        self.port = port
 
     def sendall(self, data):
         self.data += ''.join(data)
@@ -26,6 +28,9 @@ class FakeSocket:
             raise httplib.UnimplementedFileMode()
         return self.fileclass(self.text)
 
+    def close(self):
+        pass
+
 class EPipeSocket(FakeSocket):
 
     def __init__(self, text, pipe_trigger):
@@ -526,9 +531,48 @@ class HTTPSTimeoutTest(TestCase):
                 self.fail("Port incorrectly parsed: %s != %s" % (p, c.host))
 
 
+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 = httplib.HTTPConnection('proxy.com')
+        conn._create_connection = create_connection
+
+        # Once connected, we should not be able to tunnel anymore
+        conn.connect()
+        self.assertRaises(RuntimeError, conn.set_tunnel, 'destination.com')
+
+        # But if close the connection, we are 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('CONNECT destination.com' in conn.sock.data)
+        self.assertTrue('Host: destination.com' in conn.sock.data)
+
+        self.assertTrue('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('CONNECT destination.com' in conn.sock.data)
+        self.assertTrue('Host: destination.com' in conn.sock.data)
+
+
 def test_main(verbose=None):
     test_support.run_unittest(HeaderTests, OfflineTest, BasicTest, TimeoutTest,
-                              HTTPSTimeoutTest, SourceAddressTest)
+                              HTTPSTimeoutTest, SourceAddressTest, TunnelTests)
 
 if __name__ == '__main__':
     test_main()