"""The remote worker's own guards: where its token goes, and what node code gets.""" import argparse import asyncio import json import time from pathlib import Path import pytest from fluksio_worker import agent from fluksio_worker.worker_main import _Reporter @pytest.mark.parametrize( ("url", "trusted"), [ ("wss://engine.example.com/api/v1/workers/attach", True), ("https://engine.example.com/api/v1/artifacts", True), ("ws://127.0.0.1:8000/api/v1/workers/attach", True), ("ws://127.1.2.3:8000/api/v1/workers/attach", True), ("ws://[::1]:8000/api/v1/workers/attach", True), ("ws://localhost:8000/api/v1/workers/attach", True), ("ws://192.0.2.10:8000/api/v1/workers/attach", False), ("ws://engine.local:8000/api/v1/workers/attach", False), ("http://192.0.2.10:8000/api/v1/artifacts", False), ], ) def test_only_tls_or_this_machine_is_handed_the_token(url: str, trusted: bool) -> None: assert agent._encrypted_or_local(url) is trusted def test_a_cleartext_engine_elsewhere_needs_insecure( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ) -> None: url = "ws://192.0.2.10:8000/api/v1/workers/attach" with pytest.raises(SystemExit): agent.main(["--url", url, "--token", "t"]) assert "--insecure" in capsys.readouterr().err async def attached(self: agent.Agent) -> None: return None monkeypatch.setattr(agent.Agent, "serve_forever", attached) assert agent.main(["--url", url, "--token", "t", "--insecure"]) == 0 def test_node_code_does_not_get_the_hosts_credentials( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("FLUKSIO_WORKER_TOKEN", "worker-token") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "not-for-nodes") monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2") args = argparse.Namespace( url="wss://engine.example.com/api/v1/workers/attach", artifact_url="", token="worker-token", ) env = agent.Agent(args).env assert "FLUKSIO_WORKER_TOKEN" not in env assert "AWS_SECRET_ACCESS_KEY" not in env assert env["CUDA_VISIBLE_DEVICES"] == "2" assert "PATH" in env # Put back on purpose: the artifact store is reached with it. assert env["FLUKSIO_ARTIFACT_URL"] == "https://engine.example.com/api/v1/artifacts" @pytest.mark.parametrize( "digest", [ "sha256:../../../etc/passwd", "sha256:" + "a" * 63, "sha256:" + "A" * 64, "md5:" + "a" * 64, ], ) def test_only_a_whole_digest_is_opened(digest: str) -> None: """The digest is joined into a path, so anything else could read any file.""" with pytest.raises(ValueError, match="not an artifact reference"): _Reporter("fluksio").load_artifact({"digest": digest}) def _gone(pid: int) -> bool: """No such process, or one only waiting to be reaped.""" try: return Path(f"/proc/{pid}/stat").read_text().split()[2] == "Z" except FileNotFoundError: return True def test_stopping_a_service_takes_its_process_group_with_it(tmp_path: Path) -> None: """A child of the server left running would keep its port.""" args = argparse.Namespace( url="wss://engine.example.com/api/v1/workers/attach", artifact_url="", token="t", ) worker = agent.Agent(args) child = tmp_path / "child" sent: list[dict] = [] class Socket: async def send(self, text: str) -> None: sent.append(json.loads(text)) async def scenario() -> None: script = f"sleep 300 & echo $! > {child}; wait" request = {"op": "serve", "call_id": "svc", "argv": ["sh", "-c", script]} task = asyncio.create_task(worker._serve(Socket(), request)) for _ in range(500): if child.exists() and child.read_text().strip(): break await asyncio.sleep(0.01) worker._cancel("svc") await task asyncio.run(scenario()) pid = int(child.read_text()) deadline = time.monotonic() + 5 while not _gone(pid) and time.monotonic() < deadline: time.sleep(0.05) assert _gone(pid) assert not worker.serving assert sent == [ { "call_id": "svc", "ok": False, "error": { "type": "ServiceStopped", "message": "the service was stopped", "short": "the service was stopped", "traceback": "", }, } ]