def _start_slave(src_id, cmdline, router):
"""
This runs in the target context, it is invoked by _fakessh_main running in
the fakessh context immediately after startup. It starts the slave process
(the the point where it has a stdin_handle to target but not stdout_chan to
write to), and waits for main to.
"""
LOG.debug('_start_slave(%r, %r)', router, cmdline)
proc = subprocess.Popen(
cmdline,
# SSH server always uses user's shell.
shell=True,
# SSH server always executes new commands in the user's HOME.
cwd=os.path.expanduser('~'),
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
)
process = Process(
router,
proc.stdin.fileno(),
proc.stdout.fileno(),
proc,
)
return process.control_handle, process.stdin_handle
def _accept_client(self, sock):
sock.setblocking(True)
try:
pid, = struct.unpack('>L', sock.recv(4))
except socket.error:
LOG.error('%r: failed to read remote identity: %s',
self, sys.exc_info()[1])
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()))
except socket.error:
LOG.error('%r: failed to assign identity to PID %d: %s',
self, pid, sys.exc_info()[1])
return
stream.accept(sock.fileno(), sock.fileno())
self._router.register(context, stream)
def connect(path, broker=None):
LOG.debug('unix.connect(path=%r)', path)
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
sock.connect(path)
sock.send(struct.pack('>L', os.getpid()))
mitogen.context_id, remote_id, pid = struct.unpack('>LLL', sock.recv(12))
mitogen.parent_id = remote_id
mitogen.parent_ids = [remote_id]
LOG.debug('unix.connect(): local ID is %r, remote is %r',
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.name = u'unix_listener.%d' % (pid,)
context = mitogen.parent.Context(router, remote_id)
router.register(context, stream)
mitogen.core.listen(router.broker, 'shutdown',
lambda: router.disconnect_stream(stream))
sock.close()
return router, context
def _on_control(self, msg):
if not msg.is_dead:
command, arg = msg.unpickle(throw=False)
LOG.debug('%r._on_control(%r, %s)', self, command, arg)
func = getattr(self, '_on_%s' % (command,), None)
if func:
return func(msg, arg)
LOG.warning('%r: unknown command %r', self, command)
def __init__(self, router, path=None, backlog=100):
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)
def on_shutdown(self):
"""
Respond to shutdown by sending close() to every target, allowing their
receive loop to exit and clean up gracefully.
"""
LOG.debug('%r.on_shutdown()', self)
for stream, state in self._state_by_stream.items():
state.lock.acquire()
try:
for sender, fp in reversed(state.jobs):
sender.close()
fp.close()
state.jobs.pop()
finally:
state.lock.release()
def store_and_forward(self, path, data, context):
LOG.debug('%r.store_and_forward(%r, %r, %r) %r', self, path, data,
context,
threading.currentThread().getName())
self._lock.acquire()
try:
self._cache[path] = data
waiters = self._waiters.pop(path, [])
finally:
self._lock.release()
if context.context_id != mitogen.context_id:
self._forward(context, path)
for callback in waiters:
callback()
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.
"""
match = self.MAIN_RE.search(src)
if match:
return src[:match.start()]
if b('mitogen.main(') in src:
return src
LOG.error(self.main_guard_msg, path)
raise ImportError('refused')
def _on_service_call(self, recv, msg):
self._validate(msg)
service_name, method_name, kwargs = msg.unpickle()
try:
invoker = self.get_invoker(service_name, msg)
return invoker.invoke(method_name, kwargs, msg)
except mitogen.core.CallError:
e = sys.exc_info()[1]
LOG.warning('%r: call error: %s: %s', self, msg, e)
msg.reply(e)
except Exception:
LOG.exception('%r: while invoking %r of %r', self, method_name,
service_name)
e = sys.exc_info()[1]
msg.reply(mitogen.core.CallError(e))
def propagate_route(self, target, via):
self.add_route(target.context_id, via.context_id)
child = via
parent = via.via
while parent is not None:
LOG.debug('Adding route to %r for %r via %r', parent, target,
child)
parent.send(
mitogen.core.Message(
data='%s\x00%s' % (target.context_id, child.context_id),
handle=mitogen.core.ADD_ROUTE,
))
child = parent
parent = parent.via
def get(self, path):
self._lock.acquire()
try:
if path in self._cache:
return self._cache[path]
waiters = self._waiters.setdefault(path, [])
latch = mitogen.core.Latch()
waiters.append(lambda: latch.put(None))
finally:
self._lock.release()
LOG.debug('%r.get(%r) waiting for uncached file to arrive', self, path)
latch.get()
LOG.debug('%r.get(%r) -> %r', self, path, self._cache[path])
return self._cache[path]
def _run(self):
while True:
tup = self._pop()
if tup is None:
return
method_name, kwargs, msg = tup
try:
super(SerializedInvoker, self).invoke(method_name, kwargs, msg)
except mitogen.core.CallError:
e = sys.exc_info()[1]
LOG.warning('%r: call error: %s: %s', self, msg, e)
msg.reply(e)
except Exception:
LOG.exception('%r: while invoking %s()', self, method_name)
msg.reply(mitogen.core.Message.dead())
def find(self, fullname):
"""
Find `fullname` using :func:`pkgutil.find_loader`.
"""
try:
# Pre-'import spec' this returned None, in Python3.6 it raises
# ImportError.
loader = pkgutil.find_loader(fullname)
except ImportError:
e = sys.exc_info()[1]
LOG.debug('%r._get_module_via_pkgutil(%r): %s', self, fullname, e)
return None
IOLOG.debug('%r._get_module_via_pkgutil(%r) -> %r', self, fullname,
loader)
if not loader:
return
try:
path, is_special = _py_filename(loader.get_filename(fullname))
source = loader.get_source(fullname)
is_pkg = loader.is_package(fullname)
# workaround for special python modules that might only exist in memory
if is_special and is_pkg and not source:
source = '\n'
except (AttributeError, ImportError):
# - Per PEP-302, get_source() and is_package() are optional,
# calling them may throw AttributeError.
# - get_filename() may throw ImportError if pkgutil.find_loader()
# picks a "parent" package's loader for some crap that's been
# stuffed in sys.modules, for example in the case of urllib3:
# "loader for urllib3.contrib.pyopenssl cannot handle
# requests.packages.urllib3.contrib.pyopenssl"
e = sys.exc_info()[1]
LOG.debug('%r: loading %r using %r failed: %s', self, fullname,
loader, e)
return
if path is None or source is None:
return
if isinstance(source, mitogen.core.UnicodeType):
# get_source() returns "string" according to PEP-302, which was
# reinterpreted for Python 3 to mean a Unicode string.
source = source.encode('utf-8')
return path, source, is_pkg
def _on_del_route(self, msg):
if msg == mitogen.core._DEAD:
return
target_id = int(msg.data)
registered_stream = self.router.stream_by_id(target_id)
stream = self.router.stream_by_id(msg.auth_id)
if registered_stream != stream:
LOG.error('Received DEL_ROUTE for %d from %r, expected %r',
target_id, stream, registered_stream)
return
LOG.debug('Deleting route to %d via %r', target_id, stream)
stream.routes.discard(target_id)
self.router.del_route(target_id)
self.propagate(mitogen.core.DEL_ROUTE, target_id)
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 _on_forward_log(self, msg):
if msg == mitogen.core._DEAD:
return
logger = self._cache.get(msg.src_id)
if logger is None:
context = self._router.context_by_id(msg.src_id)
if context is None:
LOG.error('FORWARD_LOG received from src_id %d', msg.src_id)
return
name = '%s.%s' % (RLOG.name, context.name)
self._cache[msg.src_id] = logger = logging.getLogger(name)
name, level_s, s = msg.data.split('\x00', 2)
logger.log(int(level_s), '%s: %s', name, s)
def _send_load_module(self, stream, fullname):
if fullname not in stream.sent_modules:
LOG.debug('_send_load_module(%r, %r)', stream, fullname)
tup = self._build_tuple(fullname)
msg = mitogen.core.Message.pickled(
tup,
dst_id=stream.remote_id,
handle=mitogen.core.LOAD_MODULE,
)
self._router._async_route(msg)
stream.sent_modules.add(fullname)
if tup[2] is not None:
self.good_load_module_count += 1
self.good_load_module_size += len(msg.data)
else:
self.bad_load_module_count += 1
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)
self.buf += buf
while '\n' in self.buf:
lines = self.buf.split('\n')
self.buf = lines[-1]
for line in lines[:-1]:
LOG.debug('%r: %r', self, line.rstrip())
def _send_load_module(self, stream, fullname):
if fullname not in stream.protocol.sent_modules:
tup = self._build_tuple(fullname)
msg = mitogen.core.Message.pickled(
tup,
dst_id=stream.protocol.remote_id,
handle=mitogen.core.LOAD_MODULE,
)
LOG.debug('%s: sending %s (%.2f KiB) to %s', self, fullname,
len(msg.data) / 1024.0, stream.name)
self._router._async_route(msg)
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)
else:
self.bad_load_module_count += 1
def fetch(self, path, sender, msg):
"""
Start a transfer for a registered path.
:param str path:
File path.
:param mitogen.core.Sender sender:
Sender to receive file data.
:returns:
Dict containing the file metadata:
* ``size``: File size in bytes.
* ``mode``: Integer file mode.
* ``owner``: Owner account name on host machine.
* ``group``: Owner group name on host machine.
* ``mtime``: Floating point modification time.
* ``ctime``: Floating point change time.
:raises Error:
Unregistered path, or Sender did not match requestee context.
"""
if path not in self._metadata_by_path:
raise Error(self.unregistered_msg)
if msg.src_id != sender.context.context_id:
raise Error(self.context_mismatch_msg)
LOG.debug('Serving %r', path)
try:
fp = open(path, 'rb', self.IO_SIZE)
except IOError:
msg.reply(mitogen.core.CallError(sys.exc_info()[1]))
return
# Response must arrive first so requestee can begin receive loop,
# otherwise first ack won't arrive until all pending chunks were
# delivered. In that case max BDP would always be 128KiB, aka. max
# ~10Mbit/sec over a 100ms link.
msg.reply(self._metadata_by_path[path])
stream = self.router.stream_by_id(sender.context.context_id)
state = self._state_by_stream.setdefault(stream, FileStreamState())
state.lock.acquire()
try:
state.jobs.append((sender, fp))
self._schedule_pending_unlocked(state)
finally:
state.lock.release()
def poll(self, timeout=None):
changelist = self._changelist
self._changelist = []
events, _ = mitogen.core.io_op(self._kqueue.control, changelist, 32,
timeout)
for event in events:
fd = event.ident
if event.flags & select.KQ_EV_ERROR:
LOG.debug('ignoring stale event for fd %r: errno=%d: %s', fd,
event.data, errno.errorcode.get(event.data))
elif event.filter == select.KQ_FILTER_READ and fd in self._rfds:
# Events can still be read for an already-discarded fd.
mitogen.core._vv and IOLOG.debug('%r: POLLIN: %r', self, fd)
yield self._rfds[fd]
elif event.filter == select.KQ_FILTER_WRITE and fd in self._wfds:
mitogen.core._vv and IOLOG.debug('%r: POLLOUT: %r', self, fd)
yield self._wfds[fd]
def _on_broker_exit(self):
super(Router, self)._on_broker_exit()
dct = self.get_stats()
dct['self'] = self
dct['get_module_ms'] = 1000 * dct['get_module_secs']
dct['good_load_module_size_kb'] = dct['good_load_module_size'] / 1024.0
dct['good_load_module_size_avg'] = (
(dct['good_load_module_size'] /
(float(dct['good_load_module_count']) or 1.0)) / 1024.0)
LOG.debug('%(self)r: stats: '
'%(get_module_count)d module requests in '
'%(get_module_ms)d ms, '
'%(good_load_module_count)d sent, '
'%(bad_load_module_count)d negative responses. '
'Sent %(good_load_module_size_kb).01f kb total, '
'%(good_load_module_size_avg).01f kb avg.' % dct)
def __init__(self, router, path=None, backlog=30):
self._router = router
self.path = path or tempfile.mktemp(prefix='mitogen_unix_')
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, 0600)
self._sock.listen(backlog)
mitogen.core.set_nonblock(self._sock.fileno())
mitogen.core.set_cloexec(self._sock.fileno())
self.path = self._sock.getsockname()
self.receive_side = mitogen.core.Side(self, self._sock.fileno())
router.broker.start_receive(self)
def _on_get_module(self, msg):
if msg.is_dead:
return
LOG.debug('%r._on_get_module(%r)', self, msg.data)
self.get_module_count += 1
stream = self._router.stream_by_id(msg.src_id)
fullname = msg.data.decode()
if fullname in stream.sent_modules:
LOG.warning('_on_get_module(): dup request for %r from %r',
fullname, stream)
t0 = time.time()
try:
self._send_module_and_related(stream, fullname)
finally:
self.get_module_secs += time.time() - t0
def acknowledge(self, size, msg):
"""
Acknowledge bytes received by a transfer target, scheduling new chunks
to keep the window full. This should be called for every chunk received
by the target.
"""
stream = self.router.stream_by_id(msg.src_id)
state = self._state_by_stream[stream]
state.lock.acquire()
try:
if state.unacked < size:
LOG.error('%r.acknowledge(src_id %d): unacked=%d < size %d',
self, msg.src_id, state.unacked, size)
state.unacked -= min(state.unacked, size)
self._schedule_pending_unlocked(state)
finally:
state.lock.release()
def _send_module_and_related(self, stream, fullname):
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.protocol.sent_modules:
# Parent hasn't been sent, so don't load submodule yet.
continue
self._send_load_module(stream, name)
self._send_load_module(stream, fullname)
except Exception:
LOG.debug('While importing %r', fullname, exc_info=True)
self._send_module_load_failed(stream, fullname)
def _get_module_via_pkgutil(self, fullname):
"""
Attempt to fetch source code via pkgutil. In an ideal world, this would
be the only required implementation of get_module().
"""
try:
# Pre-'import spec' this returned None, in Python3.6 it raises
# ImportError.
loader = pkgutil.find_loader(fullname)
except ImportError:
e = sys.exc_info()[1]
LOG.debug('%r._get_module_via_pkgutil(%r): %s', self, fullname, e)
return None
IOLOG.debug('%r._get_module_via_pkgutil(%r) -> %r', self, fullname,
loader)
if not loader:
return
try:
path = self._py_filename(loader.get_filename(fullname))
source = loader.get_source(fullname)
is_pkg = loader.is_package(fullname)
except (AttributeError, ImportError):
# - Per PEP-302, get_source() and is_package() are optional,
# calling them may throw AttributeError.
# - get_filename() may throw ImportError if pkgutil.find_loader()
# picks a "parent" package's loader for some crap that's been
# stuffed in sys.modules, for example in the case of urllib3:
# "loader for urllib3.contrib.pyopenssl cannot handle
# requests.packages.urllib3.contrib.pyopenssl"
e = sys.exc_info()[1]
LOG.debug('%r: loading %r using %r failed: %s', self, fullname,
loader, e)
return
if path is None or source is None:
return
if isinstance(source, mitogen.core.UnicodeType):
# get_source() returns "string" according to PEP-302, which was
# reinterpreted for Python 3 to mean a Unicode string.
source = source.encode('utf-8')
return path, source, is_pkg
def _get_module_via_pkgutil(self, fullname):
"""Attempt to fetch source code via pkgutil. In an ideal world, this
would be the only required implementation of get_module()."""
loader = pkgutil.find_loader(fullname)
LOG.debug('pkgutil._get_module_via_pkgutil(%r) -> %r', fullname,
loader)
if not loader:
return
try:
path = self._py_filename(loader.get_filename(fullname))
source = loader.get_source(fullname)
is_pkg = loader.is_package(fullname)
except AttributeError:
return
if path is not None and source is not None:
return path, source, is_pkg
def _invoke(self, method_name, kwargs, msg):
method = getattr(self.service, method_name)
if 'msg' in func_code(method).co_varnames:
kwargs['msg'] = msg # TODO: hack
no_reply = getattr(method, 'mitogen_service__no_reply', False)
ret = None
try:
ret = method(**kwargs)
if no_reply:
return Service.NO_REPLY
return ret
except Exception:
if no_reply:
LOG.exception('While calling no-reply method %s.%s',
type(self.service).__name__, func_name(method))
else:
raise
def _on_receive_message(self, msg):
method_name, kwargs = self._validate_message(msg)
method = getattr(self, method_name)
if 'msg' in method.func_code.co_varnames:
kwargs['msg'] = msg # TODO: hack
no_reply = getattr(method, 'mitogen_service__no_reply', False)
ret = None
try:
ret = method(**kwargs)
if no_reply:
return self.NO_REPLY
return ret
except Exception:
if no_reply:
LOG.exception('While calling no-reply method %s.%s',
type(self).__name__, method.func_name)
else:
raise
def create_child(*args):
parentfp, childfp = socket.socketpair()
pid = os.fork()
if not pid:
mitogen.core.set_block(childfp.fileno())
os.dup2(childfp.fileno(), 0)
os.dup2(childfp.fileno(), 1)
childfp.close()
parentfp.close()
os.execvp(args[0], args)
childfp.close()
# Decouple the socket from the lifetime of the Python socket object.
fd = os.dup(parentfp.fileno())
parentfp.close()
LOG.debug('create_child() child %d fd %d, parent %d, cmd: %s', pid, fd,
os.getpid(), Argv(args))
return pid, fd
def get(self, path):
"""
Fetch a file from the cache.
"""
assert isinstance(path, mitogen.core.UnicodeType)
self._lock.acquire()
try:
if path in self._cache:
return self._cache[path]
latch = mitogen.core.Latch()
waiters = self._waiters.setdefault(path, [])
waiters.append(lambda: latch.put(None))
finally:
self._lock.release()
LOG.debug('%r.get(%r) waiting for uncached file to arrive', self, path)
latch.get()
LOG.debug('%r.get(%r) -> %r', self, path, self._cache[path])
return self._cache[path]
def allocate_block(self):
"""
Allocate a block of IDs for use in a child context.
This function is safe to call from any thread.
:returns:
Tuple of the form `(id, end_id)` where `id` is the first usable ID
and `end_id` is the last usable ID.
"""
self.lock.acquire()
try:
id_ = self.next_id
self.next_id += self.BLOCK_SIZE
end_id = id_ + self.BLOCK_SIZE
LOG.debug('%r: allocating [%d..%d)', self, id_, end_id)
return id_, end_id
finally:
self.lock.release()
def _on_del_route(self, msg):
if msg.is_dead:
return
target_id = int(msg.data)
registered_stream = self.router.stream_by_id(target_id)
stream = self.router.stream_by_id(msg.auth_id)
if registered_stream != stream:
LOG.error('Received DEL_ROUTE for %d from %r, expected %r',
target_id, stream, registered_stream)
return
LOG.debug('Deleting route to %d via %r', target_id, stream)
stream.routes.discard(target_id)
self.router.del_route(target_id)
self.propagate(mitogen.core.DEL_ROUTE, target_id)
context = self.router.context_by_id(target_id, create=False)
if context:
mitogen.core.fire(context, 'disconnect')
def _on_add_route(self, msg):
if msg.is_dead:
return
target_id_s, _, target_name = msg.data.partition(':')
target_id = int(target_id_s)
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:
LOG.error('Cannot add duplicate route to %r via %r, '
'already have existing route via %r',
target_id, stream, current)
return
LOG.debug('Adding route to %d via %r', target_id, stream)
stream.routes.add(target_id)
self.router.add_route(target_id, stream)
self.propagate(mitogen.core.ADD_ROUTE, target_id, target_name)
def _fakessh_main(dest_context_id, econtext):
hostname, opts, args = parse_args()
if not hostname:
die('Missing hostname')
subsystem = False
for opt, optarg in opts:
if opt == '-s':
subsystem = True
else:
LOG.debug('Warning option %s %s is ignored.', opt, optarg)
LOG.debug('hostname: %r', hostname)
LOG.debug('opts: %r', opts)
LOG.debug('args: %r', args)
if subsystem:
die('-s <subsystem> is not yet supported')
if not args:
die('fakessh: login mode not supported and no command specified')
dest = mitogen.parent.Context(econtext.router, dest_context_id)
# Even though SSH receives an argument vector, it still cats the vector
# together before sending to the server, the server just uses /bin/sh -c to
# run the command. We must remain puke-for-puke compatible.
control_handle, stdin_handle = dest.call(_start_slave,
mitogen.context_id, ' '.join(args))
LOG.debug('_fakessh_main: received control_handle=%r, stdin_handle=%r',
control_handle, stdin_handle)
process = Process(econtext.router, 1, 0)
process.start_master(
stdin=mitogen.core.Sender(dest, stdin_handle),
control=mitogen.core.Sender(dest, control_handle),
)
process.wait()
process.control.put(('exit', None))
def _on_exit(self, msg, arg):
LOG.debug('on_exit: proc = %r', self.proc)
if self.proc:
self.proc.terminate()
else:
self.router.broker.shutdown()
def _on_proc_exit(self, status):
LOG.debug('%r._on_proc_exit(%r)', self, status)
self.control.put(('exit', status))
def _on_pump_disconnect(self):
LOG.debug('%r._on_pump_disconnect()', self)
mitogen.core.fire(self, 'disconnect')
self.stdin.close()
self.wake_event.set()
