def dispatch_one(self, msg):
if not all(p.is_authorized(self, msg) for p in self.policies):
LOG.error('%r: unauthorized message %r', self, msg)
msg.reply(mitogen.core.CallError('Unauthorized'))
return
if len(msg.data) > self.max_message_size:
LOG.error('%r: larger than permitted size: %r', self, msg)
msg.reply(mitogen.core.CallError('Message size exceeded'))
return
args = msg.unpickle(throw=False)
if (args == mitogen.core._DEAD
or isinstance(args, mitogen.core.CallError)
or not self.validate_args(args)):
LOG.warning('Received junk message: %r', args)
return
try:
response = self.dispatch(args, msg)
if response is not self.NO_REPLY:
msg.reply(response)
except Exception, e:
LOG.exception('While invoking %r.dispatch()', self)
msg.reply(mitogen.core.CallError(e))
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 _on_forward_log(self, msg):
if msg.is_dead:
return
context = self._router.context_by_id(msg.src_id)
if context is None:
LOG.error('%s: dropping log from unknown context %d', self,
msg.src_id)
return
name, level_s, s = msg.data.decode('utf-8', 'replace').split('\x00', 2)
logger_name = '%s.[%s]' % (name, context.name)
logger = self._cache.get(logger_name)
if logger is None:
self._cache[logger_name] = logger = logging.getLogger(logger_name)
# See logging.Handler.makeRecord()
record = logging.LogRecord(
name=logger.name,
level=int(level_s),
pathname='(unknown file)',
lineno=0,
msg=s,
args=(),
exc_info=None,
)
record.mitogen_message = s
record.mitogen_context = self._router.context_by_id(msg.src_id)
record.mitogen_name = name
logger.handle(record)
def _accept_client(self, sock):
sock.setblocking(True)
try:
pid, = struct.unpack('>L', sock.recv(4))
except (struct.error, 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 _on_forward_log(self, msg):
if msg.is_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('%s: dropping log from unknown context ID %d', self,
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.decode('latin1').split('\x00', 2)
logger.log(int(level_s),
'%s: %s',
name,
s,
extra={
'mitogen_message': s,
'mitogen_context':
self._router.context_by_id(msg.src_id),
'mitogen_name': name,
})
def _on_forward_log(self, msg):
if msg.is_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,
extra={
'mitogen_message': s,
'mitogen_context':
self._router.context_by_id(msg.src_id),
'mitogen_name': name,
})
def _on_forward_module(self, msg):
if msg.is_dead:
return
context_id_s, _, fullname = msg.data.partition('\x00')
context_id = int(context_id_s)
stream = self.router.stream_by_id(context_id)
if stream.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:
return
LOG.debug('%r._on_forward_module() sending %r to %r via %r', self,
fullname, context_id, stream.remote_id)
self._send_module_and_related(stream, fullname)
if stream.remote_id != context_id:
stream._send(
mitogen.core.Message(
data=msg.data,
handle=mitogen.core.FORWARD_MODULE,
dst_id=stream.remote_id,
))
def _build_tuple(self, fullname):
if mitogen.core.is_blacklisted_import(self, fullname):
raise ImportError('blacklisted')
if fullname in self._cache:
return self._cache[fullname]
path, source, is_pkg = self._finder.get_module_source(fullname)
if source is None:
LOG.error('_build_tuple(%r): could not locate source', fullname)
tup = self._make_negative_response(fullname)
self._cache[fullname] = tup
return tup
if is_pkg:
pkg_present = get_child_modules(path)
LOG.debug('_build_tuple(%r, %r) -> %r', path, fullname,
pkg_present)
else:
pkg_present = None
if fullname == '__main__':
source = self.neutralize_main(path, source)
compressed = mitogen.core.Blob(zlib.compress(source, 9))
related = [
to_text(name) for name in self._finder.find_related(fullname)
if not mitogen.core.is_blacklisted_import(self, name)
]
# 0:fullname 1:pkg_present 2:path 3:compressed 4:related
tup = (to_text(fullname), pkg_present, to_text(path), compressed,
related)
self._cache[fullname] = tup
return tup
def _on_forward_log(self, msg):
if msg.is_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('%s: dropping log from unknown context ID %d', self,
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.decode('latin1').split('\x00', 2)
# See logging.Handler.makeRecord()
record = logging.LogRecord(
name=logger.name,
level=int(level_s),
pathname='(unknown file)',
lineno=0,
msg=('%s: %s' % (name, s)),
args=(),
exc_info=None,
)
record.mitogen_message = s
record.mitogen_context = self._router.context_by_id(msg.src_id)
record.mitogen_name = name
logger.handle(record)
def _build_tuple(self, fullname):
if mitogen.core.is_blacklisted_import(self, fullname):
raise ImportError('blacklisted')
if fullname in self._cache:
return self._cache[fullname]
path, source, is_pkg = self._finder.get_module_source(fullname)
if source is None:
LOG.error('_build_tuple(%r): could not locate source', fullname)
tup = fullname, None, None, None, ()
self._cache[fullname] = tup
return tup
if source is None:
raise ImportError('could not find %r' % (fullname, ))
if is_pkg:
pkg_present = get_child_modules(path)
LOG.debug('_build_tuple(%r, %r) -> %r', path, fullname,
pkg_present)
else:
pkg_present = None
if fullname == '__main__':
source = self.neutralize_main(source)
compressed = zlib.compress(source)
related = list(self._finder.find_related(fullname))
# 0:fullname 1:pkg_present 2:path 3:compressed 4:related
tup = fullname, pkg_present, path, compressed, related
self._cache[fullname] = tup
return tup
def del_route(self, target_id):
LOG.debug('%r.del_route(%r)', self, target_id)
try:
del self._stream_by_id[target_id]
except KeyError:
LOG.error('%r: cant delete route to %r: no such stream',
self, target_id)
def get(cls, context, path, out_fp):
"""
Streamily download a file from the connection multiplexer process in
the controller.
:param mitogen.core.Context context:
Reference to the context hosting the FileService that will be used
to fetch the file.
:param bytes path:
FileService registered name of the input file.
:param bytes out_path:
Name of the output path on the local disk.
:returns:
Tuple of (`ok`, `metadata`), where `ok` is :data:`True` on success,
or :data:`False` if the transfer was interrupted and the output
should be discarded.
`metadata` is a dictionary of file metadata as documented in
:meth:`fetch`.
"""
LOG.debug('get_file(): fetching %r from %r', path, context)
t0 = time.time()
recv = mitogen.core.Receiver(router=context.router)
metadata = context.call_service(
service_name=cls.name(),
method_name='fetch',
path=path,
sender=recv.to_sender(),
)
received_bytes = 0
for chunk in recv:
s = chunk.unpickle()
LOG.debug('get_file(%r): received %d bytes', path, len(s))
context.call_service_async(
service_name=cls.name(),
method_name='acknowledge',
size=len(s),
).close()
out_fp.write(s)
received_bytes += len(s)
ok = received_bytes == metadata['size']
if received_bytes < metadata['size']:
LOG.error(
'get_file(%r): receiver was closed early, controller '
'may be shutting down, or the file was truncated '
'during transfer. Expected %d bytes, received %d.', path,
metadata['size'], received_bytes)
elif received_bytes > metadata['size']:
LOG.error(
'get_file(%r): the file appears to have grown '
'while transfer was in progress. Expected %d '
'bytes, received %d.', path, metadata['size'], received_bytes)
LOG.debug('target.get_file(): fetched %d bytes of %r from %r in %dms',
metadata['size'], path, context, 1000 * (time.time() - t0))
return ok, metadata
def add_route(self, target_id, stream):
LOG.debug('%r.add_route(%r, %r)', self, target_id, stream)
assert isinstance(target_id, int)
assert isinstance(stream, Stream)
try:
self._stream_by_id[target_id] = stream
except KeyError:
LOG.error('%r: cant add route to %r via %r: no such stream',
self, target_id, stream)
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 _worker_run(self):
while True:
try:
msg = self._select.get()
except (mitogen.core.ChannelError, mitogen.core.LatchError):
e = sys.exc_info()[1]
LOG.error('%r: channel or latch closed, exitting: %s', self, e)
return
service = msg.receiver.service
try:
service.dispatch_one(msg)
except Exception:
LOG.exception('While handling %r using %r', msg, service)
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 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 get(cls, context, path, out_fp):
"""
Streamily download a file from the connection multiplexer process in
the controller.
:param mitogen.core.Context context:
Reference to the context hosting the FileService that will be used
to fetch the file.
:param bytes path:
FileService registered name of the input file.
:param bytes out_path:
Name of the output path on the local disk.
:returns:
:data:`True` on success, or :data:`False` if the transfer was
interrupted and the output should be discarded.
"""
LOG.debug('get_file(): fetching %r from %r', path, context)
t0 = time.time()
recv = mitogen.core.Receiver(router=context.router)
metadata = context.call_service(
service_name=cls.name(),
method_name='fetch',
path=path,
sender=recv.to_sender(),
)
for chunk in recv:
s = chunk.unpickle()
LOG.debug('get_file(%r): received %d bytes', path, len(s))
context.call_service_async(
service_name=cls.name(),
method_name='acknowledge',
size=len(s),
).close()
out_fp.write(s)
ok = out_fp.tell() == metadata['size']
if not ok:
LOG.error(
'get_file(%r): receiver was closed early, controller '
'is likely shutting down.', path)
LOG.debug('target.get_file(): fetched %d bytes of %r from %r in %dms',
metadata['size'], path, context, 1000 * (time.time() - t0))
return ok, metadata
def dispatch_one(self, msg):
if len(msg.data) > self.max_message_size:
LOG.error('%r: larger than permitted size: %r', self, msg)
msg.reply(mitogen.core.CallError('Message size exceeded'))
return
args = msg.unpickle(throw=False)
if (args == mitogen.core._DEAD
or isinstance(args, mitogen.core.CallError)
or not self.validate_args(args)):
LOG.warning('Received junk message: %r', args)
return
try:
msg.reply(self.dispatch(args, msg))
except Exception, e:
LOG.exception('While invoking %r.dispatch()', self)
msg.reply(mitogen.core.CallError(e))
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 _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')
class Service(object):
#: If ``None``, a handle is dynamically allocated, otherwise the fixed
#: integer handle to use.
handle = None
max_message_size = 0
def __init__(self, router):
self.router = router
self.recv = mitogen.core.Receiver(router, self.handle)
self.handle = self.recv.handle
self.running = True
def validate_args(self, args):
return True
def run_once(self):
try:
msg = self.recv.get()
except mitogen.core.ChannelError, e:
# Channel closed due to broker shutdown, exit gracefully.
LOG.debug('%r: channel closed: %s', self, e)
self.running = False
return
if len(msg.data) > self.max_message_size:
LOG.error('%r: larger than permitted size: %r', self, msg)
msg.reply(mitogen.core.CallError('Message size exceeded'))
return
args = msg.unpickle(throw=False)
if (args == mitogen.core._DEAD
or isinstance(args, mitogen.core.CallError)
or not self.validate_args(args)):
LOG.warning('Received junk message: %r', args)
return
try:
msg.reply(self.dispatch(args, msg))
except Exception, e:
LOG.exception('While invoking %r.dispatch()', self)
msg.reply(mitogen.core.CallError(e))
def forward(self, path, context):
LOG.debug('%r.forward(%r, %r)', self, path, context)
if path not in self._cache:
LOG.error('%r: %r is not in local cache', self, path)
return
self._forward(path, context)
