Patch #1627441: close sockets properly in urllib2.

5 files changed, 63 insertions, 6 deletions

diff --git a/Lib/socket.py b/Lib/socket.py
index 52fb8e3..b4969bd 100644
--- a/Lib/socket.py
+++ b/Lib/socket.py
@@ -204,9 +204,10 @@ class _fileobject(object):
 
     __slots__ = ["mode", "bufsize", "softspace",
                  # "closed" is a property, see below
-                 "_sock", "_rbufsize", "_wbufsize", "_rbuf", "_wbuf"]
+                 "_sock", "_rbufsize", "_wbufsize", "_rbuf", "_wbuf",
+                 "_close"]
 
-    def __init__(self, sock, mode='rb', bufsize=-1):
+    def __init__(self, sock, mode='rb', bufsize=-1, close=False):
         self._sock = sock
         self.mode = mode # Not actually used in this version
         if bufsize < 0:
@@ -222,6 +223,7 @@ class _fileobject(object):
         self._wbufsize = bufsize
         self._rbuf = "" # A string
         self._wbuf = [] # A list of strings
+        self._close = close
 
     def _getclosed(self):
         return self._sock is None
@@ -232,6 +234,8 @@ class _fileobject(object):
             if self._sock:
                 self.flush()
         finally:
+            if self._close:
+                self._sock.close()
             self._sock = None
 
     def __del__(self):
diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py
index 1357d54..c56c3b7 100644
--- a/Lib/test/test_socket.py
+++ b/Lib/test/test_socket.py
@@ -809,6 +809,31 @@ class SmallBufferedFileObjectClassTestCase(FileObjectClassTestCase):
 
     bufsize = 2 # Exercise the buffering code
 
+
+class Urllib2FileobjectTest(unittest.TestCase):
+
+    # urllib2.HTTPHandler has "borrowed" socket._fileobject, and requires that
+    # it close the socket if the close c'tor argument is true
+
+    def testClose(self):
+        class MockSocket:
+            closed = False
+            def flush(self): pass
+            def close(self): self.closed = True
+
+        # must not close unless we request it: the original use of _fileobject
+        # by module socket requires that the underlying socket not be closed until
+        # the _socketobject that created the _fileobject is closed
+        s = MockSocket()
+        f = socket._fileobject(s)
+        f.close()
+        self.assert_(not s.closed)
+
+        s = MockSocket()
+        f = socket._fileobject(s, close=True)
+        f.close()
+        self.assert_(s.closed)
+
 class TCPTimeoutTest(SocketTCPTest):
 
     def testTCPTimeout(self):
@@ -961,7 +986,8 @@ def test_main():
         FileObjectClassTestCase,
         UnbufferedFileObjectClassTestCase,
         LineBufferedFileObjectClassTestCase,
-        SmallBufferedFileObjectClassTestCase
+        SmallBufferedFileObjectClassTestCase,
+        Urllib2FileobjectTest,
     ])
     if hasattr(socket, "socketpair"):
         tests.append(BasicSocketPairTest)
diff --git a/Lib/test/test_urllib2net.py b/Lib/test/test_urllib2net.py
index 00cf202..b271626 100644
--- a/Lib/test/test_urllib2net.py
+++ b/Lib/test/test_urllib2net.py
@@ -64,6 +64,27 @@ class AuthTests(unittest.TestCase):
 #                          urllib2.urlopen, "http://evil:thing@example.com")
 
 
+class CloseSocketTest(unittest.TestCase):
+
+    def test_close(self):
+        import socket, httplib, gc
+
+        # calling .close() on urllib2's response objects should close the
+        # underlying socket
+
+        # delve deep into response to fetch socket._socketobject
+        response = urllib2.urlopen("http://www.python.org/")
+        abused_fileobject = response.fp
+        self.assert_(abused_fileobject.__class__ is socket._fileobject)
+        httpresponse = abused_fileobject._sock
+        self.assert_(httpresponse.__class__ is httplib.HTTPResponse)
+        fileobject = httpresponse.fp
+        self.assert_(fileobject.__class__ is socket._fileobject)
+
+        self.assert_(not fileobject.closed)
+        response.close()
+        self.assert_(fileobject.closed)
+
 class urlopenNetworkTests(unittest.TestCase):
     """Tests urllib2.urlopen using the network.
 
@@ -263,8 +284,12 @@ class OtherNetworkTests(unittest.TestCase):
 
 def test_main():
     test_support.requires("network")
-    test_support.run_unittest(URLTimeoutTest, urlopenNetworkTests,
-                              AuthTests, OtherNetworkTests)
+    test_support.run_unittest(URLTimeoutTest,
+                              urlopenNetworkTests,
+                              AuthTests,
+                              OtherNetworkTests,
+                              CloseSocketTest,
+                              )
 
 if __name__ == "__main__":
     test_main()
diff --git a/Lib/urllib2.py b/Lib/urllib2.py
index 890d3d4..1ab5c61 100644
--- a/Lib/urllib2.py
+++ b/Lib/urllib2.py
@@ -1087,7 +1087,7 @@ class AbstractHTTPHandler(BaseHandler):
         # out of socket._fileobject() and into a base class.
 
         r.recv = r.read
-        fp = socket._fileobject(r)
+        fp = socket._fileobject(r, close=True)
 
         resp = addinfourl(fp, r.msg, req.get_full_url())
         resp.code = r.status
diff --git a/Misc/NEWS b/Misc/NEWS
index 25593a3..5f4fdc1 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -111,6 +111,8 @@ Core and builtins
 Library
 -------
 
+- Patch #1627441: close sockets properly in urllib2.
+
 - Bug #494589: make ntpath.expandvars behave according to its docstring.
 
 - Changed platform module API python_version_tuple() to actually