diff --git a/.ci/debops_common_tests.py b/.ci/debops_common_tests.py index b0e2e4e8..8b35de1e 100755 --- a/.ci/debops_common_tests.py +++ b/.ci/debops_common_tests.py @@ -3,6 +3,7 @@ from __future__ import print_function import os import shutil +import sys import ci_lib @@ -68,8 +69,8 @@ with ci_lib.Fold('job_setup'): with ci_lib.Fold('first_run'): - ci_lib.run('debops common') + ci_lib.run('debops common %s', ' '.join(sys.argv[1:])) with ci_lib.Fold('second_run'): - ci_lib.run('debops common') + ci_lib.run('debops common %s', ' '.join(sys.argv[1:])) diff --git a/.gitignore b/.gitignore index 55f37f29..aa75f691 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,8 @@ venvs/** MANIFEST build/ dist/ +extra/ +tests/ansible/.*.pid docs/_build/ htmlcov/ *.egg-info diff --git a/.travis.yml b/.travis.yml index eae04cb0..b8ae0c62 100644 --- a/.travis.yml +++ b/.travis.yml @@ -1,4 +1,5 @@ sudo: required +dist: trusty notifications: email: false @@ -20,6 +21,7 @@ install: - .ci/${MODE}_install.py script: +- .ci/spawn_reverse_shell.py - .ci/${MODE}_tests.py @@ -27,6 +29,11 @@ script: # newest->oldest in various configuartions. matrix: + allow_failures: + # Python 2.4 tests are still unreliable + - language: c + env: MODE=mitogen_py24 DISTRO=centos5 + include: # Mitogen tests. # 2.4 -> 2.4 diff --git a/ansible_mitogen/affinity.py b/ansible_mitogen/affinity.py index 09a6acee..94539e21 100644 --- a/ansible_mitogen/affinity.py +++ b/ansible_mitogen/affinity.py @@ -177,9 +177,9 @@ class FixedPolicy(Policy): cores, before reusing the second hyperthread of an existing core. A hook is installed that causes :meth:`reset` to run in the child of any - process created with :func:`mitogen.parent.detach_popen`, ensuring - CPU-intensive children like SSH are not forced to share the same core as - the (otherwise potentially very busy) parent. + process created with :func:`mitogen.parent.popen`, ensuring CPU-intensive + children like SSH are not forced to share the same core as the (otherwise + potentially very busy) parent. """ def __init__(self, cpu_count=None): #: For tests. diff --git a/ansible_mitogen/module_finder.py b/ansible_mitogen/module_finder.py index 633e3cad..89aa2beb 100644 --- a/ansible_mitogen/module_finder.py +++ b/ansible_mitogen/module_finder.py @@ -57,7 +57,7 @@ def get_code(module): """ Compile and return a Module's code object. """ - fp = open(module.path) + fp = open(module.path, 'rb') try: return compile(fp.read(), str(module.name), 'exec') finally: diff --git a/ansible_mitogen/process.py b/ansible_mitogen/process.py index e4e61e8b..a8827cb1 100644 --- a/ansible_mitogen/process.py +++ b/ansible_mitogen/process.py @@ -79,8 +79,15 @@ def clean_shutdown(sock): MuxProcess, debug logs may appear on the user's terminal *after* the prompt has been printed. """ - sock.shutdown(socket.SHUT_WR) + try: + sock.shutdown(socket.SHUT_WR) + except socket.error: + # Already closed. This is possible when tests are running. + LOG.debug('clean_shutdown: ignoring duplicate call') + return + sock.recv(1) + sock.close() def getenv_int(key, default=0): @@ -154,8 +161,15 @@ class MuxProcess(object): #: forked WorkerProcesses to contact the MuxProcess unix_listener_path = None - #: Singleton. - _instance = None + @classmethod + def _reset(cls): + """ + Used to clean up in unit tests. + """ + assert cls.worker_sock is not None + cls.worker_sock.close() + cls.worker_sock = None + os.waitpid(cls.worker_pid, 0) @classmethod def start(cls, _init_logging=True): @@ -178,7 +192,7 @@ class MuxProcess(object): mitogen.utils.setup_gil() cls.unix_listener_path = mitogen.unix.make_socket_path() cls.worker_sock, cls.child_sock = socket.socketpair() - atexit.register(lambda: clean_shutdown(cls.worker_sock)) + atexit.register(clean_shutdown, cls.worker_sock) mitogen.core.set_cloexec(cls.worker_sock.fileno()) mitogen.core.set_cloexec(cls.child_sock.fileno()) @@ -189,8 +203,8 @@ class MuxProcess(object): ansible_mitogen.logging.setup() cls.original_env = dict(os.environ) - cls.child_pid = os.fork() - if cls.child_pid: + cls.worker_pid = os.fork() + if cls.worker_pid: save_pid('controller') ansible_mitogen.logging.set_process_name('top') ansible_mitogen.affinity.policy.assign_controller() @@ -308,7 +322,7 @@ class MuxProcess(object): self._setup_responder(self.router.responder) mitogen.core.listen(self.broker, 'shutdown', self.on_broker_shutdown) mitogen.core.listen(self.broker, 'exit', self.on_broker_exit) - self.listener = mitogen.unix.Listener( + self.listener = mitogen.unix.Listener.build_stream( router=self.router, path=self.unix_listener_path, backlog=C.DEFAULT_FORKS, diff --git a/docs/api.rst b/docs/api.rst index db39ad99..917fc627 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -383,6 +383,9 @@ Connection Methods the root PID of a running Docker, LXC, LXD, or systemd-nspawn container. + The setns method depends on the built-in :mod:`ctypes` module, and thus + does not support Python 2.4. + A program is required only to find the root PID, after which management of the child Python interpreter is handled directly. diff --git a/docs/changelog.rst b/docs/changelog.rst index 9135c48f..43d30456 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -102,6 +102,14 @@ Fixes potential influx of 2.8-related bug reports. +Core Library +~~~~~~~~~~~~ + +* `#170 `_: to better support child + process management and a future asynchronous connect implementation, a + :class:`mitogen.parent.TimerList` API is available. + + Thanks! ~~~~~~~ diff --git a/docs/getting_started.rst b/docs/getting_started.rst index 020760bc..945e243f 100644 --- a/docs/getting_started.rst +++ b/docs/getting_started.rst @@ -341,15 +341,13 @@ The following built-in types may be used as parameters or return values in remote procedure calls: * :class:`bool` -* :class:`bytearray` -* :func:`bytes` +* :func:`bytes` (:class:`str` on Python 2.x) * :class:`dict` * :class:`int` * :func:`list` * :class:`long` -* :class:`str` * :func:`tuple` -* :func:`unicode` +* :func:`unicode` (:class:`str` on Python 3.x) User-defined types may not be used, except for: diff --git a/docs/internals.rst b/docs/internals.rst index e1dd4a41..96f9269c 100644 --- a/docs/internals.rst +++ b/docs/internals.rst @@ -15,46 +15,49 @@ Constants .. autodata:: CHUNK_SIZE -Poller Classes -============== +Pollers +======= .. currentmodule:: mitogen.core .. autoclass:: Poller - :members: + :members: + +.. currentmodule:: mitogen.parent +.. autoclass:: KqueuePoller .. currentmodule:: mitogen.parent .. autoclass:: EpollPoller .. currentmodule:: mitogen.parent -.. autoclass:: KqueuePoller +.. autoclass:: PollPoller -Latch Class -=========== +Latch +===== .. currentmodule:: mitogen.core .. autoclass:: Latch :members: -PidfulStreamHandler Class -========================= +PidfulStreamHandler +=================== .. currentmodule:: mitogen.core .. autoclass:: PidfulStreamHandler :members: -Side Class -========== +Side +==== .. currentmodule:: mitogen.core .. autoclass:: Side :members: -Stream Classes -============== +Stream +====== .. currentmodule:: mitogen.core .. autoclass:: BasicStream @@ -79,42 +82,24 @@ Stream Classes .. autoclass:: Stream :members: - -Other Stream Subclasses -======================= - .. currentmodule:: mitogen.core - .. autoclass:: IoLogger :members: +.. currentmodule:: mitogen.core .. autoclass:: Waker :members: -Poller Class -============ - -.. currentmodule:: mitogen.core -.. autoclass:: Poller - :members: - -.. currentmodule:: mitogen.parent -.. autoclass:: KqueuePoller - -.. currentmodule:: mitogen.parent -.. autoclass:: EpollPoller - - -Importer Class -============== +Importer +======== .. currentmodule:: mitogen.core .. autoclass:: Importer :members: -Responder Class +ModuleResponder =============== .. currentmodule:: mitogen.master @@ -122,40 +107,59 @@ Responder Class :members: -RouteMonitor Class -================== +RouteMonitor +============ .. currentmodule:: mitogen.parent .. autoclass:: RouteMonitor :members: -Forwarder Class -=============== +TimerList +========= + +.. currentmodule:: mitogen.parent +.. autoclass:: TimerList + :members: + + +Timer +===== + +.. currentmodule:: mitogen.parent +.. autoclass:: Timer + :members: + + +Forwarder +========= .. currentmodule:: mitogen.parent .. autoclass:: ModuleForwarder :members: -ExternalContext Class -===================== +ExternalContext +=============== .. currentmodule:: mitogen.core .. autoclass:: ExternalContext :members: -mitogen.master -============== +Process +======= .. currentmodule:: mitogen.parent -.. autoclass:: ProcessMonitor +.. autoclass:: Process :members: -Blocking I/O Functions -====================== +Helpers +======= + +Blocking I/O +------------ These functions exist to support the blocking phase of setting up a new context. They will eventually be replaced with asynchronous equivalents. @@ -167,8 +171,8 @@ context. They will eventually be replaced with asynchronous equivalents. .. autofunction:: write_all -Subprocess Creation Functions -============================= +Subprocess Functions +------------ .. currentmodule:: mitogen.parent .. autofunction:: create_child @@ -176,8 +180,8 @@ Subprocess Creation Functions .. autofunction:: tty_create_child -Helper Functions -================ +Helpers +------- .. currentmodule:: mitogen.core .. autofunction:: to_text diff --git a/examples/mitogen-fuse.py b/examples/mitogen-fuse.py index d0cd9a3a..c1b17032 100644 --- a/examples/mitogen-fuse.py +++ b/examples/mitogen-fuse.py @@ -241,9 +241,13 @@ def main(router): print('usage: %s ' % sys.argv[0]) sys.exit(1) - blerp = fuse.FUSE( + kwargs = {} + if sys.platform == 'darwin': + kwargs['volname'] = '%s (Mitogen)' % (sys.argv[1],) + + f = fuse.FUSE( operations=Operations(sys.argv[1]), mountpoint=sys.argv[2], foreground=True, - volname='%s (Mitogen)' % (sys.argv[1],), + **kwargs ) diff --git a/mitogen/__init__.py b/mitogen/__init__.py index 47fe4d38..5e2e29b6 100644 --- a/mitogen/__init__.py +++ b/mitogen/__init__.py @@ -111,10 +111,10 @@ def main(log_level='INFO', profiling=_default_profiling): if profiling: mitogen.core.enable_profiling() mitogen.master.Router.profiling = profiling - utils.log_to_file(level=log_level) + mitogen.utils.log_to_file(level=log_level) return mitogen.core._profile_hook( 'app.main', - utils.run_with_router, + mitogen.utils.run_with_router, func, ) return wrapper diff --git a/mitogen/buildah.py b/mitogen/buildah.py index eec415f3..f850234d 100644 --- a/mitogen/buildah.py +++ b/mitogen/buildah.py @@ -37,37 +37,37 @@ import mitogen.parent LOG = logging.getLogger(__name__) -class Stream(mitogen.parent.Stream): - child_is_immediate_subprocess = False - +class Options(mitogen.parent.Options): container = None username = None buildah_path = 'buildah' - # TODO: better way of capturing errors such as "No such container." - create_child_args = { - 'merge_stdio': True - } - - def construct(self, container=None, - buildah_path=None, username=None, - **kwargs): - assert container or image - super(Stream, self).construct(**kwargs) - if container: - self.container = container + def __init__(self, container=None, buildah_path=None, username=None, + **kwargs): + super(Options, self).__init__(**kwargs) + assert container is not None + self.container = container if buildah_path: self.buildah_path = buildah_path if username: self.username = username + +class Connection(mitogen.parent.Connection): + options_class = Options + child_is_immediate_subprocess = False + + # TODO: better way of capturing errors such as "No such container." + create_child_args = { + 'merge_stdio': True + } + def _get_name(self): - return u'buildah.' + self.container + return u'buildah.' + self.options.container def get_boot_command(self): - args = [] - if self.username: - args += ['--user=' + self.username] - bits = [self.buildah_path, 'run'] + args + ['--', self.container] - - return bits + super(Stream, self).get_boot_command() + args = [self.options.buildah_path, 'run'] + if self.options.username: + args += ['--user=' + self.options.username] + args += ['--', self.options.container] + return args + super(Connection, self).get_boot_command() diff --git a/mitogen/core.py b/mitogen/core.py index ea83f961..6b182c85 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -37,6 +37,7 @@ bootstrap implementation sent to every new slave context. import binascii import collections import encodings.latin_1 +import encodings.utf_8 import errno import fcntl import itertools @@ -49,6 +50,7 @@ import signal import socket import struct import sys +import syslog import threading import time import traceback @@ -102,10 +104,9 @@ LOG = logging.getLogger('mitogen') IOLOG = logging.getLogger('mitogen.io') IOLOG.setLevel(logging.INFO) -LATIN1_CODEC = encodings.latin_1.Codec() # str.encode() may take import lock. Deadlock possible if broker calls # .encode() on behalf of thread currently waiting for module. -UTF8_CODEC = encodings.latin_1.Codec() +LATIN1_CODEC = encodings.latin_1.Codec() _v = False _vv = False @@ -214,7 +215,8 @@ else: class Error(Exception): - """Base for all exceptions raised by Mitogen. + """ + Base for all exceptions raised by Mitogen. :param str fmt: Exception text, or format string if `args` is non-empty. @@ -230,14 +232,18 @@ class Error(Exception): class LatchError(Error): - """Raised when an attempt is made to use a :class:`mitogen.core.Latch` - that has been marked closed.""" + """ + Raised when an attempt is made to use a :class:`mitogen.core.Latch` that + has been marked closed. + """ pass class Blob(BytesType): - """A serializable bytes subclass whose content is summarized in repr() - output, making it suitable for logging binary data.""" + """ + A serializable bytes subclass whose content is summarized in repr() output, + making it suitable for logging binary data. + """ def __repr__(self): return '[blob: %d bytes]' % len(self) @@ -246,8 +252,10 @@ class Blob(BytesType): class Secret(UnicodeType): - """A serializable unicode subclass whose content is masked in repr() - output, making it suitable for logging passwords.""" + """ + A serializable unicode subclass whose content is masked in repr() output, + making it suitable for logging passwords. + """ def __repr__(self): return '[secret]' @@ -281,7 +289,7 @@ class Kwargs(dict): def __init__(self, dct): for k, v in dct.iteritems(): if type(k) is unicode: - k, _ = UTF8_CODEC.encode(k) + k, _ = encodings.utf_8.encode(k) self[k] = v def __repr__(self): @@ -321,25 +329,33 @@ def _unpickle_call_error(s): class ChannelError(Error): - """Raised when a channel dies or has been closed.""" + """ + Raised when a channel dies or has been closed. + """ remote_msg = 'Channel closed by remote end.' local_msg = 'Channel closed by local end.' class StreamError(Error): - """Raised when a stream cannot be established.""" + """ + Raised when a stream cannot be established. + """ pass class TimeoutError(Error): - """Raised when a timeout occurs on a stream.""" + """ + Raised when a timeout occurs on a stream. + """ pass def to_text(o): - """Coerce `o` to Unicode by decoding it from UTF-8 if it is an instance of + """ + Coerce `o` to Unicode by decoding it from UTF-8 if it is an instance of :class:`bytes`, otherwise pass it to the :class:`str` constructor. The - returned object is always a plain :class:`str`, any subclass is removed.""" + returned object is always a plain :class:`str`, any subclass is removed. + """ if isinstance(o, BytesType): return o.decode('utf-8') return UnicodeType(o) @@ -379,11 +395,13 @@ else: def has_parent_authority(msg, _stream=None): - """Policy function for use with :class:`Receiver` and + """ + Policy function for use with :class:`Receiver` and :meth:`Router.add_handler` that requires incoming messages to originate from a parent context, or on a :class:`Stream` whose :attr:`auth_id ` has been set to that of a parent context or the current - context.""" + context. + """ return (msg.auth_id == mitogen.context_id or msg.auth_id in mitogen.parent_ids) @@ -432,35 +450,42 @@ def is_blacklisted_import(importer, fullname): def set_cloexec(fd): - """Set the file descriptor `fd` to automatically close on - :func:`os.execve`. This has no effect on file descriptors inherited across - :func:`os.fork`, they must be explicitly closed through some other means, - such as :func:`mitogen.fork.on_fork`.""" + """ + Set the file descriptor `fd` to automatically close on :func:`os.execve`. + This has no effect on file descriptors inherited across :func:`os.fork`, + they must be explicitly closed through some other means, such as + :func:`mitogen.fork.on_fork`. + """ flags = fcntl.fcntl(fd, fcntl.F_GETFD) assert fd > 2 fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC) def set_nonblock(fd): - """Set the file descriptor `fd` to non-blocking mode. For most underlying - file types, this causes :func:`os.read` or :func:`os.write` to raise + """ + Set the file descriptor `fd` to non-blocking mode. For most underlying file + types, this causes :func:`os.read` or :func:`os.write` to raise :class:`OSError` with :data:`errno.EAGAIN` rather than block the thread - when the underlying kernel buffer is exhausted.""" + when the underlying kernel buffer is exhausted. + """ flags = fcntl.fcntl(fd, fcntl.F_GETFL) fcntl.fcntl(fd, fcntl.F_SETFL, flags | os.O_NONBLOCK) def set_block(fd): - """Inverse of :func:`set_nonblock`, i.e. cause `fd` to block the thread - when the underlying kernel buffer is exhausted.""" + """ + Inverse of :func:`set_nonblock`, i.e. cause `fd` to block the thread when + the underlying kernel buffer is exhausted. + """ flags = fcntl.fcntl(fd, fcntl.F_GETFL) fcntl.fcntl(fd, fcntl.F_SETFL, flags & ~os.O_NONBLOCK) def io_op(func, *args): - """Wrap `func(*args)` that may raise :class:`select.error`, - :class:`IOError`, or :class:`OSError`, trapping UNIX error codes relating - to disconnection and retry events in various subsystems: + """ + Wrap `func(*args)` that may raise :class:`select.error`, :class:`IOError`, + or :class:`OSError`, trapping UNIX error codes relating to disconnection + and retry events in various subsystems: * When a signal is delivered to the process on Python 2, system call retry is signalled through :data:`errno.EINTR`. The invocation is automatically @@ -491,7 +516,8 @@ def io_op(func, *args): class PidfulStreamHandler(logging.StreamHandler): - """A :class:`logging.StreamHandler` subclass used when + """ + A :class:`logging.StreamHandler` subclass used when :meth:`Router.enable_debug() ` has been called, or the `debug` parameter was specified during context construction. Verifies the process ID has not changed on each call to :meth:`emit`, @@ -596,6 +622,43 @@ def import_module(modname): return __import__(modname, None, None, ['']) +def pipe(): + """ + Create a UNIX pipe pair using :func:`os.pipe`, wrapping the returned + descriptors in Python file objects in order to manage their lifetime and + ensure they are closed when their last reference is discarded and they have + not been closed explicitly. + """ + rfd, wfd = os.pipe() + return ( + os.fdopen(rfd, 'rb', 0), + os.fdopen(wfd, 'wb', 0) + ) + + +def iter_split(buf, delim, func): + """ + Invoke `func(s)` for each `delim`-delimited chunk in the potentially large + `buf`, avoiding intermediate lists and quadratic string operations. Return + the trailing undelimited portion of `buf`, or any unprocessed portion of + `buf` after `func(s)` returned :data:`False`. + + :returns: + `(trailer, cont)`, where `cont` is :data:`False` if the last call to + `func(s)` returned :data:`False`. + """ + dlen = len(delim) + start = 0 + cont = True + while cont: + nl = buf.find(delim, start) + if nl == -1: + break + cont = not func(buf[start:nl]) is False + start = nl + dlen + return buf[start:], cont + + class Py24Pickler(py_pickle.Pickler): """ Exceptions were classic classes until Python 2.5. Sadly for 2.4, cPickle @@ -687,6 +750,10 @@ class Message(object): #: the :class:`mitogen.select.Select` interface. Defaults to :data:`None`. receiver = None + HEADER_FMT = '>hLLLLLL' + HEADER_LEN = struct.calcsize(HEADER_FMT) + HEADER_MAGIC = 0x4d49 # 'MI' + def __init__(self, **kwargs): """ Construct a message from from the supplied `kwargs`. :attr:`src_id` and @@ -697,6 +764,14 @@ class Message(object): vars(self).update(kwargs) assert isinstance(self.data, BytesType) + def pack(self): + return ( + struct.pack(self.HEADER_FMT, self.HEADER_MAGIC, self.dst_id, + self.src_id, self.auth_id, self.handle, + self.reply_to or 0, len(self.data)) + + self.data + ) + def _unpickle_context(self, context_id, name): return _unpickle_context(context_id, name, router=self.router) @@ -708,8 +783,10 @@ class Message(object): return s def _find_global(self, module, func): - """Return the class implementing `module_name.class_name` or raise - `StreamError` if the module is not whitelisted.""" + """ + Return the class implementing `module_name.class_name` or raise + `StreamError` if the module is not whitelisted. + """ if module == __name__: if func == '_unpickle_call_error' or func == 'CallError': return _unpickle_call_error @@ -744,7 +821,7 @@ class Message(object): """ Syntax helper to construct a dead message. """ - kwargs['data'], _ = UTF8_CODEC.encode(reason or u'') + kwargs['data'], _ = encodings.utf_8.encode(reason or u'') return cls(reply_to=IS_DEAD, **kwargs) @classmethod @@ -893,7 +970,7 @@ def _unpickle_sender(router, context_id, dst_handle): if not (isinstance(router, Router) and isinstance(context_id, (int, long)) and context_id >= 0 and isinstance(dst_handle, (int, long)) and dst_handle > 0): - raise TypeError('cannot unpickle Sender: bad input') + raise TypeError('cannot unpickle Sender: bad input or missing router') return Sender(Context(router, context_id), dst_handle) @@ -1279,7 +1356,7 @@ class Importer(object): tup = msg.unpickle() fullname = tup[0] - _v and LOG.debug('Importer._on_load_module(%r)', fullname) + _v and LOG.debug('importer: received %s', fullname) self._lock.acquire() try: @@ -1303,10 +1380,12 @@ class Importer(object): if not present: funcs = self._callbacks.get(fullname) if funcs is not None: - _v and LOG.debug('_request_module(%r): in flight', fullname) + _v and LOG.debug('%s: existing request for %s in flight', + self, fullname) funcs.append(callback) else: - _v and LOG.debug('_request_module(%r): new request', fullname) + _v and LOG.debug('%s: requesting %s from parent', + self, fullname) self._callbacks[fullname] = [callback] self._context.send( Message(data=b(fullname), handle=GET_MODULE) @@ -1319,7 +1398,7 @@ class Importer(object): def load_module(self, fullname): fullname = to_text(fullname) - _v and LOG.debug('Importer.load_module(%r)', fullname) + _v and LOG.debug('importer: requesting %s', fullname) self._refuse_imports(fullname) event = threading.Event() @@ -1343,7 +1422,7 @@ class Importer(object): if mod.__package__ and not PY3: # 2.x requires __package__ to be exactly a string. - mod.__package__, _ = UTF8_CODEC.encode(mod.__package__) + mod.__package__, _ = encodings.utf_8.encode(mod.__package__) source = self.get_source(fullname) try: @@ -1390,6 +1469,9 @@ class LogHandler(logging.Handler): self.context = context self.local = threading.local() self._buffer = [] + # Private synchronization is needed while corked, to ensure no + # concurrent call to _send() exists during uncork(). + self._buffer_lock = threading.Lock() def uncork(self): """ @@ -1397,13 +1479,25 @@ class LogHandler(logging.Handler): possible to route messages, therefore messages are buffered until :meth:`uncork` is called by :class:`ExternalContext`. """ - self._send = self.context.send - for msg in self._buffer: - self._send(msg) - self._buffer = None + self._buffer_lock.acquire() + try: + self._send = self.context.send + for msg in self._buffer: + self._send(msg) + self._buffer = None + finally: + self._buffer_lock.release() def _send(self, msg): - self._buffer.append(msg) + self._buffer_lock.acquire() + try: + if self._buffer is None: + # uncork() may run concurrent to _send() + self._send(msg) + else: + self._buffer.append(msg) + finally: + self._buffer_lock.release() def emit(self, rec): if rec.name == 'mitogen.io' or \ @@ -1422,42 +1516,273 @@ class LogHandler(logging.Handler): self.local.in_emit = False +class Stream(object): + #: A :class:`Side` representing the stream's receive file descriptor. + receive_side = None + + #: A :class:`Side` representing the stream's transmit file descriptor. + transmit_side = None + + #: A :class:`Protocol` representing the protocol active on the stream. + protocol = None + + #: In parents, the :class:`mitogen.parent.Connection` instance. + conn = None + + name = u'default' + + def set_protocol(self, protocol): + self.protocol = protocol + self.protocol.stream = self + + def accept(self, rfp, wfp): + self.receive_side = Side(self, rfp) + self.transmit_side = Side(self, wfp) + + def __repr__(self): + return "" % (self.name,) + + def on_receive(self, broker): + """ + Called by :class:`Broker` when the stream's :attr:`receive_side` has + been marked readable using :meth:`Broker.start_receive` and the broker + has detected the associated file descriptor is ready for reading. + + Subclasses must implement this if :meth:`Broker.start_receive` is ever + called on them, and the method must call :meth:`on_disconect` if + reading produces an empty string. + """ + buf = self.receive_side.read(self.protocol.read_size) + if not buf: + LOG.debug('%r: empty read, disconnecting', self.receive_side) + return self.on_disconnect(broker) + + self.protocol.on_receive(broker, buf) + + def on_transmit(self, broker): + """ + Called by :class:`Broker` when the stream's :attr:`transmit_side` + has been marked writeable using :meth:`Broker._start_transmit` and + the broker has detected the associated file descriptor is ready for + writing. + + Subclasses must implement this if :meth:`Broker._start_transmit` is + ever called on them. + """ + self.protocol.on_transmit(broker) + + def on_shutdown(self, broker): + """ + Called by :meth:`Broker.shutdown` to allow the stream time to + gracefully shutdown. The base implementation simply called + :meth:`on_disconnect`. + """ + fire(self, 'shutdown') + self.protocol.on_shutdown(broker) + + def on_disconnect(self, broker): + """ + Called by :class:`Broker` to force disconnect the stream. The base + implementation simply closes :attr:`receive_side` and + :attr:`transmit_side` and unregisters the stream from the broker. + """ + fire(self, 'disconnect') + self.protocol.on_disconnect(broker) + + +class Protocol(object): + """ + Implement the program behaviour associated with activity on a + :class:`Stream`. The protocol in use may vary over a stream's life, for + example to allow :class:`mitogen.parent.BootstrapProtocol` to initialize + the connected child before handing it off to :class:`MitogenProtocol`. A + stream's active protocol is tracked in the :attr:`Stream.protocol` + attribute, and modified via :meth:`Stream.set_protocol`. + + Protocols do not handle IO, they are entirely reliant on the interface + provided by :class:`Stream` and :class:`Side`, allowing the underlying IO + implementation to be replaced without modifying behavioural logic. + """ + stream_class = Stream + stream = None + read_size = CHUNK_SIZE + + @classmethod + def build_stream(cls, *args, **kwargs): + stream = cls.stream_class() + stream.set_protocol(cls(*args, **kwargs)) + return stream + + def __repr__(self): + return '%s(%s)' % ( + self.__class__.__name__, + self.stream and self.stream.name, + ) + + def on_shutdown(self, broker): + _v and LOG.debug('%r: shutting down', self) + self.stream.on_disconnect(broker) + + def on_disconnect(self, broker): + LOG.debug('%r: disconnecting', self) + if self.stream.receive_side: + broker.stop_receive(self.stream) + self.stream.receive_side.close() + if self.stream.transmit_side: + broker._stop_transmit(self.stream) + self.stream.transmit_side.close() + + +class DelimitedProtocol(Protocol): + """ + Provide a :meth:`Protocol.on_receive` implementation for protocols that are + delimited by a fixed string, like text based protocols. Each message is + passed to :meth:`on_line_received` as it arrives, with incomplete messages + passed to :meth:`on_partial_line_received`. + + When emulating user input it is often necessary to respond to incomplete + lines, such as when a "Password: " prompt is sent. + :meth:`on_partial_line_received` may be called repeatedly with an + increasingly complete message. When a complete message is finally received, + :meth:`on_line_received` will be called once for it before the buffer is + discarded. + + If :func:`on_line_received` returns :data:`False`, remaining data is passed + unprocessed to the stream's current protocol's :meth:`on_receive`. This + allows switching from line-oriented to binary while the input buffer + contains both kinds of data. + """ + #: The delimiter. Defaults to newline. + delimiter = b('\n') + _trailer = b('') + + def on_receive(self, broker, buf): + IOLOG.debug('%r.on_receive()', self) + self._trailer, cont = mitogen.core.iter_split( + buf=self._trailer + buf, + delim=self.delimiter, + func=self.on_line_received, + ) + + if self._trailer: + if cont: + self.on_partial_line_received(self._trailer) + else: + assert self.stream.protocol is not self + self.stream.protocol.on_receive(broker, self._trailer) + + def on_line_received(self, line): + pass + + def on_partial_line_received(self, line): + pass + + +class BufferedWriter(object): + """ + Implement buffered output while avoiding quadratic string operations. This + is currently constructed by each protocol, in future it may become fixed + for each stream instead. + """ + def __init__(self, broker, protocol): + self._broker = broker + self._protocol = protocol + self._buf = collections.deque() + self._len = 0 + + def write(self, s): + """ + Transmit `s` immediately, falling back to enqueuing it and marking the + stream writeable if no OS buffer space is available. + """ + if not self._len: + # Modifying epoll/Kqueue state is expensive, as are needless broker + # loops. Rather than wait for writeability, just write immediately, + # and fall back to the broker loop on error or full buffer. + try: + n = self._protocol.stream.transmit_side.write(s) + if n: + if n == len(s): + return + s = s[n:] + except OSError: + pass + + self._broker._start_transmit(self._protocol.stream) + self._buf.append(s) + self._len += len(s) + + def on_transmit(self, broker): + """ + Respond to stream writeability by retrying previously buffered + :meth:`write` calls. + """ + if self._buf: + buf = self._buf.popleft() + written = self._protocol.stream.transmit_side.write(buf) + if not written: + _v and LOG.debug('%r.on_transmit(): disconnection detected', self) + self._protocol.stream.on_disconnect(broker) + return + elif written != len(buf): + self._buf.appendleft(BufferType(buf, written)) + + _vv and IOLOG.debug('%r.on_transmit() -> len %d', self, written) + self._len -= written + + if not self._buf: + broker._stop_transmit(self._protocol.stream) + + class Side(object): """ - Represent a single side of a :class:`BasicStream`. This exists to allow - streams implemented using unidirectional (e.g. UNIX pipe) and bidirectional - (e.g. UNIX socket) file descriptors to operate identically. + Represent one side of a :class:`Stream`. This allows unidirectional (e.g. + pipe) and bidirectional (e.g. socket) streams to operate identically. + + Sides are also responsible for tracking the open/closed state of the + underlying FD, preventing erroneous duplicate calls to :func:`os.close` due + to duplicate :meth:`Stream.on_disconnect` calls, which would otherwise risk + silently succeeding by closing an unrelated descriptor. For this reason, it + is crucial only one :class:`Side` exists per unique descriptor. :param mitogen.core.Stream stream: The stream this side is associated with. - - :param int fd: - Underlying file descriptor. - + :param object fp: + The file or socket object managing the underlying file descriptor. Any + object may be used that supports `fileno()` and `close()` methods. + :param bool cloexec: + If :data:`True`, the descriptor has its :data:`fcntl.FD_CLOEXEC` flag + enabled using :func:`fcntl.fcntl`. :param bool keep_alive: - Value for :attr:`keep_alive` - - During construction, the file descriptor has its :data:`os.O_NONBLOCK` flag - enabled using :func:`fcntl.fcntl`. + If :data:`True`, the continued existence of this side will extend the + shutdown grace period until it has been unregistered from the broker. + :param bool blocking: + If :data:`False`, the descriptor has its :data:`os.O_NONBLOCK` flag + enabled using :func:`fcntl.fcntl`. """ _fork_refs = weakref.WeakValueDictionary() + closed = False - def __init__(self, stream, fd, cloexec=True, keep_alive=True, blocking=False): + def __init__(self, stream, fp, cloexec=True, keep_alive=True, blocking=False): #: The :class:`Stream` for which this is a read or write side. self.stream = stream + # File or socket object responsible for the lifetime of its underlying + # file descriptor. + self.fp = fp #: Integer file descriptor to perform IO on, or :data:`None` if - #: :meth:`close` has been called. - self.fd = fd - self.closed = False + #: :meth:`close` has been called. This is saved separately from the + #: file object, since fileno() cannot be called on it after it has been + #: closed. + self.fd = fp.fileno() #: If :data:`True`, causes presence of this side in #: :class:`Broker`'s active reader set to defer shutdown until the #: side is disconnected. self.keep_alive = keep_alive self._fork_refs[id(self)] = self if cloexec: - set_cloexec(fd) + set_cloexec(self.fd) if not blocking: - set_nonblock(fd) + set_nonblock(self.fd) def __repr__(self): return '' % (self.stream, self.fd) @@ -1474,10 +1799,10 @@ class Side(object): Call :func:`os.close` on :attr:`fd` if it is not :data:`None`, then set it to :data:`None`. """ + _vv and IOLOG.debug('%r.close()', self) if not self.closed: - _vv and IOLOG.debug('%r.close()', self) self.closed = True - os.close(self.fd) + self.fp.close() def read(self, n=CHUNK_SIZE): """ @@ -1499,7 +1824,7 @@ class Side(object): return b('') s, disconnected = io_op(os.read, self.fd, n) if disconnected: - LOG.debug('%r.read(): disconnected: %s', self, disconnected) + LOG.debug('%r: disconnected during read: %s', self, disconnected) return b('') return s @@ -1513,79 +1838,21 @@ class Side(object): Number of bytes written, or :data:`None` if disconnection was detected. """ - if self.closed or self.fd is None: - # Refuse to touch the handle after closed, it may have been reused - # by another thread. + if self.closed: + # Don't touch the handle after close, it may be reused elsewhere. return None written, disconnected = io_op(os.write, self.fd, s) if disconnected: - LOG.debug('%r.write(): disconnected: %s', self, disconnected) + LOG.debug('%r: disconnected during write: %s', self, disconnected) return None return written -class BasicStream(object): - #: A :class:`Side` representing the stream's receive file descriptor. - receive_side = None - - #: A :class:`Side` representing the stream's transmit file descriptor. - transmit_side = None - - def on_receive(self, broker): - """ - Called by :class:`Broker` when the stream's :attr:`receive_side` has - been marked readable using :meth:`Broker.start_receive` and the broker - has detected the associated file descriptor is ready for reading. - - Subclasses must implement this if :meth:`Broker.start_receive` is ever - called on them, and the method must call :meth:`on_disconect` if - reading produces an empty string. - """ - pass - - def on_transmit(self, broker): - """ - Called by :class:`Broker` when the stream's :attr:`transmit_side` - has been marked writeable using :meth:`Broker._start_transmit` and - the broker has detected the associated file descriptor is ready for - writing. - - Subclasses must implement this if :meth:`Broker._start_transmit` is - ever called on them. - """ - pass - - def on_shutdown(self, broker): - """ - Called by :meth:`Broker.shutdown` to allow the stream time to - gracefully shutdown. The base implementation simply called - :meth:`on_disconnect`. - """ - _v and LOG.debug('%r.on_shutdown()', self) - fire(self, 'shutdown') - self.on_disconnect(broker) - - def on_disconnect(self, broker): - """ - Called by :class:`Broker` to force disconnect the stream. The base - implementation simply closes :attr:`receive_side` and - :attr:`transmit_side` and unregisters the stream from the broker. - """ - LOG.debug('%r.on_disconnect()', self) - if self.receive_side: - broker.stop_receive(self) - self.receive_side.close() - if self.transmit_side: - broker._stop_transmit(self) - self.transmit_side.close() - fire(self, 'disconnect') - - -class Stream(BasicStream): +class MitogenProtocol(Protocol): """ - :class:`BasicStream` subclass implementing mitogen's :ref:`stream - protocol `. + :class:`Protocol` implementing mitogen's :ref:`stream protocol + `. """ #: If not :data:`None`, :class:`Router` stamps this into #: :attr:`Message.auth_id` of every message received on this stream. @@ -1596,24 +1863,24 @@ class Stream(BasicStream): #: :data:`mitogen.parent_ids`. is_privileged = False - def __init__(self, router, remote_id, **kwargs): + def __init__(self, router, remote_id): self._router = router self.remote_id = remote_id - self.name = u'default' self.sent_modules = set(['mitogen', 'mitogen.core']) - self.construct(**kwargs) self._input_buf = collections.deque() - self._output_buf = collections.deque() self._input_buf_len = 0 - self._output_buf_len = 0 + self._writer = BufferedWriter(router.broker, self) + #: Routing records the dst_id of every message arriving from this #: stream. Any arriving DEL_ROUTE is rebroadcast for any such ID. self.egress_ids = set() - def construct(self): - pass - - def _internal_receive(self, broker, buf): + def on_receive(self, broker, buf): + """ + Handle the next complete message on the stream. Raise + :class:`StreamError` on failure. + """ + _vv and IOLOG.debug('%r.on_receive()', self) if self._input_buf and self._input_buf_len < 128: self._input_buf[0] += buf else: @@ -1623,60 +1890,45 @@ class Stream(BasicStream): while self._receive_one(broker): pass - def on_receive(self, broker): - """Handle the next complete message on the stream. Raise - :class:`StreamError` on failure.""" - _vv and IOLOG.debug('%r.on_receive()', self) - - buf = self.receive_side.read() - if not buf: - return self.on_disconnect(broker) - - self._internal_receive(broker, buf) - - HEADER_FMT = '>hLLLLLL' - HEADER_LEN = struct.calcsize(HEADER_FMT) - HEADER_MAGIC = 0x4d49 # 'MI' - corrupt_msg = ( - 'Corruption detected: frame signature incorrect. This likely means ' - 'some external process is interfering with the connection. Received:' + '%s: Corruption detected: frame signature incorrect. This likely means' + ' some external process is interfering with the connection. Received:' '\n\n' '%r' ) def _receive_one(self, broker): - if self._input_buf_len < self.HEADER_LEN: + if self._input_buf_len < Message.HEADER_LEN: return False msg = Message() msg.router = self._router (magic, msg.dst_id, msg.src_id, msg.auth_id, msg.handle, msg.reply_to, msg_len) = struct.unpack( - self.HEADER_FMT, - self._input_buf[0][:self.HEADER_LEN], + Message.HEADER_FMT, + self._input_buf[0][:Message.HEADER_LEN], ) - if magic != self.HEADER_MAGIC: - LOG.error(self.corrupt_msg, self._input_buf[0][:2048]) - self.on_disconnect(broker) + if magic != Message.HEADER_MAGIC: + LOG.error(self.corrupt_msg, self.stream.name, self._input_buf[0][:2048]) + self.stream.on_disconnect(broker) return False if msg_len > self._router.max_message_size: LOG.error('Maximum message size exceeded (got %d, max %d)', msg_len, self._router.max_message_size) - self.on_disconnect(broker) + self.stream.on_disconnect(broker) return False - total_len = msg_len + self.HEADER_LEN + total_len = msg_len + Message.HEADER_LEN if self._input_buf_len < total_len: _vv and IOLOG.debug( '%r: Input too short (want %d, got %d)', - self, msg_len, self._input_buf_len - self.HEADER_LEN + self, msg_len, self._input_buf_len - Message.HEADER_LEN ) return False - start = self.HEADER_LEN + start = Message.HEADER_LEN prev_start = start remain = total_len bits = [] @@ -1691,7 +1943,7 @@ class Stream(BasicStream): msg.data = b('').join(bits) self._input_buf.appendleft(buf[prev_start+len(bit):]) self._input_buf_len -= total_len - self._router._async_route(msg, self) + self._router._async_route(msg, self.stream) return True def pending_bytes(self): @@ -1703,68 +1955,31 @@ class Stream(BasicStream): For an accurate result, this method should be called from the Broker thread, for example by using :meth:`Broker.defer_sync`. """ - return self._output_buf_len + return self._writer._len def on_transmit(self, broker): - """Transmit buffered messages.""" + """ + Transmit buffered messages. + """ _vv and IOLOG.debug('%r.on_transmit()', self) - - if self._output_buf: - buf = self._output_buf.popleft() - written = self.transmit_side.write(buf) - if not written: - _v and LOG.debug('%r.on_transmit(): disconnection detected', self) - self.on_disconnect(broker) - return - elif written != len(buf): - self._output_buf.appendleft(BufferType(buf, written)) - - _vv and IOLOG.debug('%r.on_transmit() -> len %d', self, written) - self._output_buf_len -= written - - if not self._output_buf: - broker._stop_transmit(self) + self._writer.on_transmit(broker) def _send(self, msg): _vv and IOLOG.debug('%r._send(%r)', self, msg) - pkt = struct.pack(self.HEADER_FMT, self.HEADER_MAGIC, msg.dst_id, - msg.src_id, msg.auth_id, msg.handle, - msg.reply_to or 0, len(msg.data)) + msg.data - - if not self._output_buf_len: - # Modifying epoll/Kqueue state is expensive, as are needless broker - # loops. Rather than wait for writeability, just write immediately, - # and fall back to the broker loop on error or full buffer. - try: - n = self.transmit_side.write(pkt) - if n: - if n == len(pkt): - return - pkt = pkt[n:] - except OSError: - pass - - self._router.broker._start_transmit(self) - self._output_buf.append(pkt) - self._output_buf_len += len(pkt) + self._writer.write(msg.pack()) def send(self, msg): - """Send `data` to `handle`, and tell the broker we have output. May - be called from any thread.""" + """ + Send `data` to `handle`, and tell the broker we have output. May be + called from any thread. + """ self._router.broker.defer(self._send, msg) def on_shutdown(self, broker): - """Override BasicStream behaviour of immediately disconnecting.""" - _v and LOG.debug('%r.on_shutdown(%r)', self, broker) - - def accept(self, rfd, wfd): - # TODO: what is this os.dup for? - self.receive_side = Side(self, os.dup(rfd)) - self.transmit_side = Side(self, os.dup(wfd)) - - def __repr__(self): - cls = type(self) - return "%s.%s('%s')" % (cls.__module__, cls.__name__, self.name) + """ + Disable :class:`Protocol` immediate disconnect behaviour. + """ + _v and LOG.debug('%r: shutting down', self) class Context(object): @@ -1790,21 +2005,20 @@ class Context(object): :param str name: Context name. """ + name = None remote_name = None def __init__(self, router, context_id, name=None): self.router = router self.context_id = context_id - self.name = name + if name: + self.name = to_text(name) def __reduce__(self): - name = self.name - if name and not isinstance(name, UnicodeType): - name = UnicodeType(name, 'utf-8') - return _unpickle_context, (self.context_id, name) + return _unpickle_context, (self.context_id, self.name) def on_disconnect(self): - _v and LOG.debug('%r.on_disconnect()', self) + _v and LOG.debug('%r: disconnecting', self) fire(self, 'disconnect') def send_async(self, msg, persist=False): @@ -1946,7 +2160,7 @@ class Poller(object): self._wfds = {} def __repr__(self): - return '%s(%#x)' % (type(self).__name__, id(self)) + return '%s' % (type(self).__name__,) def _update(self, fd): """ @@ -2029,9 +2243,6 @@ class Poller(object): if gen and gen < self._generation: yield data - if timeout: - timeout *= 1000 - def poll(self, timeout=None): """ Block the calling thread until one or more FDs are ready for IO. @@ -2219,7 +2430,7 @@ class Latch(object): :meth:`put` to write a byte to our socket pair. """ _vv and IOLOG.debug( - '%r._get_sleep(timeout=%r, block=%r, rfd=%d, wfd=%d)', + '%r._get_sleep(timeout=%r, block=%r, fd=%d/%d)', self, timeout, block, rsock.fileno(), wsock.fileno() ) @@ -2297,7 +2508,7 @@ class Latch(object): ) -class Waker(BasicStream): +class Waker(Protocol): """ :class:`BasicStream` subclass implementing the `UNIX self-pipe trick`_. Used to wake the multiplexer when another thread needs to modify its state @@ -2305,22 +2516,24 @@ class Waker(BasicStream): .. _UNIX self-pipe trick: https://cr.yp.to/docs/selfpipe.html """ + read_size = 1 broker_ident = None + @classmethod + def build_stream(cls, broker): + stream = super(Waker, cls).build_stream(broker) + stream.accept(*pipe()) + return stream + def __init__(self, broker): self._broker = broker self._lock = threading.Lock() self._deferred = [] - rfd, wfd = os.pipe() - self.receive_side = Side(self, rfd) - self.transmit_side = Side(self, wfd) - def __repr__(self): - return 'Waker(%r rfd=%r, wfd=%r)' % ( - self._broker, - self.receive_side and self.receive_side.fd, - self.transmit_side and self.transmit_side.fd, + return 'Waker(fd=%r/%r)' % ( + self.stream.receive_side and self.stream.receive_side.fd, + self.stream.transmit_side and self.stream.transmit_side.fd, ) @property @@ -2334,7 +2547,7 @@ class Waker(BasicStream): finally: self._lock.release() - def on_receive(self, broker): + def on_receive(self, broker, buf): """ Drain the pipe and fire callbacks. Since :attr:`_deferred` is synchronized, :meth:`defer` and :meth:`on_receive` can conspire to @@ -2343,7 +2556,6 @@ class Waker(BasicStream): _vv and IOLOG.debug('%r.on_receive()', self) self._lock.acquire() try: - self.receive_side.read(1) deferred = self._deferred self._deferred = [] finally: @@ -2355,7 +2567,7 @@ class Waker(BasicStream): except Exception: LOG.exception('defer() crashed: %r(*%r, **%r)', func, args, kwargs) - self._broker.shutdown() + broker.shutdown() def _wake(self): """ @@ -2363,7 +2575,7 @@ class Waker(BasicStream): teardown, the FD may already be closed, so ignore EBADF. """ try: - self.transmit_side.write(b(' ')) + self.stream.transmit_side.write(b(' ')) except OSError: e = sys.exc_info()[1] if e.args[0] != errno.EBADF: @@ -2390,7 +2602,8 @@ class Waker(BasicStream): if self._broker._exitted: raise Error(self.broker_shutdown_msg) - _vv and IOLOG.debug('%r.defer() [fd=%r]', self, self.transmit_side.fd) + _vv and IOLOG.debug('%r.defer() [fd=%r]', self, + self.stream.transmit_side.fd) self._lock.acquire() try: if not self._deferred: @@ -2400,55 +2613,47 @@ class Waker(BasicStream): self._lock.release() -class IoLogger(BasicStream): +class IoLoggerProtocol(DelimitedProtocol): """ - :class:`BasicStream` subclass that sets up redirection of a standard - UNIX file descriptor back into the Python :mod:`logging` package. + Handle redirection of standard IO into the :mod:`logging` package. """ - _buf = '' - - def __init__(self, broker, name, dest_fd): - self._broker = broker - self._name = name - self._rsock, self._wsock = socket.socketpair() - os.dup2(self._wsock.fileno(), dest_fd) - set_cloexec(self._wsock.fileno()) + @classmethod + def build_stream(cls, name, dest_fd): + """ + Even though the descriptor `dest_fd` will hold the opposite end of the + socket open, we must keep a separate dup() of it (i.e. wsock) in case + some code decides to overwrite `dest_fd` later, which would thus break + :meth:`on_shutdown`. + """ + rsock, wsock = socket.socketpair() + os.dup2(wsock.fileno(), dest_fd) + stream = super(IoLoggerProtocol, cls).build_stream(name) + stream.name = name + stream.accept(rsock, wsock) + return stream + def __init__(self, name): self._log = logging.getLogger(name) # #453: prevent accidental log initialization in a child creating a # feedback loop. self._log.propagate = False self._log.handlers = logging.getLogger().handlers[:] - self.receive_side = Side(self, self._rsock.fileno()) - self.transmit_side = Side(self, dest_fd, cloexec=False, blocking=True) - self._broker.start_receive(self) - - def __repr__(self): - return '' % (self._name,) - - def _log_lines(self): - while self._buf.find('\n') != -1: - line, _, self._buf = str_partition(self._buf, '\n') - self._log.info('%s', line.rstrip('\n')) - def on_shutdown(self, broker): - """Shut down the write end of the logging socket.""" - _v and LOG.debug('%r.on_shutdown()', self) + """ + Shut down the write end of the logging socket. + """ + _v and LOG.debug('%r: shutting down', self) if not IS_WSL: - # #333: WSL generates invalid readiness indication on shutdown() - self._wsock.shutdown(socket.SHUT_WR) - self._wsock.close() - self.transmit_side.close() + # #333: WSL generates invalid readiness indication on shutdown(). + # This modifies the *kernel object* inherited by children, causing + # EPIPE on subsequent writes to any dupped FD in any process. The + # read side can then drain completely of prior buffered data. + self.stream.transmit_side.fp.shutdown(socket.SHUT_WR) + self.stream.transmit_side.close() - def on_receive(self, broker): - _vv and IOLOG.debug('%r.on_receive()', self) - buf = self.receive_side.read() - if not buf: - return self.on_disconnect(broker) - - self._buf += buf.decode('latin1') - self._log_lines() + def on_line_received(self, line): + self._log.info('%s', line.decode('utf-8', 'replace')) class Router(object): @@ -2518,12 +2723,13 @@ class Router(object): corresponding :attr:`_context_by_id` member. This is replaced by :class:`mitogen.parent.RouteMonitor` in an upgraded context. """ - LOG.error('%r._on_del_route() %r', self, msg) if msg.is_dead: return target_id_s, _, name = bytes_partition(msg.data, b(':')) target_id = int(target_id_s, 10) + LOG.error('%r: deleting route to %s (%d)', + self, to_text(name), target_id) context = self._context_by_id.get(target_id) if context: fire(context, 'disconnect') @@ -2595,7 +2801,8 @@ class Router(object): the stream's receive side to the I/O multiplexer. This method remains public while the design has not yet settled. """ - _v and LOG.debug('register(%r, %r)', context, stream) + _v and LOG.debug('%s: registering %r to stream %r', + self, context, stream) self._write_lock.acquire() try: self._stream_by_id[context.context_id] = stream @@ -2700,7 +2907,7 @@ class Router(object): return handle - duplicate_handle_msg = 'cannot register a handle that is already exists' + duplicate_handle_msg = 'cannot register a handle that already exists' refused_msg = 'refused by policy' invalid_handle_msg = 'invalid handle' too_large_msg = 'message too large (max %d bytes)' @@ -2719,9 +2926,11 @@ class Router(object): del self._handle_map[handle] def on_shutdown(self, broker): - """Called during :meth:`Broker.shutdown`, informs callbacks registered - with :meth:`add_handle_cb` the connection is dead.""" - _v and LOG.debug('%r.on_shutdown(%r)', self, broker) + """ + Called during :meth:`Broker.shutdown`, informs callbacks registered + with :meth:`add_handle_cb` the connection is dead. + """ + _v and LOG.debug('%r: shutting down', self, broker) fire(self, 'shutdown') for handle, (persist, fn) in self._handle_map.iteritems(): _v and LOG.debug('%r.on_shutdown(): killing %r: %r', self, handle, fn) @@ -2797,11 +3006,11 @@ class Router(object): self, msg.src_id, in_stream, expect, msg) return - if in_stream.auth_id is not None: - msg.auth_id = in_stream.auth_id + if in_stream.protocol.auth_id is not None: + msg.auth_id = in_stream.protocol.auth_id # Maintain a set of IDs the source ever communicated with. - in_stream.egress_ids.add(msg.dst_id) + in_stream.protocol.egress_ids.add(msg.dst_id) if msg.dst_id == mitogen.context_id: return self._invoke(msg, in_stream) @@ -2816,12 +3025,13 @@ class Router(object): return if in_stream and self.unidirectional and not \ - (in_stream.is_privileged or out_stream.is_privileged): + (in_stream.protocol.is_privileged or + out_stream.protocol.is_privileged): self._maybe_send_dead(msg, self.unidirectional_msg, - in_stream.remote_id, out_stream.remote_id) + in_stream.protocol.remote_id, out_stream.protocol.remote_id) return - out_stream._send(msg) + out_stream.protocol._send(msg) def route(self, msg): """ @@ -2836,17 +3046,26 @@ class Router(object): self.broker.defer(self._async_route, msg) +class NullTimerList(object): + def get_timeout(self): + return None + + class Broker(object): """ Responsible for handling I/O multiplexing in a private thread. - **Note:** This is the somewhat limited core version of the Broker class - used by child contexts. The master subclass is documented below. + **Note:** This somewhat limited core version is used by children. The + master subclass is documented below. """ poller_class = Poller _waker = None _thread = None + # :func:`mitogen.parent._upgrade_broker` replaces this with + # :class:`mitogen.parent.TimerList` during upgrade. + timers = NullTimerList() + #: Seconds grace to allow :class:`streams ` to shutdown gracefully #: before force-disconnecting them during :meth:`shutdown`. shutdown_timeout = 3.0 @@ -2854,11 +3073,11 @@ class Broker(object): def __init__(self, poller_class=None, activate_compat=True): self._alive = True self._exitted = False - self._waker = Waker(self) + self._waker = Waker.build_stream(self) #: Arrange for `func(\*args, \**kwargs)` to be executed on the broker #: thread, or immediately if the current thread is the broker thread. #: Safe to call from any thread. - self.defer = self._waker.defer + self.defer = self._waker.protocol.defer self.poller = self.poller_class() self.poller.start_receive( self._waker.receive_side.fd, @@ -2892,7 +3111,7 @@ class Broker(object): """ _vv and IOLOG.debug('%r.start_receive(%r)', self, stream) side = stream.receive_side - assert side and side.fd is not None + assert side and not side.closed self.defer(self.poller.start_receive, side.fd, (side, stream.on_receive)) @@ -2913,7 +3132,7 @@ class Broker(object): """ _vv and IOLOG.debug('%r._start_transmit(%r)', self, stream) side = stream.transmit_side - assert side and side.fd is not None + assert side and not side.closed self.poller.start_transmit(side.fd, (side, stream.on_transmit)) def _stop_transmit(self, stream): @@ -2932,7 +3151,7 @@ class Broker(object): progress (e.g. log draining). """ it = (side.keep_alive for (_, (side, _)) in self.poller.readers) - return sum(it, 0) + return sum(it, 0) > 0 or self.timers.get_timeout() is not None def defer_sync(self, func): """ @@ -2975,10 +3194,19 @@ class Broker(object): """ _vv and IOLOG.debug('%r._loop_once(%r, %r)', self, timeout, self.poller) + + timer_to = self.timers.get_timeout() + if timeout is None: + timeout = timer_to + elif timer_to is not None and timer_to < timeout: + timeout = timer_to + #IOLOG.debug('readers =\n%s', pformat(self.poller.readers)) #IOLOG.debug('writers =\n%s', pformat(self.poller.writers)) for side, func in self.poller.poll(timeout): self._call(side.stream, func) + if timer_to is not None: + self.timers.expire() def _broker_exit(self): """ @@ -2986,7 +3214,7 @@ class Broker(object): to shut down gracefully, then discard the :class:`Poller`. """ for _, (side, _) in self.poller.readers + self.poller.writers: - LOG.debug('_broker_main() force disconnecting %r', side) + LOG.debug('%r: force disconnecting %r', self, side) side.stream.on_disconnect(self) self.poller.close() @@ -3017,7 +3245,7 @@ class Broker(object): :meth:`shutdown` is called. """ # For Python 2.4, no way to retrieve ident except on thread. - self._waker.broker_ident = thread.get_ident() + self._waker.protocol.broker_ident = thread.get_ident() try: while self._alive: self._loop_once() @@ -3025,7 +3253,10 @@ class Broker(object): fire(self, 'shutdown') self._broker_shutdown() except Exception: - LOG.exception('_broker_main() crashed') + e = sys.exc_info()[1] + LOG.exception('broker crashed') + syslog.syslog(syslog.LOG_ERR, 'broker crashed: %s' % (e,)) + syslog.closelog() # prevent test 'fd leak'. self._alive = False # Ensure _alive is consistent on crash. self._exitted = True @@ -3040,7 +3271,7 @@ class Broker(object): Request broker gracefully disconnect streams and stop. Safe to call from any thread. """ - _v and LOG.debug('%r.shutdown()', self) + _v and LOG.debug('%r: shutting down', self) def _shutdown(): self._alive = False if self._alive and not self._exitted: @@ -3054,7 +3285,7 @@ class Broker(object): self._thread.join() def __repr__(self): - return 'Broker(%#x)' % (id(self),) + return 'Broker(%04x)' % (id(self) & 0xffff,) class Dispatcher(object): @@ -3068,6 +3299,9 @@ class Dispatcher(object): mode, any exception that occurs is recorded, and causes all subsequent calls with the same `chain_id` to fail with the same exception. """ + def __repr__(self): + return 'Dispatcher' + def __init__(self, econtext): self.econtext = econtext #: Chain ID -> CallError if prior call failed. @@ -3084,7 +3318,7 @@ class Dispatcher(object): def _parse_request(self, msg): data = msg.unpickle(throw=False) - _v and LOG.debug('_dispatch_one(%r)', data) + _v and LOG.debug('%r: dispatching %r', self, data) chain_id, modname, klass, func, args, kwargs = data obj = import_module(modname) @@ -3118,7 +3352,7 @@ class Dispatcher(object): def _dispatch_calls(self): for msg in self.recv: chain_id, ret = self._dispatch_one(msg) - _v and LOG.debug('_dispatch_calls: %r -> %r', msg, ret) + _v and LOG.debug('%r: %r -> %r', self, msg, ret) if msg.reply_to: msg.reply(ret) elif isinstance(ret, CallError) and chain_id is None: @@ -3198,8 +3432,8 @@ class ExternalContext(object): th.start() def _on_shutdown_msg(self, msg): - _v and LOG.debug('_on_shutdown_msg(%r)', msg) if not msg.is_dead: + _v and LOG.debug('shutdown request from context %d', msg.src_id) self.broker.shutdown() def _on_parent_disconnect(self): @@ -3208,7 +3442,7 @@ class ExternalContext(object): mitogen.parent_id = None LOG.info('Detachment complete') else: - _v and LOG.debug('%r: parent stream is gone, dying.', self) + _v and LOG.debug('parent stream is gone, dying.') self.broker.shutdown() def detach(self): @@ -3219,7 +3453,7 @@ class ExternalContext(object): self.parent.send_await(Message(handle=DETACHING)) LOG.info('Detaching from %r; parent is %s', stream, self.parent) for x in range(20): - pending = self.broker.defer_sync(lambda: stream.pending_bytes()) + pending = self.broker.defer_sync(stream.protocol.pending_bytes) if not pending: break time.sleep(0.05) @@ -3252,18 +3486,16 @@ class ExternalContext(object): else: self.parent = Context(self.router, parent_id, 'parent') - in_fd = self.config.get('in_fd', 100) - out_fd = self.config.get('out_fd', 1) - self.stream = Stream(self.router, parent_id) + in_fp = os.fdopen(os.dup(self.config.get('in_fd', 100)), 'rb', 0) + out_fp = os.fdopen(os.dup(self.config.get('out_fd', 1)), 'wb', 0) + self.stream = MitogenProtocol.build_stream(self.router, parent_id) + self.stream.accept(in_fp, out_fp) self.stream.name = 'parent' - self.stream.accept(in_fd, out_fd) self.stream.receive_side.keep_alive = False listen(self.stream, 'disconnect', self._on_parent_disconnect) listen(self.broker, 'exit', self._on_broker_exit) - os.close(in_fd) - def _reap_first_stage(self): try: os.wait() # Reap first stage. @@ -3331,16 +3563,13 @@ class ExternalContext(object): def _nullify_stdio(self): """ - Open /dev/null to replace stdin, and stdout/stderr temporarily. In case - of odd startup, assume we may be allocated a standard handle. + Open /dev/null to replace stdio temporarily. In case of odd startup, + assume we may be allocated a standard handle. """ - fd = os.open('/dev/null', os.O_RDWR) - try: - for stdfd in (0, 1, 2): - if fd != stdfd: - os.dup2(fd, stdfd) - finally: - if fd not in (0, 1, 2): + for stdfd, mode in ((0, os.O_RDONLY), (1, os.O_RDWR), (2, os.O_RDWR)): + fd = os.open('/dev/null', mode) + if fd != stdfd: + os.dup2(fd, stdfd) os.close(fd) def _setup_stdio(self): @@ -3352,10 +3581,11 @@ class ExternalContext(object): # around a permanent dup() to avoid receiving SIGHUP. try: if os.isatty(2): - self.reserve_tty_fd = os.dup(2) - set_cloexec(self.reserve_tty_fd) + self.reserve_tty_fp = os.fdopen(os.dup(2), 'r+b', 0) + set_cloexec(self.reserve_tty_fp.fileno()) except OSError: pass + # When sys.stdout was opened by the runtime, overwriting it will not # close FD 1. However when forking from a child that previously used # fdopen(), overwriting it /will/ close FD 1. So we must swallow the @@ -3368,8 +3598,12 @@ class ExternalContext(object): sys.stdout.close() self._nullify_stdio() - self.stdout_log = IoLogger(self.broker, 'stdout', 1) - self.stderr_log = IoLogger(self.broker, 'stderr', 2) + self.loggers = [] + for name, fd in (('stdout', 1), ('stderr', 2)): + log = IoLoggerProtocol.build_stream(name, fd) + self.broker.start_receive(log) + self.loggers.append(log) + # Reopen with line buffering. sys.stdout = os.fdopen(1, 'w', 1) @@ -3389,18 +3623,21 @@ class ExternalContext(object): self.dispatcher = Dispatcher(self) self.router.register(self.parent, self.stream) self.router._setup_logging() - self.log_handler.uncork() sys.executable = os.environ.pop('ARGV0', sys.executable) - _v and LOG.debug('Connected to context %s; my ID is %r', - self.parent, mitogen.context_id) + _v and LOG.debug('Parent is context %r (%s); my ID is %r', + self.parent.context_id, self.parent.name, + mitogen.context_id) _v and LOG.debug('pid:%r ppid:%r uid:%r/%r, gid:%r/%r host:%r', os.getpid(), os.getppid(), os.geteuid(), os.getuid(), os.getegid(), os.getgid(), socket.gethostname()) _v and LOG.debug('Recovered sys.executable: %r', sys.executable) + if self.config.get('send_ec2', True): + self.stream.transmit_side.write(b('MITO002\n')) self.broker._py24_25_compat() + self.log_handler.uncork() self.dispatcher.run() _v and LOG.debug('ExternalContext.main() normal exit') except KeyboardInterrupt: diff --git a/mitogen/debug.py b/mitogen/debug.py index 3d13347f..dbab550e 100644 --- a/mitogen/debug.py +++ b/mitogen/debug.py @@ -230,7 +230,7 @@ class ContextDebugger(object): def _handle_debug_msg(self, msg): try: method, args, kwargs = msg.unpickle() - msg.reply(getattr(cls, method)(*args, **kwargs)) + msg.reply(getattr(self, method)(*args, **kwargs)) except Exception: e = sys.exc_info()[1] msg.reply(mitogen.core.CallError(e)) diff --git a/mitogen/doas.py b/mitogen/doas.py index 1b687fb2..f3bf4c90 100644 --- a/mitogen/doas.py +++ b/mitogen/doas.py @@ -29,6 +29,7 @@ # !mitogen: minify_safe import logging +import re import mitogen.core import mitogen.parent @@ -37,77 +38,106 @@ from mitogen.core import b LOG = logging.getLogger(__name__) +password_incorrect_msg = 'doas password is incorrect' +password_required_msg = 'doas password is required' + class PasswordError(mitogen.core.StreamError): pass -class Stream(mitogen.parent.Stream): - create_child = staticmethod(mitogen.parent.hybrid_tty_create_child) - child_is_immediate_subprocess = False - - username = 'root' +class Options(mitogen.parent.Options): + username = u'root' password = None doas_path = 'doas' - password_prompt = b('Password:') + password_prompt = u'Password:' incorrect_prompts = ( - b('doas: authentication failed'), + u'doas: authentication failed', # slicer69/doas + u'doas: Authorization failed', # openbsd/src ) - def construct(self, username=None, password=None, doas_path=None, - password_prompt=None, incorrect_prompts=None, **kwargs): - super(Stream, self).construct(**kwargs) + def __init__(self, username=None, password=None, doas_path=None, + password_prompt=None, incorrect_prompts=None, **kwargs): + super(Options, self).__init__(**kwargs) if username is not None: - self.username = username + self.username = mitogen.core.to_text(username) if password is not None: - self.password = password + self.password = mitogen.core.to_text(password) if doas_path is not None: self.doas_path = doas_path if password_prompt is not None: - self.password_prompt = password_prompt.lower() + self.password_prompt = mitogen.core.to_text(password_prompt) if incorrect_prompts is not None: - self.incorrect_prompts = map(str.lower, incorrect_prompts) + self.incorrect_prompts = [ + mitogen.core.to_text(p) + for p in incorrect_prompts + ] + + +class BootstrapProtocol(mitogen.parent.RegexProtocol): + password_sent = False + + def setup_patterns(self, conn): + prompt_pattern = re.compile( + re.escape(conn.options.password_prompt).encode('utf-8'), + re.I + ) + incorrect_prompt_pattern = re.compile( + u'|'.join( + re.escape(s) + for s in conn.options.incorrect_prompts + ).encode('utf-8'), + re.I + ) + + self.PATTERNS = [ + (incorrect_prompt_pattern, type(self)._on_incorrect_password), + ] + self.PARTIAL_PATTERNS = [ + (prompt_pattern, type(self)._on_password_prompt), + ] + + def _on_incorrect_password(self, line, match): + if self.password_sent: + self.stream.conn._fail_connection( + PasswordError(password_incorrect_msg) + ) + + def _on_password_prompt(self, line, match): + if self.stream.conn.options.password is None: + self.stream.conn._fail_connection( + PasswordError(password_required_msg) + ) + return + + if self.password_sent: + self.stream.conn._fail_connection( + PasswordError(password_incorrect_msg) + ) + return + + LOG.debug('sending password') + self.stream.transmit_side.write( + (self.stream.conn.options.password + '\n').encode('utf-8') + ) + self.password_sent = True + + +class Connection(mitogen.parent.Connection): + options_class = Options + diag_protocol_class = BootstrapProtocol + + create_child = staticmethod(mitogen.parent.hybrid_tty_create_child) + child_is_immediate_subprocess = False def _get_name(self): - return u'doas.' + mitogen.core.to_text(self.username) + return u'doas.' + self.options.username + + def stderr_stream_factory(self): + stream = super(Connection, self).stderr_stream_factory() + stream.protocol.setup_patterns(self) + return stream def get_boot_command(self): - bits = [self.doas_path, '-u', self.username, '--'] - bits = bits + super(Stream, self).get_boot_command() - LOG.debug('doas command line: %r', bits) - return bits - - password_incorrect_msg = 'doas password is incorrect' - password_required_msg = 'doas password is required' - - def _connect_input_loop(self, it): - password_sent = False - for buf in it: - LOG.debug('%r: received %r', self, buf) - if buf.endswith(self.EC0_MARKER): - self._ec0_received() - return - if any(s in buf.lower() for s in self.incorrect_prompts): - if password_sent: - raise PasswordError(self.password_incorrect_msg) - elif self.password_prompt in buf.lower(): - if self.password is None: - raise PasswordError(self.password_required_msg) - if password_sent: - raise PasswordError(self.password_incorrect_msg) - LOG.debug('sending password') - self.diag_stream.transmit_side.write( - mitogen.core.to_text(self.password + '\n').encode('utf-8') - ) - 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() + bits = [self.options.doas_path, '-u', self.options.username, '--'] + return bits + super(Connection, self).get_boot_command() diff --git a/mitogen/docker.py b/mitogen/docker.py index 0c0d40e7..48848c89 100644 --- a/mitogen/docker.py +++ b/mitogen/docker.py @@ -37,45 +37,47 @@ import mitogen.parent LOG = logging.getLogger(__name__) -class Stream(mitogen.parent.Stream): - child_is_immediate_subprocess = False - +class Options(mitogen.parent.Options): container = None image = None username = None - docker_path = 'docker' - - # TODO: better way of capturing errors such as "No such container." - create_child_args = { - 'merge_stdio': True - } + docker_path = u'docker' - def construct(self, container=None, image=None, - docker_path=None, username=None, - **kwargs): + def __init__(self, container=None, image=None, docker_path=None, + username=None, **kwargs): + super(Options, self).__init__(**kwargs) assert container or image - super(Stream, self).construct(**kwargs) if container: - self.container = container + self.container = mitogen.core.to_text(container) if image: - self.image = image + self.image = mitogen.core.to_text(image) if docker_path: - self.docker_path = docker_path + self.docker_path = mitogen.core.to_text(docker_path) if username: - self.username = username + self.username = mitogen.core.to_text(username) + + +class Connection(mitogen.parent.Connection): + options_class = Options + child_is_immediate_subprocess = False + + # TODO: better way of capturing errors such as "No such container." + create_child_args = { + 'merge_stdio': True + } def _get_name(self): - return u'docker.' + (self.container or self.image) + return u'docker.' + (self.options.container or self.options.image) def get_boot_command(self): args = ['--interactive'] - if self.username: - args += ['--user=' + self.username] + if self.options.username: + args += ['--user=' + self.options.username] - bits = [self.docker_path] - if self.container: - bits += ['exec'] + args + [self.container] - elif self.image: - bits += ['run'] + args + ['--rm', self.image] + bits = [self.options.docker_path] + if self.options.container: + bits += ['exec'] + args + [self.options.container] + elif self.options.image: + bits += ['run'] + args + ['--rm', self.options.image] - return bits + super(Stream, self).get_boot_command() + return bits + super(Connection, self).get_boot_command() diff --git a/mitogen/fakessh.py b/mitogen/fakessh.py index d39a710d..e62cf84a 100644 --- a/mitogen/fakessh.py +++ b/mitogen/fakessh.py @@ -117,14 +117,12 @@ SSH_GETOPTS = ( _mitogen = None -class IoPump(mitogen.core.BasicStream): +class IoPump(mitogen.core.Protocol): _output_buf = '' _closed = False - def __init__(self, broker, stdin_fd, stdout_fd): + def __init__(self, broker): self._broker = broker - self.receive_side = mitogen.core.Side(self, stdout_fd) - self.transmit_side = mitogen.core.Side(self, stdin_fd) def write(self, s): self._output_buf += s @@ -134,13 +132,13 @@ class IoPump(mitogen.core.BasicStream): self._closed = True # If local process hasn't exitted yet, ensure its write buffer is # drained before lazily triggering disconnect in on_transmit. - if self.transmit_side.fd is not None: + if self.transmit_side.fp.fileno() is not None: self._broker._start_transmit(self) - def on_shutdown(self, broker): + def on_shutdown(self, stream, broker): self.close() - def on_transmit(self, broker): + def on_transmit(self, stream, broker): written = self.transmit_side.write(self._output_buf) IOLOG.debug('%r.on_transmit() -> len %r', self, written) if written is None: @@ -153,8 +151,8 @@ class IoPump(mitogen.core.BasicStream): if self._closed: self.on_disconnect(broker) - def on_receive(self, broker): - s = self.receive_side.read() + def on_receive(self, stream, broker): + s = stream.receive_side.read() IOLOG.debug('%r.on_receive() -> len %r', self, len(s)) if s: mitogen.core.fire(self, 'receive', s) @@ -163,8 +161,8 @@ class IoPump(mitogen.core.BasicStream): def __repr__(self): return 'IoPump(%r, %r)' % ( - self.receive_side.fd, - self.transmit_side.fd, + self.receive_side.fp.fileno(), + self.transmit_side.fp.fileno(), ) @@ -173,14 +171,15 @@ class Process(object): Manages the lifetime and pipe connections of the SSH command running in the slave. """ - def __init__(self, router, stdin_fd, stdout_fd, proc=None): + def __init__(self, router, stdin, stdout, proc=None): self.router = router - self.stdin_fd = stdin_fd - self.stdout_fd = stdout_fd + self.stdin = stdin + self.stdout = stdout self.proc = proc self.control_handle = router.add_handler(self._on_control) self.stdin_handle = router.add_handler(self._on_stdin) - self.pump = IoPump(router.broker, stdin_fd, stdout_fd) + self.pump = IoPump.build_stream(router.broker) + self.pump.accept(stdin, stdout) self.stdin = None self.control = None self.wake_event = threading.Event() @@ -193,7 +192,7 @@ class Process(object): pmon.add(proc.pid, self._on_proc_exit) def __repr__(self): - return 'Process(%r, %r)' % (self.stdin_fd, self.stdout_fd) + return 'Process(%r, %r)' % (self.stdin, self.stdout) def _on_proc_exit(self, status): LOG.debug('%r._on_proc_exit(%r)', self, status) @@ -202,12 +201,12 @@ class Process(object): def _on_stdin(self, msg): if msg.is_dead: IOLOG.debug('%r._on_stdin() -> %r', self, data) - self.pump.close() + self.pump.protocol.close() return data = msg.unpickle() IOLOG.debug('%r._on_stdin() -> len %d', self, len(data)) - self.pump.write(data) + self.pump.protocol.write(data) def _on_control(self, msg): if not msg.is_dead: @@ -279,13 +278,7 @@ def _start_slave(src_id, cmdline, router): stdout=subprocess.PIPE, ) - process = Process( - router, - proc.stdin.fileno(), - proc.stdout.fileno(), - proc, - ) - + process = Process(router, proc.stdin, proc.stdout, proc) return process.control_handle, process.stdin_handle @@ -361,7 +354,9 @@ def _fakessh_main(dest_context_id, econtext): LOG.debug('_fakessh_main: received control_handle=%r, stdin_handle=%r', control_handle, stdin_handle) - process = Process(econtext.router, 1, 0) + process = Process(econtext.router, + stdin=os.fdopen(1, 'w+b', 0), + stdout=os.fdopen(0, 'r+b', 0)) process.start_master( stdin=mitogen.core.Sender(dest, stdin_handle), control=mitogen.core.Sender(dest, control_handle), @@ -427,7 +422,7 @@ def run(dest, router, args, deadline=None, econtext=None): stream = mitogen.core.Stream(router, context_id) stream.name = u'fakessh' - stream.accept(sock1.fileno(), sock1.fileno()) + stream.accept(sock1, sock1) router.register(fakessh, stream) # Held in socket buffer until process is booted. diff --git a/mitogen/fork.py b/mitogen/fork.py index d6685d70..e2075fc3 100644 --- a/mitogen/fork.py +++ b/mitogen/fork.py @@ -28,6 +28,7 @@ # !mitogen: minify_safe +import errno import logging import os import random @@ -37,6 +38,7 @@ import traceback import mitogen.core import mitogen.parent +from mitogen.core import b LOG = logging.getLogger('mitogen') @@ -119,32 +121,45 @@ def handle_child_crash(): os._exit(1) -class Stream(mitogen.parent.Stream): - child_is_immediate_subprocess = True - +class Process(mitogen.parent.Process): + def poll(self): + try: + pid, status = os.waitpid(self.pid, os.WNOHANG) + except OSError: + e = sys.exc_info()[1] + if e.args[0] == errno.ECHILD: + LOG.warn('%r: waitpid(%r) produced ECHILD', self, self.pid) + return + raise + + if not pid: + return + if os.WIFEXITED(status): + return os.WEXITSTATUS(status) + elif os.WIFSIGNALED(status): + return -os.WTERMSIG(status) + elif os.WIFSTOPPED(status): + return -os.WSTOPSIG(status) + + +class Options(mitogen.parent.Options): #: Reference to the importer, if any, recovered from the parent. importer = None #: User-supplied function for cleaning up child process state. on_fork = None - python_version_msg = ( - "The mitogen.fork method is not supported on Python versions " - "prior to 2.6, since those versions made no attempt to repair " - "critical interpreter state following a fork. Please use the " - "local() method instead." - ) - - def construct(self, old_router, max_message_size, on_fork=None, - debug=False, profiling=False, unidirectional=False, - on_start=None): + def __init__(self, old_router, max_message_size, on_fork=None, debug=False, + profiling=False, unidirectional=False, on_start=None, + name=None): if not FORK_SUPPORTED: raise Error(self.python_version_msg) # fork method only supports a tiny subset of options. - super(Stream, self).construct(max_message_size=max_message_size, - debug=debug, profiling=profiling, - unidirectional=False) + super(Options, self).__init__( + max_message_size=max_message_size, debug=debug, + profiling=profiling, unidirectional=unidirectional, name=name, + ) self.on_fork = on_fork self.on_start = on_start @@ -152,17 +167,26 @@ class Stream(mitogen.parent.Stream): if isinstance(responder, mitogen.parent.ModuleForwarder): self.importer = responder.importer + +class Connection(mitogen.parent.Connection): + options_class = Options + child_is_immediate_subprocess = True + + python_version_msg = ( + "The mitogen.fork method is not supported on Python versions " + "prior to 2.6, since those versions made no attempt to repair " + "critical interpreter state following a fork. Please use the " + "local() method instead." + ) + name_prefix = u'fork' def start_child(self): parentfp, childfp = mitogen.parent.create_socketpair() - self.pid = os.fork() - if self.pid: + pid = os.fork() + if pid: childfp.close() - # Decouple the socket from the lifetime of the Python socket object. - fd = os.dup(parentfp.fileno()) - parentfp.close() - return self.pid, fd, None + return Process(pid, stdin=parentfp, stdout=parentfp) else: parentfp.close() self._wrap_child_main(childfp) @@ -173,12 +197,24 @@ class Stream(mitogen.parent.Stream): except BaseException: handle_child_crash() + def get_econtext_config(self): + config = super(Connection, self).get_econtext_config() + config['core_src_fd'] = None + config['importer'] = self.options.importer + config['send_ec2'] = False + config['setup_package'] = False + if self.options.on_start: + config['on_start'] = self.options.on_start + return config + def _child_main(self, childfp): on_fork() - if self.on_fork: - self.on_fork() + if self.options.on_fork: + self.options.on_fork() mitogen.core.set_block(childfp.fileno()) + childfp.send(b('MITO002\n')) + # Expected by the ExternalContext.main(). os.dup2(childfp.fileno(), 1) os.dup2(childfp.fileno(), 100) @@ -201,23 +237,12 @@ class Stream(mitogen.parent.Stream): if childfp.fileno() not in (0, 1, 100): childfp.close() - config = self.get_econtext_config() - config['core_src_fd'] = None - config['importer'] = self.importer - config['setup_package'] = False - if self.on_start: - config['on_start'] = self.on_start - try: try: - mitogen.core.ExternalContext(config).main() + mitogen.core.ExternalContext(self.get_econtext_config()).main() except Exception: # TODO: report exception somehow. os._exit(72) finally: # Don't trigger atexit handlers, they were copied from the parent. os._exit(0) - - def _connect_bootstrap(self): - # None required. - pass diff --git a/mitogen/jail.py b/mitogen/jail.py index 6e0ac68b..c7c1f0f9 100644 --- a/mitogen/jail.py +++ b/mitogen/jail.py @@ -37,29 +37,34 @@ import mitogen.parent LOG = logging.getLogger(__name__) -class Stream(mitogen.parent.Stream): - child_is_immediate_subprocess = False - create_child_args = { - 'merge_stdio': True - } - +class Options(mitogen.parent.Options): container = None username = None - jexec_path = '/usr/sbin/jexec' + jexec_path = u'/usr/sbin/jexec' - def construct(self, container, jexec_path=None, username=None, **kwargs): - super(Stream, self).construct(**kwargs) - self.container = container - self.username = username + def __init__(self, container, jexec_path=None, username=None, **kwargs): + super(Options, self).__init__(**kwargs) + self.container = mitogen.core.to_text(container) + if username: + self.username = mitogen.core.to_text(username) if jexec_path: self.jexec_path = jexec_path + +class Connection(mitogen.parent.Connection): + options_class = Options + + child_is_immediate_subprocess = False + create_child_args = { + 'merge_stdio': True + } + def _get_name(self): - return u'jail.' + self.container + return u'jail.' + self.options.container def get_boot_command(self): - bits = [self.jexec_path] - if self.username: - bits += ['-U', self.username] - bits += [self.container] - return bits + super(Stream, self).get_boot_command() + bits = [self.options.jexec_path] + if self.options.username: + bits += ['-U', self.options.username] + bits += [self.options.container] + return bits + super(Connection, self).get_boot_command() diff --git a/mitogen/kubectl.py b/mitogen/kubectl.py index ef626e1b..acc011b9 100644 --- a/mitogen/kubectl.py +++ b/mitogen/kubectl.py @@ -37,29 +37,36 @@ import mitogen.parent LOG = logging.getLogger(__name__) -class Stream(mitogen.parent.Stream): - child_is_immediate_subprocess = True - +class Options(mitogen.parent.Options): pod = None kubectl_path = 'kubectl' kubectl_args = None - # TODO: better way of capturing errors such as "No such container." - create_child_args = { - 'merge_stdio': True - } - - def construct(self, pod, kubectl_path=None, kubectl_args=None, **kwargs): - super(Stream, self).construct(**kwargs) + def __init__(self, pod, kubectl_path=None, kubectl_args=None, **kwargs): + super(Options, self).__init__(**kwargs) assert pod self.pod = pod if kubectl_path: self.kubectl_path = kubectl_path self.kubectl_args = kubectl_args or [] + +class Connection(mitogen.parent.Connection): + options_class = Options + child_is_immediate_subprocess = True + + # TODO: better way of capturing errors such as "No such container." + create_child_args = { + 'merge_stdio': True + } + def _get_name(self): - return u'kubectl.%s%s' % (self.pod, self.kubectl_args) + return u'kubectl.%s%s' % (self.options.pod, self.options.kubectl_args) def get_boot_command(self): - bits = [self.kubectl_path] + self.kubectl_args + ['exec', '-it', self.pod] - return bits + ["--"] + super(Stream, self).get_boot_command() + bits = [ + self.options.kubectl_path + ] + self.options.kubectl_args + [ + 'exec', '-it', self.options.pod + ] + return bits + ["--"] + super(Connection, self).get_boot_command() diff --git a/mitogen/lxc.py b/mitogen/lxc.py index 879d19a1..759475c1 100644 --- a/mitogen/lxc.py +++ b/mitogen/lxc.py @@ -37,7 +37,20 @@ import mitogen.parent LOG = logging.getLogger(__name__) -class Stream(mitogen.parent.Stream): +class Options(mitogen.parent.Options): + container = None + lxc_attach_path = 'lxc-attach' + + def __init__(self, container, lxc_attach_path=None, **kwargs): + super(Options, self).__init__(**kwargs) + self.container = container + if lxc_attach_path: + self.lxc_attach_path = lxc_attach_path + + +class Connection(mitogen.parent.Connection): + options_class = Options + child_is_immediate_subprocess = False create_child_args = { # If lxc-attach finds any of stdin, stdout, stderr connected to a TTY, @@ -47,29 +60,20 @@ class Stream(mitogen.parent.Stream): 'merge_stdio': True } - container = None - lxc_attach_path = 'lxc-attach' - eof_error_hint = ( 'Note: many versions of LXC do not report program execution failure ' 'meaningfully. Please check the host logs (/var/log) for more ' 'information.' ) - def construct(self, container, lxc_attach_path=None, **kwargs): - super(Stream, self).construct(**kwargs) - self.container = container - if lxc_attach_path: - self.lxc_attach_path = lxc_attach_path - def _get_name(self): - return u'lxc.' + self.container + return u'lxc.' + self.options.container def get_boot_command(self): bits = [ - self.lxc_attach_path, + self.options.lxc_attach_path, '--clear-env', - '--name', self.container, + '--name', self.options.container, '--', ] - return bits + super(Stream, self).get_boot_command() + return bits + super(Connection, self).get_boot_command() diff --git a/mitogen/lxd.py b/mitogen/lxd.py index faea2561..6fbe0694 100644 --- a/mitogen/lxd.py +++ b/mitogen/lxd.py @@ -37,7 +37,21 @@ import mitogen.parent LOG = logging.getLogger(__name__) -class Stream(mitogen.parent.Stream): +class Options(mitogen.parent.Options): + container = None + lxc_path = 'lxc' + python_path = 'python' + + def __init__(self, container, lxc_path=None, **kwargs): + super(Options, self).__init__(**kwargs) + self.container = container + if lxc_path: + self.lxc_path = lxc_path + + +class Connection(mitogen.parent.Connection): + options_class = Options + child_is_immediate_subprocess = False create_child_args = { # If lxc finds any of stdin, stdout, stderr connected to a TTY, to @@ -47,31 +61,21 @@ class Stream(mitogen.parent.Stream): 'merge_stdio': True } - container = None - lxc_path = 'lxc' - python_path = 'python' - eof_error_hint = ( 'Note: many versions of LXC do not report program execution failure ' 'meaningfully. Please check the host logs (/var/log) for more ' 'information.' ) - def construct(self, container, lxc_path=None, **kwargs): - super(Stream, self).construct(**kwargs) - self.container = container - if lxc_path: - self.lxc_path = lxc_path - def _get_name(self): - return u'lxd.' + self.container + return u'lxd.' + self.options.container def get_boot_command(self): bits = [ - self.lxc_path, + self.options.lxc_path, 'exec', '--mode=noninteractive', - self.container, + self.options.container, '--', ] - return bits + super(Stream, self).get_boot_command() + return bits + super(Connection, self).get_boot_command() diff --git a/mitogen/master.py b/mitogen/master.py index fb4f505b..909c3cef 100644 --- a/mitogen/master.py +++ b/mitogen/master.py @@ -91,7 +91,8 @@ RLOG = logging.getLogger('mitogen.ctx') def _stdlib_paths(): - """Return a set of paths from which Python imports the standard library. + """ + Return a set of paths from which Python imports the standard library. """ attr_candidates = [ 'prefix', @@ -111,8 +112,8 @@ def _stdlib_paths(): def is_stdlib_name(modname): - """Return :data:`True` if `modname` appears to come from the standard - library. + """ + Return :data:`True` if `modname` appears to come from the standard library. """ if imp.is_builtin(modname) != 0: return True @@ -139,7 +140,8 @@ def is_stdlib_path(path): def get_child_modules(path): - """Return the suffixes of submodules directly neated beneath of the package + """ + Return the suffixes of submodules directly neated beneath of the package directory at `path`. :param str path: @@ -301,8 +303,10 @@ class ThreadWatcher(object): @classmethod def _reset(cls): - """If we have forked since the watch dictionaries were initialized, all - that has is garbage, so clear it.""" + """ + If we have forked since the watch dictionaries were initialized, all + that has is garbage, so clear it. + """ if os.getpid() != cls._cls_pid: cls._cls_pid = os.getpid() cls._cls_instances_by_target.clear() @@ -394,7 +398,7 @@ class LogForwarder(object): name = '%s.%s' % (RLOG.name, context.name) self._cache[msg.src_id] = logger = logging.getLogger(name) - name, level_s, s = msg.data.decode('latin1').split('\x00', 2) + name, level_s, s = msg.data.decode('utf-8', 'replace').split('\x00', 2) # See logging.Handler.makeRecord() record = logging.LogRecord( @@ -531,14 +535,15 @@ class SysModulesMethod(FinderMethod): return if not isinstance(module, types.ModuleType): - LOG.debug('sys.modules[%r] absent or not a regular module', - fullname) + LOG.debug('%r: sys.modules[%r] absent or not a regular module', + self, fullname) return path = _py_filename(getattr(module, '__file__', '')) if not path: return + LOG.debug('%r: sys.modules[%r]: found %s', self, fullname, path) is_pkg = hasattr(module, '__path__') try: source = inspect.getsource(module) @@ -667,7 +672,8 @@ class ModuleFinder(object): ] def get_module_source(self, fullname): - """Given the name of a loaded module `fullname`, attempt to find its + """ + Given the name of a loaded module `fullname`, attempt to find its source code. :returns: @@ -691,9 +697,10 @@ class ModuleFinder(object): return tup def resolve_relpath(self, fullname, level): - """Given an ImportFrom AST node, guess the prefix that should be tacked - on to an alias name to produce a canonical name. `fullname` is the name - of the module in which the ImportFrom appears. + """ + Given an ImportFrom AST node, guess the prefix that should be tacked on + to an alias name to produce a canonical name. `fullname` is the name of + the module in which the ImportFrom appears. """ mod = sys.modules.get(fullname, None) if hasattr(mod, '__path__'): @@ -817,7 +824,7 @@ class ModuleResponder(object): ) def __repr__(self): - return 'ModuleResponder(%r)' % (self._router,) + return 'ModuleResponder' def add_source_override(self, fullname, path, source, is_pkg): """ @@ -844,9 +851,11 @@ class ModuleResponder(object): self.blacklist.append(fullname) def neutralize_main(self, path, src): - """Given the source for the __main__ module, try to find where it - begins conditional execution based on a "if __name__ == '__main__'" - guard, and remove any code after that point.""" + """ + Given the source for the __main__ module, try to find where it begins + conditional execution based on a "if __name__ == '__main__'" guard, and + remove any code after that point. + """ match = self.MAIN_RE.search(src) if match: return src[:match.start()] @@ -920,17 +929,17 @@ class ModuleResponder(object): return tup def _send_load_module(self, stream, fullname): - if fullname not in stream.sent_modules: + if fullname not in stream.protocol.sent_modules: tup = self._build_tuple(fullname) msg = mitogen.core.Message.pickled( tup, - dst_id=stream.remote_id, + dst_id=stream.protocol.remote_id, handle=mitogen.core.LOAD_MODULE, ) - LOG.debug('%s: sending module %s (%.2f KiB)', - stream.name, fullname, len(msg.data) / 1024.0) + LOG.debug('%s: sending %s (%.2f KiB) to %s', + self, fullname, len(msg.data) / 1024.0, stream.name) self._router._async_route(msg) - stream.sent_modules.add(fullname) + stream.protocol.sent_modules.add(fullname) if tup[2] is not None: self.good_load_module_count += 1 self.good_load_module_size += len(msg.data) @@ -939,23 +948,23 @@ class ModuleResponder(object): def _send_module_load_failed(self, stream, fullname): self.bad_load_module_count += 1 - stream.send( + stream.protocol.send( mitogen.core.Message.pickled( self._make_negative_response(fullname), - dst_id=stream.remote_id, + dst_id=stream.protocol.remote_id, handle=mitogen.core.LOAD_MODULE, ) ) def _send_module_and_related(self, stream, fullname): - if fullname in stream.sent_modules: + if fullname in stream.protocol.sent_modules: return try: tup = self._build_tuple(fullname) for name in tup[4]: # related parent, _, _ = str_partition(name, '.') - if parent != fullname and parent not in stream.sent_modules: + if parent != fullname and parent not in stream.protocol.sent_modules: # Parent hasn't been sent, so don't load submodule yet. continue @@ -976,7 +985,7 @@ class ModuleResponder(object): fullname = msg.data.decode() LOG.debug('%s requested module %s', stream.name, fullname) self.get_module_count += 1 - if fullname in stream.sent_modules: + if fullname in stream.protocol.sent_modules: LOG.warning('_on_get_module(): dup request for %r from %r', fullname, stream) @@ -987,12 +996,12 @@ class ModuleResponder(object): self.get_module_secs += time.time() - t0 def _send_forward_module(self, stream, context, fullname): - if stream.remote_id != context.context_id: - stream.send( + if stream.protocol.remote_id != context.context_id: + stream.protocol._send( mitogen.core.Message( data=b('%s\x00%s' % (context.context_id, fullname)), handle=mitogen.core.FORWARD_MODULE, - dst_id=stream.remote_id, + dst_id=stream.protocol.remote_id, ) ) @@ -1061,6 +1070,7 @@ class Broker(mitogen.core.Broker): on_join=self.shutdown, ) super(Broker, self).__init__() + self.timers = mitogen.parent.TimerList() def shutdown(self): super(Broker, self).shutdown() diff --git a/mitogen/minify.py b/mitogen/minify.py index dc9f517c..09fdc4eb 100644 --- a/mitogen/minify.py +++ b/mitogen/minify.py @@ -44,7 +44,8 @@ else: def minimize_source(source): - """Remove comments and docstrings from Python `source`, preserving line + """ + Remove comments and docstrings from Python `source`, preserving line numbers and syntax of empty blocks. :param str source: @@ -62,7 +63,8 @@ def minimize_source(source): def strip_comments(tokens): - """Drop comment tokens from a `tokenize` stream. + """ + Drop comment tokens from a `tokenize` stream. Comments on lines 1-2 are kept, to preserve hashbang and encoding. Trailing whitespace is remove from all lines. @@ -84,7 +86,8 @@ def strip_comments(tokens): def strip_docstrings(tokens): - """Replace docstring tokens with NL tokens in a `tokenize` stream. + """ + Replace docstring tokens with NL tokens in a `tokenize` stream. Any STRING token not part of an expression is deemed a docstring. Indented docstrings are not yet recognised. @@ -119,7 +122,8 @@ def strip_docstrings(tokens): def reindent(tokens, indent=' '): - """Replace existing indentation in a token steam, with `indent`. + """ + Replace existing indentation in a token steam, with `indent`. """ old_levels = [] old_level = 0 diff --git a/mitogen/parent.py b/mitogen/parent.py index 113fdc2e..6b3da70b 100644 --- a/mitogen/parent.py +++ b/mitogen/parent.py @@ -38,9 +38,11 @@ import codecs import errno import fcntl import getpass +import heapq import inspect import logging import os +import re import signal import socket import struct @@ -136,6 +138,9 @@ SIGNAL_BY_NUM = dict( if name.startswith('SIG') and not name.startswith('SIG_') ) +_core_source_lock = threading.Lock() +_core_source_partial = None + def get_log_level(): return (LOG.level or logging.getLogger().level or logging.INFO) @@ -157,10 +162,6 @@ def get_sys_executable(): return '/usr/bin/python' -_core_source_lock = threading.Lock() -_core_source_partial = None - - def _get_core_source(): """ In non-masters, simply fetch the cached mitogen.core source code via the @@ -208,27 +209,33 @@ def is_immediate_child(msg, stream): Handler policy that requires messages to arrive only from immediately connected children. """ - return msg.src_id == stream.remote_id + return msg.src_id == stream.protocol.remote_id def flags(names): - """Return the result of ORing a set of (space separated) :py:mod:`termios` - module constants together.""" + """ + Return the result of ORing a set of (space separated) :py:mod:`termios` + module constants together. + """ return sum(getattr(termios, name, 0) for name in names.split()) def cfmakeraw(tflags): - """Given a list returned by :py:func:`termios.tcgetattr`, return a list + """ + Given a list returned by :py:func:`termios.tcgetattr`, return a list modified in a manner similar to the `cfmakeraw()` C library function, but - additionally disabling local echo.""" - # BSD: https://github.com/freebsd/freebsd/blob/master/lib/libc/gen/termios.c#L162 - # Linux: https://github.com/lattera/glibc/blob/master/termios/cfmakeraw.c#L20 + additionally disabling local echo. + """ + # BSD: github.com/freebsd/freebsd/blob/master/lib/libc/gen/termios.c#L162 + # Linux: github.com/lattera/glibc/blob/master/termios/cfmakeraw.c#L20 iflag, oflag, cflag, lflag, ispeed, ospeed, cc = tflags - iflag &= ~flags('IMAXBEL IXOFF INPCK BRKINT PARMRK ISTRIP INLCR ICRNL IXON IGNPAR') + iflag &= ~flags('IMAXBEL IXOFF INPCK BRKINT PARMRK ' + 'ISTRIP INLCR ICRNL IXON IGNPAR') iflag &= ~flags('IGNBRK BRKINT PARMRK') oflag &= ~flags('OPOST') - lflag &= ~flags('ECHO ECHOE ECHOK ECHONL ICANON ISIG IEXTEN NOFLSH TOSTOP PENDIN') + lflag &= ~flags('ECHO ECHOE ECHOK ECHONL ICANON ISIG ' + 'IEXTEN NOFLSH TOSTOP PENDIN') cflag &= ~flags('CSIZE PARENB') cflag |= flags('CS8 CREAD') return [iflag, oflag, cflag, lflag, ispeed, ospeed, cc] @@ -245,22 +252,13 @@ def disable_echo(fd): termios.tcsetattr(fd, flags, new) -def close_nonstandard_fds(): - for fd in xrange(3, SC_OPEN_MAX): - try: - os.close(fd) - except OSError: - pass - - def create_socketpair(size=None): """ - Create a :func:`socket.socketpair` to use for use as a child process's UNIX - stdio channels. As socket pairs are bidirectional, they are economical on - file descriptor usage as the same descriptor can be used for ``stdin`` and + Create a :func:`socket.socketpair` for use as a child's UNIX stdio + channels. As socketpairs are bidirectional, they are economical on file + descriptor usage as one descriptor can be used for ``stdin`` and ``stdout``. As they are sockets their buffers are tunable, allowing large - buffers to be configured in order to improve throughput for file transfers - and reduce :class:`mitogen.core.Broker` IO loop iterations. + buffers to improve file transfer throughput and reduce IO loop iterations. """ parentfp, childfp = socket.socketpair() parentfp.setsockopt(socket.SOL_SOCKET, @@ -272,44 +270,26 @@ def create_socketpair(size=None): return parentfp, childfp -def detach_popen(**kwargs): +def popen(**kwargs): """ - Use :class:`subprocess.Popen` to construct a child process, then hack the - Popen so that it forgets the child it created, allowing it to survive a - call to Popen.__del__. - - If the child process is not detached, there is a race between it exitting - and __del__ being called. If it exits before __del__ runs, then __del__'s - call to :func:`os.waitpid` will capture the one and only exit event - delivered to this process, causing later 'legitimate' calls to fail with - ECHILD. - - :param list close_on_error: - Array of integer file descriptors to close on exception. - :returns: - Process ID of the new child. + Wrap :class:`subprocess.Popen` to ensure any global :data:`_preexec_hook` + is invoked in the child. """ - # This allows Popen() to be used for e.g. graceful post-fork error - # handling, without tying the surrounding code into managing a Popen - # object, which isn't possible for at least :mod:`mitogen.fork`. This - # should be replaced by a swappable helper class in a future version. real_preexec_fn = kwargs.pop('preexec_fn', None) def preexec_fn(): if _preexec_hook: _preexec_hook() if real_preexec_fn: real_preexec_fn() - proc = subprocess.Popen(preexec_fn=preexec_fn, **kwargs) - proc._child_created = False - return proc.pid + return subprocess.Popen(preexec_fn=preexec_fn, **kwargs) def create_child(args, merge_stdio=False, stderr_pipe=False, preexec_fn=None): """ Create a child process whose stdin/stdout is connected to a socket. - :param args: - Argument vector for execv() call. + :param list args: + Program argument vector. :param bool merge_stdio: If :data:`True`, arrange for `stderr` to be connected to the `stdout` socketpair, rather than inherited from the parent process. This may be @@ -321,52 +301,47 @@ def create_child(args, merge_stdio=False, stderr_pipe=False, preexec_fn=None): logs generated by e.g. SSH to be outpu as the session progresses, without interfering with `stdout`. :returns: - `(pid, socket_obj, :data:`None` or pipe_fd)` + :class:`Process` instance. """ parentfp, childfp = create_socketpair() # When running under a monkey patches-enabled gevent, the socket module - # yields file descriptors who already have O_NONBLOCK, which is - # persisted across fork, totally breaking Python. Therefore, drop - # O_NONBLOCK from Python's future stdin fd. + # yields descriptors who already have O_NONBLOCK, which is persisted across + # fork, totally breaking Python. Therefore, drop O_NONBLOCK from Python's + # future stdin fd. mitogen.core.set_block(childfp.fileno()) + stderr = None stderr_r = None - extra = {} if merge_stdio: - extra = {'stderr': childfp} + stderr = childfp elif stderr_pipe: - stderr_r, stderr_w = os.pipe() - mitogen.core.set_cloexec(stderr_r) - mitogen.core.set_cloexec(stderr_w) - extra = {'stderr': stderr_w} + stderr_r, stderr = mitogen.core.pipe() + mitogen.core.set_cloexec(stderr_r.fileno()) try: - pid = detach_popen( + proc = popen( args=args, stdin=childfp, stdout=childfp, + stderr=stderr, close_fds=True, preexec_fn=preexec_fn, - **extra ) - except Exception: + except: childfp.close() parentfp.close() if stderr_pipe: - os.close(stderr_r) - os.close(stderr_w) + stderr.close() + stderr_r.close() raise - if stderr_pipe: - os.close(stderr_w) childfp.close() - # Decouple the socket from the lifetime of the Python socket object. - fd = os.dup(parentfp.fileno()) - parentfp.close() + if stderr_pipe: + stderr.close() LOG.debug('create_child() child %d fd %d, parent %d, cmd: %s', - pid, fd, os.getpid(), Argv(args)) - return pid, fd, stderr_r + proc.pid, parentfp.fileno(), os.getpid(), Argv(args)) + return PopenProcess(proc, stdin=parentfp, stdout=parentfp, stderr=stderr_r) def _acquire_controlling_tty(): @@ -431,15 +406,22 @@ def openpty(): :raises mitogen.core.StreamError: Creating a PTY failed. :returns: - See :func`os.openpty`. + `(master_fp, slave_fp)` file-like objects. """ try: - return os.openpty() + master_fd, slave_fd = os.openpty() except OSError: e = sys.exc_info()[1] - if IS_LINUX and e.args[0] == errno.EPERM: - return _linux_broken_devpts_openpty() - raise mitogen.core.StreamError(OPENPTY_MSG, e) + if not (IS_LINUX and e.args[0] == errno.EPERM): + raise mitogen.core.StreamError(OPENPTY_MSG, e) + master_fd, slave_fd = _linux_broken_devpts_openpty() + + master_fp = os.fdopen(master_fd, 'r+b', 0) + slave_fp = os.fdopen(slave_fd, 'r+b', 0) + disable_echo(master_fd) + disable_echo(slave_fd) + mitogen.core.set_block(slave_fd) + return master_fp, slave_fp def tty_create_child(args): @@ -451,34 +433,29 @@ def tty_create_child(args): slave end. :param list args: - :py:func:`os.execl` argument list. - + Program argument vector. :returns: - `(pid, tty_fd, None)` + :class:`Process` instance. """ - master_fd, slave_fd = openpty() + master_fp, slave_fp = openpty() try: - mitogen.core.set_block(slave_fd) - disable_echo(master_fd) - disable_echo(slave_fd) - - pid = detach_popen( + proc = popen( args=args, - stdin=slave_fd, - stdout=slave_fd, - stderr=slave_fd, + stdin=slave_fp, + stdout=slave_fp, + stderr=slave_fp, preexec_fn=_acquire_controlling_tty, close_fds=True, ) - except Exception: - os.close(master_fd) - os.close(slave_fd) + except: + master_fp.close() + slave_fp.close() raise - os.close(slave_fd) + slave_fp.close() LOG.debug('tty_create_child() child %d fd %d, parent %d, cmd: %s', - pid, master_fd, os.getpid(), Argv(args)) - return pid, master_fd, None + proc.pid, master_fp.fileno(), os.getpid(), Argv(args)) + return PopenProcess(proc, stdin=master_fp, stdout=master_fp) def hybrid_tty_create_child(args): @@ -488,93 +465,132 @@ def hybrid_tty_create_child(args): attached to a TTY. :param list args: - :py:func:`os.execl` argument list. - + Program argument vector. :returns: - `(pid, socketpair_fd, tty_fd)` + :class:`Process` instance. """ - master_fd, slave_fd = openpty() - + master_fp, slave_fp = openpty() try: - disable_echo(master_fd) - disable_echo(slave_fd) - mitogen.core.set_block(slave_fd) - parentfp, childfp = create_socketpair() try: mitogen.core.set_block(childfp) - pid = detach_popen( + proc = popen( args=args, stdin=childfp, stdout=childfp, - stderr=slave_fd, + stderr=slave_fp, preexec_fn=_acquire_controlling_tty, close_fds=True, ) - except Exception: + except: parentfp.close() childfp.close() raise - except Exception: - os.close(master_fd) - os.close(slave_fd) + except: + master_fp.close() + slave_fp.close() raise - os.close(slave_fd) + slave_fp.close() childfp.close() - # Decouple the socket from the lifetime of the Python socket object. - stdio_fd = os.dup(parentfp.fileno()) - parentfp.close() - LOG.debug('hybrid_tty_create_child() pid=%d stdio=%d, tty=%d, cmd: %s', - pid, stdio_fd, master_fd, Argv(args)) - return pid, stdio_fd, master_fd - - -def write_all(fd, s, deadline=None): - """Arrange for all of bytestring `s` to be written to the file descriptor - `fd`. - - :param int fd: - File descriptor to write to. - :param bytes s: - Bytestring to write to file descriptor. - :param float deadline: - If not :data:`None`, absolute UNIX timestamp after which timeout should - occur. - - :raises mitogen.core.TimeoutError: - Bytestring could not be written entirely before deadline was exceeded. - :raises mitogen.parent.EofError: - Stream indicated EOF, suggesting the child process has exitted. - :raises mitogen.core.StreamError: - File descriptor was disconnected before write could complete. + proc.pid, parentfp.fileno(), master_fp.fileno(), Argv(args)) + return PopenProcess(proc, stdin=parentfp, stdout=parentfp, stderr=master_fp) + + +class Timer(object): """ - timeout = None - written = 0 - poller = PREFERRED_POLLER() - poller.start_transmit(fd) + Represents a future event. + """ + cancelled = False - try: - while written < len(s): - if deadline is not None: - timeout = max(0, deadline - time.time()) - if timeout == 0: - raise mitogen.core.TimeoutError('write timed out') - - if mitogen.core.PY3: - window = memoryview(s)[written:] - else: - window = buffer(s, written) + def __init__(self, when, func): + self.when = when + self.func = func - for fd in poller.poll(timeout): - n, disconnected = mitogen.core.io_op(os.write, fd, window) - if disconnected: - raise EofError('EOF on stream during write') + def __repr__(self): + return 'Timer(%r, %r)' % (self.when, self.func) - written += n - finally: - poller.close() + def __eq__(self, other): + return self.when == other.when + + def __lt__(self, other): + return self.when < other.when + + def __le__(self, other): + return self.when <= other.when + + def cancel(self): + """ + Cancel this event. If it has not yet executed, it will not execute + during any subsequent :meth:`TimerList.expire` call. + """ + self.cancelled = True + + +class TimerList(object): + """ + Efficiently manage a list of cancellable future events relative to wall + clock time. An instance of this class is installed as + :attr:`mitogen.master.Broker.timers` by default, and as + :attr:`mitogen.core.Broker.timers` in children after a call to + :func:`mitogen.parent.upgrade_router`. + + You can use :class:`TimerList` to cause the broker to wake at arbitrary + future moments, useful for implementing timeouts and polling in an + asynchronous context. + + :class:`TimerList` methods can only be called from asynchronous context, + for example via :meth:`mitogen.core.Broker.defer`. + + The broker automatically adjusts its sleep delay according to the installed + timer list, and arranges for timers to expire via automatic calls to + :meth:`expire`. The main user interface to :class:`TimerList` is + :meth:`schedule`. + """ + _now = time.time + + def __init__(self): + self._lst = [] + + def get_timeout(self): + """ + Return the floating point seconds until the next event is due. + + :returns: + Floating point delay, or 0.0, or :data:`None` if no events are + scheduled. + """ + while self._lst and self._lst[0].cancelled: + heapq.heappop(self._lst) + if self._lst: + return max(0, self._lst[0].when - self._now()) + + def schedule(self, when, func): + """ + Schedule a future event. + + :param float when: + UNIX time in seconds when event should occur. + :param callable func: + Callable to invoke on expiry. + :returns: + A :class:`Timer` instance, exposing :meth:`Timer.cancel`, which may + be used to cancel the future invocation. + """ + timer = Timer(when, func) + heapq.heappush(self._lst, timer) + return timer + + def expire(self): + """ + Invoke callbacks for any events in the past. + """ + now = self._now() + while self._lst and self._lst[0].when <= now: + timer = heapq.heappop(self._lst) + if not timer.cancelled: + timer.func() class PartialZlib(object): @@ -614,103 +630,6 @@ class PartialZlib(object): return out + compressor.flush() -class IteratingRead(object): - def __init__(self, fds, deadline=None): - self.deadline = deadline - self.timeout = None - self.poller = PREFERRED_POLLER() - for fd in fds: - self.poller.start_receive(fd) - - self.bits = [] - self.timeout = None - - def close(self): - self.poller.close() - - def __iter__(self): - return self - - def next(self): - while self.poller.readers: - if self.deadline is not None: - self.timeout = max(0, self.deadline - time.time()) - if self.timeout == 0: - break - - for fd in self.poller.poll(self.timeout): - s, disconnected = mitogen.core.io_op(os.read, fd, 4096) - if disconnected or not s: - LOG.debug('iter_read(%r) -> disconnected: %s', - fd, disconnected) - self.poller.stop_receive(fd) - else: - IOLOG.debug('iter_read(%r) -> %r', fd, s) - self.bits.append(s) - return s - - if not self.poller.readers: - raise EofError(u'EOF on stream; last 300 bytes received: %r' % - (b('').join(self.bits)[-300:].decode('latin1'),)) - - raise mitogen.core.TimeoutError('read timed out') - - __next__ = next - - -def iter_read(fds, deadline=None): - """Return a generator that arranges for up to 4096-byte chunks to be read - at a time from the file descriptor `fd` until the generator is destroyed. - - :param int fd: - File descriptor to read from. - :param float deadline: - If not :data:`None`, an absolute UNIX timestamp after which timeout - should occur. - - :raises mitogen.core.TimeoutError: - Attempt to read beyond deadline. - :raises mitogen.parent.EofError: - All streams indicated EOF, suggesting the child process has exitted. - :raises mitogen.core.StreamError: - Attempt to read past end of file. - """ - return IteratingRead(fds=fds, deadline=deadline) - - -def discard_until(fd, s, deadline): - """Read chunks from `fd` until one is encountered that ends with `s`. This - is used to skip output produced by ``/etc/profile``, ``/etc/motd`` and - mandatory SSH banners while waiting for :attr:`Stream.EC0_MARKER` to - appear, indicating the first stage is ready to receive the compressed - :mod:`mitogen.core` source. - - :param int fd: - File descriptor to read from. - :param bytes s: - Marker string to discard until encountered. - :param float deadline: - Absolute UNIX timestamp after which timeout should occur. - - :raises mitogen.core.TimeoutError: - Attempt to read beyond deadline. - :raises mitogen.parent.EofError: - All streams indicated EOF, suggesting the child process has exitted. - :raises mitogen.core.StreamError: - Attempt to read past end of file. - """ - 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): """ Extract the poller state from Broker and replace it with the industrial @@ -726,17 +645,20 @@ def _upgrade_broker(broker): root = logging.getLogger() old_level = root.level root.setLevel(logging.CRITICAL) + try: + old = broker.poller + new = PREFERRED_POLLER() + for fd, data in old.readers: + new.start_receive(fd, data) + for fd, data in old.writers: + new.start_transmit(fd, data) + + old.close() + broker.poller = new + finally: + root.setLevel(old_level) - old = broker.poller - new = PREFERRED_POLLER() - for fd, data in old.readers: - new.start_receive(fd, data) - for fd, data in old.writers: - new.start_transmit(fd, data) - - old.close() - broker.poller = new - root.setLevel(old_level) + broker.timers = TimerList() LOG.debug('replaced %r with %r (new: %d readers, %d writers; ' 'old: %d readers, %d writers)', old, new, len(new.readers), len(new.writers), @@ -754,7 +676,7 @@ def upgrade_router(econtext): ) -def stream_by_method_name(name): +def get_connection_class(name): """ Given the name of a Mitogen connection method, import its implementation module and return its Stream subclass. @@ -762,7 +684,7 @@ def stream_by_method_name(name): if name == u'local': name = u'parent' module = mitogen.core.import_module(u'mitogen.' + name) - return module.Stream + return module.Connection @mitogen.core.takes_econtext @@ -783,7 +705,7 @@ def _proxy_connect(name, method_name, kwargs, econtext): try: context = econtext.router._connect( - klass=stream_by_method_name(method_name), + klass=get_connection_class(method_name), name=name, **kwargs ) @@ -804,19 +726,13 @@ def _proxy_connect(name, method_name, kwargs, econtext): } -def wstatus_to_str(status): +def returncode_to_str(n): """ Parse and format a :func:`os.waitpid` exit status. """ - if os.WIFEXITED(status): - return 'exited with return code %d' % (os.WEXITSTATUS(status),) - if os.WIFSIGNALED(status): - n = os.WTERMSIG(status) - return 'exited due to signal %d (%s)' % (n, SIGNAL_BY_NUM.get(n)) - if os.WIFSTOPPED(status): - n = os.WSTOPSIG(status) - return 'stopped due to signal %d (%s)' % (n, SIGNAL_BY_NUM.get(n)) - return 'unknown wait status (%d)' % (status,) + if n < 0: + return 'exited due to signal %d (%s)' % (-n, SIGNAL_BY_NUM.get(-n)) + return 'exited with return code %d' % (n,) class EofError(mitogen.core.StreamError): @@ -1096,123 +1012,188 @@ for _klass in mitogen.core.Poller, PollPoller, KqueuePoller, EpollPoller: if _klass.SUPPORTED: PREFERRED_POLLER = _klass -# For apps that start threads dynamically, it's possible Latch will also get -# very high-numbered wait fds when there are many connections, and so select() -# becomes useless there too. So swap in our favourite poller. +# For processes that start many threads or connections, it's possible Latch +# will also get high-numbered FDs, and so select() becomes useless there too. +# So swap in our favourite poller. if PollPoller.SUPPORTED: mitogen.core.Latch.poller_class = PollPoller else: mitogen.core.Latch.poller_class = PREFERRED_POLLER -class DiagLogStream(mitogen.core.BasicStream): +class LineLoggingProtocolMixin(object): + def __init__(self, **kwargs): + super(LineLoggingProtocolMixin, self).__init__(**kwargs) + self.logged_lines = [] + self.logged_partial = None + + def on_line_received(self, line): + self.logged_partial = None + self.logged_lines.append((time.time(), line)) + self.logged_lines[:] = self.logged_lines[-100:] + return super(LineLoggingProtocolMixin, self).on_line_received(line) + + def on_partial_line_received(self, line): + self.logged_partial = line + return super(LineLoggingProtocolMixin, self).on_partial_line_received(line) + + def on_disconnect(self, broker): + if self.logged_partial: + self.logged_lines.append((time.time(), self.logged_partial)) + self.logged_partial = None + super(LineLoggingProtocolMixin, self).on_disconnect(broker) + + +def get_history(streams): + history = [] + for stream in streams: + if stream: + history.extend(getattr(stream.protocol, 'logged_lines', [])) + history.sort() + + s = b('\n').join(h[1] for h in history) + return mitogen.core.to_text(s) + + +class RegexProtocol(LineLoggingProtocolMixin, mitogen.core.DelimitedProtocol): + """ + Implement a delimited protocol where messages matching a set of regular + expressions are dispatched to individual handler methods. Input is + dispatches using :attr:`PATTERNS` and :attr:`PARTIAL_PATTERNS`, before + falling back to :meth:`on_unrecognized_line_received` and + :meth:`on_unrecognized_partial_line_received`. """ - For "hybrid TTY/socketpair" mode, after a connection has been setup, a - spare TTY file descriptor will exist that cannot be closed, and to which - SSH or sudo may continue writing log messages. - The descriptor cannot be closed since the UNIX TTY layer will send a - termination signal to any processes whose controlling TTY is the TTY that - has been closed. + #: A sequence of 2-tuples of the form `(compiled pattern, method)` for + #: patterns that should be matched against complete (delimited) messages, + #: i.e. full lines. + PATTERNS = [] + + #: Like :attr:`PATTERNS`, but patterns that are matched against incomplete + #: lines. + PARTIAL_PATTERNS = [] + + def on_line_received(self, line): + super(RegexProtocol, self).on_line_received(line) + for pattern, func in self.PATTERNS: + match = pattern.search(line) + if match is not None: + return func(self, line, match) + + return self.on_unrecognized_line_received(line) + + def on_unrecognized_line_received(self, line): + LOG.debug('%s: (unrecognized): %s', + self.stream.name, line.decode('utf-8', 'replace')) + + def on_partial_line_received(self, line): + super(RegexProtocol, self).on_partial_line_received(line) + LOG.debug('%s: (partial): %s', + self.stream.name, line.decode('utf-8', 'replace')) + for pattern, func in self.PARTIAL_PATTERNS: + match = pattern.search(line) + if match is not None: + return func(self, line, match) + + return self.on_unrecognized_partial_line_received(line) + + def on_unrecognized_partial_line_received(self, line): + LOG.debug('%s: (unrecognized partial): %s', + self.stream.name, line.decode('utf-8', 'replace')) + - DiagLogStream takes over this descriptor and creates corresponding log - messages for anything written to it. +class BootstrapProtocol(RegexProtocol): """ + Respond to stdout of a child during bootstrap. Wait for EC0_MARKER to be + written by the first stage to indicate it can receive the bootstrap, then + await EC1_MARKER to indicate success, and + :class:`mitogen.core.MitogenProtocol` can be enabled. + """ + #: Sentinel value emitted by the first stage to indicate it is ready to + #: receive the compressed bootstrap. For :mod:`mitogen.ssh` this must have + #: length of at least `max(len('password'), len('debug1:'))` + EC0_MARKER = b('MITO000') + EC1_MARKER = b('MITO001') + EC2_MARKER = b('MITO002') - def __init__(self, fd, stream): - self.receive_side = mitogen.core.Side(self, fd) - self.transmit_side = self.receive_side - self.stream = stream - self.buf = '' + def __init__(self, broker): + super(BootstrapProtocol, self).__init__() + self._writer = mitogen.core.BufferedWriter(broker, self) - def __repr__(self): - return "mitogen.parent.DiagLogStream(fd=%r, '%s')" % ( - self.receive_side.fd, - self.stream.name, - ) + def on_transmit(self, broker): + self._writer.on_transmit(broker) - def on_receive(self, broker): - """ - This handler is only called after the stream is registered with the IO - loop, the descriptor is manually read/written by _connect_bootstrap() - prior to that. - """ - buf = self.receive_side.read() - if not buf: - return self.on_disconnect(broker) + def _on_ec0_received(self, line, match): + LOG.debug('%r: first stage started succcessfully', self) + self._writer.write(self.stream.conn.get_preamble()) + + def _on_ec1_received(self, line, match): + LOG.debug('%r: first stage received bootstrap', self) - self.buf += buf.decode('utf-8', 'replace') - while u'\n' in self.buf: - lines = self.buf.split('\n') - self.buf = lines[-1] - for line in lines[:-1]: - LOG.debug('%s: %s', self.stream.name, line.rstrip()) + def _on_ec2_received(self, line, match): + LOG.debug('%r: new child booted successfully', self) + self.stream.conn._complete_connection() + return False + def on_unrecognized_line_received(self, line): + LOG.debug('%s: stdout: %s', self.stream.name, line) -class Stream(mitogen.core.Stream): + PATTERNS = [ + (re.compile(EC0_MARKER), _on_ec0_received), + (re.compile(EC1_MARKER), _on_ec1_received), + (re.compile(EC2_MARKER), _on_ec2_received), + ] + + +class LogProtocol(LineLoggingProtocolMixin, mitogen.core.DelimitedProtocol): """ - Base for streams capable of starting new slaves. + For "hybrid TTY/socketpair" mode, after connection setup a spare TTY master + FD exists that cannot be closed, and to which SSH or sudo may continue + writing log messages. + + The descriptor cannot be closed since the UNIX TTY layer sends SIGHUP to + processes whose controlling TTY is the slave whose master side was closed. + LogProtocol takes over this FD and creates log messages for anything + written to it. """ + def on_line_received(self, line): + super(LogProtocol, self).on_line_received(line) + LOG.info(u'%s: %s', self.stream.name, line.decode('utf-8', 'replace')) + + +class Options(object): + name = None + #: The path to the remote Python interpreter. python_path = get_sys_executable() #: Maximum time to wait for a connection attempt. connect_timeout = 30.0 - #: Derived from :py:attr:`connect_timeout`; absolute floating point - #: UNIX timestamp after which the connection attempt should be abandoned. - connect_deadline = None - #: True to cause context to write verbose /tmp/mitogen..log. debug = False #: True to cause context to write /tmp/mitogen.stats...log. profiling = False - #: Set to the child's PID by connect(). - pid = None + #: True if unidirectional routing is enabled in the new child. + unidirectional = False #: Passed via Router wrapper methods, must eventually be passed to #: ExternalContext.main(). max_message_size = None - #: If :attr:`create_child` supplied a diag_fd, references the corresponding - #: :class:`DiagLogStream`, allowing it to be disconnected when this stream - #: is disconnected. Set to :data:`None` if no `diag_fd` was present. - diag_stream = None - - #: Function with the semantics of :func:`create_child` used to create the - #: child process. - create_child = staticmethod(create_child) - - #: Dictionary of extra kwargs passed to :attr:`create_child`. - create_child_args = {} + #: Remote name. + remote_name = None - #: :data:`True` if the remote has indicated that it intends to detach, and - #: should not be killed on disconnect. - detached = False - - #: If :data:`True`, indicates the child should not be killed during - #: graceful detachment, as it the actual process implementing the child - #: context. In all other cases, the subprocess is SSH, sudo, or a similar - #: tool that should be reminded to quit during disconnection. - child_is_immediate_subprocess = True - - #: Prefix given to default names generated by :meth:`connect`. - name_prefix = u'local' - - _reaped = False + #: Derived from :py:attr:`connect_timeout`; absolute floating point + #: UNIX timestamp after which the connection attempt should be abandoned. + connect_deadline = None - def __init__(self, *args, **kwargs): - super(Stream, self).__init__(*args, **kwargs) - self.sent_modules = set(['mitogen', 'mitogen.core']) - - def construct(self, max_message_size, remote_name=None, python_path=None, - debug=False, connect_timeout=None, profiling=False, - unidirectional=False, old_router=None, **kwargs): - """Get the named context running on the local machine, creating it if - it does not exist.""" - super(Stream, self).construct(**kwargs) + def __init__(self, max_message_size, name=None, remote_name=None, + python_path=None, debug=False, connect_timeout=None, + profiling=False, unidirectional=False, old_router=None): + self.name = name self.max_message_size = max_message_size if python_path: self.python_path = python_path @@ -1222,72 +1203,68 @@ class Stream(mitogen.core.Stream): remote_name = get_default_remote_name() if '/' in remote_name or '\\' in remote_name: raise ValueError('remote_name= cannot contain slashes') - self.remote_name = remote_name + if remote_name: + self.remote_name = mitogen.core.to_text(remote_name) self.debug = debug self.profiling = profiling self.unidirectional = unidirectional self.max_message_size = max_message_size self.connect_deadline = time.time() + self.connect_timeout - def on_shutdown(self, broker): - """Request the slave gracefully shut itself down.""" - LOG.debug('%r closing CALL_FUNCTION channel', self) - self._send( - mitogen.core.Message( - src_id=mitogen.context_id, - dst_id=self.remote_id, - handle=mitogen.core.SHUTDOWN, - ) - ) - def _reap_child(self): - """ - Reap the child process during disconnection. - """ - if self.detached and self.child_is_immediate_subprocess: - LOG.debug('%r: immediate child is detached, won\'t reap it', self) - return +class Connection(object): + """ + Base for streams capable of starting children. + """ + options_class = Options - if self.profiling: - LOG.info('%r: wont kill child because profiling=True', self) - return + #: The protocol attached to stdio of the child. + stream_protocol_class = BootstrapProtocol - if self._reaped: - # on_disconnect() may be invoked more than once, for example, if - # there is still a pending message to be sent after the first - # on_disconnect() call. - return + #: The protocol attached to stderr of the child. + diag_protocol_class = LogProtocol - try: - pid, status = os.waitpid(self.pid, os.WNOHANG) - except OSError: - e = sys.exc_info()[1] - if e.args[0] == errno.ECHILD: - LOG.warn('%r: waitpid(%r) produced ECHILD', self, self.pid) - return - raise + #: :class:`Process` + proc = None - self._reaped = True - if pid: - LOG.debug('%r: PID %d %s', self, pid, wstatus_to_str(status)) - return + #: :class:`mitogen.core.Stream` with sides connected to stdin/stdout. + stdio_stream = None - if not self._router.profiling: - # For processes like sudo we cannot actually send sudo a signal, - # because it is setuid, so this is best-effort only. - LOG.debug('%r: child process still alive, sending SIGTERM', self) - try: - os.kill(self.pid, signal.SIGTERM) - except OSError: - e = sys.exc_info()[1] - if e.args[0] != errno.EPERM: - raise + #: If `proc.stderr` is set, referencing either a plain pipe or the + #: controlling TTY, this references the corresponding + #: :class:`LogProtocol`'s stream, allowing it to be disconnected when this + #: stream is disconnected. + stderr_stream = None - def on_disconnect(self, broker): - super(Stream, self).on_disconnect(broker) - if self.diag_stream is not None: - self.diag_stream.on_disconnect(broker) - self._reap_child() + #: Function with the semantics of :func:`create_child` used to create the + #: child process. + create_child = staticmethod(create_child) + + #: Dictionary of extra kwargs passed to :attr:`create_child`. + create_child_args = {} + + #: :data:`True` if the remote has indicated that it intends to detach, and + #: should not be killed on disconnect. + detached = False + + #: If :data:`True`, indicates the child should not be killed during + #: graceful detachment, as it the actual process implementing the child + #: context. In all other cases, the subprocess is SSH, sudo, or a similar + #: tool that should be reminded to quit during disconnection. + child_is_immediate_subprocess = True + + #: Prefix given to default names generated by :meth:`connect`. + name_prefix = u'local' + + timer = None + + def __init__(self, options, router): + #: :class:`Options` + self.options = options + self._router = router + + def __repr__(self): + return 'Connection(%r)' % (self.stdio_stream,) # Minimised, gzipped, base64'd and passed to 'python -c'. It forks, dups # file descriptor 0 as 100, creates a pipe, then execs a new interpreter @@ -1346,15 +1323,15 @@ class Stream(mitogen.core.Stream): This allows emulation of existing tools where the Python invocation may be set to e.g. `['/usr/bin/env', 'python']`. """ - if isinstance(self.python_path, list): - return self.python_path - return [self.python_path] + if isinstance(self.options.python_path, list): + return self.options.python_path + return [self.options.python_path] def get_boot_command(self): source = inspect.getsource(self._first_stage) source = textwrap.dedent('\n'.join(source.strip().split('\n')[2:])) source = source.replace(' ', '\t') - source = source.replace('CONTEXT_NAME', self.remote_name) + source = source.replace('CONTEXT_NAME', self.options.remote_name) preamble_compressed = self.get_preamble() source = source.replace('PREAMBLE_COMPRESSED_LEN', str(len(preamble_compressed))) @@ -1372,19 +1349,19 @@ class Stream(mitogen.core.Stream): ] def get_econtext_config(self): - assert self.max_message_size is not None + assert self.options.max_message_size is not None parent_ids = mitogen.parent_ids[:] parent_ids.insert(0, mitogen.context_id) return { 'parent_ids': parent_ids, - 'context_id': self.remote_id, - 'debug': self.debug, - 'profiling': self.profiling, - 'unidirectional': self.unidirectional, + 'context_id': self.context.context_id, + 'debug': self.options.debug, + 'profiling': self.options.profiling, + 'unidirectional': self.options.unidirectional, 'log_level': get_log_level(), 'whitelist': self._router.get_module_whitelist(), 'blacklist': self._router.get_module_blacklist(), - 'max_message_size': self.max_message_size, + 'max_message_size': self.options.max_message_size, 'version': mitogen.__version__, } @@ -1396,10 +1373,18 @@ class Stream(mitogen.core.Stream): partial = get_core_source_partial() return partial.append(suffix.encode('utf-8')) + def _get_name(self): + """ + Called by :meth:`connect` after :attr:`pid` is known. Subclasses can + override it to specify a default stream name, or set + :attr:`name_prefix` to generate a default format. + """ + return u'%s.%s' % (self.name_prefix, self.proc.pid) + def start_child(self): args = self.get_boot_command() try: - return self.create_child(args, **self.create_child_args) + return self.create_child(args=args, **self.create_child_args) except OSError: e = sys.exc_info()[1] msg = 'Child start failed: %s. Command was: %s' % (e, Argv(args)) @@ -1409,65 +1394,162 @@ class Stream(mitogen.core.Stream): def _adorn_eof_error(self, e): """ - Used by subclasses to provide additional information in the case of a - failed connection. + Subclasses may provide additional information in the case of a failed + connection. """ if self.eof_error_hint: e.args = ('%s\n\n%s' % (e.args[0], self.eof_error_hint),) - def _get_name(self): + exception = None + + def _complete_connection(self): + self.timer.cancel() + if not self.exception: + self._router.register(self.context, self.stdio_stream) + self.stdio_stream.set_protocol( + mitogen.core.MitogenProtocol( + router=self._router, + remote_id=self.context.context_id, + ) + ) + self.latch.put() + + def _fail_connection(self, exc): """ - Called by :meth:`connect` after :attr:`pid` is known. Subclasses can - override it to specify a default stream name, or set - :attr:`name_prefix` to generate a default format. + Fail the connection attempt. + """ + LOG.debug('%s: failing connection due to %r', + self.stdio_stream.name, exc) + if self.exception is None: + self._adorn_eof_error(exc) + self.exception = exc + for stream in self.stdio_stream, self.stderr_stream: + if stream and not stream.receive_side.closed: + stream.on_disconnect(self._router.broker) + self._complete_connection() + + def on_stream_shutdown(self): + """ + Request the slave gracefully shut itself down. """ - return u'%s.%s' % (self.name_prefix, self.pid) + LOG.debug('%r: requesting child shutdown', self) + self.stdio_stream.protocol._send( + mitogen.core.Message( + src_id=mitogen.context_id, + dst_id=self.stdio_stream.protocol.remote_id, + handle=mitogen.core.SHUTDOWN, + ) + ) - def connect(self): - LOG.debug('%r.connect()', self) - self.pid, fd, diag_fd = self.start_child() - self.name = self._get_name() - self.receive_side = mitogen.core.Side(self, fd) - self.transmit_side = mitogen.core.Side(self, os.dup(fd)) - if diag_fd is not None: - self.diag_stream = DiagLogStream(diag_fd, self) - else: - self.diag_stream = None + eof_error_msg = 'EOF on stream; last 100 lines received:\n' - LOG.debug('%r.connect(): pid:%r stdin:%r, stdout:%r, diag:%r', - self, self.pid, self.receive_side.fd, self.transmit_side.fd, - self.diag_stream and self.diag_stream.receive_side.fd) + def on_stdio_disconnect(self): + """ + Handle stdio stream disconnection by failing the Connection if the + stderr stream has already been closed. Otherwise, wait for it to close + (or timeout), to allow buffered diagnostic logs to be consumed. + + It is normal that when a subprocess aborts, stdio has nothing buffered + when it is closed, thus signalling readability, causing an empty read + (interpreted as indicating disconnection) on the next loop iteration, + even if its stderr pipe has lots of diagnostic logs still buffered in + the kernel. Therefore we must wait for both pipes to indicate they are + empty before triggering connection failure. + """ + stderr = self.stderr_stream + if stderr is None or stderr.receive_side.closed: + self._on_streams_disconnected() - try: - self._connect_bootstrap() - except EofError: - self.on_disconnect(self._router.broker) - e = sys.exc_info()[1] - self._adorn_eof_error(e) - raise - except Exception: - self.on_disconnect(self._router.broker) - self._reap_child() - raise + def on_stderr_disconnect(self): + """ + Inverse of :func:`on_stdio_disconnect`. + """ + if self.stdio_stream.receive_side.closed: + self._on_streams_disconnected() - #: Sentinel value emitted by the first stage to indicate it is ready to - #: receive the compressed bootstrap. For :mod:`mitogen.ssh` this must have - #: length of at least `max(len('password'), len('debug1:'))` - EC0_MARKER = mitogen.core.b('MITO000\n') - EC1_MARKER = mitogen.core.b('MITO001\n') + def _on_streams_disconnected(self): + """ + When disconnection has been detected for both our streams, cancel the + connection timer, mark the connection failed, and reap the child + process. Do nothing if the timer has already been cancelled, indicating + some existing failure has already been noticed. + """ + if not self.timer.cancelled: + self.timer.cancel() + self._fail_connection(EofError( + self.eof_error_msg + get_history( + [self.stdio_stream, self.stderr_stream] + ) + )) + self.proc._async_reap(self, self._router) - def _ec0_received(self): - LOG.debug('%r._ec0_received()', self) - write_all(self.transmit_side.fd, self.get_preamble()) - discard_until(self.receive_side.fd, self.EC1_MARKER, - self.connect_deadline) - if self.diag_stream: - self._router.broker.start_receive(self.diag_stream) + def _start_timer(self): + self.timer = self._router.broker.timers.schedule( + when=self.options.connect_deadline, + func=self._on_timer_expired, + ) + + def _on_timer_expired(self): + self._fail_connection( + mitogen.core.TimeoutError( + 'Failed to setup connection after %.2f seconds', + self.options.connect_timeout, + ) + ) + + def stream_factory(self): + return self.stream_protocol_class.build_stream( + broker=self._router.broker, + ) + + def stderr_stream_factory(self): + return self.diag_protocol_class.build_stream() + + def _setup_stdio_stream(self): + stream = self.stream_factory() + stream.conn = self + stream.name = self.options.name or self._get_name() + stream.accept(self.proc.stdout, self.proc.stdin) + + mitogen.core.listen(stream, 'shutdown', self.on_stream_shutdown) + mitogen.core.listen(stream, 'disconnect', self.on_stdio_disconnect) + self._router.broker.start_receive(stream) + return stream + + def _setup_stderr_stream(self): + stream = self.stderr_stream_factory() + stream.conn = self + stream.name = self.options.name or self._get_name() + stream.accept(self.proc.stderr, self.proc.stderr) + + mitogen.core.listen(stream, 'disconnect', self.on_stderr_disconnect) + self._router.broker.start_receive(stream) + return stream + + def _async_connect(self): + self._start_timer() + self.stdio_stream = self._setup_stdio_stream() + if self.context.name is None: + self.context.name = self.stdio_stream.name + self.proc.name = self.stdio_stream.name + if self.proc.stderr: + self.stderr_stream = self._setup_stderr_stream() + + def connect(self, context): + LOG.debug('%r.connect()', self) + self.context = context + self.proc = self.start_child() + LOG.debug('%r.connect(): pid:%r stdin:%r stdout:%r stderr:%r', + self, self.proc.pid, + self.proc.stdin.fileno(), + self.proc.stdout.fileno(), + self.proc.stderr and self.proc.stderr.fileno()) - def _connect_bootstrap(self): - discard_until(self.receive_side.fd, self.EC0_MARKER, - self.connect_deadline) - self._ec0_received() + self.latch = mitogen.core.Latch() + self._router.broker.defer(self._async_connect) + self.latch.get() + if self.exception: + raise self.exception class ChildIdAllocator(object): @@ -1482,7 +1564,7 @@ class ChildIdAllocator(object): for id_ in self.it: return id_ - master = mitogen.core.Context(self.router, 0) + master = self.router.context_by_id(0) start, end = master.send_await( mitogen.core.Message(dst_id=0, handle=mitogen.core.ALLOCATE_ID) ) @@ -1739,9 +1821,11 @@ class Context(mitogen.core.Context): return not (self == other) def __eq__(self, other): - return (isinstance(other, mitogen.core.Context) and - (other.context_id == self.context_id) and - (other.router == self.router)) + return ( + isinstance(other, mitogen.core.Context) and + (other.context_id == self.context_id) and + (other.router == self.router) + ) def __hash__(self): return hash((self.router, self.context_id)) @@ -1869,11 +1953,11 @@ class RouteMonitor(object): data = str(target_id) if name: data = '%s:%s' % (target_id, name) - stream.send( + stream.protocol.send( mitogen.core.Message( handle=handle, data=data.encode('utf-8'), - dst_id=stream.remote_id, + dst_id=stream.protocol.remote_id, ) ) @@ -1907,9 +1991,9 @@ class RouteMonitor(object): ID of the connecting or disconnecting context. """ for stream in self.router.get_streams(): - if target_id in stream.egress_ids and ( + if target_id in stream.protocol.egress_ids and ( (self.parent is None) or - (self.parent.context_id != stream.remote_id) + (self.parent.context_id != stream.protocol.remote_id) ): self._send_one(stream, mitogen.core.DEL_ROUTE, target_id, None) @@ -1919,8 +2003,8 @@ class RouteMonitor(object): stream, we're also responsible for broadcasting DEL_ROUTE upstream if/when that child disconnects. """ - self._routes_by_stream[stream] = set([stream.remote_id]) - self._propagate_up(mitogen.core.ADD_ROUTE, stream.remote_id, + self._routes_by_stream[stream] = set([stream.protocol.remote_id]) + self._propagate_up(mitogen.core.ADD_ROUTE, stream.protocol.remote_id, stream.name) mitogen.core.listen( obj=stream, @@ -1974,7 +2058,7 @@ class RouteMonitor(object): self.router.context_by_id(target_id).name = target_name stream = self.router.stream_by_id(msg.auth_id) current = self.router.stream_by_id(target_id) - if current and current.remote_id != mitogen.parent_id: + if current and current.protocol.remote_id != mitogen.parent_id: LOG.error('Cannot add duplicate route to %r via %r, ' 'already have existing route via %r', target_id, stream, current) @@ -2017,7 +2101,7 @@ class RouteMonitor(object): routes.discard(target_id) self.router.del_route(target_id) - if stream.remote_id != mitogen.parent_id: + if stream.protocol.remote_id != mitogen.parent_id: self._propagate_up(mitogen.core.DEL_ROUTE, target_id) self._propagate_down(mitogen.core.DEL_ROUTE, target_id) @@ -2051,11 +2135,11 @@ class Router(mitogen.core.Router): if msg.is_dead: return stream = self.stream_by_id(msg.src_id) - if stream.remote_id != msg.src_id or stream.detached: + if stream.protocol.remote_id != msg.src_id or stream.conn.detached: LOG.warning('bad DETACHING received on %r: %r', stream, msg) return LOG.debug('%r: marking as detached', stream) - stream.detached = True + stream.conn.detached = True msg.reply(None) def get_streams(self): @@ -2078,7 +2162,7 @@ class Router(mitogen.core.Router): """ LOG.debug('%r.add_route(%r, %r)', self, target_id, stream) assert isinstance(target_id, int) - assert isinstance(stream, Stream) + assert isinstance(stream, mitogen.core.Stream) self._write_lock.acquire() try: @@ -2087,7 +2171,7 @@ class Router(mitogen.core.Router): self._write_lock.release() def del_route(self, target_id): - LOG.debug('%r.del_route(%r)', self, target_id) + LOG.debug('%r: deleting route to %r', self, target_id) # DEL_ROUTE may be sent by a parent if it knows this context sent # messages to a peer that has now disconnected, to let us raise # 'disconnect' event on the appropriate Context instance. In that case, @@ -2114,35 +2198,37 @@ class Router(mitogen.core.Router): connection_timeout_msg = u"Connection timed out." - def _connect(self, klass, name=None, **kwargs): + def _connect(self, klass, **kwargs): context_id = self.allocate_id() context = self.context_class(self, context_id) + context.name = kwargs.get('name') + kwargs['old_router'] = self kwargs['max_message_size'] = self.max_message_size - stream = klass(self, context_id, **kwargs) - if name is not None: - stream.name = name + conn = klass(klass.options_class(**kwargs), self) try: - stream.connect() + conn.connect(context=context) except mitogen.core.TimeoutError: raise mitogen.core.StreamError(self.connection_timeout_msg) - context.name = stream.name - self.route_monitor.notice_stream(stream) - self.register(context, stream) + + self.route_monitor.notice_stream(conn.stdio_stream) return context def connect(self, method_name, name=None, **kwargs): - klass = stream_by_method_name(method_name) + if name: + name = mitogen.core.to_text(name) + + klass = get_connection_class(method_name) kwargs.setdefault(u'debug', self.debug) kwargs.setdefault(u'profiling', self.profiling) kwargs.setdefault(u'unidirectional', self.unidirectional) + kwargs.setdefault(u'name', name) via = kwargs.pop(u'via', None) if via is not None: - return self.proxy_connect(via, method_name, name=name, - **mitogen.core.Kwargs(kwargs)) - return self._connect(klass, name=name, - **mitogen.core.Kwargs(kwargs)) + return self.proxy_connect(via, method_name, + **mitogen.core.Kwargs(kwargs)) + return self._connect(klass, **mitogen.core.Kwargs(kwargs)) def proxy_connect(self, via_context, method_name, name=None, **kwargs): resp = via_context.call(_proxy_connect, @@ -2203,43 +2289,86 @@ class Router(mitogen.core.Router): return self.connect(u'ssh', **kwargs) -class ProcessMonitor(object): - """ - Install a :data:`signal.SIGCHLD` handler that generates callbacks when a - specific child process has exitted. This class is obsolete, do not use. - """ - def __init__(self): - # pid -> callback() - self.callback_by_pid = {} - signal.signal(signal.SIGCHLD, self._on_sigchld) +class Process(object): + _delays = [0.05, 0.15, 0.3, 1.0, 5.0, 10.0] + name = None - def _on_sigchld(self, _signum, _frame): - for pid, callback in self.callback_by_pid.items(): - pid, status = os.waitpid(pid, os.WNOHANG) - if pid: - callback(status) - del self.callback_by_pid[pid] + def __init__(self, pid, stdin, stdout, stderr=None): + self.pid = pid + self.stdin = stdin + self.stdout = stdout + self.stderr = stderr + self._returncode = None + self._reap_count = 0 - def add(self, pid, callback): - """ - Add a callback function to be notified of the exit status of a process. + def __repr__(self): + return '%s %s pid %d' % ( + type(self).__name__, + self.name, + self.pid, + ) - :param int pid: - Process ID to be notified of. + def poll(self): + raise NotImplementedError() + + def _signal_child(self, signum): + # For processes like sudo we cannot actually send sudo a signal, + # because it is setuid, so this is best-effort only. + LOG.debug('%r: child process still alive, sending %s', + self, SIGNAL_BY_NUM[signum]) + try: + os.kill(self.pid, signum) + except OSError: + e = sys.exc_info()[1] + if e.args[0] != errno.EPERM: + raise - :param callback: - Function invoked as `callback(status)`, where `status` is the raw - exit status of the child process. + def _async_reap(self, conn, router): + """ + Reap the child process during disconnection. """ - self.callback_by_pid[pid] = callback + if self._returncode is not None: + # on_disconnect() may be invoked more than once, for example, if + # there is still a pending message to be sent after the first + # on_disconnect() call. + return - _instance = None + if conn.detached and conn.child_is_immediate_subprocess: + LOG.debug('%r: immediate child is detached, won\'t reap it', self) + return - @classmethod - def instance(cls): - if cls._instance is None: - cls._instance = cls() - return cls._instance + if router.profiling: + LOG.info('%r: wont kill child because profiling=True', self) + return + + self._reap_count += 1 + status = self.poll() + if status is not None: + LOG.debug('%r: %s', self, returncode_to_str(status)) + return + + i = self._reap_count - 1 + if i >= len(self._delays): + LOG.warning('%r: child will not die, abandoning it', self) + return + elif i == 0: + self._signal_child(signal.SIGTERM) + elif i == 1: + self._signal_child(signal.SIGKILL) + + router.broker.timers.schedule( + when=time.time() + self._delays[i], + func=lambda: self._async_reap(conn, router), + ) + + +class PopenProcess(Process): + def __init__(self, proc, stdin, stdout, stderr=None): + super(PopenProcess, self).__init__(proc.pid, stdin, stdout, stderr) + self.proc = proc + + def poll(self): + return self.proc.poll() class ModuleForwarder(object): @@ -2265,7 +2394,7 @@ class ModuleForwarder(object): ) def __repr__(self): - return 'ModuleForwarder(%r)' % (self.router,) + return 'ModuleForwarder' def _on_forward_module(self, msg): if msg.is_dead: @@ -2275,38 +2404,38 @@ class ModuleForwarder(object): fullname = mitogen.core.to_text(fullname) context_id = int(context_id_s) stream = self.router.stream_by_id(context_id) - if stream.remote_id == mitogen.parent_id: + if stream.protocol.remote_id == mitogen.parent_id: LOG.error('%r: dropping FORWARD_MODULE(%d, %r): no route to child', self, context_id, fullname) return - if fullname in stream.sent_modules: + if fullname in stream.protocol.sent_modules: return LOG.debug('%r._on_forward_module() sending %r to %r via %r', - self, fullname, context_id, stream.remote_id) + self, fullname, context_id, stream.protocol.remote_id) self._send_module_and_related(stream, fullname) - if stream.remote_id != context_id: - stream._send( + if stream.protocol.remote_id != context_id: + stream.protocol._send( mitogen.core.Message( data=msg.data, handle=mitogen.core.FORWARD_MODULE, - dst_id=stream.remote_id, + dst_id=stream.protocol.remote_id, ) ) def _on_get_module(self, msg): - LOG.debug('%r._on_get_module(%r)', self, msg) if msg.is_dead: return fullname = msg.data.decode('utf-8') + LOG.debug('%r: %s requested by %d', self, fullname, msg.src_id) callback = lambda: self._on_cache_callback(msg, fullname) self.importer._request_module(fullname, callback) def _on_cache_callback(self, msg, fullname): - LOG.debug('%r._on_get_module(): sending %r', self, fullname) stream = self.router.stream_by_id(msg.src_id) + LOG.debug('%r: sending %s to %r', self, fullname, stream) self._send_module_and_related(stream, fullname) def _send_module_and_related(self, stream, fullname): @@ -2316,18 +2445,18 @@ class ModuleForwarder(object): if rtup: self._send_one_module(stream, rtup) else: - LOG.debug('%r._send_module_and_related(%r): absent: %r', - self, fullname, related) + LOG.debug('%r: %s not in cache (for %s)', + self, related, fullname) self._send_one_module(stream, tup) def _send_one_module(self, stream, tup): - if tup[0] not in stream.sent_modules: - stream.sent_modules.add(tup[0]) + if tup[0] not in stream.protocol.sent_modules: + stream.protocol.sent_modules.add(tup[0]) self.router._async_route( mitogen.core.Message.pickled( tup, - dst_id=stream.remote_id, + dst_id=stream.protocol.remote_id, handle=mitogen.core.LOAD_MODULE, ) ) diff --git a/mitogen/profiler.py b/mitogen/profiler.py index 74bbdb23..e697d599 100644 --- a/mitogen/profiler.py +++ b/mitogen/profiler.py @@ -28,7 +28,8 @@ # !mitogen: minify_safe -"""mitogen.profiler +""" +mitogen.profiler Record and report cProfile statistics from a run. Creates one aggregated output file, one aggregate containing only workers, and one for the top-level process. @@ -152,7 +153,7 @@ def do_stat(tmpdir, sort, *args): def main(): if len(sys.argv) < 2 or sys.argv[1] not in ('record', 'report', 'stat'): - sys.stderr.write(__doc__) + sys.stderr.write(__doc__.lstrip()) sys.exit(1) func = globals()['do_' + sys.argv[1]] diff --git a/mitogen/service.py b/mitogen/service.py index 942ed4f7..886012e8 100644 --- a/mitogen/service.py +++ b/mitogen/service.py @@ -485,7 +485,6 @@ class Pool(object): ) thread.start() self._threads.append(thread) - LOG.debug('%r: initialized', self) def _py_24_25_compat(self): @@ -658,7 +657,7 @@ class PushFileService(Service): def _forward(self, context, path): stream = self.router.stream_by_id(context.context_id) - child = mitogen.core.Context(self.router, stream.remote_id) + child = mitogen.core.Context(self.router, stream.protocol.remote_id) sent = self._sent_by_stream.setdefault(stream, set()) if path in sent: if child.context_id != context.context_id: @@ -891,7 +890,7 @@ class FileService(Service): # The IO loop pumps 128KiB chunks. An ideal message is a multiple of this, # odd-sized messages waste one tiny write() per message on the trailer. # Therefore subtract 10 bytes pickle overhead + 24 bytes header. - IO_SIZE = mitogen.core.CHUNK_SIZE - (mitogen.core.Stream.HEADER_LEN + ( + IO_SIZE = mitogen.core.CHUNK_SIZE - (mitogen.core.Message.HEADER_LEN + ( len( mitogen.core.Message.pickled( mitogen.core.Blob(b(' ') * mitogen.core.CHUNK_SIZE) diff --git a/mitogen/setns.py b/mitogen/setns.py index b1d69783..46a50301 100644 --- a/mitogen/setns.py +++ b/mitogen/setns.py @@ -116,9 +116,15 @@ def get_machinectl_pid(path, name): raise Error("could not find PID from machinectl output.\n%s", output) -class Stream(mitogen.parent.Stream): - child_is_immediate_subprocess = False +GET_LEADER_BY_KIND = { + 'docker': ('docker_path', get_docker_pid), + 'lxc': ('lxc_info_path', get_lxc_pid), + 'lxd': ('lxc_path', get_lxd_pid), + 'machinectl': ('machinectl_path', get_machinectl_pid), +} + +class Options(mitogen.parent.Options): container = None username = 'root' kind = None @@ -128,24 +134,17 @@ class Stream(mitogen.parent.Stream): lxc_info_path = 'lxc-info' machinectl_path = 'machinectl' - GET_LEADER_BY_KIND = { - 'docker': ('docker_path', get_docker_pid), - 'lxc': ('lxc_info_path', get_lxc_pid), - 'lxd': ('lxc_path', get_lxd_pid), - 'machinectl': ('machinectl_path', get_machinectl_pid), - } - - def construct(self, container, kind, username=None, docker_path=None, - lxc_path=None, lxc_info_path=None, machinectl_path=None, - **kwargs): - super(Stream, self).construct(**kwargs) - if kind not in self.GET_LEADER_BY_KIND: + def __init__(self, container, kind, username=None, docker_path=None, + lxc_path=None, lxc_info_path=None, machinectl_path=None, + **kwargs): + super(Options, self).__init__(**kwargs) + if kind not in GET_LEADER_BY_KIND: raise Error('unsupported container kind: %r', kind) - self.container = container + self.container = mitogen.core.to_text(container) self.kind = kind if username: - self.username = username + self.username = mitogen.core.to_text(username) if docker_path: self.docker_path = docker_path if lxc_path: @@ -155,6 +154,11 @@ class Stream(mitogen.parent.Stream): if machinectl_path: self.machinectl_path = machinectl_path + +class Connection(mitogen.parent.Connection): + options_class = Options + child_is_immediate_subprocess = False + # Order matters. https://github.com/karelzak/util-linux/commit/854d0fe/ NS_ORDER = ('ipc', 'uts', 'net', 'pid', 'mnt', 'user') @@ -189,15 +193,15 @@ class Stream(mitogen.parent.Stream): try: os.setgroups([grent.gr_gid for grent in grp.getgrall() - if self.username in grent.gr_mem]) - pwent = pwd.getpwnam(self.username) + if self.options.username in grent.gr_mem]) + pwent = pwd.getpwnam(self.options.username) os.setreuid(pwent.pw_uid, pwent.pw_uid) # shadow-4.4/libmisc/setupenv.c. Not done: MAIL, PATH os.environ.update({ 'HOME': pwent.pw_dir, 'SHELL': pwent.pw_shell or '/bin/sh', - 'LOGNAME': self.username, - 'USER': self.username, + 'LOGNAME': self.options.username, + 'USER': self.options.username, }) if ((os.path.exists(pwent.pw_dir) and os.access(pwent.pw_dir, os.X_OK))): @@ -217,7 +221,7 @@ class Stream(mitogen.parent.Stream): # namespaces, meaning starting new threads in the exec'd program will # fail. The solution is forking, so inject a /bin/sh call to achieve # this. - argv = super(Stream, self).get_boot_command() + argv = super(Connection, self).get_boot_command() # bash will exec() if a single command was specified and the shell has # nothing left to do, so "; exit $?" gives bash a reason to live. return ['/bin/sh', '-c', '%s; exit $?' % (mitogen.parent.Argv(argv),)] @@ -226,13 +230,12 @@ class Stream(mitogen.parent.Stream): return mitogen.parent.create_child(args, preexec_fn=self.preexec_fn) def _get_name(self): - return u'setns.' + self.container + return u'setns.' + self.options.container - def connect(self): - self.name = self._get_name() - attr, func = self.GET_LEADER_BY_KIND[self.kind] - tool_path = getattr(self, attr) - self.leader_pid = func(tool_path, self.container) + def connect(self, **kwargs): + attr, func = GET_LEADER_BY_KIND[self.options.kind] + tool_path = getattr(self.options, attr) + self.leader_pid = func(tool_path, self.options.container) LOG.debug('Leader PID for %s container %r: %d', - self.kind, self.container, self.leader_pid) - super(Stream, self).connect() + self.options.kind, self.options.container, self.leader_pid) + return super(Connection, self).connect(**kwargs) diff --git a/mitogen/ssh.py b/mitogen/ssh.py index 11b74c1b..b4c247c1 100644 --- a/mitogen/ssh.py +++ b/mitogen/ssh.py @@ -29,7 +29,7 @@ # !mitogen: minify_safe """ -Functionality to allow establishing new slave contexts over an SSH connection. +Construct new children via the OpenSSH client. """ import logging @@ -52,82 +52,122 @@ except NameError: LOG = logging.getLogger('mitogen') +auth_incorrect_msg = 'SSH authentication is incorrect' +password_incorrect_msg = 'SSH password is incorrect' +password_required_msg = 'SSH password was requested, but none specified' +hostkey_config_msg = ( + 'SSH requested permission to accept unknown host key, but ' + 'check_host_keys=ignore. This is likely due to ssh_args= ' + 'conflicting with check_host_keys=. Please correct your ' + 'configuration.' +) +hostkey_failed_msg = ( + 'Host key checking is enabled, and SSH reported an unrecognized or ' + 'mismatching host key.' +) + # sshpass uses 'assword' because it doesn't lowercase the input. -PASSWORD_PROMPT = b('password') -HOSTKEY_REQ_PROMPT = b('are you sure you want to continue connecting (yes/no)?') -HOSTKEY_FAIL = b('host key verification failed.') +PASSWORD_PROMPT_PATTERN = re.compile( + b('password'), + re.I +) + +HOSTKEY_REQ_PATTERN = re.compile( + b(r'are you sure you want to continue connecting \(yes/no\)\?'), + re.I +) + +HOSTKEY_FAIL_PATTERN = re.compile( + b(r'host key verification failed\.'), + re.I +) # [user@host: ] permission denied -PERMDENIED_RE = re.compile( - ('(?:[^@]+@[^:]+: )?' # Absent in OpenSSH <7.5 - 'Permission denied').encode(), +# issue #271: work around conflict with user shell reporting 'permission +# denied' e.g. during chdir($HOME) by only matching it at the start of the +# line. +PERMDENIED_PATTERN = re.compile( + b('^(?:[^@]+@[^:]+: )?' # Absent in OpenSSH <7.5 + 'Permission denied'), re.I ) +DEBUG_PATTERN = re.compile(b('^debug[123]:')) + -DEBUG_PREFIXES = (b('debug1:'), b('debug2:'), b('debug3:')) +class PasswordError(mitogen.core.StreamError): + pass -def filter_debug(stream, it): - """ - Read line chunks from it, either yielding them directly, or building up and - logging individual lines if they look like SSH debug output. +class HostKeyError(mitogen.core.StreamError): + pass - This contains the mess of dealing with both line-oriented input, and partial - lines such as the password prompt. - Yields `(line, partial)` tuples, where `line` is the line, `partial` is - :data:`True` if no terminating newline character was present and no more - data exists in the read buffer. Consuming code can use this to unreliably - detect the presence of an interactive prompt. +class SetupProtocol(mitogen.parent.RegexProtocol): + """ + This protocol is attached to stderr of the SSH client. It responds to + various interactive prompts as required. """ - # The `partial` test is unreliable, but is only problematic when verbosity - # is enabled: it's possible for a combination of SSH banner, password - # prompt, verbose output, timing and OS buffering specifics to create a - # situation where an otherwise newline-terminated line appears to not be - # terminated, due to a partial read(). If something is broken when - # ssh_debug_level>0, this is the first place to look. - state = 'start_of_line' - buf = b('') - for chunk in it: - buf += chunk - while buf: - if state == 'start_of_line': - if len(buf) < 8: - # short read near buffer limit, block awaiting at least 8 - # bytes so we can discern a debug line, or the minimum - # interesting token from above or the bootstrap - # ('password', 'MITO000\n'). - break - elif any(buf.startswith(p) for p in DEBUG_PREFIXES): - state = 'in_debug' - else: - state = 'in_plain' - elif state == 'in_debug': - if b('\n') not in buf: - break - line, _, buf = bytes_partition(buf, b('\n')) - LOG.debug('%s: %s', stream.name, - mitogen.core.to_text(line.rstrip())) - state = 'start_of_line' - elif state == 'in_plain': - line, nl, buf = bytes_partition(buf, b('\n')) - yield line + nl, not (nl or buf) - if nl: - state = 'start_of_line' + password_sent = False + def _on_host_key_request(self, line, match): + if self.stream.conn.options.check_host_keys == 'accept': + LOG.debug('%s: accepting host key', self.stream.name) + self.stream.transmit_side.write(b('yes\n')) + return -class PasswordError(mitogen.core.StreamError): - pass + # _host_key_prompt() should never be reached with ignore or enforce + # mode, SSH should have handled that. User's ssh_args= is conflicting + # with ours. + self.stream.conn._fail_connection(HostKeyError(hostkey_config_msg)) + + def _on_host_key_failed(self, line, match): + self.stream.conn._fail_connection(HostKeyError(hostkey_failed_msg)) + + def _on_permission_denied(self, line, match): + if self.stream.conn.options.password is not None and \ + self.password_sent: + self.stream.conn._fail_connection( + PasswordError(password_incorrect_msg) + ) + elif PASSWORD_PROMPT_PATTERN.search(line) and \ + self.stream.conn.options.password is None: + # Permission denied (password,pubkey) + self.stream.conn._fail_connection( + PasswordError(password_required_msg) + ) + else: + self.stream.conn._fail_connection( + PasswordError(auth_incorrect_msg) + ) + def _on_password_prompt(self, line, match): + LOG.debug('%s: (password prompt): %s', self.stream.name, line) + if self.stream.conn.options.password is None: + self.stream.conn._fail(PasswordError(password_required_msg)) -class HostKeyError(mitogen.core.StreamError): - pass + self.stream.transmit_side.write( + (self.stream.conn.options.password + '\n').encode('utf-8') + ) + self.password_sent = True + def _on_debug_line(self, line, match): + text = mitogen.core.to_text(line.rstrip()) + LOG.debug('%s: %s', self.stream.name, text) + + PATTERNS = [ + (DEBUG_PATTERN, _on_debug_line), + (HOSTKEY_FAIL_PATTERN, _on_host_key_failed), + (PERMDENIED_PATTERN, _on_permission_denied), + ] + + PARTIAL_PATTERNS = [ + (PASSWORD_PROMPT_PATTERN, _on_password_prompt), + (HOSTKEY_REQ_PATTERN, _on_host_key_request), + ] -class Stream(mitogen.parent.Stream): - child_is_immediate_subprocess = False +class Options(mitogen.parent.Options): #: Default to whatever is available as 'python' on the remote machine, #: overriding sys.executable use. python_path = 'python' @@ -141,19 +181,19 @@ class Stream(mitogen.parent.Stream): hostname = None username = None port = None - identity_file = None password = None ssh_args = None check_host_keys_msg = 'check_host_keys= must be set to accept, enforce or ignore' - def construct(self, hostname, username=None, ssh_path=None, port=None, - check_host_keys='enforce', password=None, identity_file=None, - compression=True, ssh_args=None, keepalive_enabled=True, - keepalive_count=3, keepalive_interval=15, - identities_only=True, ssh_debug_level=None, **kwargs): - super(Stream, self).construct(**kwargs) + def __init__(self, hostname, username=None, ssh_path=None, port=None, + check_host_keys='enforce', password=None, identity_file=None, + compression=True, ssh_args=None, keepalive_enabled=True, + keepalive_count=3, keepalive_interval=15, + identities_only=True, ssh_debug_level=None, **kwargs): + super(Options, self).__init__(**kwargs) + if check_host_keys not in ('accept', 'enforce', 'ignore'): raise ValueError(self.check_host_keys_msg) @@ -175,143 +215,81 @@ class Stream(mitogen.parent.Stream): if ssh_debug_level: self.ssh_debug_level = ssh_debug_level - self._init_create_child() + +class Connection(mitogen.parent.Connection): + options_class = Options + diag_protocol_class = SetupProtocol + + child_is_immediate_subprocess = False + + def _get_name(self): + s = u'ssh.' + mitogen.core.to_text(self.options.hostname) + if self.options.port and self.options.port != 22: + s += u':%s' % (self.options.port,) + return s def _requires_pty(self): """ - Return :data:`True` if the configuration requires a PTY to be - allocated. This is only true if we must interactively accept host keys, - or type a password. + Return :data:`True` if a PTY to is required for this configuration, + because it must interactively accept host keys or type a password. """ - return (self.check_host_keys == 'accept' or - self.password is not None) + return ( + self.options.check_host_keys == 'accept' or + self.options.password is not None + ) - def _init_create_child(self): + def create_child(self, **kwargs): """ - Initialize the base class :attr:`create_child` and - :attr:`create_child_args` according to whether we need a PTY or not. + Avoid PTY use when possible to avoid a scaling limitation. """ if self._requires_pty(): - self.create_child = mitogen.parent.hybrid_tty_create_child + return mitogen.parent.hybrid_tty_create_child(**kwargs) else: - self.create_child = mitogen.parent.create_child - self.create_child_args = { - 'stderr_pipe': True, - } + return mitogen.parent.create_child(stderr_pipe=True, **kwargs) def get_boot_command(self): - bits = [self.ssh_path] - if self.ssh_debug_level: - bits += ['-' + ('v' * min(3, self.ssh_debug_level))] + bits = [self.options.ssh_path] + if self.options.ssh_debug_level: + bits += ['-' + ('v' * min(3, self.options.ssh_debug_level))] else: # issue #307: suppress any login banner, as it may contain the # password prompt, and there is no robust way to tell the # difference. bits += ['-o', 'LogLevel ERROR'] - if self.username: - bits += ['-l', self.username] - if self.port is not None: - bits += ['-p', str(self.port)] - if self.identities_only and (self.identity_file or self.password): + if self.options.username: + bits += ['-l', self.options.username] + if self.options.port is not None: + bits += ['-p', str(self.options.port)] + if self.options.identities_only and (self.options.identity_file or + self.options.password): bits += ['-o', 'IdentitiesOnly yes'] - if self.identity_file: - bits += ['-i', self.identity_file] - if self.compression: + if self.options.identity_file: + bits += ['-i', self.options.identity_file] + if self.options.compression: bits += ['-o', 'Compression yes'] - if self.keepalive_enabled: + if self.options.keepalive_enabled: bits += [ - '-o', 'ServerAliveInterval %s' % (self.keepalive_interval,), - '-o', 'ServerAliveCountMax %s' % (self.keepalive_count,), + '-o', 'ServerAliveInterval %s' % ( + self.options.keepalive_interval, + ), + '-o', 'ServerAliveCountMax %s' % ( + self.options.keepalive_count, + ), ] if not self._requires_pty(): bits += ['-o', 'BatchMode yes'] - if self.check_host_keys == 'enforce': + if self.options.check_host_keys == 'enforce': bits += ['-o', 'StrictHostKeyChecking yes'] - if self.check_host_keys == 'accept': + if self.options.check_host_keys == 'accept': bits += ['-o', 'StrictHostKeyChecking ask'] - elif self.check_host_keys == 'ignore': + elif self.options.check_host_keys == 'ignore': bits += [ '-o', 'StrictHostKeyChecking no', '-o', 'UserKnownHostsFile /dev/null', '-o', 'GlobalKnownHostsFile /dev/null', ] - if self.ssh_args: - bits += self.ssh_args - bits.append(self.hostname) - base = super(Stream, self).get_boot_command() + if self.options.ssh_args: + bits += self.options.ssh_args + bits.append(self.options.hostname) + base = super(Connection, self).get_boot_command() return bits + [shlex_quote(s).strip() for s in base] - - def _get_name(self): - s = u'ssh.' + mitogen.core.to_text(self.hostname) - if self.port: - s += u':%s' % (self.port,) - return s - - auth_incorrect_msg = 'SSH authentication is incorrect' - password_incorrect_msg = 'SSH password is incorrect' - password_required_msg = 'SSH password was requested, but none specified' - hostkey_config_msg = ( - 'SSH requested permission to accept unknown host key, but ' - 'check_host_keys=ignore. This is likely due to ssh_args= ' - 'conflicting with check_host_keys=. Please correct your ' - 'configuration.' - ) - hostkey_failed_msg = ( - 'Host key checking is enabled, and SSH reported an unrecognized or ' - 'mismatching host key.' - ) - - def _host_key_prompt(self): - if self.check_host_keys == 'accept': - LOG.debug('%s: accepting host key', self.name) - self.diag_stream.transmit_side.write(b('yes\n')) - return - - # _host_key_prompt() should never be reached with ignore or enforce - # mode, SSH should have handled that. User's ssh_args= is conflicting - # with ours. - raise HostKeyError(self.hostkey_config_msg) - - def _connect_input_loop(self, it): - password_sent = False - for buf, partial in filter_debug(self, it): - LOG.debug('%s: stdout: %s', self.name, buf.rstrip()) - if buf.endswith(self.EC0_MARKER): - self._ec0_received() - return - elif HOSTKEY_REQ_PROMPT in buf.lower(): - self._host_key_prompt() - elif HOSTKEY_FAIL in buf.lower(): - raise HostKeyError(self.hostkey_failed_msg) - elif PERMDENIED_RE.match(buf): - # issue #271: work around conflict with user shell reporting - # 'permission denied' e.g. during chdir($HOME) by only matching - # it at the start of the line. - if self.password is not None and password_sent: - raise PasswordError(self.password_incorrect_msg) - elif PASSWORD_PROMPT in buf and self.password is None: - # Permission denied (password,pubkey) - raise PasswordError(self.password_required_msg) - else: - raise PasswordError(self.auth_incorrect_msg) - elif partial and PASSWORD_PROMPT in buf.lower(): - if self.password is None: - raise PasswordError(self.password_required_msg) - LOG.debug('%s: sending password', self.name) - self.diag_stream.transmit_side.write( - (self.password + '\n').encode() - ) - 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 5ff9e177..5e9a237a 100644 --- a/mitogen/su.py +++ b/mitogen/su.py @@ -29,6 +29,7 @@ # !mitogen: minify_safe import logging +import re import mitogen.core import mitogen.parent @@ -42,87 +43,119 @@ except NameError: LOG = logging.getLogger(__name__) +password_incorrect_msg = 'su password is incorrect' +password_required_msg = 'su password is required' + class PasswordError(mitogen.core.StreamError): pass -class Stream(mitogen.parent.Stream): - # TODO: BSD su cannot handle stdin being a socketpair, but it does let the - # child inherit fds from the parent. So we can still pass a socketpair in - # for hybrid_tty_create_child(), there just needs to be either a shell - # snippet or bootstrap support for fixing things up afterwards. - create_child = staticmethod(mitogen.parent.tty_create_child) - child_is_immediate_subprocess = False +class SetupBootstrapProtocol(mitogen.parent.BootstrapProtocol): + password_sent = False + + def setup_patterns(self, conn): + """ + su options cause the regexes used to vary. This is a mess, requires + reworking. + """ + incorrect_pattern = re.compile( + mitogen.core.b('|').join( + re.escape(s.encode('utf-8')) + for s in conn.options.incorrect_prompts + ), + re.I + ) + prompt_pattern = re.compile( + re.escape( + conn.options.password_prompt.encode('utf-8') + ), + re.I + ) + + self.PATTERNS = mitogen.parent.BootstrapProtocol.PATTERNS + [ + (incorrect_pattern, type(self)._on_password_incorrect), + ] + self.PARTIAL_PATTERNS = mitogen.parent.BootstrapProtocol.PARTIAL_PATTERNS + [ + (prompt_pattern, type(self)._on_password_prompt), + ] + + def _on_password_prompt(self, line, match): + LOG.debug('%r: (password prompt): %r', + self.stream.name, line.decode('utf-8', 'replace')) + + if self.stream.conn.options.password is None: + self.stream.conn._fail_connection( + PasswordError(password_required_msg) + ) + return + + if self.password_sent: + self.stream.conn._fail_connection( + PasswordError(password_incorrect_msg) + ) + return + + self.stream.transmit_side.write( + (self.stream.conn.options.password + '\n').encode('utf-8') + ) + self.password_sent = True + + def _on_password_incorrect(self, line, match): + self.stream.conn._fail_connection( + PasswordError(password_incorrect_msg) + ) - #: Once connected, points to the corresponding DiagLogStream, allowing it to - #: be disconnected at the same time this stream is being torn down. - username = 'root' +class Options(mitogen.parent.Options): + username = u'root' password = None su_path = 'su' - password_prompt = b('password:') + password_prompt = u'password:' incorrect_prompts = ( - b('su: sorry'), # BSD - b('su: authentication failure'), # Linux - b('su: incorrect password'), # CentOS 6 - b('authentication is denied'), # AIX + u'su: sorry', # BSD + u'su: authentication failure', # Linux + u'su: incorrect password', # CentOS 6 + u'authentication is denied', # AIX ) - def construct(self, username=None, password=None, su_path=None, - password_prompt=None, incorrect_prompts=None, **kwargs): - super(Stream, self).construct(**kwargs) + def __init__(self, username=None, password=None, su_path=None, + password_prompt=None, incorrect_prompts=None, **kwargs): + super(Options, self).__init__(**kwargs) if username is not None: - self.username = username + self.username = mitogen.core.to_text(username) if password is not None: - self.password = password + self.password = mitogen.core.to_text(password) if su_path is not None: self.su_path = su_path if password_prompt is not None: - self.password_prompt = password_prompt.lower() + self.password_prompt = password_prompt if incorrect_prompts is not None: - self.incorrect_prompts = map(str.lower, incorrect_prompts) + self.incorrect_prompts = [ + mitogen.core.to_text(p) + for p in incorrect_prompts + ] + + +class Connection(mitogen.parent.Connection): + options_class = Options + stream_protocol_class = SetupBootstrapProtocol + + # TODO: BSD su cannot handle stdin being a socketpair, but it does let the + # child inherit fds from the parent. So we can still pass a socketpair in + # for hybrid_tty_create_child(), there just needs to be either a shell + # snippet or bootstrap support for fixing things up afterwards. + create_child = staticmethod(mitogen.parent.tty_create_child) + child_is_immediate_subprocess = False def _get_name(self): - return u'su.' + mitogen.core.to_text(self.username) + return u'su.' + self.options.username + + def stream_factory(self): + stream = super(Connection, self).stream_factory() + stream.protocol.setup_patterns(self) + return stream def get_boot_command(self): - argv = mitogen.parent.Argv(super(Stream, self).get_boot_command()) - return [self.su_path, self.username, '-c', str(argv)] - - password_incorrect_msg = 'su password is incorrect' - password_required_msg = 'su password is required' - - def _connect_input_loop(self, it): - password_sent = False - - for buf in it: - LOG.debug('%r: received %r', self, buf) - if buf.endswith(self.EC0_MARKER): - self._ec0_received() - return - if any(s in buf.lower() for s in self.incorrect_prompts): - if password_sent: - raise PasswordError(self.password_incorrect_msg) - elif self.password_prompt in buf.lower(): - if self.password is None: - raise PasswordError(self.password_required_msg) - if password_sent: - raise PasswordError(self.password_incorrect_msg) - LOG.debug('sending password') - self.transmit_side.write( - mitogen.core.to_text(self.password + '\n').encode('utf-8') - ) - password_sent = True - - raise mitogen.core.StreamError('bootstrap failed') - - def _connect_bootstrap(self): - it = mitogen.parent.iter_read( - fds=[self.receive_side.fd], - deadline=self.connect_deadline, - ) - try: - self._connect_input_loop(it) - finally: - it.close() + argv = mitogen.parent.Argv(super(Connection, self).get_boot_command()) + return [self.options.su_path, self.options.username, '-c', str(argv)] diff --git a/mitogen/sudo.py b/mitogen/sudo.py index 868d4d76..725e6aff 100644 --- a/mitogen/sudo.py +++ b/mitogen/sudo.py @@ -40,6 +40,9 @@ from mitogen.core import b LOG = logging.getLogger(__name__) +password_incorrect_msg = 'sudo password is incorrect' +password_required_msg = 'sudo password is required' + # These are base64-encoded UTF-8 as our existing minifier/module server # struggles with Unicode Python source in some (forgotten) circumstances. PASSWORD_PROMPTS = [ @@ -99,14 +102,13 @@ PASSWORD_PROMPTS = [ PASSWORD_PROMPT_RE = re.compile( - u'|'.join( - base64.b64decode(s).decode('utf-8') + mitogen.core.b('|').join( + base64.b64decode(s) for s in PASSWORD_PROMPTS - ) + ), + re.I ) - -PASSWORD_PROMPT = b('password') SUDO_OPTIONS = [ #(False, 'bool', '--askpass', '-A') #(False, 'str', '--auth-type', '-a') @@ -181,10 +183,7 @@ def option(default, *args): return default -class Stream(mitogen.parent.Stream): - create_child = staticmethod(mitogen.parent.hybrid_tty_create_child) - child_is_immediate_subprocess = False - +class Options(mitogen.parent.Options): sudo_path = 'sudo' username = 'root' password = None @@ -195,15 +194,16 @@ class Stream(mitogen.parent.Stream): selinux_role = None selinux_type = None - def construct(self, username=None, sudo_path=None, password=None, - preserve_env=None, set_home=None, sudo_args=None, - login=None, selinux_role=None, selinux_type=None, **kwargs): - super(Stream, self).construct(**kwargs) + def __init__(self, username=None, sudo_path=None, password=None, + preserve_env=None, set_home=None, sudo_args=None, + login=None, selinux_role=None, selinux_type=None, **kwargs): + super(Options, self).__init__(**kwargs) opts = parse_sudo_flags(sudo_args or []) self.username = option(self.username, username, opts.user) self.sudo_path = option(self.sudo_path, sudo_path) - self.password = password or None + if password: + self.password = mitogen.core.to_text(password) self.preserve_env = option(self.preserve_env, preserve_env, opts.preserve_env) self.set_home = option(self.set_home, set_home, opts.set_home) @@ -211,67 +211,59 @@ class Stream(mitogen.parent.Stream): self.selinux_role = option(self.selinux_role, selinux_role, opts.role) self.selinux_type = option(self.selinux_type, selinux_type, opts.type) + +class SetupProtocol(mitogen.parent.RegexProtocol): + password_sent = False + + def _on_password_prompt(self, line, match): + LOG.debug('%s: (password prompt): %s', + self.stream.name, line.decode('utf-8', 'replace')) + + if self.stream.conn.options.password is None: + self.stream.conn._fail_connection( + PasswordError(password_required_msg) + ) + return + + if self.password_sent: + self.stream.conn._fail_connection( + PasswordError(password_incorrect_msg) + ) + return + + self.stream.transmit_side.write( + (self.stream.conn.options.password + '\n').encode('utf-8') + ) + self.password_sent = True + + PARTIAL_PATTERNS = [ + (PASSWORD_PROMPT_RE, _on_password_prompt), + ] + + +class Connection(mitogen.parent.Connection): + diag_protocol_class = SetupProtocol + options_class = Options + create_child = staticmethod(mitogen.parent.hybrid_tty_create_child) + child_is_immediate_subprocess = False + def _get_name(self): - return u'sudo.' + mitogen.core.to_text(self.username) + return u'sudo.' + mitogen.core.to_text(self.options.username) def get_boot_command(self): # Note: sudo did not introduce long-format option processing until July # 2013, so even though we parse long-format options, supply short-form # to the sudo command. - bits = [self.sudo_path, '-u', self.username] - if self.preserve_env: + bits = [self.options.sudo_path, '-u', self.options.username] + if self.options.preserve_env: bits += ['-E'] - if self.set_home: + if self.options.set_home: bits += ['-H'] - if self.login: + if self.options.login: bits += ['-i'] - if self.selinux_role: - bits += ['-r', self.selinux_role] - if self.selinux_type: - bits += ['-t', self.selinux_type] - - bits = bits + ['--'] + super(Stream, self).get_boot_command() - LOG.debug('sudo command line: %r', bits) - return bits - - password_incorrect_msg = 'sudo password is incorrect' - password_required_msg = 'sudo password is required' - - def _connect_input_loop(self, it): - password_sent = False - - for buf in it: - LOG.debug('%s: received %r', self.name, buf) - if buf.endswith(self.EC0_MARKER): - self._ec0_received() - return - - match = PASSWORD_PROMPT_RE.search(buf.decode('utf-8').lower()) - if match is not None: - LOG.debug('%s: matched password prompt %r', - self.name, match.group(0)) - if self.password is None: - raise PasswordError(self.password_required_msg) - if password_sent: - raise PasswordError(self.password_incorrect_msg) - self.diag_stream.transmit_side.write( - (mitogen.core.to_text(self.password) + '\n').encode('utf-8') - ) - 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, - ) + if self.options.selinux_role: + bits += ['-r', self.options.selinux_role] + if self.options.selinux_type: + bits += ['-t', self.options.selinux_type] - try: - self._connect_input_loop(it) - finally: - it.close() + return bits + ['--'] + super(Connection, self).get_boot_command() diff --git a/mitogen/unix.py b/mitogen/unix.py index 66141eec..c0d2bb9c 100644 --- a/mitogen/unix.py +++ b/mitogen/unix.py @@ -65,9 +65,38 @@ def make_socket_path(): return tempfile.mktemp(prefix='mitogen_unix_', suffix='.sock') -class Listener(mitogen.core.BasicStream): +class ListenerStream(mitogen.core.Stream): + def on_receive(self, broker): + sock, _ = self.receive_side.fp.accept() + try: + self.protocol.on_accept_client(sock) + except: + sock.close() + raise + + +class Listener(mitogen.core.Protocol): + stream_class = ListenerStream keep_alive = True + @classmethod + def build_stream(cls, router, path=None, backlog=100): + if not path: + path = make_socket_path() + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + if os.path.exists(path) and is_path_dead(path): + LOG.debug('%r: deleting stale %r', cls.__name__, path) + os.unlink(path) + + sock.bind(path) + os.chmod(path, int('0600', 8)) + sock.listen(backlog) + + stream = super(Listener, cls).build_stream(router, path) + stream.accept(sock, sock) + router.broker.start_receive(stream) + return stream + def __repr__(self): return '%s.%s(%r)' % ( __name__, @@ -75,20 +104,9 @@ class Listener(mitogen.core.BasicStream): self.path, ) - def __init__(self, router, path=None, backlog=100): + def __init__(self, router, path): self._router = router - self.path = path or make_socket_path() - self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - - if os.path.exists(self.path) and is_path_dead(self.path): - LOG.debug('%r: deleting stale %r', self, self.path) - os.unlink(self.path) - - self._sock.bind(self.path) - os.chmod(self.path, int('0600', 8)) - self._sock.listen(backlog) - self.receive_side = mitogen.core.Side(self, self._sock.fileno()) - router.broker.start_receive(self) + self.path = path def _unlink_socket(self): try: @@ -100,12 +118,11 @@ class Listener(mitogen.core.BasicStream): raise def on_shutdown(self, broker): - broker.stop_receive(self) + broker.stop_receive(self.stream) self._unlink_socket() - self._sock.close() - self.receive_side.closed = True + self.stream.receive_side.close() - def _accept_client(self, sock): + def on_accept_client(self, sock): sock.setblocking(True) try: pid, = struct.unpack('>L', sock.recv(4)) @@ -115,12 +132,6 @@ class Listener(mitogen.core.BasicStream): return context_id = self._router.id_allocator.allocate() - context = mitogen.parent.Context(self._router, context_id) - stream = mitogen.core.Stream(self._router, context_id) - stream.name = u'unix_client.%d' % (pid,) - stream.auth_id = mitogen.context_id - stream.is_privileged = True - try: sock.send(struct.pack('>LLL', context_id, mitogen.context_id, os.getpid())) @@ -129,21 +140,20 @@ class Listener(mitogen.core.BasicStream): self, pid, sys.exc_info()[1]) return + context = mitogen.parent.Context(self._router, context_id) + stream = mitogen.core.MitogenProtocol.build_stream( + router=self._router, + remote_id=context_id, + ) + stream.name = u'unix_client.%d' % (pid,) + stream.protocol.auth_id = mitogen.context_id + stream.protocol.is_privileged = True + stream.accept(sock, sock) LOG.debug('%r: accepted %r', self, stream) - stream.accept(sock.fileno(), sock.fileno()) self._router.register(context, stream) - def on_receive(self, broker): - sock, _ = self._sock.accept() - try: - self._accept_client(sock) - finally: - sock.close() - -def connect(path, broker=None): - LOG.debug('unix.connect(path=%r)', path) - sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) +def _connect(path, broker, sock): sock.connect(path) sock.send(struct.pack('>L', os.getpid())) mitogen.context_id, remote_id, pid = struct.unpack('>LLL', sock.recv(12)) @@ -154,15 +164,35 @@ def connect(path, broker=None): mitogen.context_id, remote_id) router = mitogen.master.Router(broker=broker) - stream = mitogen.core.Stream(router, remote_id) - stream.accept(sock.fileno(), sock.fileno()) + stream = mitogen.core.MitogenProtocol.build_stream(router, remote_id) + stream.accept(sock, sock) stream.name = u'unix_listener.%d' % (pid,) + mitogen.core.listen(stream, 'disconnect', _cleanup) + mitogen.core.listen(router.broker, 'shutdown', + lambda: router.disconnect_stream(stream)) + context = mitogen.parent.Context(router, remote_id) router.register(context, stream) + return router, context - mitogen.core.listen(router.broker, 'shutdown', - lambda: router.disconnect_stream(stream)) - sock.close() - return router, context +def connect(path, broker=None): + LOG.debug('unix.connect(path=%r)', path) + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + return _connect(path, broker, sock) + except: + sock.close() + raise + + +def _cleanup(): + """ + Reset mitogen.context_id and friends when our connection to the parent is + lost. Per comments on #91, these globals need to move to the Router so + fix-ups like this become unnecessary. + """ + mitogen.context_id = 0 + mitogen.parent_id = None + mitogen.parent_ids = [] diff --git a/preamble_size.py b/preamble_size.py index f5f1adc1..692ad7b1 100644 --- a/preamble_size.py +++ b/preamble_size.py @@ -19,15 +19,17 @@ import mitogen.sudo router = mitogen.master.Router() context = mitogen.parent.Context(router, 0) -stream = mitogen.ssh.Stream(router, 0, max_message_size=0, hostname='foo') +options = mitogen.ssh.Options(max_message_size=0, hostname='foo') +conn = mitogen.ssh.Connection(options, router) +conn.context = context -print('SSH command size: %s' % (len(' '.join(stream.get_boot_command())),)) +print('SSH command size: %s' % (len(' '.join(conn.get_boot_command())),)) print('Preamble size: %s (%.2fKiB)' % ( - len(stream.get_preamble()), - len(stream.get_preamble()) / 1024.0, + len(conn.get_preamble()), + len(conn.get_preamble()) / 1024.0, )) if '--dump' in sys.argv: - print(zlib.decompress(stream.get_preamble())) + print(zlib.decompress(conn.get_preamble())) exit() diff --git a/scripts/release-notes.py b/scripts/release-notes.py new file mode 100644 index 00000000..1444d7a3 --- /dev/null +++ b/scripts/release-notes.py @@ -0,0 +1,48 @@ +# coding=UTF-8 + +# Generate the fragment used to make email release announcements +# usage: release-notes.py 0.2.6 + +import os +import sys +import urllib +import lxml.html + +import subprocess + + +response = urllib.urlopen('https://mitogen.networkgenomics.com/changelog.html') +tree = lxml.html.parse(response) + +prefix = 'v' + sys.argv[1].replace('.', '-') + +for elem in tree.getroot().cssselect('div.section[id]'): + if elem.attrib['id'].startswith(prefix): + break +else: + print('cant find') + + + +for child in tree.getroot().cssselect('body > *'): + child.getparent().remove(child) + +body, = tree.getroot().cssselect('body') +body.append(elem) + +proc = subprocess.Popen( + args=['w3m', '-T', 'text/html', '-dump', '-cols', '72'], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, +) + +stdout, _ = proc.communicate(input=(lxml.html.tostring(tree))) +stdout = stdout.decode('UTF-8') +stdout = stdout.translate({ + ord(u'¶'): None, + ord(u'•'): ord(u'*'), + ord(u'’'): ord(u"'"), + ord(u'“'): ord(u'"'), + ord(u'”'): ord(u'"'), +}) +print(stdout) diff --git a/tests/ansible/tests/affinity_test.py b/tests/ansible/tests/affinity_test.py index 8ee96085..641455bd 100644 --- a/tests/ansible/tests/affinity_test.py +++ b/tests/ansible/tests/affinity_test.py @@ -199,10 +199,10 @@ class LinuxPolicyTest(testlib.TestCase): self.policy._set_cpu(3) my_cpu = self._get_cpus() - pid = mitogen.parent.detach_popen( + proc = mitogen.parent.popen( args=['cp', '/proc/self/status', tf.name] ) - os.waitpid(pid, 0) + proc.wait() his_cpu = self._get_cpus(tf.name) self.assertNotEquals(my_cpu, his_cpu) diff --git a/tests/ansible/tests/connection_test.py b/tests/ansible/tests/connection_test.py index 401cbe9e..d663ecc5 100644 --- a/tests/ansible/tests/connection_test.py +++ b/tests/ansible/tests/connection_test.py @@ -13,42 +13,29 @@ import ansible.errors import ansible.playbook.play_context import mitogen.core +import mitogen.utils + import ansible_mitogen.connection import ansible_mitogen.plugins.connection.mitogen_local import ansible_mitogen.process -import testlib - -LOGGER_NAME = ansible_mitogen.target.LOG.name - - -# TODO: fixtureize -import mitogen.utils -mitogen.utils.log_to_file() -ansible_mitogen.process.MuxProcess.start(_init_logging=False) - - -class OptionalIntTest(unittest2.TestCase): - func = staticmethod(ansible_mitogen.connection.optional_int) - - def test_already_int(self): - self.assertEquals(0, self.func(0)) - self.assertEquals(1, self.func(1)) - self.assertEquals(-1, self.func(-1)) +import testlib - def test_is_string(self): - self.assertEquals(0, self.func("0")) - self.assertEquals(1, self.func("1")) - self.assertEquals(-1, self.func("-1")) - def test_is_none(self): - self.assertEquals(None, self.func(None)) +class MuxProcessMixin(object): + @classmethod + def setUpClass(cls): + #mitogen.utils.log_to_file() + ansible_mitogen.process.MuxProcess.start(_init_logging=False) + super(MuxProcessMixin, cls).setUpClass() - def test_is_junk(self): - self.assertEquals(None, self.func({1:2})) + @classmethod + def tearDownClass(cls): + super(MuxProcessMixin, cls).tearDownClass() + ansible_mitogen.process.MuxProcess._reset() -class ConnectionMixin(object): +class ConnectionMixin(MuxProcessMixin): klass = ansible_mitogen.plugins.connection.mitogen_local.Connection def make_connection(self): @@ -70,6 +57,26 @@ class ConnectionMixin(object): super(ConnectionMixin, self).tearDown() +class OptionalIntTest(unittest2.TestCase): + func = staticmethod(ansible_mitogen.connection.optional_int) + + def test_already_int(self): + self.assertEquals(0, self.func(0)) + self.assertEquals(1, self.func(1)) + self.assertEquals(-1, self.func(-1)) + + def test_is_string(self): + self.assertEquals(0, self.func("0")) + self.assertEquals(1, self.func("1")) + self.assertEquals(-1, self.func("-1")) + + def test_is_none(self): + self.assertEquals(None, self.func(None)) + + def test_is_junk(self): + self.assertEquals(None, self.func({1:2})) + + class PutDataTest(ConnectionMixin, unittest2.TestCase): def test_out_path(self): path = tempfile.mktemp(prefix='mitotest') diff --git a/tests/bench/ssh-roundtrip.py b/tests/bench/ssh-roundtrip.py new file mode 100644 index 00000000..8745505d --- /dev/null +++ b/tests/bench/ssh-roundtrip.py @@ -0,0 +1,35 @@ +""" +Measure latency of SSH RPC. +""" + +import sys +import time + +import mitogen +import mitogen.utils +import ansible_mitogen.affinity + +mitogen.utils.setup_gil() +ansible_mitogen.affinity.policy.assign_worker() + +try: + xrange +except NameError: + xrange = range + +def do_nothing(): + pass + +@mitogen.main() +def main(router): + f = router.ssh(hostname=sys.argv[1]) + f.call(do_nothing) + t0 = time.time() + end = time.time() + 5.0 + i = 0 + while time.time() < end: + f.call(do_nothing) + i += 1 + t1 = time.time() + + print('++', float(1e3 * (t1 - t0) / (1.0+i)), 'ms') diff --git a/tests/broker_test.py b/tests/broker_test.py index 23839a54..2212d8aa 100644 --- a/tests/broker_test.py +++ b/tests/broker_test.py @@ -1,4 +1,5 @@ +import time import threading import mock diff --git a/tests/buildah_test.py b/tests/buildah_test.py index dad2534f..874205cd 100644 --- a/tests/buildah_test.py +++ b/tests/buildah_test.py @@ -21,7 +21,7 @@ class ConstructorTest(testlib.RouterMixin, testlib.TestCase): self.assertEquals(argv[1], 'run') self.assertEquals(argv[2], '--') self.assertEquals(argv[3], 'container_name') - self.assertEquals(argv[4], stream.python_path) + self.assertEquals(argv[4], stream.conn.options.python_path) if __name__ == '__main__': diff --git a/tests/serialization_test.py b/tests/context_test.py similarity index 50% rename from tests/serialization_test.py rename to tests/context_test.py index 6cf5f8b7..4bc4bd2e 100644 --- a/tests/serialization_test.py +++ b/tests/context_test.py @@ -8,40 +8,7 @@ from mitogen.core import b import testlib -class EvilObject(object): - pass - - -def roundtrip(v): - msg = mitogen.core.Message.pickled(v) - return mitogen.core.Message(data=msg.data).unpickle() - - -class EvilObjectTest(testlib.TestCase): - def test_deserialization_fails(self): - msg = mitogen.core.Message.pickled(EvilObject()) - e = self.assertRaises(mitogen.core.StreamError, - lambda: msg.unpickle() - ) - - -class BlobTest(testlib.TestCase): - klass = mitogen.core.Blob - - # Python 3 pickle protocol 2 does weird stuff depending on whether an empty - # or nonempty bytes is being serialized. For non-empty, it yields a - # _codecs.encode() call. For empty, it yields a bytes() call. - - def test_nonempty_bytes(self): - v = mitogen.core.Blob(b('dave')) - self.assertEquals(b('dave'), roundtrip(v)) - - def test_empty_bytes(self): - v = mitogen.core.Blob(b('')) - self.assertEquals(b(''), roundtrip(v)) - - -class ContextTest(testlib.RouterMixin, testlib.TestCase): +class PickleTest(testlib.RouterMixin, testlib.TestCase): klass = mitogen.core.Context # Ensure Context can be round-tripped by regular pickle in addition to diff --git a/tests/create_child_test.py b/tests/create_child_test.py new file mode 100644 index 00000000..21591fe8 --- /dev/null +++ b/tests/create_child_test.py @@ -0,0 +1,284 @@ + +import fcntl +import os +import stat +import sys +import time +import tempfile + +import mock +import unittest2 + +import mitogen.parent +from mitogen.core import b + +import testlib + + +def run_fd_check(func, fd, mode, on_start=None): + tf = tempfile.NamedTemporaryFile() + args = [ + sys.executable, + testlib.data_path('fd_check.py'), + tf.name, + str(fd), + mode, + ] + + proc = func(args=args) + os = None + if on_start: + os = on_start(proc) + proc.proc.wait() + try: + return proc, eval(tf.read()), os + finally: + tf.close() + + +def close_proc(proc): + proc.stdin.close() + proc.stdout.close() + if proc.stderr: + prco.stderr.close() + + +def wait_read(fp, n): + poller = mitogen.core.Poller() + try: + poller.start_receive(fp.fileno()) + for _ in poller.poll(): + return os.read(fp.fileno(), n) + assert False + finally: + poller.close() + + +class StdinSockMixin(object): + def test_stdin(self): + proc, info, _ = run_fd_check(self.func, 0, 'read', + lambda proc: proc.stdin.send(b('TEST'))) + st = os.fstat(proc.stdin.fileno()) + self.assertTrue(stat.S_ISSOCK(st.st_mode)) + self.assertEquals(st.st_dev, info['st_dev']) + self.assertEquals(st.st_mode, info['st_mode']) + flags = fcntl.fcntl(proc.stdin.fileno(), fcntl.F_GETFL) + self.assertTrue(flags & os.O_RDWR) + self.assertTrue(info['buf'], 'TEST') + self.assertTrue(info['flags'] & os.O_RDWR) + + +class StdoutSockMixin(object): + def test_stdout(self): + proc, info, buf = run_fd_check(self.func, 1, 'write', + lambda proc: wait_read(proc.stdout, 4)) + st = os.fstat(proc.stdout.fileno()) + self.assertTrue(stat.S_ISSOCK(st.st_mode)) + self.assertEquals(st.st_dev, info['st_dev']) + self.assertEquals(st.st_mode, info['st_mode']) + flags = fcntl.fcntl(proc.stdout.fileno(), fcntl.F_GETFL) + self.assertTrue(flags & os.O_RDWR) + self.assertTrue(buf, 'TEST') + self.assertTrue(info['flags'] & os.O_RDWR) + + +class CreateChildTest(StdinSockMixin, StdoutSockMixin, testlib.TestCase): + func = staticmethod(mitogen.parent.create_child) + + def test_stderr(self): + proc, info, _ = run_fd_check(self.func, 2, 'write') + st = os.fstat(sys.stderr.fileno()) + self.assertEquals(st.st_dev, info['st_dev']) + self.assertEquals(st.st_mode, info['st_mode']) + self.assertEquals(st.st_ino, info['st_ino']) + + +class CreateChildMergedTest(StdinSockMixin, StdoutSockMixin, + testlib.TestCase): + def func(self, *args, **kwargs): + kwargs['merge_stdio'] = True + return mitogen.parent.create_child(*args, **kwargs) + + def test_stderr(self): + proc, info, buf = run_fd_check(self.func, 2, 'write', + lambda proc: wait_read(proc.stdout, 4)) + self.assertEquals(None, proc.stderr) + st = os.fstat(proc.stdout.fileno()) + self.assertTrue(stat.S_ISSOCK(st.st_mode)) + self.assertEquals(st.st_dev, info['st_dev']) + self.assertEquals(st.st_mode, info['st_mode']) + flags = fcntl.fcntl(proc.stdout.fileno(), fcntl.F_GETFL) + self.assertTrue(flags & os.O_RDWR) + self.assertTrue(buf, 'TEST') + self.assertTrue(info['flags'] & os.O_RDWR) + + +class CreateChildStderrPipeTest(StdinSockMixin, StdoutSockMixin, + testlib.TestCase): + def func(self, *args, **kwargs): + kwargs['stderr_pipe'] = True + return mitogen.parent.create_child(*args, **kwargs) + + def test_stderr(self): + proc, info, buf = run_fd_check(self.func, 2, 'write', + lambda proc: wait_read(proc.stderr, 4)) + st = os.fstat(proc.stderr.fileno()) + self.assertTrue(stat.S_ISFIFO(st.st_mode)) + self.assertEquals(st.st_dev, info['st_dev']) + self.assertEquals(st.st_mode, info['st_mode']) + flags = fcntl.fcntl(proc.stderr.fileno(), fcntl.F_GETFL) + self.assertFalse(flags & os.O_WRONLY) + self.assertFalse(flags & os.O_RDWR) + self.assertTrue(buf, 'TEST') + self.assertTrue(info['flags'] & os.O_WRONLY) + + +class TtyCreateChildTest(testlib.TestCase): + func = staticmethod(mitogen.parent.tty_create_child) + + def test_stdin(self): + proc, info, _ = run_fd_check(self.func, 0, 'read', + lambda proc: proc.stdin.write(b('TEST'))) + st = os.fstat(proc.stdin.fileno()) + self.assertTrue(stat.S_ISCHR(st.st_mode)) + self.assertTrue(stat.S_ISCHR(info['st_mode'])) + + self.assertTrue(isinstance(info['ttyname'], + mitogen.core.UnicodeType)) + os.ttyname(proc.stdin.fileno()) # crashes if not TTY + + flags = fcntl.fcntl(proc.stdin.fileno(), fcntl.F_GETFL) + self.assertTrue(flags & os.O_RDWR) + self.assertTrue(info['flags'] & os.O_RDWR) + + self.assertNotEquals(st.st_dev, info['st_dev']) + self.assertTrue(info['buf'], 'TEST') + + def test_stdout(self): + proc, info, buf = run_fd_check(self.func, 1, 'write', + lambda proc: wait_read(proc.stdout, 4)) + + st = os.fstat(proc.stdout.fileno()) + self.assertTrue(stat.S_ISCHR(st.st_mode)) + self.assertTrue(stat.S_ISCHR(info['st_mode'])) + + self.assertTrue(isinstance(info['ttyname'], + mitogen.core.UnicodeType)) + os.ttyname(proc.stdout.fileno()) # crashes if wrong + + flags = fcntl.fcntl(proc.stdout.fileno(), fcntl.F_GETFL) + self.assertTrue(flags & os.O_RDWR) + self.assertTrue(info['flags'] & os.O_RDWR) + + self.assertNotEquals(st.st_dev, info['st_dev']) + self.assertTrue(flags & os.O_RDWR) + self.assertTrue(buf, 'TEST') + + def test_stderr(self): + proc, info, buf = run_fd_check(self.func, 2, 'write', + lambda proc: wait_read(proc.stdout, 4)) + + st = os.fstat(proc.stdout.fileno()) + self.assertTrue(stat.S_ISCHR(st.st_mode)) + self.assertTrue(stat.S_ISCHR(info['st_mode'])) + + self.assertTrue(isinstance(info['ttyname'], + mitogen.core.UnicodeType)) + os.ttyname(proc.stdin.fileno()) # crashes if not TTY + + flags = fcntl.fcntl(proc.stdout.fileno(), fcntl.F_GETFL) + self.assertTrue(flags & os.O_RDWR) + self.assertTrue(info['flags'] & os.O_RDWR) + + self.assertNotEquals(st.st_dev, info['st_dev']) + self.assertTrue(flags & os.O_RDWR) + self.assertTrue(buf, 'TEST') + + def test_dev_tty_open_succeeds(self): + # In the early days of UNIX, a process that lacked a controlling TTY + # would acquire one simply by opening an existing TTY. Linux and OS X + # continue to follow this behaviour, however at least FreeBSD moved to + # requiring an explicit ioctl(). Linux supports it, but we don't yet + # use it there and anyway the behaviour will never change, so no point + # in fixing things that aren't broken. Below we test that + # getpass-loving apps like sudo and ssh get our slave PTY when they + # attempt to open /dev/tty, which is what they both do on attempting to + # read a password. + tf = tempfile.NamedTemporaryFile() + try: + proc = self.func([ + 'bash', '-c', 'exec 2>%s; echo hi > /dev/tty' % (tf.name,) + ]) + deadline = time.time() + 5.0 + self.assertEquals(mitogen.core.b('hi\n'), wait_read(proc.stdout, 3)) + waited_pid, status = os.waitpid(proc.pid, 0) + self.assertEquals(proc.pid, waited_pid) + self.assertEquals(0, status) + self.assertEquals(mitogen.core.b(''), tf.read()) + proc.stdout.close() + finally: + tf.close() + + +class StderrDiagTtyMixin(object): + def test_stderr(self): + proc, info, buf = run_fd_check(self.func, 2, 'write', + lambda proc: wait_read(proc.stderr, 4)) + + st = os.fstat(proc.stderr.fileno()) + self.assertTrue(stat.S_ISCHR(st.st_mode)) + self.assertTrue(stat.S_ISCHR(info['st_mode'])) + + self.assertTrue(isinstance(info['ttyname'], + mitogen.core.UnicodeType)) + os.ttyname(proc.stderr.fileno()) # crashes if wrong + + flags = fcntl.fcntl(proc.stderr.fileno(), fcntl.F_GETFL) + self.assertTrue(flags & os.O_RDWR) + self.assertTrue(info['flags'] & os.O_RDWR) + + self.assertNotEquals(st.st_dev, info['st_dev']) + self.assertTrue(flags & os.O_RDWR) + self.assertTrue(buf, 'TEST') + + +class HybridTtyCreateChildTest(StdinSockMixin, StdoutSockMixin, + StderrDiagTtyMixin, testlib.TestCase): + func = staticmethod(mitogen.parent.hybrid_tty_create_child) + + + +if 0: + # issue #410 + class SelinuxHybridTtyCreateChildTest(StderrDiagTtyMixin, testlib.TestCase): + func = staticmethod(mitogen.parent.selinux_hybrid_tty_create_child) + + def test_stdin(self): + proc, info, buf = run_fd_check(self.func, 0, 'read', + lambda proc: proc.transmit_side.write('TEST')) + st = os.fstat(proc.transmit_side.fd) + self.assertTrue(stat.S_ISFIFO(st.st_mode)) + self.assertEquals(st.st_dev, info['st_dev']) + self.assertEquals(st.st_mode, info['st_mode']) + flags = fcntl.fcntl(proc.transmit_side.fd, fcntl.F_GETFL) + self.assertTrue(flags & os.O_WRONLY) + self.assertTrue(buf, 'TEST') + self.assertFalse(info['flags'] & os.O_WRONLY) + self.assertFalse(info['flags'] & os.O_RDWR) + + def test_stdout(self): + proc, info, buf = run_fd_check(self.func, 1, 'write', + lambda proc: wait_read(proc.receive_side, 4)) + st = os.fstat(proc.receive_side.fd) + self.assertTrue(stat.S_ISFIFO(st.st_mode)) + self.assertEquals(st.st_dev, info['st_dev']) + self.assertEquals(st.st_mode, info['st_mode']) + flags = fcntl.fcntl(proc.receive_side.fd, fcntl.F_GETFL) + self.assertFalse(flags & os.O_WRONLY) + self.assertFalse(flags & os.O_RDWR) + self.assertTrue(info['flags'] & os.O_WRONLY) + self.assertTrue(buf, 'TEST') + + +if __name__ == '__main__': + unittest2.main() diff --git a/tests/data/docker/README.md b/tests/data/docker/README.md new file mode 100644 index 00000000..d3d37d52 --- /dev/null +++ b/tests/data/docker/README.md @@ -0,0 +1,7 @@ + + +# doas-debian.tar.gz + +A dynamically linked copy of the OpenBSD ``doas`` tool for Debian, port is from +https://github.com/multiplexd/doas (the slicer69 port is broken, it reads the +password from stdin). diff --git a/tests/data/docker/doas-debian.tar.gz b/tests/data/docker/doas-debian.tar.gz new file mode 100644 index 00000000..2deb72ff Binary files /dev/null and b/tests/data/docker/doas-debian.tar.gz differ diff --git a/tests/data/docker/mitogen__permdenied.profile b/tests/data/docker/mitogen__permdenied.profile new file mode 100644 index 00000000..4a2be07e --- /dev/null +++ b/tests/data/docker/mitogen__permdenied.profile @@ -0,0 +1,4 @@ + +mkdir -p bad +chmod 0 bad +cd bad diff --git a/tests/data/docker/ssh_login_banner.txt b/tests/data/docker/ssh_login_banner.txt index 1ae4cd03..8a03fbe4 100644 --- a/tests/data/docker/ssh_login_banner.txt +++ b/tests/data/docker/ssh_login_banner.txt @@ -19,3 +19,5 @@ incidents to law enforcement officials. ************************************************************** NOTE: This system is connected to DOMAIN.COM, please use your password. + +ستتم محاكمة المعتدين. هذا يختبر التدويل diff --git a/tests/data/fd_check.py b/tests/data/fd_check.py new file mode 100755 index 00000000..0a87a95e --- /dev/null +++ b/tests/data/fd_check.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python + +import fcntl +import os +import sys + + +def ttyname(fd): + try: + t = os.ttyname(fd) + if hasattr(t, 'decode'): + t = t.decode() + return t + except OSError: + return None + + +def controlling_tty(): + try: + fp = open('/dev/tty') + try: + return ttyname(fp.fileno()) + finally: + fp.close() + except (IOError, OSError): + return None + + +fd = int(sys.argv[2]) +st = os.fstat(fd) + +if sys.argv[3] == 'write': + os.write(fd, u'TEST'.encode()) + buf = u'' +else: + buf = os.read(fd, 4).decode() + +open(sys.argv[1], 'w').write(repr({ + 'buf': buf, + 'flags': fcntl.fcntl(fd, fcntl.F_GETFL), + 'st_mode': st.st_mode, + 'st_dev': st.st_dev, + 'st_ino': st.st_ino, + 'ttyname': ttyname(fd), + 'controlling_tty': controlling_tty(), +})) diff --git a/tests/data/iter_read_generator.py b/tests/data/iter_read_generator.py deleted file mode 100755 index 3fd3c08c..00000000 --- a/tests/data/iter_read_generator.py +++ /dev/null @@ -1,13 +0,0 @@ -#!/usr/bin/env python -# I produce text every 100ms, for testing mitogen.core.iter_read() - -import sys -import time - - -i = 0 -while True: - i += 1 - sys.stdout.write(str(i)) - sys.stdout.flush() - time.sleep(0.1) diff --git a/tests/data/stubs/stub-jexec.py b/tests/data/stubs/stub-jexec.py new file mode 100755 index 00000000..3f3e3bdc --- /dev/null +++ b/tests/data/stubs/stub-jexec.py @@ -0,0 +1,15 @@ +#!/usr/bin/env python + +import json +import os +import subprocess +import sys + +os.environ['ORIGINAL_ARGV'] = json.dumps(sys.argv) +os.environ['THIS_IS_STUB_JEXEC'] = '1' + +# This must be a child process and not exec() since Mitogen replaces its stderr +# descriptor, causing the last user of the slave PTY to close it, resulting in +# the master side indicating EIO. +subprocess.call(sys.argv[sys.argv.index('somejail') + 1:]) +os._exit(0) diff --git a/tests/data/stubs/stub-su.py b/tests/data/stubs/stub-su.py index c32c91de..1f5e512d 100755 --- a/tests/data/stubs/stub-su.py +++ b/tests/data/stubs/stub-su.py @@ -4,6 +4,16 @@ import json import os import subprocess import sys +import time + +# #363: old input loop would fail to spot auth failure because of scheduling +# vs. su calling write() twice. +if 'DO_SLOW_AUTH_FAILURE' in os.environ: + os.write(2, u'su: '.encode()) + time.sleep(0.5) + os.write(2, u'incorrect password\n'.encode()) + os._exit(1) + os.environ['ORIGINAL_ARGV'] = json.dumps(sys.argv) os.environ['THIS_IS_STUB_SU'] = '1' diff --git a/tests/data/write_all_consumer.py b/tests/data/write_all_consumer.py deleted file mode 100755 index 4013ccdd..00000000 --- a/tests/data/write_all_consumer.py +++ /dev/null @@ -1,9 +0,0 @@ -#!/usr/bin/env python -# I consume 65535 bytes every 10ms, for testing mitogen.core.write_all() - -import os -import time - -while True: - os.read(0, 65535) - time.sleep(0.01) diff --git a/tests/doas_test.py b/tests/doas_test.py index 0e27c2ab..560ada99 100644 --- a/tests/doas_test.py +++ b/tests/doas_test.py @@ -2,6 +2,7 @@ import os import mitogen +import mitogen.doas import mitogen.parent import unittest2 @@ -27,5 +28,38 @@ class ConstructorTest(testlib.RouterMixin, testlib.TestCase): self.assertEquals('1', context.call(os.getenv, 'THIS_IS_STUB_DOAS')) +class DoasTest(testlib.DockerMixin, testlib.TestCase): + # Only mitogen/debian-test has doas. + mitogen_test_distro = 'debian' + + def test_password_required(self): + ssh = self.docker_ssh( + username='mitogen__has_sudo', + password='has_sudo_password', + ) + e = self.assertRaises(mitogen.core.StreamError, + lambda: self.router.doas(via=ssh) + ) + self.assertTrue(mitogen.doas.password_required_msg in str(e)) + + def test_password_incorrect(self): + ssh = self.docker_ssh( + username='mitogen__has_sudo', + password='has_sudo_password', + ) + e = self.assertRaises(mitogen.core.StreamError, + lambda: self.router.doas(via=ssh, password='x') + ) + self.assertTrue(mitogen.doas.password_incorrect_msg in str(e)) + + def test_password_okay(self): + ssh = self.docker_ssh( + username='mitogen__has_sudo', + password='has_sudo_password', + ) + context = self.router.su(via=ssh, password='rootpassword') + self.assertEquals(0, context.call(os.getuid)) + + if __name__ == '__main__': unittest2.main() diff --git a/tests/docker_test.py b/tests/docker_test.py index 49c742ee..b5d15707 100644 --- a/tests/docker_test.py +++ b/tests/docker_test.py @@ -21,7 +21,7 @@ class ConstructorTest(testlib.RouterMixin, testlib.TestCase): self.assertEquals(argv[1], 'exec') self.assertEquals(argv[2], '--interactive') self.assertEquals(argv[3], 'container_name') - self.assertEquals(argv[4], stream.python_path) + self.assertEquals(argv[4], stream.conn.options.python_path) if __name__ == '__main__': diff --git a/tests/first_stage_test.py b/tests/first_stage_test.py index 470afc7a..53f98373 100644 --- a/tests/first_stage_test.py +++ b/tests/first_stage_test.py @@ -19,8 +19,10 @@ class CommandLineTest(testlib.RouterMixin, testlib.TestCase): # * 3.x starting 2.7 def test_valid_syntax(self): - stream = mitogen.parent.Stream(self.router, 0, max_message_size=123) - args = stream.get_boot_command() + options = mitogen.parent.Options(max_message_size=123) + conn = mitogen.parent.Connection(options, self.router) + conn.context = mitogen.core.Context(None, 123) + args = conn.get_boot_command() # Executing the boot command will print "EC0" and expect to read from # stdin, which will fail because it's pointing at /dev/null, causing @@ -38,7 +40,8 @@ class CommandLineTest(testlib.RouterMixin, testlib.TestCase): ) stdout, stderr = proc.communicate() self.assertEquals(0, proc.returncode) - self.assertEquals(mitogen.parent.Stream.EC0_MARKER, stdout) + self.assertEquals(stdout, + mitogen.parent.BootstrapProtocol.EC0_MARKER+b('\n')) self.assertIn(b("Error -5 while decompressing data"), stderr) finally: fp.close() diff --git a/tests/image_prep/README.md b/tests/image_prep/README.md index d275672f..a970b319 100644 --- a/tests/image_prep/README.md +++ b/tests/image_prep/README.md @@ -11,10 +11,13 @@ code, the OS X config just has the user accounts. See ../README.md for a (mostly) description of the accounts created. + ## Building the containers ``./build_docker_images.sh`` +Requires Ansible 2.3.x.x in order to target CentOS 5 + ## Preparing an OS X box diff --git a/tests/image_prep/_container_setup.yml b/tests/image_prep/_container_setup.yml index dc0bbf53..9d001f48 100644 --- a/tests/image_prep/_container_setup.yml +++ b/tests/image_prep/_container_setup.yml @@ -46,65 +46,98 @@ - when: ansible_virtualization_type != "docker" meta: end_play - - apt: + - name: Ensure requisite Debian packages are installed + apt: name: "{{packages.common + packages[distro][ver]}}" state: installed update_cache: true when: distro == "Debian" - - yum: + - name: Ensure requisite Red Hat packaed are installed + yum: name: "{{packages.common + packages[distro][ver]}}" state: installed update_cache: true when: distro == "CentOS" - - command: apt-get clean + - name: Clean up apt cache + command: apt-get clean when: distro == "Debian" - - command: yum clean all - when: distro == "CentOS" - - - shell: rm -rf {{item}}/* + - name: Clean up apt package lists + shell: rm -rf {{item}}/* with_items: - /var/cache/apt - /var/lib/apt/lists + when: distro == "Debian" - - copy: + - name: Clean up yum cache + command: yum clean all + when: distro == "CentOS" + + - name: Enable UTF-8 locale on Debian + copy: dest: /etc/locale.gen content: | en_US.UTF-8 UTF-8 fr_FR.UTF-8 UTF-8 when: distro == "Debian" - - shell: locale-gen + - name: Generate UTF-8 locale on Debian + shell: locale-gen when: distro == "Debian" - # Vanilla Ansible needs simplejson on CentOS 5. - - shell: mkdir -p /usr/lib/python2.4/site-packages/simplejson/ + - name: Install prebuilt 'doas' binary + unarchive: + dest: / + src: ../data/docker/doas-debian.tar.gz + + - name: Make prebuilt 'doas' binary executable + file: + path: /usr/local/bin/doas + mode: 'u=rwxs,go=rx' + owner: root + group: root + + - name: Install doas.conf + copy: + dest: /etc/doas.conf + content: | + permit :mitogen__group + permit :root + + - name: Vanilla Ansible needs simplejson on CentOS 5. + shell: mkdir -p /usr/lib/python2.4/site-packages/simplejson/ when: distro == "CentOS" and ver == "5" - - synchronize: + - name: Vanilla Ansible needs simplejson on CentOS 5. + synchronize: dest: /usr/lib/python2.4/site-packages/simplejson/ src: ../../ansible_mitogen/compat/simplejson/ when: distro == "CentOS" and ver == "5" - - user: + - name: Set root user password and shell + user: name: root password: "{{ 'rootpassword' | password_hash('sha256') }}" shell: /bin/bash - - file: + - name: Ensure /var/run/sshd exists + file: path: /var/run/sshd state: directory - - command: ssh-keygen -t rsa -f /etc/ssh/ssh_host_rsa_key + - name: Generate SSH host key + command: ssh-keygen -t rsa -f /etc/ssh/ssh_host_rsa_key args: creates: /etc/ssh/ssh_host_rsa_key - - group: + - name: Ensure correct sudo group exists + group: name: "{{sudo_group[distro]}}" - - copy: + - name: Ensure /etc/sentinel exists + copy: dest: /etc/sentinel content: | i-am-mitogen-test-docker-image @@ -119,7 +152,8 @@ path: /etc/sudoers.d mode: 'u=rwx,go=' - - blockinfile: + - name: Install test-related sudo rules + blockinfile: path: /etc/sudoers block: | # https://www.toofishes.net/blog/trouble-sudoers-or-last-entry-wins/ @@ -131,31 +165,36 @@ Defaults>mitogen__require_tty requiretty Defaults>mitogen__require_tty_pw_required requiretty,targetpw - # Prevent permission denied errors. - - file: + - name: Prevent permission denied errors. + file: path: /etc/sudoers.d/README state: absent - - lineinfile: + - name: Install CentOS wheel sudo rule + lineinfile: path: /etc/sudoers line: "%wheel ALL=(ALL) ALL" when: distro == "CentOS" - - lineinfile: + - name: Enable SSH banner + lineinfile: path: /etc/ssh/sshd_config line: Banner /etc/ssh/banner.txt - - lineinfile: + - name: Allow remote SSH root login + lineinfile: path: /etc/ssh/sshd_config line: PermitRootLogin yes regexp: '.*PermitRootLogin.*' - - lineinfile: + - name: Allow remote SSH root login + lineinfile: path: /etc/pam.d/sshd regexp: '.*session.*required.*pam_loginuid.so' line: session optional pam_loginuid.so - - copy: + - name: Install convenience script for running an straced Python + copy: mode: 'u+rwx,go=rx' dest: /usr/local/bin/pywrap content: | diff --git a/tests/image_prep/_user_accounts.yml b/tests/image_prep/_user_accounts.yml index a5b63c13..5f1bf0dc 100644 --- a/tests/image_prep/_user_accounts.yml +++ b/tests/image_prep/_user_accounts.yml @@ -20,6 +20,7 @@ - readonly_homedir - require_tty - require_tty_pw_required + - permdenied - slow_user - webapp - sudo1 @@ -98,6 +99,14 @@ - bashrc - profile + - name: "Login throws permission denied errors (issue #271)" + copy: + dest: ~mitogen__permdenied/.{{item}} + src: ../data/docker/mitogen__permdenied.profile + with_items: + - bashrc + - profile + - name: Install pubkey for mitogen__has_sudo_pubkey block: - file: diff --git a/tests/image_prep/build_docker_images.py b/tests/image_prep/build_docker_images.py index 9fc89c05..76564297 100755 --- a/tests/image_prep/build_docker_images.py +++ b/tests/image_prep/build_docker_images.py @@ -26,7 +26,13 @@ label_by_id = {} for base_image, label in [ ('astj/centos5-vault', 'centos5'), # Python 2.4.3 - ('debian:stretch', 'debian'), # Python 2.7.13, 3.5.3 + # Debian containers later than debuerreotype/debuerreotype#48 no longer + # ship a stub 'initctl', causing (apparently) the Ansible service + # module run at the end of DebOps to trigger a full stop/start of SSHd. + # When SSHd is killed, Docker responds by destroying the container. + # Proper solution is to include a full /bin/init; Docker --init doesn't + # help. In the meantime, just use a fixed older version. + ('debian:stretch-20181112', 'debian'), # Python 2.7.13, 3.5.3 ('centos:6', 'centos6'), # Python 2.6.6 ('centos:7', 'centos7') # Python 2.7.5 ]: diff --git a/tests/iter_split_test.py b/tests/iter_split_test.py new file mode 100644 index 00000000..ee5e97d9 --- /dev/null +++ b/tests/iter_split_test.py @@ -0,0 +1,66 @@ + +import mock +import unittest2 + +import mitogen.core + +import testlib + +try: + next +except NameError: + def next(it): + return it.next() + + +class IterSplitTest(unittest2.TestCase): + func = staticmethod(mitogen.core.iter_split) + + def test_empty_buffer(self): + lst = [] + trailer, cont = self.func(buf='', delim='\n', func=lst.append) + self.assertTrue(cont) + self.assertEquals('', trailer) + self.assertEquals([], lst) + + def test_empty_line(self): + lst = [] + trailer, cont = self.func(buf='\n', delim='\n', func=lst.append) + self.assertTrue(cont) + self.assertEquals('', trailer) + self.assertEquals([''], lst) + + def test_one_line(self): + buf = 'xxxx\n' + lst = [] + trailer, cont = self.func(buf=buf, delim='\n', func=lst.append) + self.assertTrue(cont) + self.assertEquals('', trailer) + self.assertEquals(lst, ['xxxx']) + + def test_one_incomplete(self): + buf = 'xxxx\nyy' + lst = [] + trailer, cont = self.func(buf=buf, delim='\n', func=lst.append) + self.assertTrue(cont) + self.assertEquals('yy', trailer) + self.assertEquals(lst, ['xxxx']) + + def test_returns_false_immediately(self): + buf = 'xxxx\nyy' + func = lambda buf: False + trailer, cont = self.func(buf=buf, delim='\n', func=func) + self.assertFalse(cont) + self.assertEquals('yy', trailer) + + def test_returns_false_second_call(self): + buf = 'xxxx\nyy\nzz' + it = iter([True, False]) + func = lambda buf: next(it) + trailer, cont = self.func(buf=buf, delim='\n', func=func) + self.assertFalse(cont) + self.assertEquals('zz', trailer) + + +if __name__ == '__main__': + unittest2.main() diff --git a/tests/jail_test.py b/tests/jail_test.py new file mode 100644 index 00000000..7239d32f --- /dev/null +++ b/tests/jail_test.py @@ -0,0 +1,33 @@ + +import os + +import mitogen +import mitogen.parent + +import unittest2 + +import testlib + + +class ConstructorTest(testlib.RouterMixin, testlib.TestCase): + jexec_path = testlib.data_path('stubs/stub-jexec.py') + + def test_okay(self): + context = self.router.jail( + jexec_path=self.jexec_path, + container='somejail', + ) + stream = self.router.stream_by_id(context.context_id) + + argv = eval(context.call(os.getenv, 'ORIGINAL_ARGV')) + self.assertEquals(argv[:4], [ + self.jexec_path, + 'somejail', + stream.conn.options.python_path, + '-c', + ]) + self.assertEquals('1', context.call(os.getenv, 'THIS_IS_STUB_JEXEC')) + + +if __name__ == '__main__': + unittest2.main() diff --git a/tests/lxc_test.py b/tests/lxc_test.py index ae5990f6..f78846ff 100644 --- a/tests/lxc_test.py +++ b/tests/lxc_test.py @@ -38,7 +38,7 @@ class ConstructorTest(testlib.RouterMixin, testlib.TestCase): lxc_attach_path='true', ) ) - self.assertTrue(str(e).endswith(mitogen.lxc.Stream.eof_error_hint)) + self.assertTrue(str(e).endswith(mitogen.lxc.Connection.eof_error_hint)) if __name__ == '__main__': diff --git a/tests/lxd_test.py b/tests/lxd_test.py index e59da43c..c80f8251 100644 --- a/tests/lxd_test.py +++ b/tests/lxd_test.py @@ -30,7 +30,7 @@ class ConstructorTest(testlib.RouterMixin, testlib.TestCase): lxc_path='true', ) ) - self.assertTrue(str(e).endswith(mitogen.lxd.Stream.eof_error_hint)) + self.assertTrue(str(e).endswith(mitogen.lxd.Connection.eof_error_hint)) if __name__ == '__main__': diff --git a/tests/message_test.py b/tests/message_test.py new file mode 100644 index 00000000..79deb2c6 --- /dev/null +++ b/tests/message_test.py @@ -0,0 +1,545 @@ + +import sys +import struct + +import mock +import unittest2 + +import mitogen.core +import mitogen.master +import testlib + +from mitogen.core import b + + +class ConstructorTest(testlib.TestCase): + klass = mitogen.core.Message + + def test_dst_id_default(self): + self.assertEquals(self.klass().dst_id, None) + + def test_dst_id_explicit(self): + self.assertEquals(self.klass(dst_id=1111).dst_id, 1111) + + @mock.patch('mitogen.context_id', 1234) + def test_src_id_default(self): + self.assertEquals(self.klass().src_id, 1234) + + def test_src_id_explicit(self): + self.assertEquals(self.klass(src_id=4321).src_id, 4321) + + @mock.patch('mitogen.context_id', 5555) + def test_auth_id_default(self): + self.assertEquals(self.klass().auth_id, 5555) + + def test_auth_id_explicit(self): + self.assertEquals(self.klass(auth_id=2222).auth_id, 2222) + + def test_handle_default(self): + self.assertEquals(self.klass().handle, None) + + def test_handle_explicit(self): + self.assertEquals(self.klass(handle=1234).handle, 1234) + + def test_reply_to_default(self): + self.assertEquals(self.klass().reply_to, None) + + def test_reply_to_explicit(self): + self.assertEquals(self.klass(reply_to=8888).reply_to, 8888) + + def test_data_default(self): + m = self.klass() + self.assertEquals(m.data, b('')) + self.assertTrue(isinstance(m.data, mitogen.core.BytesType)) + + def test_data_explicit(self): + m = self.klass(data=b('asdf')) + self.assertEquals(m.data, b('asdf')) + self.assertTrue(isinstance(m.data, mitogen.core.BytesType)) + + def test_data_hates_unicode(self): + self.assertRaises(Exception, + lambda: self.klass(data=u'asdf')) + + +class PackTest(testlib.TestCase): + klass = mitogen.core.Message + + def test_header_format_sanity(self): + self.assertEquals(self.klass.HEADER_LEN, + struct.calcsize(self.klass.HEADER_FMT)) + + def test_header_length_correct(self): + s = self.klass(dst_id=123, handle=123).pack() + self.assertEquals(len(s), self.klass.HEADER_LEN) + + def test_magic(self): + s = self.klass(dst_id=123, handle=123).pack() + magic, = struct.unpack('>h', s[:2]) + self.assertEquals(self.klass.HEADER_MAGIC, magic) + + def test_dst_id(self): + s = self.klass(dst_id=123, handle=123).pack() + dst_id, = struct.unpack('>L', s[2:6]) + self.assertEquals(123, dst_id) + + def test_src_id(self): + s = self.klass(src_id=5432, dst_id=123, handle=123).pack() + src_id, = struct.unpack('>L', s[6:10]) + self.assertEquals(5432, src_id) + + def test_auth_id(self): + s = self.klass(auth_id=1919, src_id=5432, dst_id=123, handle=123).pack() + auth_id, = struct.unpack('>L', s[10:14]) + self.assertEquals(1919, auth_id) + + def test_handle(self): + s = self.klass(dst_id=123, handle=9999).pack() + handle, = struct.unpack('>L', s[14:18]) + self.assertEquals(9999, handle) + + def test_reply_to(self): + s = self.klass(dst_id=1231, handle=7777, reply_to=9132).pack() + reply_to, = struct.unpack('>L', s[18:22]) + self.assertEquals(9132, reply_to) + + def test_data_length_empty(self): + s = self.klass(dst_id=1231, handle=7777).pack() + data_length, = struct.unpack('>L', s[22:26]) + self.assertEquals(0, data_length) + + def test_data_length_present(self): + s = self.klass(dst_id=1231, handle=7777, data=b('hello')).pack() + data_length, = struct.unpack('>L', s[22:26]) + self.assertEquals(5, data_length) + + def test_data_empty(self): + s = self.klass(dst_id=1231, handle=7777).pack() + data = s[26:] + self.assertEquals(b(''), data) + + def test_data_present(self): + s = self.klass(dst_id=11, handle=77, data=b('hello')).pack() + data = s[26:] + self.assertEquals(b('hello'), data) + + +class IsDeadTest(testlib.TestCase): + klass = mitogen.core.Message + + def test_is_dead(self): + msg = self.klass(reply_to=mitogen.core.IS_DEAD) + self.assertTrue(msg.is_dead) + + def test_is_not_dead(self): + msg = self.klass(reply_to=5555) + self.assertFalse(msg.is_dead) + + +class DeadTest(testlib.TestCase): + klass = mitogen.core.Message + + def test_no_reason(self): + msg = self.klass.dead() + self.assertEquals(msg.reply_to, mitogen.core.IS_DEAD) + self.assertTrue(msg.is_dead) + self.assertEquals(msg.data, b('')) + + def test_with_reason(self): + msg = self.klass.dead(reason=u'oh no') + self.assertEquals(msg.reply_to, mitogen.core.IS_DEAD) + self.assertTrue(msg.is_dead) + self.assertEquals(msg.data, b('oh no')) + + +class EvilObject(object): + pass + + +class PickledTest(testlib.TestCase): + # getting_started.html#rpc-serialization-rules + klass = mitogen.core.Message + + def roundtrip(self, v, router=None): + msg = self.klass.pickled(v) + msg2 = self.klass(data=msg.data) + msg2.router = router + return msg2.unpickle() + + def test_bool(self): + for b in True, False: + self.assertEquals(b, self.roundtrip(b)) + + @unittest2.skipIf(condition=sys.version_info < (2, 6), + reason='bytearray missing on <2.6') + def test_bytearray(self): + ba = bytearray(b('123')) + self.assertRaises(mitogen.core.StreamError, + lambda: self.roundtrip(ba) + ) + + def test_bytes(self): + by = b('123') + self.assertEquals(by, self.roundtrip(by)) + + def test_dict(self): + d = {1: 2, u'a': 3, b('b'): 4, 'c': {}} + roundtrip = self.roundtrip(d) + self.assertEquals(d, roundtrip) + self.assertTrue(isinstance(roundtrip, dict)) + for k in d: + self.assertTrue(isinstance(roundtrip[k], type(d[k]))) + + def test_int(self): + self.assertEquals(123, self.klass.pickled(123).unpickle()) + + def test_list(self): + l = [1, u'b', b('c')] + roundtrip = self.roundtrip(l) + self.assertTrue(isinstance(roundtrip, list)) + self.assertEquals(l, roundtrip) + for k in range(len(l)): + self.assertTrue(isinstance(roundtrip[k], type(l[k]))) + + @unittest2.skipIf(condition=sys.version_info > (3, 0), + reason='long missing in >3.x') + def test_long(self): + l = long(0xffffffffffff) + roundtrip = self.roundtrip(l) + self.assertEquals(l, roundtrip) + self.assertTrue(isinstance(roundtrip, long)) + + def test_tuple(self): + l = (1, u'b', b('c')) + roundtrip = self.roundtrip(l) + self.assertEquals(l, roundtrip) + self.assertTrue(isinstance(roundtrip, tuple)) + for k in range(len(l)): + self.assertTrue(isinstance(roundtrip[k], type(l[k]))) + + def test_unicode(self): + u = u'abcd' + roundtrip = self.roundtrip(u) + self.assertEquals(u, roundtrip) + self.assertTrue(isinstance(roundtrip, mitogen.core.UnicodeType)) + + #### custom types. see also: types_test.py, call_error_test.py + + # Python 3 pickle protocol 2 does weird stuff depending on whether an empty + # or nonempty bytes is being serialized. For non-empty, it yields a + # _codecs.encode() call. For empty, it yields a bytes() call. + + def test_blob_nonempty(self): + v = mitogen.core.Blob(b('dave')) + roundtrip = self.roundtrip(v) + self.assertTrue(isinstance(roundtrip, mitogen.core.Blob)) + self.assertEquals(b('dave'), roundtrip) + + def test_blob_empty(self): + v = mitogen.core.Blob(b('')) + roundtrip = self.roundtrip(v) + self.assertTrue(isinstance(roundtrip, mitogen.core.Blob)) + self.assertEquals(b(''), v) + + def test_secret_nonempty(self): + s = mitogen.core.Secret(u'dave') + roundtrip = self.roundtrip(s) + self.assertTrue(isinstance(roundtrip, mitogen.core.Secret)) + self.assertEquals(u'dave', roundtrip) + + def test_secret_empty(self): + s = mitogen.core.Secret(u'') + roundtrip = self.roundtrip(s) + self.assertTrue(isinstance(roundtrip, mitogen.core.Secret)) + self.assertEquals(u'', roundtrip) + + def test_call_error(self): + ce = mitogen.core.CallError('nope') + ce2 = self.assertRaises(mitogen.core.CallError, + lambda: self.roundtrip(ce)) + self.assertEquals(ce.args[0], ce2.args[0]) + + def test_context(self): + router = mitogen.master.Router() + try: + c = router.context_by_id(1234) + roundtrip = self.roundtrip(c) + self.assertTrue(isinstance(roundtrip, mitogen.core.Context)) + self.assertEquals(c.context_id, 1234) + finally: + router.broker.shutdown() + router.broker.join() + + def test_sender(self): + router = mitogen.master.Router() + try: + recv = mitogen.core.Receiver(router) + sender = recv.to_sender() + roundtrip = self.roundtrip(sender, router=router) + self.assertTrue(isinstance(roundtrip, mitogen.core.Sender)) + self.assertEquals(roundtrip.context.context_id, mitogen.context_id) + self.assertEquals(roundtrip.dst_handle, sender.dst_handle) + finally: + router.broker.shutdown() + router.broker.join() + + #### + + def test_custom_object_deserialization_fails(self): + self.assertRaises(mitogen.core.StreamError, + lambda: self.roundtrip(EvilObject()) + ) + + +class ReplyTest(testlib.TestCase): + # getting_started.html#rpc-serialization-rules + klass = mitogen.core.Message + + def test_reply_calls_router_route(self): + msg = self.klass(src_id=1234, reply_to=9191) + router = mock.Mock() + msg.reply(123, router=router) + self.assertEquals(1, router.route.call_count) + + def test_reply_pickles_object(self): + msg = self.klass(src_id=1234, reply_to=9191) + router = mock.Mock() + msg.reply(123, router=router) + _, (reply,), _ = router.route.mock_calls[0] + self.assertEquals(reply.dst_id, 1234) + self.assertEquals(reply.unpickle(), 123) + + def test_reply_uses_preformatted_message(self): + msg = self.klass(src_id=1234, reply_to=9191) + router = mock.Mock() + my_reply = mitogen.core.Message.pickled(4444) + msg.reply(my_reply, router=router) + _, (reply,), _ = router.route.mock_calls[0] + self.assertTrue(my_reply is reply) + self.assertEquals(reply.dst_id, 1234) + self.assertEquals(reply.unpickle(), 4444) + + def test_reply_sets_dst_id(self): + msg = self.klass(src_id=1234, reply_to=9191) + router = mock.Mock() + msg.reply(123, router=router) + _, (reply,), _ = router.route.mock_calls[0] + self.assertEquals(reply.dst_id, 1234) + + def test_reply_sets_handle(self): + msg = self.klass(src_id=1234, reply_to=9191) + router = mock.Mock() + msg.reply(123, router=router) + _, (reply,), _ = router.route.mock_calls[0] + self.assertEquals(reply.handle, 9191) + + +class UnpickleTest(testlib.TestCase): + # mostly done by PickleTest, just check behaviour of parameters + klass = mitogen.core.Message + + def test_throw(self): + ce = mitogen.core.CallError('nope') + m = self.klass.pickled(ce) + ce2 = self.assertRaises(mitogen.core.CallError, + lambda: m.unpickle()) + self.assertEquals(ce.args[0], ce2.args[0]) + + def test_no_throw(self): + ce = mitogen.core.CallError('nope') + m = self.klass.pickled(ce) + ce2 = m.unpickle(throw=False) + self.assertEquals(ce.args[0], ce2.args[0]) + + def test_throw_dead(self): + m = self.klass.pickled('derp', reply_to=mitogen.core.IS_DEAD) + self.assertRaises(mitogen.core.ChannelError, + lambda: m.unpickle()) + + def test_no_throw_dead(self): + m = self.klass.pickled('derp', reply_to=mitogen.core.IS_DEAD) + self.assertEquals('derp', m.unpickle(throw_dead=False)) + + +class UnpickleCompatTest(testlib.TestCase): + # try weird variations of pickles from different Python versions. + klass = mitogen.core.Message + + def check(self, value, encoded, **kwargs): + if isinstance(encoded, mitogen.core.UnicodeType): + encoded = encoded.encode('latin1') + m = self.klass(data=encoded) + m.router = mitogen.master.Router() + try: + return m.unpickle(**kwargs) + finally: + m.router.broker.shutdown() + m.router.broker.join() + + def test_py24_bytes(self): + self.check('test', + ('\x80\x02U\x04testq\x00.')) + + def test_py24_unicode(self): + self.check(u'test', + ('\x80\x02X\x04\x00\x00\x00testq\x00.')) + + def test_py24_int(self): + self.check(123, + ('\x80\x02K{.')) + + def test_py24_long(self): + self.check(17592186044415, + ('\x80\x02\x8a\x06\xff\xff\xff\xff\xff\x0f.')) + + def test_py24_dict(self): + self.check({}, + ('\x80\x02}q\x00.')) + + def test_py24_tuple(self): + self.check((1, 2, u'b'), + ('\x80\x02K\x01K\x02X\x01\x00\x00\x00bq\x00\x87q\x01.')) + + def test_py24_bool(self): + self.check(True, + ('\x80\x02\x88.')) + + def test_py24_list(self): + self.check([1, 2, u'b'], + ('\x80\x02]q\x00(K\x01K\x02X\x01\x00\x00\x00bq\x01e.')) + + def test_py24_blob(self): + self.check(mitogen.core.mitogen.core.Blob(b('bigblob')), + ('\x80\x02cmitogen.core\nBlob\nq\x00U\x07bigblobq\x01\x85q\x02Rq\x03.')) + + def test_py24_secret(self): + self.check(mitogen.core.Secret(u'mypassword'), + ('\x80\x02cmitogen.core\nSecret\nq\x00X\n\x00\x00\x00mypasswordq\x01\x85q\x02Rq\x03.')) + + def test_py24_call_error(self): + self.check(mitogen.core.CallError('big error'), + ('\x80\x02cmitogen.core\n_unpickle_call_error\nq\x00X\t\x00\x00\x00big errorq\x01\x85q\x02R.'), throw=False) + + def test_py24_context(self): + self.check(mitogen.core.Context(1234, None), + ('\x80\x02cmitogen.core\n_unpickle_context\nq\x00M\xd2\x04N\x86q\x01Rq\x02.')) + + def test_py24_sender(self): + self.check(mitogen.core.Sender(mitogen.core.Context(55555, None), 4444), + ('\x80\x02cmitogen.core\n_unpickle_sender\nq\x00M\x03\xd9M\\\x11\x86q\x01Rq\x02.')) + + def test_py27_bytes(self): + self.check(b('test'), + ('\x80\x02U\x04testq\x01.')) + + def test_py27_unicode(self): + self.check(u'test', + ('\x80\x02X\x04\x00\x00\x00testq\x01.')) + + def test_py27_int(self): + self.check(123, + ('\x80\x02K{.')) + + def test_py27_long(self): + self.check(17592186044415, + ('\x80\x02\x8a\x06\xff\xff\xff\xff\xff\x0f.')) + + def test_py27_dict(self): + self.check({}, + ('\x80\x02}q\x01.')) + + def test_py27_tuple(self): + self.check((1, 2, u'b'), + ('\x80\x02K\x01K\x02X\x01\x00\x00\x00b\x87q\x01.')) + + def test_py27_bool(self): + self.check(True, + ('\x80\x02\x88.')) + + def test_py27_list(self): + self.check([1, 2, u'b'], + ('\x80\x02]q\x01(K\x01K\x02X\x01\x00\x00\x00be.')) + + def test_py27_blob(self): + self.check(mitogen.core.mitogen.core.Blob(b('bigblob')), + ('\x80\x02cmitogen.core\nBlob\nq\x01U\x07bigblob\x85Rq\x02.')) + + def test_py27_secret(self): + self.check(mitogen.core.Secret(u'mypassword'), + ('\x80\x02cmitogen.core\nSecret\nq\x01X\n\x00\x00\x00mypassword\x85Rq\x02.')) + + def test_py27_call_error(self): + self.check(mitogen.core.CallError(u'big error',), + ('\x80\x02cmitogen.core\n_unpickle_call_error\nq\x01X\t\x00\x00\x00big errorq\x02\x85Rq\x03.'), throw=False) + + def test_py27_context(self): + self.check(mitogen.core.Context(1234, None), + ('\x80\x02cmitogen.core\n_unpickle_context\nq\x01M\xd2\x04N\x86Rq\x02.')) + + def test_py27_sender(self): + self.check(mitogen.core.Sender(mitogen.core.Context(55555, None), 4444), + ('\x80\x02cmitogen.core\n_unpickle_sender\nq\x01M\x03\xd9M\\\x11\x86Rq\x02.')) + + def test_py36_bytes(self): + self.check(b('test'), + ('\x80\x02c_codecs\nencode\nq\x00X\x04\x00\x00\x00testq\x01X\x06\x00\x00\x00latin1q\x02\x86q\x03Rq\x04.')) + + def test_py36_unicode(self): + self.check('test', + ('\x80\x02X\x04\x00\x00\x00testq\x00.')) + + def test_py36_int(self): + self.check(123, + ('\x80\x02K{.')) + + def test_py36_long(self): + self.check(17592186044415, + ('\x80\x02\x8a\x06\xff\xff\xff\xff\xff\x0f.')) + + def test_py36_dict(self): + self.check({}, + ('\x80\x02}q\x00.')) + + def test_py36_tuple(self): + self.check((1, 2, u'b'), + ('\x80\x02K\x01K\x02X\x01\x00\x00\x00bq\x00\x87q\x01.')) + + def test_py36_bool(self): + self.check(True, + ('\x80\x02\x88.')) + + def test_py36_list(self): + self.check([1, 2, u'b'], + ('\x80\x02]q\x00(K\x01K\x02X\x01\x00\x00\x00bq\x01e.')) + + def test_py36_blob(self): + self.check(mitogen.core.mitogen.core.Blob(b('bigblob')), + ('\x80\x02cmitogen.core\nBlob\nq\x00c_codecs\nencode\nq\x01X\x07\x00\x00\x00bigblobq\x02X\x06\x00\x00\x00latin1q\x03\x86q\x04Rq\x05\x85q\x06Rq\x07.')) + + def test_py36_secret(self): + self.check(mitogen.core.Secret('mypassword'), + ('\x80\x02cmitogen.core\nSecret\nq\x00X\n\x00\x00\x00mypasswordq\x01\x85q\x02Rq\x03.')) + + def test_py36_call_error(self): + self.check(mitogen.core.CallError('big error'), + ('\x80\x02cmitogen.core\n_unpickle_call_error\nq\x00X\t\x00\x00\x00big errorq\x01\x85q\x02Rq\x03.'), throw=False) + + def test_py36_context(self): + self.check(mitogen.core.Context(1234, None), + ('\x80\x02cmitogen.core\n_unpickle_context\nq\x00M\xd2\x04N\x86q\x01Rq\x02.')) + + def test_py36_sender(self): + self.check(mitogen.core.Sender(mitogen.core.Context(55555, None), 4444), + ('\x80\x02cmitogen.core\n_unpickle_sender\nq\x00M\x03\xd9M\\\x11\x86q\x01Rq\x02.')) + + +class ReprTest(testlib.TestCase): + klass = mitogen.core.Message + + def test_repr(self): + # doesn't crash + repr(self.klass.pickled('test')) + + +if __name__ == '__main__': + unittest2.main() diff --git a/tests/mitogen_protocol_test.py b/tests/mitogen_protocol_test.py new file mode 100644 index 00000000..834fb437 --- /dev/null +++ b/tests/mitogen_protocol_test.py @@ -0,0 +1,34 @@ + +import unittest2 +import mock + +import mitogen.core + +import testlib + + +class ReceiveOneTest(testlib.TestCase): + klass = mitogen.core.MitogenProtocol + + def test_corruption(self): + broker = mock.Mock() + router = mock.Mock() + stream = mock.Mock() + + protocol = self.klass(router, 1) + protocol.stream = stream + + junk = mitogen.core.b('x') * mitogen.core.Message.HEADER_LEN + + capture = testlib.LogCapturer() + capture.start() + protocol.on_receive(broker, junk) + capture.stop() + + self.assertEquals(1, stream.on_disconnect.call_count) + expect = self.klass.corrupt_msg % (stream.name, junk) + self.assertTrue(expect in capture.raw()) + + +if __name__ == '__main__': + unittest2.main() diff --git a/tests/parent_test.py b/tests/parent_test.py index 00bddb4d..7ac482c5 100644 --- a/tests/parent_test.py +++ b/tests/parent_test.py @@ -14,6 +14,11 @@ from testlib import Popen__terminate import mitogen.parent +try: + file +except NameError: + from io import FileIO as file + def wait_for_child(pid, timeout=1.0): deadline = time.time() + timeout @@ -49,7 +54,7 @@ def wait_for_empty_output_queue(sync_recv, context): while True: # Now wait for the RPC to exit the output queue. stream = router.stream_by_id(context.context_id) - if broker.defer_sync(lambda: stream.pending_bytes()) == 0: + if broker.defer_sync(lambda: stream.protocol.pending_bytes()) == 0: return time.sleep(0.1) @@ -69,35 +74,17 @@ class GetDefaultRemoteNameTest(testlib.TestCase): self.assertEquals("ECORP_Administrator@box:123", self.func()) -class WstatusToStrTest(testlib.TestCase): - func = staticmethod(mitogen.parent.wstatus_to_str) +class ReturncodeToStrTest(testlib.TestCase): + func = staticmethod(mitogen.parent.returncode_to_str) def test_return_zero(self): - pid = os.fork() - if not pid: - os._exit(0) - (pid, status), _ = mitogen.core.io_op(os.waitpid, pid, 0) - self.assertEquals(self.func(status), - 'exited with return code 0') + self.assertEquals(self.func(0), 'exited with return code 0') def test_return_one(self): - pid = os.fork() - if not pid: - os._exit(1) - (pid, status), _ = mitogen.core.io_op(os.waitpid, pid, 0) - self.assertEquals( - self.func(status), - 'exited with return code 1' - ) + self.assertEquals(self.func(1), 'exited with return code 1') def test_sigkill(self): - pid = os.fork() - if not pid: - time.sleep(600) - os.kill(pid, signal.SIGKILL) - (pid, status), _ = mitogen.core.io_op(os.waitpid, pid, 0) - self.assertEquals( - self.func(status), + self.assertEquals(self.func(-signal.SIGKILL), 'exited due to signal %s (SIGKILL)' % (int(signal.SIGKILL),) ) @@ -107,20 +94,20 @@ class WstatusToStrTest(testlib.TestCase): class ReapChildTest(testlib.RouterMixin, testlib.TestCase): def test_connect_timeout(self): # Ensure the child process is reaped if the connection times out. - stream = mitogen.parent.Stream( - router=self.router, - remote_id=1234, + options = mitogen.parent.Options( old_router=self.router, max_message_size=self.router.max_message_size, python_path=testlib.data_path('python_never_responds.py'), connect_timeout=0.5, ) + + conn = mitogen.parent.Connection(options, router=self.router) self.assertRaises(mitogen.core.TimeoutError, - lambda: stream.connect() + lambda: conn.connect(context=mitogen.core.Context(None, 1234)) ) - wait_for_child(stream.pid) + wait_for_child(conn.proc.pid) e = self.assertRaises(OSError, - lambda: os.kill(stream.pid, 0) + lambda: os.kill(conn.proc.pid, 0) ) self.assertEquals(e.args[0], errno.ESRCH) @@ -133,7 +120,7 @@ class StreamErrorTest(testlib.RouterMixin, testlib.TestCase): connect_timeout=3, ) ) - prefix = "EOF on stream; last 300 bytes received: " + prefix = mitogen.parent.Connection.eof_error_msg self.assertTrue(e.args[0].startswith(prefix)) def test_via_eof(self): @@ -142,12 +129,12 @@ class StreamErrorTest(testlib.RouterMixin, testlib.TestCase): e = self.assertRaises(mitogen.core.StreamError, lambda: self.router.local( via=local, - python_path='true', + python_path='echo', connect_timeout=3, ) ) - s = "EOF on stream; last 300 bytes received: " - self.assertTrue(s in e.args[0]) + expect = mitogen.parent.Connection.eof_error_msg + self.assertTrue(expect in e.args[0]) def test_direct_enoent(self): e = self.assertRaises(mitogen.core.StreamError, @@ -185,11 +172,15 @@ class OpenPtyTest(testlib.TestCase): func = staticmethod(mitogen.parent.openpty) def test_pty_returned(self): - master_fd, slave_fd = self.func() - self.assertTrue(isinstance(master_fd, int)) - self.assertTrue(isinstance(slave_fd, int)) - os.close(master_fd) - os.close(slave_fd) + master_fp, slave_fp = self.func() + try: + self.assertTrue(master_fp.isatty()) + self.assertTrue(isinstance(master_fp, file)) + self.assertTrue(slave_fp.isatty()) + self.assertTrue(isinstance(slave_fp, file)) + finally: + master_fp.close() + slave_fp.close() @mock.patch('os.openpty') def test_max_reached(self, openpty): @@ -204,20 +195,20 @@ class OpenPtyTest(testlib.TestCase): @mock.patch('os.openpty') def test_broken_linux_fallback(self, openpty): openpty.side_effect = OSError(errno.EPERM) - master_fd, slave_fd = self.func() + master_fp, slave_fp = self.func() try: - st = os.fstat(master_fd) + st = os.fstat(master_fp.fileno()) self.assertEquals(5, os.major(st.st_rdev)) - flags = fcntl.fcntl(master_fd, fcntl.F_GETFL) + flags = fcntl.fcntl(master_fp.fileno(), fcntl.F_GETFL) self.assertTrue(flags & os.O_RDWR) - st = os.fstat(slave_fd) + st = os.fstat(slave_fp.fileno()) self.assertEquals(136, os.major(st.st_rdev)) - flags = fcntl.fcntl(slave_fd, fcntl.F_GETFL) + flags = fcntl.fcntl(slave_fp.fileno(), fcntl.F_GETFL) self.assertTrue(flags & os.O_RDWR) finally: - os.close(master_fd) - os.close(slave_fd) + master_fp.close() + slave_fp.close() class TtyCreateChildTest(testlib.TestCase): @@ -235,123 +226,20 @@ class TtyCreateChildTest(testlib.TestCase): # read a password. tf = tempfile.NamedTemporaryFile() try: - pid, fd, _ = self.func([ + proc = self.func([ 'bash', '-c', 'exec 2>%s; echo hi > /dev/tty' % (tf.name,) ]) deadline = time.time() + 5.0 - for line in mitogen.parent.iter_read([fd], deadline): - self.assertEquals(mitogen.core.b('hi\n'), line) - break - waited_pid, status = os.waitpid(pid, 0) - self.assertEquals(pid, waited_pid) + mitogen.core.set_block(proc.stdin.fileno()) + # read(3) below due to https://bugs.python.org/issue37696 + self.assertEquals(mitogen.core.b('hi\n'), proc.stdin.read(3)) + waited_pid, status = os.waitpid(proc.pid, 0) + self.assertEquals(proc.pid, waited_pid) self.assertEquals(0, status) self.assertEquals(mitogen.core.b(''), tf.read()) - os.close(fd) - finally: - tf.close() - - -class IterReadTest(testlib.TestCase): - func = staticmethod(mitogen.parent.iter_read) - - def make_proc(self): - # I produce text every 100ms. - args = [testlib.data_path('iter_read_generator.py')] - proc = subprocess.Popen(args, stdout=subprocess.PIPE) - mitogen.core.set_nonblock(proc.stdout.fileno()) - return proc - - def test_no_deadline(self): - proc = self.make_proc() - try: - reader = self.func([proc.stdout.fileno()]) - for i, chunk in enumerate(reader): - self.assertEqual(1+i, int(chunk)) - if i > 2: - break - finally: - Popen__terminate(proc) proc.stdout.close() - - def test_deadline_exceeded_before_call(self): - proc = self.make_proc() - reader = self.func([proc.stdout.fileno()], 0) - try: - got = [] - try: - for chunk in reader: - got.append(chunk) - assert 0, 'TimeoutError not raised' - except mitogen.core.TimeoutError: - self.assertEqual(len(got), 0) finally: - Popen__terminate(proc) - proc.stdout.close() - - def test_deadline_exceeded_during_call(self): - proc = self.make_proc() - deadline = time.time() + 0.4 - - reader = self.func([proc.stdout.fileno()], deadline) - try: - got = [] - try: - for chunk in reader: - if time.time() > (deadline + 1.0): - assert 0, 'TimeoutError not raised' - got.append(chunk) - except mitogen.core.TimeoutError: - # Give a little wiggle room in case of imperfect scheduling. - # Ideal number should be 9. - self.assertLess(deadline, time.time()) - self.assertLess(1, len(got)) - self.assertLess(len(got), 20) - finally: - Popen__terminate(proc) - proc.stdout.close() - - -class WriteAllTest(testlib.TestCase): - func = staticmethod(mitogen.parent.write_all) - - def make_proc(self): - args = [testlib.data_path('write_all_consumer.py')] - proc = subprocess.Popen(args, stdin=subprocess.PIPE) - mitogen.core.set_nonblock(proc.stdin.fileno()) - return proc - - ten_ms_chunk = (mitogen.core.b('x') * 65535) - - def test_no_deadline(self): - proc = self.make_proc() - try: - self.func(proc.stdin.fileno(), self.ten_ms_chunk) - finally: - Popen__terminate(proc) - proc.stdin.close() - - def test_deadline_exceeded_before_call(self): - proc = self.make_proc() - try: - self.assertRaises(mitogen.core.TimeoutError, ( - lambda: self.func(proc.stdin.fileno(), self.ten_ms_chunk, 0) - )) - finally: - Popen__terminate(proc) - proc.stdin.close() - - def test_deadline_exceeded_during_call(self): - proc = self.make_proc() - try: - deadline = time.time() + 0.1 # 100ms deadline - self.assertRaises(mitogen.core.TimeoutError, ( - lambda: self.func(proc.stdin.fileno(), - self.ten_ms_chunk * 100, # 1s of data - deadline) - )) - finally: - Popen__terminate(proc) - proc.stdin.close() + tf.close() class DisconnectTest(testlib.RouterMixin, testlib.TestCase): @@ -394,7 +282,7 @@ class DisconnectTest(testlib.RouterMixin, testlib.TestCase): c2 = self.router.local() # Let c1 call functions in c2. - self.router.stream_by_id(c1.context_id).auth_id = mitogen.context_id + self.router.stream_by_id(c1.context_id).protocol.auth_id = mitogen.context_id c1.call(mitogen.parent.upgrade_router) sync_recv = mitogen.core.Receiver(self.router) @@ -412,14 +300,14 @@ class DisconnectTest(testlib.RouterMixin, testlib.TestCase): def test_far_sibling_disconnected(self): # God mode: child of child notices child of child of parent has # disconnected. - c1 = self.router.local() - c11 = self.router.local(via=c1) + c1 = self.router.local(name='c1') + c11 = self.router.local(name='c11', via=c1) - c2 = self.router.local() - c22 = self.router.local(via=c2) + c2 = self.router.local(name='c2') + c22 = self.router.local(name='c22', via=c2) # Let c1 call functions in c2. - self.router.stream_by_id(c1.context_id).auth_id = mitogen.context_id + self.router.stream_by_id(c1.context_id).protocol.auth_id = mitogen.context_id c11.call(mitogen.parent.upgrade_router) sync_recv = mitogen.core.Receiver(self.router) diff --git a/tests/poller_test.py b/tests/poller_test.py index e2e3cdd7..b05a9b94 100644 --- a/tests/poller_test.py +++ b/tests/poller_test.py @@ -42,8 +42,8 @@ class SockMixin(object): self.l2_sock, self.r2_sock = socket.socketpair() self.l2 = self.l2_sock.fileno() self.r2 = self.r2_sock.fileno() - for fd in self.l1, self.r1, self.l2, self.r2: - mitogen.core.set_nonblock(fd) + for fp in self.l1, self.r1, self.l2, self.r2: + mitogen.core.set_nonblock(fp) def fill(self, fd): """Make `fd` unwriteable.""" @@ -354,17 +354,17 @@ class FileClosedMixin(PollerMixin, SockMixin): class TtyHangupMixin(PollerMixin): def test_tty_hangup_detected(self): # bug in initial select.poll() implementation failed to detect POLLHUP. - master_fd, slave_fd = mitogen.parent.openpty() + master_fp, slave_fp = mitogen.parent.openpty() try: - self.p.start_receive(master_fd) + self.p.start_receive(master_fp.fileno()) self.assertEquals([], list(self.p.poll(0))) - os.close(slave_fd) - slave_fd = None - self.assertEquals([master_fd], list(self.p.poll(0))) + slave_fp.close() + slave_fp = None + self.assertEquals([master_fp.fileno()], list(self.p.poll(0))) finally: - if slave_fd is not None: - os.close(slave_fd) - os.close(master_fd) + if slave_fp is not None: + slave_fp.close() + master_fp.close() class DistinctDataMixin(PollerMixin, SockMixin): diff --git a/tests/requirements.txt b/tests/requirements.txt index 327f563a..bbcdc7cc 100644 --- a/tests/requirements.txt +++ b/tests/requirements.txt @@ -13,3 +13,5 @@ unittest2==1.1.0 # Fix InsecurePlatformWarning while creating py26 tox environment # https://urllib3.readthedocs.io/en/latest/advanced-usage.html#ssl-warnings urllib3[secure]; python_version < '2.7.9' +# Last idna compatible with Python 2.6 was idna 2.7. +idna==2.7; python_version < '2.7' diff --git a/tests/responder_test.py b/tests/responder_test.py index dbc68a3c..285acd6f 100644 --- a/tests/responder_test.py +++ b/tests/responder_test.py @@ -105,7 +105,7 @@ class BrokenModulesTest(testlib.TestCase): # unavailable. Should never happen in the real world. stream = mock.Mock() - stream.sent_modules = set() + stream.protocol.sent_modules = set() router = mock.Mock() router.stream_by_id = lambda n: stream @@ -143,7 +143,7 @@ class BrokenModulesTest(testlib.TestCase): import six_brokenpkg stream = mock.Mock() - stream.sent_modules = set() + stream.protocol.sent_modules = set() router = mock.Mock() router.stream_by_id = lambda n: stream diff --git a/tests/router_test.py b/tests/router_test.py index 80169e34..1bd6c26a 100644 --- a/tests/router_test.py +++ b/tests/router_test.py @@ -171,7 +171,7 @@ class CrashTest(testlib.BrokerMixin, testlib.TestCase): self.assertTrue(sem.get().is_dead) # Ensure it was logged. - expect = '_broker_main() crashed' + expect = 'broker crashed' self.assertTrue(expect in log.stop()) self.broker.join() @@ -364,8 +364,8 @@ class UnidirectionalTest(testlib.RouterMixin, testlib.TestCase): # treated like a parent. l1 = self.router.local() l1s = self.router.stream_by_id(l1.context_id) - l1s.auth_id = mitogen.context_id - l1s.is_privileged = True + l1s.protocol.auth_id = mitogen.context_id + l1s.protocol.is_privileged = True l2 = self.router.local() e = self.assertRaises(mitogen.core.CallError, @@ -378,12 +378,21 @@ class UnidirectionalTest(testlib.RouterMixin, testlib.TestCase): class EgressIdsTest(testlib.RouterMixin, testlib.TestCase): def test_egress_ids_populated(self): # Ensure Stream.egress_ids is populated on message reception. - c1 = self.router.local() - stream = self.router.stream_by_id(c1.context_id) - self.assertEquals(set(), stream.egress_ids) + c1 = self.router.local(name='c1') + c2 = self.router.local(name='c2') - c1.call(time.sleep, 0) - self.assertEquals(set([mitogen.context_id]), stream.egress_ids) + c1s = self.router.stream_by_id(c1.context_id) + try: + c1.call(ping_context, c2) + except mitogen.core.CallError: + # Fails because siblings cant call funcs in each other, but this + # causes messages to be sent. + pass + + self.assertEquals(c1s.protocol.egress_ids, set([ + mitogen.context_id, + c2.context_id, + ])) if __name__ == '__main__': diff --git a/tests/service_test.py b/tests/service_test.py index 3869f713..438766f7 100644 --- a/tests/service_test.py +++ b/tests/service_test.py @@ -44,8 +44,8 @@ class ActivationTest(testlib.RouterMixin, testlib.TestCase): self.assertTrue(isinstance(id_, int)) def test_sibling_cannot_activate_framework(self): - l1 = self.router.local() - l2 = self.router.local() + l1 = self.router.local(name='l1') + l2 = self.router.local(name='l2') exc = self.assertRaises(mitogen.core.CallError, lambda: l2.call(call_service_in, l1, MyService2.name(), 'get_id')) self.assertTrue(mitogen.core.Router.refused_msg in exc.args[0]) diff --git a/tests/setns_test.py b/tests/setns_test.py new file mode 100644 index 00000000..d48179b1 --- /dev/null +++ b/tests/setns_test.py @@ -0,0 +1,46 @@ + +import os +import socket +import sys + +import mitogen +import mitogen.parent + +import unittest2 + +import testlib + + +class DockerTest(testlib.DockerMixin, testlib.TestCase): + def test_okay(self): + # Magic calls must happen as root. + try: + root = self.router.sudo() + except mitogen.core.StreamError: + raise unittest2.SkipTest("requires sudo to localhost root") + + via_ssh = self.docker_ssh( + username='mitogen__has_sudo', + password='has_sudo_password', + ) + + via_setns = self.router.setns( + kind='docker', + container=self.dockerized_ssh.container_name, + via=root, + ) + + self.assertEquals( + via_ssh.call(socket.gethostname), + via_setns.call(socket.gethostname), + ) + + +DockerTest = unittest2.skipIf( + condition=sys.version_info < (2, 5), + reason="mitogen.setns unsupported on Python <2.4" +)(DockerTest) + + +if __name__ == '__main__': + unittest2.main() diff --git a/tests/ssh_test.py b/tests/ssh_test.py index 496710b8..273412e8 100644 --- a/tests/ssh_test.py +++ b/tests/ssh_test.py @@ -42,8 +42,6 @@ class ConstructorTest(testlib.RouterMixin, testlib.TestCase): class SshTest(testlib.DockerMixin, testlib.TestCase): - stream_class = mitogen.ssh.Stream - def test_debug_decoding(self): # ensure filter_debug_logs() decodes the logged string. capture = testlib.LogCapturer() @@ -60,6 +58,14 @@ class SshTest(testlib.DockerMixin, testlib.TestCase): expect = "%s: debug1: Reading configuration data" % (context.name,) self.assertTrue(expect in s) + def test_bash_permission_denied(self): + # issue #271: only match Permission Denied at start of line. + context = self.docker_ssh( + username='mitogen__permdenied', + password='permdenied_password', + ssh_debug_level=3, + ) + def test_stream_name(self): context = self.docker_ssh( username='mitogen__has_sudo', @@ -85,27 +91,21 @@ class SshTest(testlib.DockerMixin, testlib.TestCase): self.assertEquals(name, sudo.name) def test_password_required(self): - try: - context = self.docker_ssh( + e = self.assertRaises(mitogen.ssh.PasswordError, + lambda: self.docker_ssh( username='mitogen__has_sudo', ) - assert 0, 'exception not thrown' - except mitogen.ssh.PasswordError: - e = sys.exc_info()[1] - - self.assertEqual(e.args[0], self.stream_class.password_required_msg) + ) + self.assertEqual(e.args[0], mitogen.ssh.password_required_msg) def test_password_incorrect(self): - try: - context = self.docker_ssh( + e = self.assertRaises(mitogen.ssh.PasswordError, + lambda: self.docker_ssh( username='mitogen__has_sudo', password='badpw', ) - assert 0, 'exception not thrown' - except mitogen.ssh.PasswordError: - e = sys.exc_info()[1] - - self.assertEqual(e.args[0], self.stream_class.password_incorrect_msg) + ) + self.assertEqual(e.args[0], mitogen.ssh.password_incorrect_msg) def test_password_specified(self): context = self.docker_ssh( @@ -119,15 +119,12 @@ class SshTest(testlib.DockerMixin, testlib.TestCase): ) def test_pubkey_required(self): - try: - context = self.docker_ssh( + e = self.assertRaises(mitogen.ssh.PasswordError, + lambda: self.docker_ssh( username='mitogen__has_sudo_pubkey', ) - assert 0, 'exception not thrown' - except mitogen.ssh.PasswordError: - e = sys.exc_info()[1] - - self.assertEqual(e.args[0], self.stream_class.password_required_msg) + ) + self.assertEqual(e.args[0], mitogen.ssh.password_required_msg) def test_pubkey_specified(self): context = self.docker_ssh( @@ -150,7 +147,7 @@ class SshTest(testlib.DockerMixin, testlib.TestCase): check_host_keys='enforce', ) ) - self.assertEquals(e.args[0], mitogen.ssh.Stream.hostkey_failed_msg) + self.assertEquals(e.args[0], mitogen.ssh.hostkey_failed_msg) finally: fp.close() @@ -184,8 +181,6 @@ class SshTest(testlib.DockerMixin, testlib.TestCase): class BannerTest(testlib.DockerMixin, testlib.TestCase): # Verify the ability to disambiguate random spam appearing in the SSHd's # login banner from a legitimate password prompt. - stream_class = mitogen.ssh.Stream - def test_verbose_enabled(self): context = self.docker_ssh( username='mitogen__has_sudo', @@ -210,8 +205,6 @@ class StubPermissionDeniedTest(StubSshMixin, testlib.TestCase): class StubCheckHostKeysTest(StubSshMixin, testlib.TestCase): - stream_class = mitogen.ssh.Stream - def test_check_host_keys_accept(self): # required=true, host_key_checking=accept context = self.stub_ssh(STUBSSH_MODE='ask', check_host_keys='accept') diff --git a/tests/stream_test.py b/tests/stream_test.py deleted file mode 100644 index d844e610..00000000 --- a/tests/stream_test.py +++ /dev/null @@ -1,33 +0,0 @@ - -import unittest2 -import mock - -import mitogen.core - -import testlib - - -class ReceiveOneTest(testlib.TestCase): - klass = mitogen.core.Stream - - def test_corruption(self): - broker = mock.Mock() - router = mock.Mock() - - stream = self.klass(router, 1) - junk = mitogen.core.b('x') * stream.HEADER_LEN - stream._input_buf = [junk] - stream._input_buf_len = len(junk) - - capture = testlib.LogCapturer() - capture.start() - ret = stream._receive_one(broker) - #self.assertEquals(1, broker.stop_receive.mock_calls) - capture.stop() - - self.assertFalse(ret) - self.assertTrue((self.klass.corrupt_msg % (junk,)) in capture.raw()) - - -if __name__ == '__main__': - unittest2.main() diff --git a/tests/su_test.py b/tests/su_test.py index 2af17c6e..320f9cef 100644 --- a/tests/su_test.py +++ b/tests/su_test.py @@ -2,8 +2,7 @@ import os import mitogen -import mitogen.lxd -import mitogen.parent +import mitogen.su import unittest2 @@ -11,22 +10,64 @@ import testlib class ConstructorTest(testlib.RouterMixin, testlib.TestCase): - su_path = testlib.data_path('stubs/stub-su.py') + stub_su_path = testlib.data_path('stubs/stub-su.py') def run_su(self, **kwargs): context = self.router.su( - su_path=self.su_path, + su_path=self.stub_su_path, **kwargs ) argv = eval(context.call(os.getenv, 'ORIGINAL_ARGV')) return context, argv - def test_basic(self): context, argv = self.run_su() self.assertEquals(argv[1], 'root') self.assertEquals(argv[2], '-c') +class SuTest(testlib.DockerMixin, testlib.TestCase): + stub_su_path = testlib.data_path('stubs/stub-su.py') + + def test_slow_auth_failure(self): + # #363: old input loop would fail to spot auth failure because of + # scheduling vs. su calling write() twice. + os.environ['DO_SLOW_AUTH_FAILURE'] = '1' + try: + self.assertRaises(mitogen.su.PasswordError, + lambda: self.router.su(su_path=self.stub_su_path) + ) + finally: + del os.environ['DO_SLOW_AUTH_FAILURE'] + + def test_password_required(self): + ssh = self.docker_ssh( + username='mitogen__has_sudo', + password='has_sudo_password', + ) + e = self.assertRaises(mitogen.core.StreamError, + lambda: self.router.su(via=ssh) + ) + self.assertTrue(mitogen.su.password_required_msg in str(e)) + + def test_password_incorrect(self): + ssh = self.docker_ssh( + username='mitogen__has_sudo', + password='has_sudo_password', + ) + e = self.assertRaises(mitogen.core.StreamError, + lambda: self.router.su(via=ssh, password='x') + ) + self.assertTrue(mitogen.su.password_incorrect_msg in str(e)) + + def test_password_okay(self): + ssh = self.docker_ssh( + username='mitogen__has_sudo', + password='has_sudo_password', + ) + context = self.router.su(via=ssh, password='rootpassword') + self.assertEquals(0, context.call(os.getuid)) + + if __name__ == '__main__': unittest2.main() diff --git a/tests/sudo_test.py b/tests/sudo_test.py index 1d10ba9a..9ecf103d 100644 --- a/tests/sudo_test.py +++ b/tests/sudo_test.py @@ -2,8 +2,7 @@ import os import mitogen -import mitogen.lxd -import mitogen.parent +import mitogen.sudo import unittest2 @@ -79,7 +78,7 @@ class NonEnglishPromptTest(testlib.DockerMixin, testlib.TestCase): e = self.assertRaises(mitogen.core.StreamError, lambda: self.router.sudo(via=ssh) ) - self.assertTrue(mitogen.sudo.Stream.password_required_msg in str(e)) + self.assertTrue(mitogen.sudo.password_required_msg in str(e)) def test_password_incorrect(self): ssh = self.docker_ssh( @@ -91,7 +90,7 @@ class NonEnglishPromptTest(testlib.DockerMixin, testlib.TestCase): e = self.assertRaises(mitogen.core.StreamError, lambda: self.router.sudo(via=ssh, password='x') ) - self.assertTrue(mitogen.sudo.Stream.password_incorrect_msg in str(e)) + self.assertTrue(mitogen.sudo.password_incorrect_msg in str(e)) def test_password_okay(self): ssh = self.docker_ssh( @@ -103,7 +102,7 @@ class NonEnglishPromptTest(testlib.DockerMixin, testlib.TestCase): e = self.assertRaises(mitogen.core.StreamError, lambda: self.router.sudo(via=ssh, password='rootpassword') ) - self.assertTrue(mitogen.sudo.Stream.password_incorrect_msg in str(e)) + self.assertTrue(mitogen.sudo.password_incorrect_msg in str(e)) if __name__ == '__main__': diff --git a/tests/testlib.py b/tests/testlib.py index 04a48d84..3eeaa461 100644 --- a/tests/testlib.py +++ b/tests/testlib.py @@ -283,7 +283,11 @@ class LogCapturer(object): self.logger.level = logging.DEBUG def raw(self): - return self.sio.getvalue() + s = self.sio.getvalue() + # Python 2.x logging package hard-wires UTF-8 output. + if isinstance(s, mitogen.core.BytesType): + s = s.decode('utf-8') + return s def msgs(self): return self.handler.msgs @@ -327,17 +331,36 @@ class TestCase(unittest2.TestCase): for name in counts: assert counts[name] == 1, \ - 'Found %d copies of thread %r running after tests.' % (name,) + 'Found %d copies of thread %r running after tests.' % ( + counts[name], name + ) def _teardown_check_fds(self): mitogen.core.Latch._on_fork() if get_fd_count() != self._fd_count_before: - import os; os.system('lsof -p %s' % (os.getpid(),)) + import os; os.system('lsof -w -p %s' % (os.getpid(),)) assert 0, "%s leaked FDs. Count before: %s, after: %s" % ( self, self._fd_count_before, get_fd_count(), ) + def _teardown_check_zombies(self): + try: + pid, status = os.waitpid(0, os.WNOHANG) + except OSError: + return # ECHILD + + if pid: + assert 0, "%s failed to reap subprocess %d (status %d)." % ( + self, pid, status + ) + + print() + print('Children of unit test process:') + os.system('ps uww --ppid ' + str(os.getpid())) + assert 0, "%s leaked still-running subprocesses." % (self,) + def tearDown(self): + self._teardown_check_zombies() self._teardown_check_threads() self._teardown_check_fds() super(TestCase, self).tearDown() diff --git a/tests/timer_test.py b/tests/timer_test.py new file mode 100644 index 00000000..14a9c080 --- /dev/null +++ b/tests/timer_test.py @@ -0,0 +1,189 @@ + +import time + +import mock +import unittest2 + +import mitogen.core +import mitogen.parent + +import testlib + + +class TimerListMixin(object): + klass = mitogen.parent.TimerList + + def setUp(self): + self.list = self.klass() + + +class GetTimeoutTest(TimerListMixin, testlib.TestCase): + def test_empty(self): + self.assertEquals(None, self.list.get_timeout()) + + def test_one_event(self): + self.list.schedule(2, lambda: None) + self.list._now = lambda: 1 + self.assertEquals(1, self.list.get_timeout()) + + def test_two_events_same_moment(self): + self.list.schedule(2, lambda: None) + self.list.schedule(2, lambda: None) + self.list._now = lambda: 1 + self.assertEquals(1, self.list.get_timeout()) + + def test_two_events(self): + self.list.schedule(2, lambda: None) + self.list.schedule(3, lambda: None) + self.list._now = lambda: 1 + self.assertEquals(1, self.list.get_timeout()) + + def test_two_events_expired(self): + self.list.schedule(2, lambda: None) + self.list.schedule(3, lambda: None) + self.list._now = lambda: 3 + self.assertEquals(0, self.list.get_timeout()) + + def test_two_events_in_past(self): + self.list.schedule(2, lambda: None) + self.list.schedule(3, lambda: None) + self.list._now = lambda: 30 + self.assertEquals(0, self.list.get_timeout()) + + def test_two_events_in_past(self): + self.list.schedule(2, lambda: None) + self.list.schedule(3, lambda: None) + self.list._now = lambda: 30 + self.assertEquals(0, self.list.get_timeout()) + + def test_one_cancelled(self): + t1 = self.list.schedule(2, lambda: None) + t2 = self.list.schedule(3, lambda: None) + self.list._now = lambda: 0 + t1.cancel() + self.assertEquals(3, self.list.get_timeout()) + + def test_two_cancelled(self): + t1 = self.list.schedule(2, lambda: None) + t2 = self.list.schedule(3, lambda: None) + self.list._now = lambda: 0 + t1.cancel() + t2.cancel() + self.assertEquals(None, self.list.get_timeout()) + + +class ScheduleTest(TimerListMixin, testlib.TestCase): + def test_in_past(self): + self.list._now = lambda: 30 + timer = self.list.schedule(29, lambda: None) + self.assertEquals(29, timer.when) + self.assertEquals(0, self.list.get_timeout()) + + def test_in_future(self): + self.list._now = lambda: 30 + timer = self.list.schedule(31, lambda: None) + self.assertEquals(31, timer.when) + self.assertEquals(1, self.list.get_timeout()) + + def test_same_moment(self): + self.list._now = lambda: 30 + timer = self.list.schedule(31, lambda: None) + timer2 = self.list.schedule(31, lambda: None) + self.assertEquals(31, timer.when) + self.assertEquals(31, timer2.when) + self.assertTrue(timer is not timer2) + self.assertEquals(1, self.list.get_timeout()) + + +class ExpireTest(TimerListMixin, testlib.TestCase): + def test_in_past(self): + timer = self.list.schedule(29, mock.Mock()) + self.list._now = lambda: 30 + self.list.expire() + self.assertEquals(1, len(timer.func.mock_calls)) + + def test_in_future(self): + timer = self.list.schedule(29, mock.Mock()) + self.list._now = lambda: 28 + self.list.expire() + self.assertEquals(0, len(timer.func.mock_calls)) + + def test_same_moment(self): + timer = self.list.schedule(29, mock.Mock()) + timer2 = self.list.schedule(29, mock.Mock()) + self.list._now = lambda: 29 + self.list.expire() + self.assertEquals(1, len(timer.func.mock_calls)) + self.assertEquals(1, len(timer2.func.mock_calls)) + + def test_cancelled(self): + self.list._now = lambda: 29 + timer = self.list.schedule(29, mock.Mock()) + timer.cancel() + self.assertEquals(None, self.list.get_timeout()) + self.list._now = lambda: 29 + self.list.expire() + self.assertEquals(0, len(timer.func.mock_calls)) + self.assertEquals(None, self.list.get_timeout()) + + +class CancelTest(TimerListMixin, testlib.TestCase): + def test_single_cancel(self): + self.list._now = lambda: 29 + timer = self.list.schedule(29, mock.Mock()) + timer.cancel() + self.list.expire() + self.assertEquals(0, len(timer.func.mock_calls)) + + def test_double_cancel(self): + self.list._now = lambda: 29 + timer = self.list.schedule(29, mock.Mock()) + timer.cancel() + timer.cancel() + self.list.expire() + self.assertEquals(0, len(timer.func.mock_calls)) + + +@mitogen.core.takes_econtext +def do_timer_test_econtext(econtext): + do_timer_test(econtext.broker) + + +def do_timer_test(broker): + now = time.time() + latch = mitogen.core.Latch() + broker.defer(lambda: + broker.timers.schedule( + now + 0.250, + lambda: latch.put('hi'), + ) + ) + + assert 'hi' == latch.get() + assert time.time() > (now + 0.250) + + +class BrokerTimerTest(testlib.TestCase): + klass = mitogen.master.Broker + + def test_call_later(self): + broker = self.klass() + try: + do_timer_test(broker) + finally: + broker.shutdown() + broker.join() + + def test_child_upgrade(self): + router = mitogen.master.Router() + try: + c = router.local() + c.call(mitogen.parent.upgrade_router) + c.call(do_timer_test_econtext) + finally: + router.broker.shutdown() + router.broker.join() + + +if __name__ == '__main__': + unittest2.main() diff --git a/tests/types_test.py b/tests/types_test.py index 8f120931..8e441c65 100644 --- a/tests/types_test.py +++ b/tests/types_test.py @@ -16,6 +16,11 @@ from mitogen.core import b import testlib +#### +#### see also message_test.py / PickledTest +#### + + class BlobTest(testlib.TestCase): klass = mitogen.core.Blob diff --git a/tests/unix_test.py b/tests/unix_test.py index 02dc11a4..cb8c08f5 100644 --- a/tests/unix_test.py +++ b/tests/unix_test.py @@ -67,12 +67,12 @@ class ListenerTest(testlib.RouterMixin, testlib.TestCase): klass = mitogen.unix.Listener def test_constructor_basic(self): - listener = self.klass(router=self.router) + listener = self.klass.build_stream(router=self.router) capture = testlib.LogCapturer() capture.start() try: - self.assertFalse(mitogen.unix.is_path_dead(listener.path)) - os.unlink(listener.path) + self.assertFalse(mitogen.unix.is_path_dead(listener.protocol.path)) + os.unlink(listener.protocol.path) # ensure we catch 0 byte read error log message self.broker.shutdown() self.broker.join() @@ -86,25 +86,28 @@ class ClientTest(testlib.TestCase): def _try_connect(self, path): # give server a chance to setup listener - for x in range(10): + timeout = time.time() + 30.0 + while True: try: return mitogen.unix.connect(path) except socket.error: - if x == 9: + if time.time() > timeout: raise time.sleep(0.1) def _test_simple_client(self, path): router, context = self._try_connect(path) - self.assertEquals(0, context.context_id) - self.assertEquals(1, mitogen.context_id) - self.assertEquals(0, mitogen.parent_id) - resp = context.call_service(service_name=MyService, method_name='ping') - self.assertEquals(mitogen.context_id, resp['src_id']) - self.assertEquals(0, resp['auth_id']) - router.broker.shutdown() - router.broker.join() - os.unlink(path) + try: + self.assertEquals(0, context.context_id) + self.assertEquals(1, mitogen.context_id) + self.assertEquals(0, mitogen.parent_id) + resp = context.call_service(service_name=MyService, method_name='ping') + self.assertEquals(mitogen.context_id, resp['src_id']) + self.assertEquals(0, resp['auth_id']) + finally: + router.broker.shutdown() + router.broker.join() + os.unlink(path) @classmethod def _test_simple_server(cls, path): @@ -112,7 +115,7 @@ class ClientTest(testlib.TestCase): latch = mitogen.core.Latch() try: try: - listener = cls.klass(path=path, router=router) + listener = cls.klass.build_stream(path=path, router=router) pool = mitogen.service.Pool(router=router, services=[ MyService(latch=latch, router=router), ])