import asyncio
import logging
import os
from typing import Any, Callable, Sequence
from .jsonutil import JsonObject, JsonValue
from .packages import BridgeConfig
from .protocol import CockpitProblem, CockpitProtocol, CockpitProtocolError
from .router import Endpoint, Router, RoutingRule
from .transports import SubprocessProtocol, SubprocessTransport
logger = logging.getLogger(__name__)
class PeerError(CockpitProblem):
pass
class PeerExited(Exception):
def __init__(self, exit_code: int):
self.exit_code = exit_code
class Peer(CockpitProtocol, SubprocessProtocol, Endpoint):
done_callbacks: 'list[Callable[[], None]]'
init_future: 'asyncio.Future[JsonObject] | None'
def __init__(self, router: Router):
super().__init__(router)
self.freeze_endpoint()
self.init_future = asyncio.get_running_loop().create_future()
self.done_callbacks = []
async def do_connect_transport(self) -> None:
raise NotImplementedError
async def spawn(self, argv: Sequence[str], env: Sequence[str], **kwargs: Any) -> asyncio.Transport:
loop = asyncio.get_running_loop()
user_env = dict(e.split('=', 1) for e in env)
return SubprocessTransport(loop, self, argv, env=dict(os.environ, **user_env), **kwargs)
async def start(self, init_host: 'str | None' = None, **kwargs: JsonValue) -> JsonObject:
"""Request that the Peer is started and connected to the router.
Creates the transport, connects it to the protocol, and participates in
exchanging of init messages. If anything goes wrong, the connection
will be closed and an exception will be raised.
The Peer starts out in a frozen state (ie: attempts to send messages to
it will initially be queued). If init_host is not None then an init
message is sent with the given 'host' field, plus any extra kwargs, and
the queue is thawed. Otherwise, the caller is responsible for sending
the init message and thawing the peer.
In any case, the return value is the init message from the peer.
"""
assert self.init_future is not None
def _connect_task_done(task: 'asyncio.Task[None]') -> None:
assert task is connect_task
try:
task.result()
except asyncio.CancelledError:
pass
except Exception as exc:
self.close(exc)
connect_task = asyncio.create_task(self.do_connect_transport())
connect_task.add_done_callback(_connect_task_done)
try:
init_message = await self.init_future
except (PeerExited, BrokenPipeError):
await connect_task
raise
finally:
self.init_future = None
if not connect_task.done():
connect_task.cancel()
if init_host is not None:
logger.debug(' sending init message back, host %s', init_host)
self.write_control(None, command='init', version=1, host=init_host, **kwargs)
self.thaw_endpoint()
return init_message
def start_in_background(self, init_host: 'str | None' = None, **kwargs: JsonValue) -> None:
def _start_task_done(task: 'asyncio.Task[JsonObject]') -> None:
assert task is start_task
try:
task.result()
except (OSError, PeerExited, CockpitProblem, asyncio.CancelledError):
pass
start_task = asyncio.create_task(self.start(init_host, **kwargs))
start_task.add_done_callback(_start_task_done)
def add_done_callback(self, callback: Callable[[], None]) -> None:
self.done_callbacks.append(callback)
def do_superuser_init_done(self) -> None:
pass
def do_authorize(self, message: JsonObject) -> None:
pass
def transport_control_received(self, command: str, message: JsonObject) -> None:
if command == 'init' and self.init_future is not None:
logger.debug('Got init message with active init_future. Setting result.')
self.init_future.set_result(message)
elif command == 'authorize':
self.do_authorize(message)
elif command == 'superuser-init-done':
self.do_superuser_init_done()
else:
raise CockpitProtocolError(f'Received unexpected control message {command}')
def eof_received(self) -> bool:
logger.debug('Peer %s received unexpected EOF', self.__class__.__name__)
return True
def do_closed(self, exc: 'Exception | None') -> None:
logger.debug('Peer %s connection lost %s %s', self.__class__.__name__, type(exc), exc)
if exc is None:
self.shutdown_endpoint(problem='terminated')
elif isinstance(exc, PeerExited):
if exc.exit_code == 127:
self.shutdown_endpoint(problem='no-cockpit')
else:
self.shutdown_endpoint(problem='terminated', message=f'Peer exited with status {exc.exit_code}')
elif isinstance(exc, CockpitProblem):
self.shutdown_endpoint(exc.attrs)
else:
self.shutdown_endpoint(problem='internal-error',
message=f"[{exc.__class__.__name__}] {exc!s}")
if self.init_future is not None and not self.init_future.done():
if exc is not None:
self.init_future.set_exception(exc)
else:
self.init_future.cancel()
for callback in self.done_callbacks:
callback()
def process_exited(self) -> None:
assert isinstance(self.transport, SubprocessTransport)
logger.debug('Peer %s exited, status %d', self.__class__.__name__, self.transport.get_returncode())
returncode = self.transport.get_returncode()
assert isinstance(returncode, int)
self.close(PeerExited(returncode))
def channel_control_received(self, channel: str, command: str, message: JsonObject) -> None:
if self.init_future is not None:
raise CockpitProtocolError('Received unexpected channel control message before init')
self.send_channel_control(channel, command, message)
def channel_data_received(self, channel: str, data: bytes) -> None:
if self.init_future is not None:
raise CockpitProtocolError('Received unexpected channel data before init')
self.send_channel_data(channel, data)
def do_channel_control(self, channel: str, command: str, message: JsonObject) -> None:
assert self.init_future is None
self.write_control(message)
def do_channel_data(self, channel: str, data: bytes) -> None:
assert self.init_future is None
self.write_channel_data(channel, data)
def do_kill(self, host: 'str | None', group: 'str | None', message: JsonObject) -> None:
assert self.init_future is None
self.write_control(message)
def do_close(self) -> None:
self.close()
class ConfiguredPeer(Peer):
config: BridgeConfig
method: str
args: Sequence[str]
env: Sequence[str]
def __init__(self, router: Router, config: BridgeConfig):
self.config = config
self.method = config.method
self.args = config.spawn
self.env = config.environ
super().__init__(router)
async def do_connect_transport(self) -> None:
assert self.args is not None
await self.spawn(self.args, self.env)
class PeerRoutingRule(RoutingRule):
config: BridgeConfig
match: JsonObject
peer: 'Peer | None'
def __init__(self, router: Router, config: BridgeConfig):
super().__init__(router)
self.config = config
self.match = config.match
self.peer = None
def apply_rule(self, options: JsonObject) -> 'Peer | None':
for key, value in self.match.items():
if key not in options:
logger.debug(' rejecting because key %s is missing', key)
return None
if value is not None and options[key] != value:
logger.debug(' rejecting because key %s has wrong value %s (vs %s)', key, options[key], value)
return None
if self.peer is None:
self.peer = ConfiguredPeer(self.router, self.config)
self.peer.add_done_callback(self.peer_closed)
assert self.router.init_host
self.peer.start_in_background(init_host=self.router.init_host)
return self.peer
def peer_closed(self) -> None:
self.peer = None
def shutdown(self) -> None:
if self.peer is not None:
self.peer.close()
class PeersRoutingRule(RoutingRule):
rules: 'list[PeerRoutingRule]' = []
def apply_rule(self, options: JsonObject) -> 'Endpoint | None':
logger.debug(' considering %d rules', len(self.rules))
for rule in self.rules:
logger.debug(' considering %s', rule.config.name)
endpoint = rule.apply_rule(options)
if endpoint is not None:
logger.debug(' selected')
return endpoint
logger.debug(' no peer rules matched')
return None
def set_configs(self, bridge_configs: Sequence[BridgeConfig]) -> None:
old_rules = self.rules
self.rules = []
for config in bridge_configs:
if config.privileged or 'host' in config.match:
continue
for rule in list(old_rules):
if rule.config == config:
old_rules.remove(rule)
break
else:
rule = PeerRoutingRule(self.router, config)
self.rules.append(rule)
for rule in old_rules:
rule.shutdown()
def shutdown(self) -> None:
for rule in self.rules:
rule.shutdown()
