From 08f5f7aa81321eb667609b2f096b2fc0c092cad4 Mon Sep 17 00:00:00 2001
From: Nadeem Vawda <nadeem.vawda@gmail.com>
Date: Sat, 23 Jul 2011 14:03:00 +0200
Subject: Issue #10883: Fix socket leaks in urllib.request.

* ftpwrapper now uses reference counting to ensure that the underlying socket
  is closed when the ftpwrapper object is no longer in use
* ftplib.FTP.ntransfercmd() now closes the socket if an error occurs

Initial patch by Victor Stinner.
---
 Lib/urllib/request.py | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

(limited to 'Lib/urllib/request.py')

diff --git a/Lib/urllib/request.py b/Lib/urllib/request.py
index 35fd1f136f..a09a35348c 100644
--- a/Lib/urllib/request.py
+++ b/Lib/urllib/request.py
@@ -1362,8 +1362,8 @@ class FTPHandler(BaseHandler):
             raise exc.with_traceback(sys.exc_info()[2])
 
     def connect_ftp(self, user, passwd, host, port, dirs, timeout):
-        fw = ftpwrapper(user, passwd, host, port, dirs, timeout)
-        return fw
+        return ftpwrapper(user, passwd, host, port, dirs, timeout,
+                          persistent=False)
 
 class CacheFTPHandler(FTPHandler):
     # XXX would be nice to have pluggable cache strategies
@@ -1412,6 +1412,13 @@ class CacheFTPHandler(FTPHandler):
                     break
             self.soonest = min(list(self.timeout.values()))
 
+    def clear_cache(self):
+        for conn in self.cache.values():
+            conn.close()
+        self.cache.clear()
+        self.timeout.clear()
+
+
 # Code move from the old urllib module
 
 MAXFTPCACHE = 10        # Trim the ftp cache beyond this size
@@ -2135,13 +2142,16 @@ def noheaders():
 class ftpwrapper:
     """Class used by open_ftp() for cache of open FTP connections."""
 
-    def __init__(self, user, passwd, host, port, dirs, timeout=None):
+    def __init__(self, user, passwd, host, port, dirs, timeout=None,
+                 persistent=True):
         self.user = user
         self.passwd = passwd
         self.host = host
         self.port = port
         self.dirs = dirs
         self.timeout = timeout
+        self.refcount = 0
+        self.keepalive = persistent
         self.init()
 
     def init(self):
@@ -2192,7 +2202,8 @@ class ftpwrapper:
             conn, retrlen = self.ftp.ntransfercmd(cmd)
         self.busy = 1
 
-        ftpobj = addclosehook(conn.makefile('rb'), self.endtransfer)
+        ftpobj = addclosehook(conn.makefile('rb'), self.file_close)
+        self.refcount += 1
         conn.close()
         # Pass back both a suitably decorated object and a retrieval length
         return (ftpobj, retrlen)
@@ -2207,6 +2218,17 @@ class ftpwrapper:
             pass
 
     def close(self):
+        self.keepalive = False
+        if self.refcount <= 0:
+            self.real_close()
+
+    def file_close(self):
+        self.endtransfer()
+        self.refcount -= 1
+        if self.refcount <= 0 and not self.keepalive:
+            self.real_close()
+
+    def real_close(self):
         self.endtransfer()
         try:
             self.ftp.close()
-- 
cgit v1.2.1