Stuff

colony/core/dispatcher.py

@@ -8,14 +8,31 @@ class Dispatcher:
         self.network = network
         self.role_registry = role_registry
 
-    def dispatch(self, payload):
-        self.logger.debug(f"Dispatching payload for command {payload['command']['role']}.{payload["command"]['name']}")
+    @staticmethod
+    def _parse(payload):
+        command = payload.get("command") if isinstance(payload, dict) else None
+        if not isinstance(command, dict):
+            raise ValueError("Payload has no command")
 
-        c = payload["command"]
-        role_name = c["role"]
-        command_name = c["name"]
-        args = tuple(c.get("args", []))
-        kwargs = dict(c.get("kwargs", {}))
+        role_name = command.get("role")
+        command_name = command.get("name")
+        if not isinstance(role_name, str) or not role_name or not isinstance(command_name, str) or not command_name:
+            raise ValueError("Command needs a role and a name")
+
+        args = command.get("args")
+        if args is None:
+            args = []
+        kwargs = command.get("kwargs")
+        if kwargs is None:
+            kwargs = {}
+        if not isinstance(args, (list, tuple)) or not isinstance(kwargs, dict):
+            raise ValueError("Command args must be a list and kwargs a dict")
+
+        return role_name, command_name, tuple(args), kwargs
+
+    def dispatch(self, payload):
+        role_name, command_name, args, kwargs = self._parse(payload)
+        self.logger.debug(f"Dispatching payload for command {role_name}.{command_name}")
 
         role_entry = self.role_registry.get_role_class_by_name(role_name)
         if role_entry is None:

colony/core/network/core.py

@@ -22,10 +22,10 @@ class Network:
         )
         self.logger.info(f"Self node created: {self.self_node.id}")
 
-        self.nodes = self.load_network_from_snapshot()
-
-        if not self.nodes:
-            self.nodes.append(self.self_node)
+        self.nodes = [self.self_node] + [
+            node for node in self.load_network_from_snapshot()
+            if (node.ip, node.port) != (self.self_node.ip, self.self_node.port)
+        ]
 
     def get_self(self):
         return self.self_node
@@ -70,21 +70,42 @@ class Network:
                 snapshot = json.load(f)
         except FileNotFoundError:
             self.logger.warning(f"Snapshot file {file_path} not found. Starting with empty network")
-            snapshot = {"nodes": []}
+            return nodes
+        except (json.JSONDecodeError, OSError) as e:
+            self.logger.warning(f"Could not read snapshot {file_path}: {e}")
+            return nodes
+
+        if not isinstance(snapshot, dict):
+            self.logger.warning(f"Snapshot {file_path} is not an object. Ignoring it")
+            return nodes
 
         for node_ip, node_data in snapshot.items():
-            result = node_data.get('result')
+            if not isinstance(node_data, dict):
+                self.logger.warning(f"Skipping node {node_ip}: entry is not a dict")
+                continue
+
+            if node_data.get("status") == "unreachable":
+                self.logger.info(f"Skipping node {node_ip}: last seen unreachable")
+                continue
+
+            result = node_data.get("result")
 
             if not isinstance(result, dict):
                 self.logger.warning(f"Skipping node {node_ip}: result is not a dict")
                 continue
 
-            node = Node(
-                hostname=result["hostname"],
-                ip=result["ip"],
-                port=result["port"],
-                roles=result.get("roles", [])
-            )
+            roles = result.get("roles", [])
+            try:
+                node = Node(
+                    hostname=result["hostname"],
+                    ip=result["ip"],
+                    port=result["port"],
+                    roles=roles if isinstance(roles, list) else []
+                )
+            except KeyError as e:
+                self.logger.warning(f"Skipping node {node_ip}: result has no {e}")
+                continue
+
             self.logger.info(f"Loaded node from snapshot: {node.id}")
             nodes.append(node)
 

colony/core/network/hardware.py

@@ -42,7 +42,7 @@ class HardwareStatistics:
         return json.dumps(self.to_dict(), indent=2)
 
     @staticmethod
-    def from_dict(hw_dict: str):
+    def from_dict(hw_dict: dict):
         hw = HardwareStatistics()
         hw.cpu = hw_dict.get("cpu")
         hw.mem = hw_dict.get("memory")
@@ -59,7 +59,7 @@ class HardwareStatistics:
         return {
             "cores": psutil.cpu_count(logical=False),
             "threads": psutil.cpu_count(logical=True),
-            "usage_percent": psutil.cpu_percent(interval=1)
+            "usage_percent": psutil.cpu_percent(interval=0.2)
         }
 
     def _get_memory_stats(self):
@@ -72,7 +72,10 @@ class HardwareStatistics:
         }
 
     def _get_gpu_stats(self):
-        gpus = GPUtil.getGPUs()
+        try:
+            gpus = GPUtil.getGPUs()
+        except Exception:
+            return []
         if not gpus:
             return []
         return [
@@ -91,7 +94,10 @@ class HardwareStatistics:
     def _get_disk_stats(self):
         disk_info = []
         for d in psutil.disk_partitions():
-            usage = psutil.disk_usage(d.mountpoint)
+            try:
+                usage = psutil.disk_usage(d.mountpoint)
+            except OSError:
+                continue
             disk_info.append({
                 "device": d.device,
                 "mountpoint": d.mountpoint,
@@ -125,9 +131,9 @@ class HardwareScorer:
         cpu = hw_obj.cpu
         if not cpu:
             return 0.0
-        cores = cpu.get("cores", 1)
-        threads = cpu.get("threads", 1)
-        usage = cpu.get("usage_percent", 0)  # current CPU utilization in percent
+        cores = cpu.get("cores") or cpu.get("threads") or 1
+        threads = cpu.get("threads") or cores
+        usage = cpu.get("usage_percent") or 0  # current CPU utilization in percent
         available = (cores + threads) * (1 - usage / 100)
         return round(available, 2)
 
@@ -183,15 +189,18 @@ class HardwareScorer:
     def longevity_score(hw_obj) -> float:
         """
         Compute uptime in hours.
-        Converts boot_time from ISO format to datetime, then calculates
+        Accepts boot_time as a datetime or an ISO string, then calculates
         the difference from the current time.
         """
-        boot_time_iso = hw_obj.boot_time
-        try:
-            boot_time = datetime.datetime.fromisoformat(boot_time_iso)
-        except Exception:
+        boot_time = hw_obj.boot_time
+        if isinstance(boot_time, str):
+            try:
+                boot_time = datetime.datetime.fromisoformat(boot_time)
+            except ValueError:
+                return 0.0
+        if not isinstance(boot_time, datetime.datetime):
             return 0.0
-        uptime_hours = (datetime.datetime.now() - boot_time).total_seconds() / 3600
+        uptime_hours = (datetime.datetime.now(tz=boot_time.tzinfo) - boot_time).total_seconds() / 3600
         return round(uptime_hours, 2)
 
     @classmethod

colony/core/network/nodes.py

@@ -37,11 +37,10 @@ class LocalNode(Node):
         self.__roles = dict()
 
     def get_role_instances(self):
-        return self.__roles
+        return dict(self.__roles)
 
     def set_role_instances(self, roles):
-        for role_name, role_instance in roles.items():
-            self.__roles[role_name] = role_instance
+        self.__roles = dict(roles)
 
     def execute_local(self, command: Enum, *args, **kwargs):
         request = CommandRequest.form(command, *args, **kwargs)

colony/core/roles/explorer/role.py

@@ -1,19 +1,23 @@
 import datetime
+from concurrent.futures import ThreadPoolExecutor
 from enum import Enum
 
-from colony.core.network.requests import CommandRequest
 from colony.core.storage import DEFAULT_STORAGE_FILES
 from colony.core.roles.role import Role
 from colony.core.roles.unit.role import Unit
 from colony.core.roles.vaultwarden.role import VaultWarden
+from colony.settings import configuration
 from colony.transport.discover import Discover
-from colony.transport.whisperer import Whisperer
+from colony.utils.network import get_local_subnet
 
 
 class Explorer(Role):
     CONTINUOUS = True
     LOOP_SECONDS_WAIT = 60
 
+    SCAN_WORKERS = 64
+    SCAN_TIMEOUT = 1
+
     class COMMANDS(Enum):
         pass
 
@@ -26,29 +30,77 @@ class Explorer(Role):
 
     def loop(self):
         self.logger.info("Getting network snapshot")
-        snapshot = self.get_snapshot()['result'] or {}
+        snapshot = self.load_snapshot()
+        if snapshot is None:
+            return
         self.logger.debug(f"Loaded snapshot")
 
         discovered_peers = self.discover_peers()
         self.logger.info(f"Discovered peers: {list(discovered_peers.keys())}")
 
+        reachable = set()
         for ip, port in discovered_peers.items():
             self.logger.info(f"Sending SHOW_STATS command to {ip}:{port}")
 
             stats = self.fetch_stats_from_peer(ip, port)
+            if stats is None:
+                continue
 
             self.logger.debug(f"Received stats from {ip}")
-            snapshot[ip] = stats | {'last_update': datetime.datetime.now().isoformat(), 'last_seen': datetime.datetime.now().isoformat(), 'status': 'reachable'}
+            now = datetime.datetime.now().isoformat()
+            snapshot[ip] = {
+                'status': 'reachable',
+                'result': stats,
+                'last_update': now,
+                'last_seen': now,
+            }
+            reachable.add(ip)
 
-        missing_ips = set(snapshot.keys()) - set(discovered_peers.keys())
-        for ip in missing_ips:
+        for ip in snapshot.keys() - reachable:
             self.logger.warning(f"Peer {ip} was in snapshot but not reachable. Marking as unreachable")
             snapshot[ip]['last_update'] = datetime.datetime.now().isoformat()
             snapshot[ip]['status'] = 'unreachable'
 
-        self.set_snapshot(snapshot)
-        self.logger.info("Snapshot updated successfully")
+        if self._stop_event.is_set():
+            self.logger.info("Role was removed during the scan, not writing the snapshot")
+            return
 
+        self.save_snapshot(snapshot)
+
+    def load_snapshot(self):
+        try:
+            resp = self.get_snapshot()
+        except Exception as e:
+            self.logger.warning(f"Could not read the snapshot: {e}")
+            return None
+
+        if not isinstance(resp, dict) or resp.get('status') != 'ok':
+            self.logger.warning("Could not read the snapshot: VaultWarden returned an error")
+            return None
+
+        snapshot = resp.get('result')
+        if snapshot is None:
+            return {}
+        if not isinstance(snapshot, dict):
+            self.logger.warning("Stored snapshot is not an object. Refusing to overwrite it")
+            return None
+
+        for ip in [ip for ip, entry in snapshot.items() if not isinstance(entry, dict)]:
+            self.logger.warning(f"Dropping malformed snapshot entry for {ip}")
+            del snapshot[ip]
+        return snapshot
+
+    def save_snapshot(self, snapshot):
+        try:
+            resp = self.set_snapshot(snapshot)
+        except Exception as e:
+            self.logger.error(f"Could not write the snapshot: {e}")
+            return
+
+        if not isinstance(resp, dict) or resp.get('status') != 'ok':
+            self.logger.error("Could not write the snapshot: VaultWarden returned an error")
+            return
+        self.logger.info("Snapshot updated successfully")
 
     def get_snapshot(self):
         return self.context.execute(VaultWarden.COMMANDS.OPEN, DEFAULT_STORAGE_FILES.COLONY_SNAPSHOT.value)
@@ -56,17 +108,32 @@ class Explorer(Role):
     def set_snapshot(self, snapshot):
         return self.context.execute(VaultWarden.COMMANDS.OVERWRITE, DEFAULT_STORAGE_FILES.COLONY_SNAPSHOT.value, snapshot)
 
-    def discover_peers(self):
-        peers = {}
+    def _hosts_to_scan(self):
+        configured = configuration.get("explorer.hosts")
+        if isinstance(configured, (list, tuple)) and configured:
+            return [str(host) for host in configured]
+        return [str(host) for host in get_local_subnet().hosts()]
 
-        # for ip in get_local_subnet().hosts():
-        for ip in ['192.168.65.101', '192.168.65.102']:
-            try:
-                port = Discover.request(str(ip))
-                peers[str(ip)] = port
-            except (TimeoutError, ConnectionRefusedError, OSError):
-                continue
-        return peers
+    def _probe(self, ip):
+        try:
+            return ip, Discover.request(ip, timeout=self.SCAN_TIMEOUT)
+        except Exception:
+            return ip, None
+
+    def discover_peers(self):
+        with ThreadPoolExecutor(max_workers=self.SCAN_WORKERS) as pool:
+            results = list(pool.map(self._probe, self._hosts_to_scan()))
+
+        return {ip: port for ip, port in results if port is not None}
 
     def fetch_stats_from_peer(self, ip, port):
-        return self.context.execute_to_unregistered_node(ip, port, Unit.COMMANDS.SHOW_STATS)
+        try:
+            resp = self.context.execute_to_unregistered_node(ip, port, Unit.COMMANDS.SHOW_STATS)
+        except Exception as e:
+            self.logger.warning(f"Could not fetch stats from {ip}:{port}: {e}")
+            return None
+
+        if not isinstance(resp, dict) or resp.get('status') != 'ok' or not isinstance(resp.get('result'), dict):
+            self.logger.warning(f"Peer {ip}:{port} answered SHOW_STATS with an error")
+            return None
+        return resp['result']

colony/core/roles/foreman/role.py

@@ -36,21 +36,30 @@ class Foreman(Role):
         k = (self_node.ip, self_node.port)
         if k not in by_key:
             by_key[k] = self_node
-        return list(by_key.values())
+        return [by_key[key] for key in sorted(by_key)]
+
+    def _call(self, peer, command, *args):
+        try:
+            resp = peer.execute(command, *args)
+        except Exception as e:
+            self.logger.warning("%s failed for %s:%s: %s", command.value, peer.ip, peer.port, e)
+            return None
+
+        if not isinstance(resp, dict) or resp.get("status") != "ok":
+            self.logger.warning("%s failed for %s:%s", command.value, peer.ip, peer.port)
+            return None
+        return resp
 
     def _fetch_live_roles(self, peers):
         live = {}
         for peer in peers:
-            resp = peer.execute(Unit.COMMANDS.SHOW_STATS)
-            if resp.get("status") != "ok":
-                self.logger.warning(
-                    "SHOW_STATS failed for %s:%s, assuming no roles",
-                    peer.ip,
-                    peer.port,
-                )
-                live[peer] = []
+            resp = self._call(peer, Unit.COMMANDS.SHOW_ROLES)
+            if resp is None:
                 continue
-            live[peer] = list(resp["result"].get("roles", []))
+            if not isinstance(resp.get("result"), list):
+                self.logger.warning("Bad role list from %s:%s, leaving it out of this round", peer.ip, peer.port)
+                continue
+            live[peer] = list(resp["result"])
         return live
 
     def _assign_critical_holders(self, peers, critical_roles, live_roles):
@@ -68,15 +77,23 @@ class Foreman(Role):
                 holders[critical_role] = peer_list[i % n]
         return holders
 
-    def _best_peer_for_optional(self, peers, role_stat):
+    def _fetch_hardware_scores(self, peers):
+        scores = {}
+        for peer in peers:
+            resp = self._call(peer, Unit.COMMANDS.SHOW_HARDWARE)
+            if resp is None:
+                continue
+            hw_obj = HardwareStatistics.from_dict(resp["result"])
+            scores[peer] = HardwareScorer.from_object(hw_obj)
+        return scores
+
+    def _best_peer_for_optional(self, peers, role_stat, hardware_scores):
         best_score = float("-inf")
         best_peer = None
         for peer in peers:
-            hw_resp = peer.execute(Unit.COMMANDS.SHOW_HARDWARE)
-            if hw_resp.get("status") != "ok":
+            hs = hardware_scores.get(peer)
+            if hs is None:
                 continue
-            hw_obj = HardwareStatistics.from_dict(hw_resp["result"])
-            hs = HardwareScorer.from_object(hw_obj)
             score = 0.0
             if role_stat.hardware_preferences:
                 for hw_type, weight in role_stat.hardware_preferences.priorities.items():
@@ -119,6 +136,10 @@ class Foreman(Role):
             return
 
         live_roles = self._fetch_live_roles(peers)
+        peers = [peer for peer in peers if peer in live_roles]
+        if not peers:
+            return
+
         critical_holders = self._assign_critical_holders(
             peers, critical_roles, live_roles
         )
@@ -127,14 +148,16 @@ class Foreman(Role):
         for critical_role, holder in critical_holders.items():
             desired[holder].append(critical_role)
 
+        hardware_scores = self._fetch_hardware_scores(peers) if optional_roles else {}
         for role_name in optional_roles:
             role_stat = roles[role_name]
-            best_peer = self._best_peer_for_optional(peers, role_stat)
+            best_peer = self._best_peer_for_optional(peers, role_stat, hardware_scores)
             if best_peer is not None:
                 desired[best_peer].append(role_name)
 
         for peer in peers:
-            peer.execute(
+            self._call(
+                peer,
                 Unit.COMMANDS.SET_ROLES,
                 self._dedupe_preserve_order(desired[peer]),
             )

colony/core/roles/manager.py

@@ -12,20 +12,25 @@ class RoleManager:
         self.state_manager = state_manager
 
         self.logger = get_logger(self.__class__.__name__)
+        self._lock = threading.RLock()
 
     def initialize_roles(self):
         self.logger.info("Initializing roles...")
         context = RoleContext(self.network, self)
 
-        roles = {}
-        for role_name in self.state_manager.get(StateManager.Data.Roles):
-            role_cls = self.role_registry.get_role_class_by_name(role_name)
-            self.logger.debug(f"Creating instance of role {role_name}. ")
-            role_instance = role_cls(context=context)
-            roles[role_name] = role_instance
+        with self._lock:
+            roles = {}
+            for role_name in self.state_manager.get(StateManager.Data.Roles):
+                role_cls = self.role_registry.get_role_class_by_name(role_name)
+                if role_cls is None:
+                    self.logger.warning(f"Skipping unknown role {role_name}")
+                    continue
 
-        self.network.self_node.roles = [role for role in roles.keys()]
-        self.network.self_node.set_role_instances(roles)
+                self.logger.debug(f"Creating instance of role {role_name}. ")
+                roles[role_name] = role_cls(context=context)
+
+            self.network.self_node.roles = list(roles)
+            self.network.self_node.set_role_instances(roles)
         self.logger.info(f"Roles initialized: {list(roles.keys())}. ")
 
     def get_role_instances(self):
@@ -34,39 +39,54 @@ class RoleManager:
     def start_all(self):
         for role_name, role in self.get_role_instances().items():
             self.logger.info(f"Starting role {role_name}")
-            t = threading.Thread(target=role.start, daemon=True, name=role_name)
-            t.start()
-            self.logger.debug(f"Thread started for role {role_name}: {t.name}")
+            role.start()
+
+    def _validate_roles(self, new_roles):
+        if not isinstance(new_roles, list):
+            raise ValueError("Roles must be a list of role names")
+
+        for role_name in new_roles:
+            if self.role_registry.get_role_class_by_name(role_name) is None:
+                raise ValueError(f"Unknown role {role_name!r}")
+
+        return list(dict.fromkeys(new_roles))
 
     def set_roles(self, new_roles):
         self.logger.info("Updating roles...")
 
-        current_roles = self.get_role_instances()
-        roles_to_remove = set(current_roles) - set(new_roles)
-        roles_to_add = set(new_roles) - set(current_roles)
+        new_roles = self._validate_roles(new_roles)
 
-        for role_name in roles_to_remove:
-            self.logger.debug(f"Stopping and removing role {role_name}")
-            current_roles[role_name].stop()
-            del current_roles[role_name]
+        with self._lock:
+            current_roles = self.get_role_instances()
+            roles_to_remove = [name for name in current_roles if name not in new_roles]
+            roles_to_add = [name for name in new_roles if name not in current_roles]
 
-        context = RoleContext(self.network, self)
-        for role_name in roles_to_add:
-            self.logger.debug(f"Creating and adding new role {role_name}")
-            role_cls = self.role_registry.get_role_class_by_name(role_name)
-            current_roles[role_name] = role_cls(context=context)
+            if not roles_to_remove and not roles_to_add:
+                self.logger.info("Roles already up to date")
+                return
 
-        self.state_manager.set(StateManager.Data.Roles, list(current_roles.keys()))
+            context = RoleContext(self.network, self)
+            added = {}
+            for role_name in roles_to_add:
+                self.logger.debug(f"Creating new role {role_name}")
+                role_cls = self.role_registry.get_role_class_by_name(role_name)
+                added[role_name] = role_cls(context=context)
 
-        self.network.self_node.roles = list(current_roles.keys())
-        self.network.self_node.set_role_instances(current_roles)
+            final_roles = {
+                name: current_roles[name] if name in current_roles else added[name]
+                for name in new_roles
+            }
 
-        for role_name in roles_to_add:
-            role = current_roles[role_name]
-            self.logger.info(f"Starting role {role_name}")
-            t = threading.Thread(target=role.start, daemon=True, name=role_name)
-            t.start()
-            self.logger.debug(f"Thread started for role {role_name}: {t.name}")
+            self.state_manager.set(StateManager.Data.Roles, list(final_roles.keys()))
+            self.network.self_node.set_role_instances(final_roles)
+            self.network.self_node.roles = list(final_roles.keys())
 
-        self.logger.info(f"Roles updated to: {list(current_roles.keys())}")
+            for role_name in roles_to_remove:
+                self.logger.debug(f"Stopping removed role {role_name}")
+                current_roles[role_name].stop()
 
+            for role_name, role in added.items():
+                self.logger.info(f"Starting role {role_name}")
+                role.start()
+
+        self.logger.info(f"Roles updated to: {list(final_roles.keys())}")

colony/core/roles/role.py

@@ -31,7 +31,7 @@ class Role:
         return self.__class__.__name__
 
     def start(self):
-        if not self.CONTINUOUS:
+        if not self.CONTINUOUS or self._thread is not None:
             return
 
         self._thread = threading.Thread(target=self._run_loop, daemon=True, name=str(self))
@@ -39,7 +39,10 @@ class Role:
 
     def _run_loop(self):
         while not self._stop_event.is_set():
-            self.loop()
+            try:
+                self.loop()
+            except Exception:
+                self.logger.exception("Loop failed, retrying after the usual wait")
             self._stop_event.wait(self.LOOP_SECONDS_WAIT)
 
     def loop(self):

colony/core/roles/unit/role.py

@@ -6,12 +6,14 @@ from colony.core.roles.role import Role, Command
 class Unit(Role):
     class COMMANDS(Enum):
         SHOW_STATS = 'Unit.ShowStats'
+        SHOW_ROLES = 'Unit.ShowRoles'
         SHOW_HARDWARE = 'Unit.ShowHardware'
         SET_ROLES = 'Unit.SetRoles'
 
     def __init__(self, context):
         commands = {
             self.COMMANDS.SHOW_STATS.value: Command(self.show_stats),
+            self.COMMANDS.SHOW_ROLES.value: Command(self.show_roles),
             self.COMMANDS.SHOW_HARDWARE.value: Command(self.show_hardware),
             self.COMMANDS.SET_ROLES.value: Command(self.set_roles),
         }
@@ -21,6 +23,9 @@ class Unit(Role):
     def show_hardware(self):
         return HardwareStatistics.collect().to_dict()
 
+    def show_roles(self):
+        return list(self.context.self().roles)
+
     def show_stats(self):
         self_node = self.context.self()
         return self_node.to_dict()

colony/core/state_manager.py

@@ -1,41 +1,73 @@
+import copy
 import json
+import shutil
+import threading
 from enum import Enum
 from pathlib import Path
 from typing import Any
 
+from colony.core.logger.logging import get_logger
+from colony.utils.files import atomic_write_json
+
 
 class StateManager:
     class Data(Enum):
         Roles = "roles"
 
+    DEFAULTS = {
+        Data.Roles: [],
+    }
+
     STATE_FILE = Path("state.json")
 
     def __init__(self):
+        self.logger = get_logger(self.__class__.__name__)
+        self._lock = threading.RLock()
         self._data = {}
         self._load()
 
     def _load(self):
-        if self.STATE_FILE.exists():
-            with open(self.STATE_FILE, "r", encoding="utf-8") as f:
+        with self._lock:
+            if self.STATE_FILE.exists():
                 try:
-                    self._data = json.load(f)
-                except json.decoder.JSONDecodeError:
-                    self._data = {}
+                    with open(self.STATE_FILE, "r", encoding="utf-8") as f:
+                        data = json.load(f)
+                    if not isinstance(data, dict):
+                        raise ValueError("top level is not an object")
+                    for key, default in self.DEFAULTS.items():
+                        if key.value in data and not isinstance(data[key.value], type(default)):
+                            raise ValueError(f"'{key.value}' has the wrong type")
+                    self._data = data
+                except (ValueError, OSError) as e:
+                    self._recover_from_corrupt_file(e)
+            else:
+                self._data = {}
+                self._apply_defaults()
+                self._save()
 
-        else:
-            self._data = {}
-            self._save()
+            self._apply_defaults()
 
-        for key in self.Data:
-            self._data.setdefault(key.value, {})
+    def _recover_from_corrupt_file(self, error):
+        backup = self.STATE_FILE.with_name(self.STATE_FILE.name + ".corrupt")
+        self.logger.warning(f"Could not use {self.STATE_FILE} ({error}). Keeping a copy as {backup} and resetting to defaults")
+        try:
+            shutil.copy2(self.STATE_FILE, backup)
+        except OSError as copy_error:
+            self.logger.warning(f"Could not back up {self.STATE_FILE}: {copy_error}")
+        self._data = {}
+
+    def _apply_defaults(self):
+        for key, default in self.DEFAULTS.items():
+            self._data.setdefault(key.value, copy.deepcopy(default))
 
     def _save(self):
-        with open(self.STATE_FILE, "w", encoding="utf-8") as f:
-            json.dump(self._data, f, indent=4)
+        atomic_write_json(self.STATE_FILE, self._data, indent=4)
 
     def get(self, key: Data) -> Any:
-        return self._data[key.value]
+        with self._lock:
+            return copy.deepcopy(self._data[key.value])
 
     def set(self, key: Data, value: Any):
-        self._data[key.value] = value
-        self._save()
+        with self._lock:
+            self._data[key.value] = value
+            self._save()

colony/core/storage.py

@@ -1,31 +1,50 @@
 import os
 import json
+import threading
 from enum import Enum
 
+from colony.utils.files import atomic_write_json
+
 STORAGE_DIR = "storage"
 
+_write_lock = threading.Lock()
+
+
 class Storage:
     def __init__(self, storage_dir):
         self.storage_dir = storage_dir
 
         os.makedirs(self.storage_dir, exist_ok=True)
+        self._root = os.path.realpath(self.storage_dir)
+
+    def _resolve(self, name, allow_root=False):
+        if not isinstance(name, str) or "\0" in name:
+            raise ValueError("Invalid storage path")
+
+        path = os.path.realpath(os.path.join(self._root, name))
+        if os.path.commonpath([self._root, path]) != self._root:
+            raise ValueError("Storage path escapes the storage directory")
+        if path == self._root and not allow_root:
+            raise ValueError("Storage path must name a file")
+        return path
 
     def get_path(self, filename):
-        return os.path.join(self.storage_dir, filename)
+        return self._resolve(filename)
 
     def ls(self, subdir):
-        path = os.path.join(self.storage_dir, subdir)
+        path = self._resolve(subdir, allow_root=True)
 
-        if not os.path.exists(path):
+        if not os.path.isdir(path):
             return []
 
-        return os.listdir(path)
+        return [name for name in os.listdir(path) if not name.startswith(".tmp-")]
 
     def save(self, filename, data):
         path = self.get_path(filename)
 
-        with open(path, 'w', encoding='utf-8') as f:
-            json.dump(data, f, indent=2)
+        with _write_lock:
+            os.makedirs(os.path.dirname(path), exist_ok=True)
+            atomic_write_json(path, data, indent=2)
 
         return path
 
@@ -44,9 +63,10 @@ class Storage:
     def delete(self, filename):
         path = self.get_path(filename)
 
-        if os.path.exists(path):
-            os.remove(path)
-            return True
+        with _write_lock:
+            if os.path.exists(path):
+                os.remove(path)
+                return True
 
         return False
 

colony/role_registry.py

@@ -47,4 +47,5 @@ class RoleRegistry:
         self.roles = roles
 
     def get_role_class_by_name(self, name):
-        return self.roles.get(name).cls
+        role_stat = self.roles.get(name) if isinstance(name, str) else None
+        return role_stat.cls if role_stat is not None else None

colony/transport/discover.py

@@ -1,32 +1,50 @@
 import socket
 import struct
+import time
+
+from colony.core.logger.logging import get_logger
+
 
 class Discover:
     DISCOVER_PORT = 28149
     MAGIC = b"OK"
+    CLIENT_TIMEOUT = 2
 
     @staticmethod
     def listen(real_port, host="0.0.0.0"):
+        logger = get_logger(Discover.__name__)
+
         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         s.bind((host, Discover.DISCOVER_PORT))
         s.listen()
 
         while True:
-            conn, _ = s.accept()
             try:
-                data = conn.recv(2)
-                if data == Discover.MAGIC:
-                    conn.sendall(struct.pack("!H", real_port))
-            finally:
-                conn.close()
+                conn, _ = s.accept()
+            except OSError as e:
+                logger.error(f"Accept failed: {e}")
+                time.sleep(1)
+                continue
+
+            with conn:
+                try:
+                    conn.settimeout(Discover.CLIENT_TIMEOUT)
+                    if conn.recv(len(Discover.MAGIC)) == Discover.MAGIC:
+                        conn.sendall(struct.pack("!H", real_port))
+                except OSError as e:
+                    logger.debug(f"Discovery request failed: {e}")
 
     @staticmethod
     def request(target_host, timeout=3):
-        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        s.settimeout(timeout)
-        s.connect((target_host, Discover.DISCOVER_PORT))
-        s.sendall(Discover.MAGIC)
-        data = s.recv(2)
-        s.close()
+        with socket.create_connection((target_host, Discover.DISCOVER_PORT), timeout=timeout) as s:
+            s.sendall(Discover.MAGIC)
+
+            data = b""
+            while len(data) < 2:
+                chunk = s.recv(2 - len(data))
+                if not chunk:
+                    raise ConnectionError("Discovery response was cut short")
+                data += chunk
+
         return struct.unpack("!H", data)[0]

colony/transport/listener.py

@@ -2,63 +2,92 @@ import json
 import socket
 import ssl
 import threading
-from asyncio import Queue
+import time
 
+from colony.core.logger.logging import get_logger
 from colony.settings import configuration
+from colony.transport.framing import send_frame, recv_frame
 
 
 class Listener:
+    CLIENT_TIMEOUT = 10
+
     def __init__(self, dispatcher, host=None, port=None, cert_file=None, keyfile=None):
+        self.logger = get_logger(self.__class__.__name__)
+
         self.dispatcher = dispatcher
-        self.host = host or getattr(configuration.listener, "host", "0.0.0.0")
-        self.port = port or getattr(configuration.listener, "port", 9000)
+        self.host = host if host is not None else getattr(configuration.listener, "host", "0.0.0.0")
+        self.port = port if port is not None else getattr(configuration.listener, "port", 9000)
         self.cert_file = cert_file or getattr(configuration.listener, "cert_file", "certs/server.crt")
         self.keyfile = keyfile or getattr(configuration.listener, "keyfile", "certs/server.key")
-        self.queue = Queue()
         self.running = False
+        self._sock = None
+        self._ssl_context = None
+
+    def _build_ssl_context(self):
+        context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        context.minimum_version = ssl.TLSVersion.TLSv1_2
+        context.load_cert_chain(certfile=self.cert_file, keyfile=self.keyfile)
+        context.load_verify_locations(cafile=self.cert_file)
+        context.verify_mode = ssl.CERT_REQUIRED
+        return context
 
     def handle_client(self, conn):
-        context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
-        context.load_cert_chain(certfile=self.cert_file, keyfile=self.keyfile)
-
-        tls_conn = context.wrap_socket(conn, server_side=True)
-
         try:
-            data = tls_conn.recv(65536)
-            payload = json.loads(data.decode())
+            conn.settimeout(self.CLIENT_TIMEOUT)
+            tls_conn = self._ssl_context.wrap_socket(conn, server_side=True)
+        except OSError as e:
+            self.logger.warning(f"TLS handshake failed: {e}")
+            conn.close()
+            return
 
-            result = self.dispatcher.dispatch(payload)
+        with tls_conn:
+            try:
+                payload = json.loads(recv_frame(tls_conn).decode())
+                result = self.dispatcher.dispatch(payload)
+                body = json.dumps({"status": "ok", "result": result}).encode()
+            except Exception as e:
+                self.logger.exception("Request failed")
+                body = json.dumps({"status": "error", "error": type(e).__name__}).encode()
 
-            response = json.dumps({
-                "status": "ok",
-                "result": result
-            }).encode()
+            try:
+                send_frame(tls_conn, body)
+            except (OSError, ValueError) as e:
+                self.logger.warning(f"Could not send response: {e}")
 
-            tls_conn.sendall(response)
-        except Exception as e:
-            error = json.dumps({
-                "status": "error",
-                "error": str(e)
-            }).encode()
-            tls_conn.sendall(error)
-            raise
-        finally:
-            tls_conn.close()
+    def bind(self):
+        if self._sock is not None:
+            return
 
-    def start(self):
-        self.running = True
+        self._ssl_context = self._build_ssl_context()
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        sock.bind((self.host, self.port))
-        sock.listen(5)
+        try:
+            sock.bind((self.host, self.port))
+            sock.listen(5)
+        except OSError:
+            sock.close()
+            raise
+        self.port = sock.getsockname()[1]
+        self._sock = sock
+
+    def start(self):
+        self.bind()
+        self.running = True
+        self.logger.info(f"Listening on {self.host}:{self.port}")
 
         while self.running:
             try:
-                conn, addr = sock.accept()
-                t = threading.Thread(target=self.handle_client, args=(conn,), name='HandleClient')
-                t.daemon = True
-                t.start()
-            except KeyboardInterrupt:
-                break
+                conn, _ = self._sock.accept()
+            except OSError as e:
+                if not self.running:
+                    break
+                self.logger.error(f"Accept failed: {e}")
+                time.sleep(1)
+                continue
 
-        sock.close()
+            t = threading.Thread(target=self.handle_client, args=(conn,), name='HandleClient')
+            t.daemon = True
+            t.start()
+
+        self._sock.close()

colony/transport/whisperer.py

@@ -3,36 +3,37 @@ import socket
 import ssl
 
 from colony.settings import configuration
+from colony.transport.framing import send_frame, recv_frame
 
 
 class Whisperer:
+    DEFAULT_TIMEOUT = 10
+
     def __init__(self, host=None, port=None, cert_file=None, keyfile=None):
         self.host = host or getattr(configuration.whisperer, "host", "127.0.0.1")
         self.port = port or getattr(configuration.whisperer, "port", 9000)
-        self.cert_file = cert_file or getattr(configuration.whisperer, "cert_file", None)
-        self.keyfile = keyfile or getattr(configuration.whisperer, "keyfile", None)
+        self.cert_file = cert_file or getattr(configuration.whisperer, "cert_file", "certs/server.crt")
+        self.keyfile = keyfile or getattr(configuration.whisperer, "keyfile", "certs/server.key")
 
-    def send(self, payload, timeout=500):
-        data = json.dumps(payload).encode()
-
-        context = ssl.create_default_context()
+    def _build_ssl_context(self):
+        context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+        context.minimum_version = ssl.TLSVersion.TLSv1_2
         context.check_hostname = False
-        context.verify_mode = ssl.CERT_NONE
+        context.verify_mode = ssl.CERT_REQUIRED
+        context.load_verify_locations(cafile=self.cert_file)
+        context.load_cert_chain(certfile=self.cert_file, keyfile=self.keyfile)
+        return context
 
-        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        sock.settimeout(timeout)
-        tls_sock = context.wrap_socket(sock)
+    def send(self, payload, timeout=DEFAULT_TIMEOUT):
+        data = json.dumps(payload).encode()
+        context = self._build_ssl_context()
 
         try:
-            tls_sock.connect((self.host, self.port))
-            tls_sock.sendall(data)
-
-            response = tls_sock.recv(65536)
-            if not response:
-                raise RuntimeError("Connection closed without response. ")
-
-            return json.loads(response.decode())
+            with socket.create_connection((self.host, self.port), timeout=timeout) as sock:
+                with context.wrap_socket(sock) as tls_sock:
+                    send_frame(tls_sock, data)
+                    response = recv_frame(tls_sock)
         except socket.timeout:
             raise TimeoutError("No response before timeout. ")
-        finally:
-            tls_sock.close()
+
+        return json.loads(response.decode())

colony/utils/network.py

@@ -3,15 +3,20 @@ import socket
 
 
 def get_local_subnet():
-    hostname = socket.gethostname()
-    local_ip = socket.gethostbyname(hostname)
-
-    subnet = ipaddress.ip_network(f"{local_ip}/24", strict=False)
-
-    return subnet
+    return ipaddress.ip_network(f"{get_local_ip()}/24", strict=False)
 
 def get_local_ip():
-    return socket.gethostbyname(socket.gethostname())
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+            s.connect(("192.0.2.1", 9))
+            return s.getsockname()[0]
+    except OSError:
+        pass
+
+    try:
+        return socket.gethostbyname(socket.gethostname())
+    except OSError:
+        return "127.0.0.1"
 
 def get_hostname():
     return socket.gethostname()