Add automatic agent updates with verified rollback

This commit is contained in:
2026-10-04 10:40:45 -04:00
parent 020894b816
commit 8d1331c67d
26 changed files with 2568 additions and 110 deletions
+418
View File
@@ -0,0 +1,418 @@
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()