# HG changeset patch # Parent dd2522058ce8043cb7e4caa968afe932bfef723e Issue #22350: Close NNTP connection after failures to avoid attempting QUIT diff -r dd2522058ce8 Lib/nntplib.py --- a/Lib/nntplib.py Thu Dec 18 01:20:29 2014 +0100 +++ b/Lib/nntplib.py Fri Dec 19 03:56:06 2014 +0000 @@ -78,6 +78,7 @@ from email.header import decode_header as _email_decode_header from socket import _GLOBAL_DEFAULT_TIMEOUT +from contextlib import contextmanager __all__ = ["NNTP", "NNTPError", "NNTPReplyError", "NNTPTemporaryError", @@ -370,6 +371,15 @@ if is_connected(): self._close() + @contextmanager + def _error_cleanup(self): + """Aborts the connection in case of an unexpected error""" + try: + yield + except: + self._close() + raise + def getwelcome(self): """Get the welcome message from the server (this is read and squirreled away by __init__()). @@ -423,7 +433,8 @@ The `line` must be an unicode string.""" if self.debugging: print('*cmd*', repr(line)) line = line.encode(self.encoding, self.errors) - self._putline(line) + with self._error_cleanup(): + self._putline(line) def _getline(self, strip_crlf=True): """Internal: return one line from the server, stripping _CRLF. @@ -446,15 +457,17 @@ """Internal: get a response from the server. Raise various errors if the response indicates an error. Returns an unicode string.""" - resp = self._getline() - if self.debugging: print('*resp*', repr(resp)) - resp = resp.decode(self.encoding, self.errors) + with self._error_cleanup(): + resp = self._getline() + if self.debugging: print('*resp*', repr(resp)) + resp = resp.decode(self.encoding, self.errors) c = resp[:1] if c == '4': raise NNTPTemporaryError(resp) if c == '5': raise NNTPPermanentError(resp) if c not in '123': + self._close() raise NNTPProtocolError(resp) return resp @@ -467,36 +480,37 @@ If `file` is a file-like object, it must be open in binary mode. """ + resp = self._getresp() openedFile = None try: - # If a string was passed then open a file with that name - if isinstance(file, (str, bytes)): - openedFile = file = open(file, "wb") + with self._error_cleanup(): + # If a string was passed then open a file with that name + if isinstance(file, (str, bytes)): + openedFile = file = open(file, "wb") - resp = self._getresp() - if resp[:3] not in _LONGRESP: - raise NNTPReplyError(resp) + if resp[:3] not in _LONGRESP: + raise NNTPReplyError(resp) - lines = [] - if file is not None: - # XXX lines = None instead? - terminators = (b'.' + _CRLF, b'.\n') - while 1: - line = self._getline(False) - if line in terminators: - break - if line.startswith(b'..'): - line = line[1:] - file.write(line) - else: - terminator = b'.' - while 1: - line = self._getline() - if line == terminator: - break - if line.startswith(b'..'): - line = line[1:] - lines.append(line) + lines = [] + if file is not None: + # XXX lines = None instead? + terminators = (b'.' + _CRLF, b'.\n') + while 1: + line = self._getline(False) + if line in terminators: + break + if line.startswith(b'..'): + line = line[1:] + file.write(line) + else: + terminator = b'.' + while 1: + line = self._getline() + if line == terminator: + break + if line.startswith(b'..'): + line = line[1:] + lines.append(line) finally: # If this method created the file, then it must close it if openedFile: @@ -890,22 +904,23 @@ def _post(self, command, f): resp = self._shortcmd(command) # Raises a specific exception if posting is not allowed - if not resp.startswith('3'): - raise NNTPReplyError(resp) - if isinstance(f, (bytes, bytearray)): - f = f.splitlines() - # We don't use _putline() because: - # - we don't want additional CRLF if the file or iterable is already - # in the right format - # - we don't want a spurious flush() after each line is written - for line in f: - if not line.endswith(_CRLF): - line = line.rstrip(b"\r\n") + _CRLF - if line.startswith(b'.'): - line = b'.' + line - self.file.write(line) - self.file.write(b".\r\n") - self.file.flush() + with self._error_cleanup(): + if not resp.startswith('3'): + raise NNTPReplyError(resp) + if isinstance(f, (bytes, bytearray)): + f = f.splitlines() + # We don't use _putline() because: + # - we don't want additional CRLF if + # the file or iterable is already in the right format + # - we don't want a spurious flush() after each line is written + for line in f: + if not line.endswith(_CRLF): + line = line.rstrip(b"\r\n") + _CRLF + if line.startswith(b'.'): + line = b'.' + line + self.file.write(line) + self.file.write(b".\r\n") + self.file.flush() return self._getresp() def post(self, data): @@ -933,8 +948,10 @@ - resp: server response if successful""" try: resp = self._shortcmd('QUIT') - finally: + except (NNTPTemporaryError, NNTPPermanentError): self._close() + raise + self._close() return resp def login(self, user=None, password=None, usenetrc=True): @@ -1004,10 +1021,11 @@ raise ValueError("TLS cannot be started after authentication.") resp = self._shortcmd('STARTTLS') if resp.startswith('382'): - self.file.close() - self.sock = _encrypt_on(self.sock, context, self.host) - self.file = self.sock.makefile("rwb") - self.tls_on = True + with self._error_cleanup(): + self.file.close() + self.sock = _encrypt_on(self.sock, context, self.host) + self.file = self.sock.makefile("rwb") + self.tls_on = True # Capabilities may change after TLS starts up, so ask for them # again. self._caps = None diff -r dd2522058ce8 Lib/test/test_nntplib.py --- a/Lib/test/test_nntplib.py Thu Dec 18 01:20:29 2014 +0100 +++ b/Lib/test/test_nntplib.py Fri Dec 19 03:56:06 2014 +0000 @@ -1190,6 +1190,17 @@ self.assertRaises(nntplib.NNTPDataError, self.server.newnews, "comp.lang.python", dt) + def test_unrecoverable_cleanup(self): + """Check cleanup after unrecoverable error (Issue 22350)""" + class Aborter(io.BufferedIOBase): + def write(self, data): + raise EnvironmentError() + with self.server: + # Break protocol state by aborting in middle of response + self.assertRaises(EnvironmentError, + self.server.body, file=Aborter()) + self.assertTrue(self.sio.closed, "Connection should be aborted") + class NNTPv1Tests(NNTPv1v2TestsMixin, MockedNNTPTestsMixin, unittest.TestCase): """Tests an NNTP v1 server (no capabilities)."""