419 lines
22 KiB
Python
419 lines
22 KiB
Python
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()
|