diff --git a/agent/job.py b/agent/job.py index 61e8c5f3..d479d3f8 100644 --- a/agent/job.py +++ b/agent/job.py @@ -136,8 +136,15 @@ def enqueue(self, name, function, args, kwargs, agent_job_id=None): @save def cancel(self): self.job.cancel() + self._cleanup_cancelled_workload_environment() self.model.status = "Failure" + def _cleanup_cancelled_workload_environment(self): + from agent.server import Server + from agent.workload import Workload + + Workload(Server()).cleanup_cancelled_environment(self.model.data) + @save def stop(self): send_stop_job_command(self.redis, self.job.get_id()) diff --git a/agent/tests/test_workload.py b/agent/tests/test_workload.py new file mode 100644 index 00000000..489d75c7 --- /dev/null +++ b/agent/tests/test_workload.py @@ -0,0 +1,123 @@ +import json +import os +import tempfile +import unittest +from pathlib import Path + +from agent.workload import ENVIRONMENT_FILE_MAX_AGE, InvalidWorkload, Workload, WorkloadConfig + + +class TestWorkloadConfig(unittest.TestCase): + def valid_config(self): + return { + "name": "xassida-search", + "image": "ghcr.io/jvrlc/xassida-search:abc123", + "container_port": 3000, + "host_port": 13000, + "health_path": "/api/health", + } + + def test_valid_config_builds_loopback_only_docker_command(self): + config = WorkloadConfig(self.valid_config(), {"SUPABASE_URL": "safe"}) + config.validate() + + command = config.run_command("xassida-search-candidate", Path("/secure/environment")) + + self.assertIn("127.0.0.1:13000:3000", command) + self.assertNotIn("safe", command) + + def test_public_config_contains_secret_names_but_not_values(self): + config = WorkloadConfig(self.valid_config(), {"SUPABASE_SERVICE_ROLE_KEY": "secret"}) + config.validate() + + public = config.public_config() + + self.assertEqual(public["environment_keys"], ["SUPABASE_SERVICE_ROLE_KEY"]) + self.assertNotIn("secret", str(public)) + + def test_queued_config_does_not_contain_environment_values(self): + config = WorkloadConfig(self.valid_config(), {"SUPABASE_SERVICE_ROLE_KEY": "secret"}) + config.validate() + + queued = config.public_config() + + self.assertNotIn("secret", str(queued)) + + def test_rejects_shell_metacharacters_in_workload_name(self): + config = self.valid_config() + config["name"] = "xassida;rm" + + with self.assertRaisesRegex(InvalidWorkload, "Workload name"): + WorkloadConfig(config, {}).validate() + + def test_rejects_environment_values_with_newlines(self): + with self.assertRaisesRegex(InvalidWorkload, "SUPABASE_URL"): + WorkloadConfig(self.valid_config(), {"SUPABASE_URL": "safe\nINJECTED=yes"}).validate() + + def test_rejects_public_privileged_port(self): + config = self.valid_config() + config["host_port"] = 443 + + with self.assertRaisesRegex(InvalidWorkload, "host_port"): + WorkloadConfig(config, {}).validate() + + +class TestWorkloadEnvironmentFiles(unittest.TestCase): + def test_queued_environment_files_are_distinct_and_encrypted(self): + workload = object.__new__(Workload) + with tempfile.TemporaryDirectory() as temporary_directory: + directory = Path(temporary_directory) + workload.directory = temporary_directory + + first = workload._write_environment(directory, {"TOKEN": "first"}) + second = workload._write_environment(directory, {"TOKEN": "second"}) + + self.assertNotEqual(first, second) + self.assertNotIn(b"first", first.read_bytes()) + self.assertNotIn(b"second", second.read_bytes()) + self.assertEqual(first.stat().st_mode & 0o777, 0o600) + self.assertEqual(second.stat().st_mode & 0o777, 0o600) + + def test_plaintext_exists_only_while_worker_uses_it(self): + workload = object.__new__(Workload) + with tempfile.TemporaryDirectory() as temporary_directory: + directory = Path(temporary_directory) + workload.directory = temporary_directory + queued = workload._write_environment(directory, {"TOKEN": "secret"}) + + with workload._decrypted_environment(queued) as runtime: + self.assertEqual(runtime.read_text(), "TOKEN=secret\n") + self.assertEqual(runtime.stat().st_mode & 0o777, 0o600) + + self.assertFalse(runtime.exists()) + + def test_expired_encrypted_file_from_cancelled_job_is_removed(self): + workload = object.__new__(Workload) + with tempfile.TemporaryDirectory() as temporary_directory: + directory = Path(temporary_directory) + workload.directory = temporary_directory + expired = workload._write_environment(directory, {"TOKEN": "expired"}) + fresh = workload._write_environment(directory, {"TOKEN": "fresh"}) + expired_time = expired.stat().st_mtime - ENVIRONMENT_FILE_MAX_AGE - 1 + os.utime(expired, (expired_time, expired_time)) + + workload._remove_expired_environments(directory) + + self.assertFalse(expired.exists()) + self.assertTrue(fresh.exists()) + + def test_cancelled_job_removes_its_encrypted_environment(self): + workload = object.__new__(Workload) + with tempfile.TemporaryDirectory() as temporary_directory: + directory = Path(temporary_directory) + workload.directory = temporary_directory + queued = workload._write_environment(directory, {"TOKEN": "secret"}) + job_data = json.dumps({"function": "deploy", "args": [{}, str(queued)], "kwargs": {}}) + + workload.cleanup_cancelled_environment(job_data) + + self.assertFalse(queued.exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/agent/web.py b/agent/web.py index 7f0fb383..be7fa77e 100644 --- a/agent/web.py +++ b/agent/web.py @@ -36,6 +36,7 @@ from agent.snapshot_recovery import SnapshotRecovery from agent.ssh import SSHProxy from agent.utils import check_installed_pyspy +from agent.workload import InvalidWorkload, Workload if TYPE_CHECKING: from datetime import datetime, timedelta @@ -234,6 +235,41 @@ def patch_build_image(): return {"job": job} +@application.route("/workloads", methods=["POST"]) +def deploy_workload(): + data = request.json or {} + try: + job = Workload(Server()).prepare_deploy(data.get("config", {}), data.get("environment", {})) + except InvalidWorkload as error: + return {"message": str(error)}, 400 + return {"job": job} + + +@application.route("/workloads//status", methods=["GET"]) +def workload_status(name: str): + try: + return Workload(Server()).status(name) + except InvalidWorkload as error: + return {"message": str(error)}, 400 + + +@application.route("/workloads//logs", methods=["GET"]) +def workload_logs(name: str): + try: + lines = request.args.get("lines", 200, type=int) + return Workload(Server()).logs(name, lines) + except InvalidWorkload as error: + return {"message": str(error)}, 400 + + +@application.route("/workloads//rollback", methods=["POST"]) +def rollback_workload(name: str): + try: + return {"job": Workload(Server()).rollback(name)} + except InvalidWorkload as error: + return {"message": str(error)}, 400 + + @application.route("/server") def get_server(): return Server().dump() diff --git a/agent/workload.py b/agent/workload.py new file mode 100644 index 00000000..20765985 --- /dev/null +++ b/agent/workload.py @@ -0,0 +1,329 @@ +from __future__ import annotations + +import json +import os +import re +import tempfile +import time +from contextlib import contextmanager +from pathlib import Path +from urllib.error import URLError +from urllib.request import urlopen + +from cryptography.fernet import Fernet +from filelock import FileLock + +from agent.base import Base +from agent.job import job, step + +WORKLOAD_NAME = re.compile(r"^[a-z0-9][a-z0-9-]{0,62}$") +IMAGE_REFERENCE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._/:@-]{0,254}$") +ENVIRONMENT_FILE_MAX_AGE = 24 * 60 * 60 + + +class InvalidWorkload(ValueError): + pass + + +class Workload(Base): + def __init__(self, server): + super().__init__() + self.server = server + self.directory = self._workloads_directory() + + @property + def job_record(self): + return self.server.job_record + + @property + def step_record(self): + return self.server.step_record + + @step_record.setter + def step_record(self, value): + self.server.step_record = value + + def prepare_deploy(self, config: dict, environment: dict[str, str]): + deployment = WorkloadConfig(config, environment) + deployment.validate() + directory = self._create_directory(deployment.name) + self._remove_expired_environments(directory) + environment_file = self._write_environment(directory, deployment.environment) + try: + return self.deploy(deployment.public_config(), str(environment_file)) + except Exception: + environment_file.unlink(missing_ok=True) + raise + + @job("Deploy Workload", priority="low") + def deploy(self, config: dict, environment_file: str): + deployment = WorkloadConfig(config, {}) + deployment.validate() + try: + lock_path = Path(environment_file).parent / ".deploy.lock" + with FileLock(lock_path, timeout=300), self._decrypted_environment( + Path(environment_file) + ) as decrypted: + self._deploy(deployment, decrypted) + finally: + Path(environment_file).unlink(missing_ok=True) + + @step("Deploy Workload") + def _deploy(self, deployment: WorkloadConfig, environment_file: Path): + directory = self._create_directory(deployment.name) + self.server.execute(["docker", "pull", deployment.image]) + self._replace_container(deployment, environment_file) + self._write_config(directory, deployment.config_without_secrets()) + + def status(self, name: str): + validate_name(name) + result = self.server.execute( + ["docker", "inspect", "--format", "{{json .State}}", name], + non_zero_throw=False, + ) + if result["returncode"]: + return {"name": name, "status": "missing"} + return {"name": name, "status": json.loads(result["output"])} + + def logs(self, name: str, lines: int = 200): + validate_name(name) + if not 1 <= lines <= 1000: + raise InvalidWorkload("Log line count must be between 1 and 1000") + result = self.server.execute( + ["docker", "logs", "--tail", str(lines), name], + skip_output_log=True, + ) + return {"name": name, "logs": result["output"]} + + def cleanup_cancelled_environment(self, job_data: str): + data = json.loads(job_data) + if data.get("function") != "deploy": + return + arguments = data.get("args", []) + if len(arguments) < 2: + return + self._remove_queued_environment(arguments[1]) + + @job("Rollback Workload", priority="low") + def rollback(self, name: str): + validate_name(name) + previous = f"{name}-previous" + failed = f"{name}-failed" + if not self._container_exists(previous): + raise InvalidWorkload("No previous workload deployment is available") + self.server.execute(["docker", "stop", name]) + try: + self.server.execute(["docker", "rename", name, failed]) + except Exception: + self.server.execute(["docker", "start", name]) + raise + try: + self.server.execute(["docker", "rename", previous, name]) + self.server.execute(["docker", "start", name]) + except Exception: + if self._container_exists(name): + self.server.execute(["docker", "rename", name, previous]) + self.server.execute(["docker", "rename", failed, name]) + self.server.execute(["docker", "start", name]) + raise + self._remove_container(failed) + + def _workloads_directory(self) -> str: + configured = self.server.config.get("workloads_directory") + if configured: + return configured + return os.path.join(os.path.dirname(self.server.benches_directory), "workloads") + + def _create_directory(self, name: str) -> Path: + root = Path(self.directory) + root.mkdir(mode=0o700, parents=True, exist_ok=True) + root.chmod(0o700) + directory = root / name + directory.mkdir(mode=0o700, exist_ok=True) + directory.chmod(0o700) + return directory + + def _write_environment(self, directory: Path, environment: dict[str, str]) -> Path: + content = "".join(f"{key}={value}\n" for key, value in sorted(environment.items())) + encrypted = self._environment_cipher().encrypt(content.encode()) + return self._write_protected_file(directory, "environment-queued-", encrypted) + + @contextmanager + def _decrypted_environment(self, encrypted_path: Path): + content = self._environment_cipher().decrypt(encrypted_path.read_bytes()) + path = self._write_protected_file(encrypted_path.parent, "environment-runtime-", content) + try: + yield path + finally: + path.unlink(missing_ok=True) + + def _environment_cipher(self) -> Fernet: + key_path = Path(self.directory) / ".environment.key" + with FileLock(f"{key_path}.lock", timeout=30): + if not key_path.exists(): + self._write_exclusive_file(key_path, Fernet.generate_key()) + key_path.chmod(0o600) + return Fernet(key_path.read_bytes()) + + def _write_protected_file(self, directory: Path, prefix: str, content: bytes) -> Path: + descriptor, name = tempfile.mkstemp(prefix=prefix, dir=directory) + path = Path(name) + try: + os.fchmod(descriptor, 0o600) + with os.fdopen(descriptor, "wb") as file: + descriptor = -1 + file.write(content) + except Exception: + if descriptor != -1: + os.close(descriptor) + path.unlink(missing_ok=True) + raise + return path + + def _write_exclusive_file(self, path: Path, content: bytes): + descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + try: + with os.fdopen(descriptor, "wb") as file: + descriptor = -1 + file.write(content) + finally: + if descriptor != -1: + os.close(descriptor) + + def _remove_expired_environments(self, directory: Path): + cutoff = time.time() - ENVIRONMENT_FILE_MAX_AGE + for path in directory.glob("environment-queued-*"): + if path.stat().st_mtime < cutoff: + path.unlink(missing_ok=True) + + def _remove_queued_environment(self, name: str): + root = Path(self.directory).resolve() + path = Path(name).resolve() + if os.path.commonpath((root, path)) != str(root): + raise InvalidWorkload("Queued environment path is outside the workloads directory") + if not path.name.startswith("environment-queued-"): + raise InvalidWorkload("Invalid queued environment path") + path.unlink(missing_ok=True) + + def _write_config(self, directory: Path, config: dict): + path = directory / "config.json" + temporary = path.with_suffix(".tmp") + temporary.write_text(json.dumps(config, indent=2, sort_keys=True)) + temporary.chmod(0o600) + temporary.replace(path) + + def _replace_container(self, deployment: WorkloadConfig, environment_file: Path): + previous = f"{deployment.name}-previous" + candidate = f"{deployment.name}-candidate" + self._remove_container(candidate) + self._remove_container(previous) + current_exists = self._container_exists(deployment.name) + if current_exists: + self.server.execute(["docker", "rename", deployment.name, previous]) + try: + self.server.execute(["docker", "stop", previous]) + except Exception: + self.server.execute(["docker", "rename", previous, deployment.name]) + raise + try: + self.server.execute(deployment.run_command(candidate, environment_file)) + self._wait_until_healthy(deployment) + except Exception: + self._remove_container(candidate) + if current_exists: + self.server.execute(["docker", "rename", previous, deployment.name]) + self.server.execute(["docker", "start", deployment.name]) + raise + self.server.execute(["docker", "rename", candidate, deployment.name]) + + def _wait_until_healthy(self, deployment: WorkloadConfig): + deadline = time.monotonic() + deployment.health_timeout + url = f"http://127.0.0.1:{deployment.host_port}{deployment.health_path}" + while time.monotonic() < deadline: + try: + with urlopen(url, timeout=2) as response: + if response.status < 400: + return + except (OSError, URLError): + time.sleep(1) + raise RuntimeError(f"Workload health check failed on {deployment.health_path}") + + def _container_exists(self, name: str) -> bool: + result = self.server.execute(["docker", "inspect", name], non_zero_throw=False) + return result["returncode"] == 0 + + def _remove_container(self, name: str): + if self._container_exists(name): + self.server.execute(["docker", "rm", "--force", name]) + + +class WorkloadConfig: + def __init__(self, config: dict, environment: dict[str, str]): + self.name = config.get("name", "") + self.image = config.get("image", "") + self.container_port = config.get("container_port") + self.host_port = config.get("host_port") + self.health_path = config.get("health_path", "/") + self.health_timeout = config.get("health_timeout", 60) + self.environment = environment + self.environment_keys = config.get("environment_keys", sorted(environment)) + + def validate(self): + validate_name(self.name) + if not IMAGE_REFERENCE.fullmatch(self.image): + raise InvalidWorkload("Invalid Docker image reference") + self._validate_port(self.container_port, "container_port") + self._validate_port(self.host_port, "host_port") + if not isinstance(self.health_path, str) or not self.health_path.startswith("/"): + raise InvalidWorkload("Health path must start with /") + if not isinstance(self.health_timeout, int) or not 1 <= self.health_timeout <= 300: + raise InvalidWorkload("Health timeout must be between 1 and 300 seconds") + self._validate_environment() + + def _validate_port(self, port, field: str): + if not isinstance(port, int) or not 1024 <= port <= 65535: + raise InvalidWorkload(f"{field} must be between 1024 and 65535") + + def _validate_environment(self): + if not isinstance(self.environment, dict): + raise InvalidWorkload("Environment must be an object") + for key, value in self.environment.items(): + if not re.fullmatch(r"[A-Z_][A-Z0-9_]*", key): + raise InvalidWorkload(f"Invalid environment variable name: {key}") + if not isinstance(value, str) or "\n" in value or "\r" in value: + raise InvalidWorkload(f"Invalid environment variable value: {key}") + + def run_command(self, container_name: str, environment_file: Path) -> list[str]: + return [ + "docker", + "run", + "--detach", + "--name", + container_name, + "--restart", + "unless-stopped", + "--env-file", + str(environment_file), + "--publish", + f"127.0.0.1:{self.host_port}:{self.container_port}", + self.image, + ] + + def public_config(self) -> dict: + return self.config_without_secrets() + + def config_without_secrets(self) -> dict: + return { + "name": self.name, + "image": self.image, + "container_port": self.container_port, + "host_port": self.host_port, + "health_path": self.health_path, + "health_timeout": self.health_timeout, + "environment_keys": self.environment_keys, + } + + +def validate_name(name: str): + if not isinstance(name, str) or not WORKLOAD_NAME.fullmatch(name): + raise InvalidWorkload("Workload name must contain lowercase letters, numbers, or hyphens") diff --git a/requirements.txt b/requirements.txt index 9f405e22..e9bbb927 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ boto3==1.18.0 certifi==2019.11.28 chardet==3.0.4 click==8.1.3 +cryptography==46.0.7 filewarmer==0.0.17 Flask==1.1.1 gunicorn==20.0.4 @@ -31,4 +32,4 @@ mariadb-binlog-indexer==0.0.19 mariadb-table-usage==0.0.3 psutil==7.0.0 tomli==2.4.1 -semantic-version==2.10.0 \ No newline at end of file +semantic-version==2.10.0