From 08f5f7aa81321eb667609b2f096b2fc0c092cad4 Mon Sep 17 00:00:00 2001 From: Nadeem Vawda 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/ftplib.py | 56 +++++++++++++++++++++++++-------------------- Lib/test/test_urllib2.py | 1 + Lib/test/test_urllib2net.py | 1 + Lib/urllib/request.py | 30 ++++++++++++++++++++---- 4 files changed, 59 insertions(+), 29 deletions(-) diff --git a/Lib/ftplib.py b/Lib/ftplib.py index be1797f..8e53023 100644 --- a/Lib/ftplib.py +++ b/Lib/ftplib.py @@ -336,33 +336,39 @@ class FTP: if self.passiveserver: host, port = self.makepasv() conn = socket.create_connection((host, port), self.timeout) - if rest is not None: - self.sendcmd("REST %s" % rest) - resp = self.sendcmd(cmd) - # Some servers apparently send a 200 reply to - # a LIST or STOR command, before the 150 reply - # (and way before the 226 reply). This seems to - # be in violation of the protocol (which only allows - # 1xx or error messages for LIST), so we just discard - # this response. - if resp[0] == '2': - resp = self.getresp() - if resp[0] != '1': - raise error_reply(resp) + try: + if rest is not None: + self.sendcmd("REST %s" % rest) + resp = self.sendcmd(cmd) + # Some servers apparently send a 200 reply to + # a LIST or STOR command, before the 150 reply + # (and way before the 226 reply). This seems to + # be in violation of the protocol (which only allows + # 1xx or error messages for LIST), so we just discard + # this response. + if resp[0] == '2': + resp = self.getresp() + if resp[0] != '1': + raise error_reply(resp) + except: + conn.close() + raise else: sock = self.makeport() - if rest is not None: - self.sendcmd("REST %s" % rest) - resp = self.sendcmd(cmd) - # See above. - if resp[0] == '2': - resp = self.getresp() - if resp[0] != '1': - raise error_reply(resp) - conn, sockaddr = sock.accept() - if self.timeout is not _GLOBAL_DEFAULT_TIMEOUT: - conn.settimeout(self.timeout) - sock.close() + try: + if rest is not None: + self.sendcmd("REST %s" % rest) + resp = self.sendcmd(cmd) + # See above. + if resp[0] == '2': + resp = self.getresp() + if resp[0] != '1': + raise error_reply(resp) + conn, sockaddr = sock.accept() + if self.timeout is not _GLOBAL_DEFAULT_TIMEOUT: + conn.settimeout(self.timeout) + finally: + sock.close() if resp[:3] == '150': # this is conditional in case we received a 125 size = parse150(resp) diff --git a/Lib/test/test_urllib2.py b/Lib/test/test_urllib2.py index 58ef836..03cd927 100644 --- a/Lib/test/test_urllib2.py +++ b/Lib/test/test_urllib2.py @@ -622,6 +622,7 @@ class HandlerTests(unittest.TestCase): def retrfile(self, filename, filetype): self.filename, self.filetype = filename, filetype return io.StringIO(self.data), len(self.data) + def close(self): pass class NullFTPHandler(urllib.request.FTPHandler): def __init__(self, data): self.data = data diff --git a/Lib/test/test_urllib2net.py b/Lib/test/test_urllib2net.py index eaad325..a475f56 100644 --- a/Lib/test/test_urllib2net.py +++ b/Lib/test/test_urllib2net.py @@ -222,6 +222,7 @@ class OtherNetworkTests(unittest.TestCase): handlers = [] cfh = urllib.request.CacheFTPHandler() + self.addCleanup(cfh.clear_cache) cfh.setTimeout(1) handlers.append(cfh) diff --git a/Lib/urllib/request.py b/Lib/urllib/request.py index 35fd1f1..a09a353 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 v0.12