import contextlib import fcntl import importlib.util import io import json import os from pathlib import Path import stat import tempfile import time import unittest from unittest.mock import patch spec = importlib.util.spec_from_file_location("ec_agent_update", Path(__file__).resolve().parents[1] / "ec_agent_update.py") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) PREVIOUS = b'VERSION = "1.0.0"\n# previous agent\n' CANDIDATE = b'VERSION = "1.1.0"\n# next agent\n' class PowerLoss(BaseException): pass class Response(io.BytesIO): def __init__(self, value, headers=None): super().__init__(value) self.headers = headers or {} self.status = 200 class UpdateTests(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.base = Path(self.temp.name).resolve() self.config_dir = self.base / "etc" / "electric-cherry" self.bin_dir = self.base / "libexec" self.root = self.base / "data" for directory in (self.config_dir, self.bin_dir, self.root / "agent"): directory.mkdir(parents=True) paths = { "AGENT_CONFIG": self.config_dir / "agent.json", "AGENT_EXECUTABLE": self.bin_dir / "electric-cherry-agent.py", "UNINSTALL_EXECUTABLE": self.bin_dir / "electric-cherry-uninstall.py", "UNINSTALL_REQUEST": self.config_dir / "decommission.json", "UPDATE_LOCK": self.config_dir / "update.lock", "UPDATE_DIR": self.config_dir / "update", "UPDATE_JOURNAL": self.config_dir / "update" / "journal.json", "UPDATE_STATUS": self.config_dir / "update-status.json", "AGENT_HEALTH": self.config_dir / "agent-health.json", } self.patches = [patch.object(module, name, value) for name, value in paths.items()] def trusted_for_test(value, missing=False, directory=False): path = Path(value) if any(item.is_symlink() for item in (path, *path.parents)): raise module.UpdateError("Unsafe symbolic link") if not path.exists(): if missing: return path raise module.UpdateError("File unavailable") if directory and not path.is_dir() or not directory and not path.is_file(): raise module.UpdateError("Wrong file type") return path self.patches.extend([patch.object(module, "trusted", side_effect=trusted_for_test), patch.object(module, "sync_parent"), patch.object(module, "systemctl", side_effect=self.systemctl)]) for item in self.patches: item.start() self.config = {"hostId": "host-one", "root": str(self.root), "controller": "https://controller.example", "credential": "secret-enrollment-credential"} self.config_bytes = json.dumps(self.config).encode() module.AGENT_CONFIG.write_bytes(self.config_bytes) (self.root / ".electric-cherry.json").write_text(json.dumps({"version": 1, "root": str(self.root), "hostId": "host-one", "controller": "https://controller.example"})) module.AGENT_EXECUTABLE.write_bytes(PREVIOUS) module.UNINSTALL_EXECUTABLE.write_bytes(PREVIOUS) (self.root / "games.ff").write_bytes(b"large game files unchanged") self.commands = [] self.active = True self.pid = 123 self.write_health = True self.downloads = [] self.source = CANDIDATE self.release = {"version": "1.1.0", "sha256": module.sha256(CANDIDATE), "protocolVersion": 1, "updaterProtocolVersion": 1, "platform": "debian-amd64"} def tearDown(self): for item in reversed(self.patches): item.stop() self.temp.cleanup() def systemctl(self, *args): self.commands.append(args) if args == ("stop", module.AGENT_SERVICE): self.active = False return "" if args == ("start", module.AGENT_SERVICE): # The candidate must be able to take the operation lock at startup. with module.exclusive_lock(self.root / "agent" / "activity.lock") as activity: self.assertIsNotNone(activity, "Updater held the operation lock while starting the agent") self.active = True self.pid += 1 if self.write_health and module.UPDATE_JOURNAL.exists(): journal = json.loads(module.UPDATE_JOURNAL.read_text()) module.AGENT_HEALTH.write_text(json.dumps({"version": module.source_version(module.AGENT_EXECUTABLE.read_bytes()), "pid": self.pid, "checkedAt": time.time(), "updateId": journal["id"]})) return "" if args == ("show", module.AGENT_SERVICE, "--property=ActiveState", "--property=MainPID"): return "ActiveState=" + ("active" if self.active else "inactive") + "\nMainPID=" + str(self.pid if self.active else 0) + "\n" self.fail("Unexpected service command " + str(args)) def fetch(self, path, limit): self.downloads.append(path) if path.endswith("manifest.json"): return json.dumps(self.release).encode() if path.endswith("ec-agent.py"): return self.source self.fail("Unexpected download " + path) def updater(self): updater = module.Updater() updater.fetch = self.fetch return updater def status(self): return json.loads(module.UPDATE_STATUS.read_text()) def assert_unchanged(self): self.assertEqual(module.AGENT_EXECUTABLE.read_bytes(), PREVIOUS) self.assertEqual(module.UNINSTALL_EXECUTABLE.read_bytes(), PREVIOUS) self.assertEqual(module.AGENT_CONFIG.read_bytes(), self.config_bytes) self.assertEqual((self.root / "games.ff").read_bytes(), b"large game files unchanged") def test_verified_upgrade_waits_for_authenticated_candidate_health(self): self.assertEqual(self.updater().run(), "updated") self.assertEqual(module.AGENT_EXECUTABLE.read_bytes(), CANDIDATE) self.assertEqual(module.UNINSTALL_EXECUTABLE.read_bytes(), CANDIDATE) self.assertEqual(module.AGENT_CONFIG.read_bytes(), self.config_bytes) self.assertEqual((self.root / "games.ff").read_bytes(), b"large game files unchanged") self.assertEqual(self.status()["currentVersion"], "1.1.0") self.assertFalse(module.UPDATE_JOURNAL.exists()) self.assertFalse((module.UPDATE_DIR / "agent.backup.py").exists()) self.assertEqual(next(command for command in self.commands if command[0] != "show"), ("stop", module.AGENT_SERVICE)) def test_equal_or_older_versions_do_not_download_or_restart(self): for target in ("1.0.0", "0.9.9"): with self.subTest(target=target): self.release["version"] = target self.assertEqual(self.updater().run(), "current") self.assert_unchanged() self.assertFalse(any(path.endswith("ec-agent.py") for path in self.downloads)) self.assertEqual(self.commands, []) def test_busy_operation_defers_update_without_stopping_agent(self): with module.exclusive_lock(self.root / "agent" / "activity.lock"): self.assertEqual(self.updater().run(), "deferred") self.assertIn("executing an operation", self.status()["error"]) self.assert_unchanged() self.assertEqual(self.commands, []) self.assertFalse(module.UPDATE_JOURNAL.exists()) def test_deliberately_stopped_agent_is_not_updated_or_started(self): self.active = False self.assertEqual(self.updater().run(), "deferred") self.assertIn("Agent service is not running", self.status()["error"]) self.assert_unchanged() self.assertFalse(module.UPDATE_JOURNAL.exists()) self.assertFalse(self.active) self.assertFalse(any(command[0] in ("start", "stop") for command in self.commands)) def test_concurrent_worker_does_not_overwrite_active_worker_status(self): module.UPDATE_STATUS.write_text('{"state":"updating"}') with module.exclusive_lock(module.UPDATE_LOCK): self.assertEqual(self.updater().run(), "deferred") self.assertEqual(self.status(), {"state": "updating"}) self.assertEqual(self.downloads, []) def test_uninstall_request_prevents_any_download_or_mutation(self): module.UNINSTALL_REQUEST.write_text("{}") self.assertEqual(self.updater().run(), "deferred") self.assertIn("deletion", self.status()["error"]) self.assertEqual(self.downloads, []) self.assertEqual(self.commands, []) self.assert_unchanged() def test_uninstall_starting_during_download_wins_before_code_changes(self): updater = self.updater() def fetch(path, limit): data = self.fetch(path, limit) if path.endswith("ec-agent.py"): module.UNINSTALL_REQUEST.write_text("{}") return data updater.fetch = fetch self.assertEqual(updater.run(), "deferred") self.assert_unchanged() self.assertEqual(self.commands, []) def test_controller_unavailable_keeps_running_agent(self): updater = self.updater() updater.fetch = lambda *_: (_ for _ in ()).throw(module.ContactError("Controller unavailable")) self.assertEqual(updater.run(), "deferred") self.assert_unchanged() self.assertTrue(self.active) self.assertEqual(self.commands, []) def test_checksum_failure_holds_artifact_without_executing_it(self): self.source = CANDIDATE + b"# changed in transit" self.assertEqual(self.updater().run(), "failed") self.assertIn("checksum", self.status()["error"]) self.assert_unchanged() self.assertEqual(self.commands, []) self.downloads.clear() self.source = CANDIDATE self.assertEqual(self.updater().run(), "failed") self.assertEqual(self.downloads, ["/api/agent/release/manifest.json"]) def test_invalid_or_mismatched_python_version_is_never_installed(self): for source in (b'VERSION = "1.1.0"\ninvalid syntax ?\n', b'VERSION = "5.0.0"\n'): with self.subTest(source=source): self.source = source self.release["sha256"] = module.sha256(source) self.assertEqual(self.updater().run(), "failed") self.assert_unchanged() self.assertEqual(self.commands, []) def test_incompatible_protocol_platform_or_version_never_downloads_code(self): original = dict(self.release) for key, value in (("protocolVersion", 2), ("protocolVersion", True), ("platform", "linux-arm64"), ("version", "1.1.0-rc1"), ("version", "01.1.0"), ("updaterProtocolVersion", 2)): with self.subTest(key=key, value=value): self.release = {**original, key: value} self.assertIn(self.updater().run(), ("failed", "deferred")) self.assert_unchanged() self.assertFalse(any(path.endswith("ec-agent.py") for path in self.downloads)) self.assertEqual(self.commands, []) def test_failed_health_rolls_back_both_files_then_holds_release(self): updater = self.updater() updater.wait_healthy = lambda journal: journal["targetVersion"] == "1.0.0" self.assertEqual(updater.run(), "failed") self.assert_unchanged() self.assertTrue(self.active) self.assertFalse(module.UPDATE_JOURNAL.exists()) self.assertEqual(self.status()["currentVersion"], "1.0.0") self.assertEqual(sum(command == ("start", module.AGENT_SERVICE) for command in self.commands), 2) self.downloads.clear() self.assertEqual(self.updater().run(), "failed") self.assertEqual(self.downloads, ["/api/agent/release/manifest.json"]) def test_changed_artifact_of_failed_target_can_be_retried(self): updater = self.updater() updater.wait_healthy = lambda journal: journal["targetVersion"] == "1.0.0" self.assertEqual(updater.run(), "failed") self.source = CANDIDATE + b"# corrected release\n" self.release["sha256"] = module.sha256(self.source) self.assertEqual(self.updater().run(), "updated") self.assertEqual(module.AGENT_EXECUTABLE.read_bytes(), self.source) def crash_during_second_replacement(self): original = module.atomic_bytes def interrupt(path, value, mode=0o600): if Path(path) == module.UNINSTALL_EXECUTABLE and value == CANDIDATE: raise PowerLoss() return original(path, value, mode) with patch.object(module, "atomic_bytes", side_effect=interrupt): with self.assertRaises(PowerLoss): self.updater().run() self.assertEqual(module.AGENT_EXECUTABLE.read_bytes(), CANDIDATE) self.assertEqual(module.UNINSTALL_EXECUTABLE.read_bytes(), PREVIOUS) self.assertTrue(module.UPDATE_JOURNAL.exists()) def test_crash_between_executable_replacements_recovers_from_backups(self): self.crash_during_second_replacement() self.downloads.clear() self.assertEqual(self.updater().run(), "failed") self.assert_unchanged() self.assertEqual(self.downloads, []) self.assertTrue(self.active) self.assertFalse(module.UPDATE_JOURNAL.exists()) def test_crash_after_verification_still_recovers_conservatively(self): updater = self.updater() updater.finish_journal = lambda: (_ for _ in ()).throw(PowerLoss()) with self.assertRaises(PowerLoss): updater.run() self.assertEqual(json.loads(module.UPDATE_JOURNAL.read_text())["phase"], "verifying") self.assertEqual(self.updater().run(), "failed") self.assert_unchanged() self.assertFalse(module.UPDATE_JOURNAL.exists()) def test_corrupt_backup_keeps_journal_and_never_overwrites_code(self): self.crash_during_second_replacement() (module.UPDATE_DIR / "agent.backup.py").write_bytes(b"corrupt") before = len(self.commands) self.assertEqual(self.updater().run(), "failed") self.assertTrue(module.UPDATE_JOURNAL.exists()) self.assertEqual(module.AGENT_EXECUTABLE.read_bytes(), CANDIDATE) self.assertEqual(len(self.commands), before) def test_changed_enrollment_prevents_update_and_rollback(self): self.crash_during_second_replacement() changed = {**self.config, "hostId": "different-host"} module.AGENT_CONFIG.write_text(json.dumps(changed)) (self.root / ".electric-cherry.json").write_text(json.dumps({"version": 1, "root": str(self.root), "hostId": "different-host", "controller": "https://controller.example"})) before = len(self.commands) self.assertEqual(self.updater().run(), "failed") self.assertTrue(module.UPDATE_JOURNAL.exists()) self.assertEqual(module.AGENT_EXECUTABLE.read_bytes(), CANDIDATE) self.assertEqual(len(self.commands), before) def test_rollback_busy_lock_keeps_recovery_journal_for_next_tick(self): self.crash_during_second_replacement() with module.exclusive_lock(self.root / "agent" / "activity.lock"): self.assertEqual(self.updater().run(), "deferred") self.assertEqual(json.loads(module.UPDATE_JOURNAL.read_text())["phase"], "rollback") self.assertEqual(self.updater().run(), "failed") self.assert_unchanged() def test_unconfirmed_rollback_health_retains_backups_and_journal(self): updater = self.updater() updater.wait_healthy = lambda _: False self.assertEqual(updater.run(), "failed") self.assert_unchanged() self.assertEqual(json.loads(module.UPDATE_JOURNAL.read_text())["phase"], "rollback") self.assertTrue((module.UPDATE_DIR / "agent.backup.py").exists()) self.assertIn("recovery is pending", self.status()["error"]) # A subsequent worker can safely confirm the restored release. self.assertEqual(self.updater().run(), "failed") self.assertFalse(module.UPDATE_JOURNAL.exists()) def test_stale_wrong_pid_wrong_release_or_wrong_update_health_is_rejected(self): updater = self.updater() journal = {"id": "this-update", "targetVersion": "1.1.0", "switchedAt": time.time()} valid = {"version": "1.1.0", "pid": self.pid, "checkedAt": time.time(), "updateId": "this-update"} for key, value in (("pid", self.pid + 1), ("version", "1.0.0"), ("updateId", "previous-update"), ("checkedAt", journal["switchedAt"] - 1), ("checkedAt", float("nan")), ("checkedAt", time.time() + 100)): with self.subTest(key=key, value=value): module.AGENT_HEALTH.write_text(json.dumps({**valid, key: value})) with patch.object(module, "HEALTH_TIMEOUT", 1), patch.object(module.time, "monotonic", side_effect=[0, 0, 2]), patch.object(module.time, "sleep"): self.assertFalse(updater.wait_healthy(journal)) module.AGENT_HEALTH.write_text(json.dumps(valid)) self.assertTrue(updater.wait_healthy(journal)) def test_symlinked_executable_is_rejected_without_touching_target(self): external = self.base / "external.py" external.write_bytes(PREVIOUS) module.AGENT_EXECUTABLE.unlink() module.AGENT_EXECUTABLE.symlink_to(external) self.assertEqual(self.updater().run(), "failed") self.assertEqual(external.read_bytes(), PREVIOUS) self.assertEqual(self.commands, []) def test_network_reader_enforces_header_and_stream_bounds(self): updater = module.Updater() updater.load_installation() for response in (Response(b"abc", {"Content-Length": "100"}), Response(b"abcdef")): with self.subTest(response=response), patch.object(module.urllib.request, "build_opener") as opener: opener.return_value.open.return_value = response with self.assertRaises(module.IntegrityError): updater.fetch("/api/agent/release/ec-agent.py", 5) def test_trickling_or_truncated_download_defers_without_holding_release(self): updater = module.Updater() updater.load_installation() with patch.object(module.urllib.request, "build_opener") as opener: opener.return_value.open.return_value = Response(b"source") with patch.object(module.time, "monotonic", side_effect=[0, 0, 121]): with self.assertRaisesRegex(module.ContactError, "timed out"): updater.fetch("/api/agent/release/ec-agent.py", 100) with patch.object(module.urllib.request, "build_opener") as opener: response = Response(b"source") response.read1 = lambda *_: (_ for _ in ()).throw(module.http.client.IncompleteRead(b"partial")) opener.return_value.open.return_value = response with self.assertRaises(module.ContactError): updater.fetch("/api/agent/release/ec-agent.py", 100) def test_stale_code_staging_is_recovered_without_leaving_extra_executables(self): temporary = module.AGENT_EXECUTABLE.with_name(module.AGENT_EXECUTABLE.name + ".update") temporary.write_bytes(b"partial previous download") self.assertEqual(self.updater().run(), "updated") self.assertFalse(temporary.exists()) self.assertEqual(module.AGENT_EXECUTABLE.read_bytes(), CANDIDATE) def test_download_uses_fixed_https_controller_and_no_enrollment_credentials(self): updater = module.Updater() updater.load_installation() with patch.object(module.urllib.request, "build_opener") as opener: opener.return_value.open.return_value = Response(b"source", {"Content-Length": "6"}) self.assertEqual(updater.fetch("/api/agent/release/ec-agent.py", 100), b"source") request = opener.return_value.open.call_args.args[0] self.assertEqual(request.full_url, "https://controller.example/api/agent/release/ec-agent.py") self.assertIsNone(request.get_header("Authorization")) redirect = opener.call_args.args[0] with self.assertRaises(module.ContactError): redirect.redirect_request(None, None, 302, "redirect", {}, "https://other.example") def test_insecure_or_credential_embedded_controller_is_rejected(self): for controller in ("http://controller.example", "https://token:secret@controller.example"): with self.subTest(controller=controller): module.AGENT_CONFIG.write_text(json.dumps({**self.config, "controller": controller})) self.assertEqual(self.updater().run(), "failed") self.assertEqual(self.downloads, []) self.assertEqual(self.commands, []) class UpdateSecurityTests(unittest.TestCase): def test_untrusted_ownership_writable_paths_and_special_files_are_rejected(self): for mode, uid in ((stat.S_IFREG | 0o666, 0), (stat.S_IFDIR | 0o755, 1000), (stat.S_IFLNK | 0o777, 0), (stat.S_IFIFO | 0o600, 0)): with self.subTest(mode=mode, uid=uid), patch.object(Path, "lstat", return_value=os.stat_result((mode, 1, 1, 1, uid, 0, 0, 0, 0, 0))): with self.assertRaises(module.UpdateError): module.trusted("/etc/electric-cherry/update/journal.json") def test_source_is_parsed_and_compiled_without_execution(self): source = b'VERSION="1.2.3"\nraise RuntimeError("Never execute downloaded code during verification")\n' self.assertEqual(module.source_version(source), "1.2.3") def test_numeric_version_order_and_strict_format(self): self.assertGreater(module.version("1.10.0"), module.version("1.9.9")) for value in ("v1.0.0", "1.0", "1.2.03", "1.0.0+local", "1.0.0-rc1", True, None, "1.0.000000000000"): with self.subTest(value=value), self.assertRaises(module.IntegrityError): module.version(value) if __name__ == "__main__": unittest.main()