diff --git a/mitogen/doas.py b/mitogen/doas.py index 09b2be9e..3a6f881d 100644 --- a/mitogen/doas.py +++ b/mitogen/doas.py @@ -80,12 +80,7 @@ class Stream(mitogen.parent.Stream): password_incorrect_msg = 'doas password is incorrect' password_required_msg = 'doas password is required' - def _connect_bootstrap(self): - it = mitogen.parent.iter_read( - fds=[self.receive_side.fd, self.diag_stream.receive_side.fd], - deadline=self.connect_deadline, - ) - + def _connect_input_loop(self, it): password_sent = False for buf in it: LOG.debug('%r: received %r', self, buf) @@ -106,3 +101,13 @@ class Stream(mitogen.parent.Stream): ) password_sent = True raise mitogen.core.StreamError('bootstrap failed') + + def _connect_bootstrap(self): + it = mitogen.parent.iter_read( + fds=[self.receive_side.fd, self.diag_stream.receive_side.fd], + deadline=self.connect_deadline, + ) + try: + self._connect_input_loop(it) + finally: + it.close() diff --git a/mitogen/parent.py b/mitogen/parent.py index 04f7784e..865d8fd2 100644 --- a/mitogen/parent.py +++ b/mitogen/parent.py @@ -532,12 +532,16 @@ def discard_until(fd, s, deadline): :raises mitogen.core.StreamError: Attempt to read past end of file. """ - for buf in iter_read([fd], deadline): - if IOLOG.level == logging.DEBUG: - for line in buf.splitlines(): - IOLOG.debug('discard_until: discarding %r', line) - if buf.endswith(s): - return + it = iter_read([fd], deadline) + try: + for buf in it: + if IOLOG.level == logging.DEBUG: + for line in buf.splitlines(): + IOLOG.debug('discard_until: discarding %r', line) + if buf.endswith(s): + return + finally: + it.close() # ensure Poller.close() is called. def _upgrade_broker(broker): diff --git a/mitogen/ssh.py b/mitogen/ssh.py index e3891f9c..106dfd56 100644 --- a/mitogen/ssh.py +++ b/mitogen/ssh.py @@ -264,13 +264,7 @@ class Stream(mitogen.parent.Stream): # with ours. raise HostKeyError(self.hostkey_config_msg) - def _connect_bootstrap(self): - fds = [self.receive_side.fd] - if self.diag_stream is not None: - fds.append(self.diag_stream.receive_side.fd) - - it = mitogen.parent.iter_read(fds=fds, deadline=self.connect_deadline) - + def _connect_input_loop(self, it): password_sent = False for buf, partial in filter_debug(self, it): LOG.debug('%r: received %r', self, buf) @@ -302,3 +296,14 @@ class Stream(mitogen.parent.Stream): password_sent = True raise mitogen.core.StreamError('bootstrap failed') + + def _connect_bootstrap(self): + fds = [self.receive_side.fd] + if self.diag_stream is not None: + fds.append(self.diag_stream.receive_side.fd) + + it = mitogen.parent.iter_read(fds=fds, deadline=self.connect_deadline) + try: + self._connect_input_loop(it) + finally: + it.close() diff --git a/mitogen/su.py b/mitogen/su.py index 9b0172c8..65357a3e 100644 --- a/mitogen/su.py +++ b/mitogen/su.py @@ -87,13 +87,7 @@ class Stream(mitogen.parent.Stream): password_incorrect_msg = 'su password is incorrect' password_required_msg = 'su password is required' - def _connect_bootstrap(self): - password_sent = False - it = mitogen.parent.iter_read( - fds=[self.receive_side.fd], - deadline=self.connect_deadline, - ) - + def _connect_input_loop(self, it): for buf in it: LOG.debug('%r: received %r', self, buf) if buf.endswith(self.EC0_MARKER): @@ -113,3 +107,14 @@ class Stream(mitogen.parent.Stream): ) password_sent = True raise mitogen.core.StreamError('bootstrap failed') + + def _connect_bootstrap(self): + password_sent = False + it = mitogen.parent.iter_read( + fds=[self.receive_side.fd], + deadline=self.connect_deadline, + ) + try: + self._connect_input_loop(it) + finally: + it.close() diff --git a/mitogen/sudo.py b/mitogen/sudo.py index b2eaabce..68b27fec 100644 --- a/mitogen/sudo.py +++ b/mitogen/sudo.py @@ -173,17 +173,7 @@ class Stream(mitogen.parent.Stream): password_incorrect_msg = 'sudo password is incorrect' password_required_msg = 'sudo password is required' - def _connect_bootstrap(self): - fds = [self.receive_side.fd] - if self.diag_stream is not None: - fds.append(self.diag_stream.receive_side.fd) - - password_sent = False - it = mitogen.parent.iter_read( - fds=fds, - deadline=self.connect_deadline, - ) - + def _connect_input_loop(self, it): for buf in it: LOG.debug('%r: received %r', self, buf) if buf.endswith(self.EC0_MARKER): @@ -199,3 +189,19 @@ class Stream(mitogen.parent.Stream): ) password_sent = True raise mitogen.core.StreamError('bootstrap failed') + + def _connect_bootstrap(self): + fds = [self.receive_side.fd] + if self.diag_stream is not None: + fds.append(self.diag_stream.receive_side.fd) + + password_sent = False + it = mitogen.parent.iter_read( + fds=fds, + deadline=self.connect_deadline, + ) + + try: + self._connect_input_loop(it) + finally: + it.close()