Add shared managed lock broker
This commit is contained in:
@@ -0,0 +1,261 @@
|
||||
import json
|
||||
import os
|
||||
import secrets
|
||||
import stat
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from harness.lib.managed_lock import ManagedLock, probe_free
|
||||
from harness.lib import managed_lock_broker as broker
|
||||
|
||||
BROKER = Path(__file__).parents[2] / "lib" / "managed_lock_broker.py"
|
||||
|
||||
|
||||
def wait_ready(proc, timeout=3):
|
||||
line = proc.stdout.readline().decode()
|
||||
if not line:
|
||||
raise AssertionError(proc.stderr.read().decode())
|
||||
return json.loads(line)
|
||||
|
||||
|
||||
class ManagedLockBrokerTests(unittest.TestCase):
|
||||
def start(self, root, owner="alice", repo="ledger", issue=7, ttl=5, role="gitea-impl-agent"):
|
||||
locks = root / "repositories" / owner / repo / "locks"
|
||||
locks.mkdir(parents=True)
|
||||
os.chmod(locks, 0o700)
|
||||
(locks / "repository.lock").touch(mode=0o600)
|
||||
runtime = root / "runtime"
|
||||
runtime.mkdir()
|
||||
os.chmod(runtime, 0o700)
|
||||
handle = runtime / "handle.json"
|
||||
socket = runtime / "broker.sock"
|
||||
cmd = [sys.executable, str(BROKER), "serve", "--root", str(root), "--owner", owner,
|
||||
"--repository", repo, "--issue", str(issue), "--role", role, "--run-id", "run-1",
|
||||
"--handle", str(handle), "--socket", str(socket), "--ttl-seconds", str(ttl),
|
||||
"--acquire-timeout", "0.3"]
|
||||
proc = subprocess.Popen(cmd, cwd=Path.cwd(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
ready = wait_ready(proc)
|
||||
self.assertEqual(ready["status"], "held")
|
||||
self.assertNotIn("token", json.dumps(ready))
|
||||
return proc, handle, socket
|
||||
|
||||
def request(self, handle, action, ttl=None, token=None):
|
||||
data = json.loads(handle.read_text())
|
||||
if token is not None:
|
||||
data["token"] = token
|
||||
import socket
|
||||
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
|
||||
s.connect(data["socket"])
|
||||
payload = {"action": action, "token": data["token"] if token is None else token}
|
||||
if ttl is not None:
|
||||
payload["ttl_seconds"] = ttl
|
||||
s.sendall((json.dumps(payload) + "\n").encode())
|
||||
with s.makefile("rb") as response_file:
|
||||
return json.loads(response_file.readline())
|
||||
|
||||
def test_cross_process_contention_and_authenticated_release(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
root = Path(d)
|
||||
proc, handle, _ = self.start(root)
|
||||
try:
|
||||
self.assertFalse(probe_free(root / "repositories/alice/ledger/locks/issue-7.lock"))
|
||||
self.assertEqual(self.request(handle, "release")["status"], "released")
|
||||
proc.wait(timeout=2)
|
||||
self.assertTrue(probe_free(root / "repositories/alice/ledger/locks/issue-7.lock"))
|
||||
self.assertNotIn("token", json.loads(handle.read_text()))
|
||||
finally:
|
||||
if proc.poll() is None:
|
||||
proc.kill(); proc.wait()
|
||||
proc.stdout.close(); proc.stderr.close()
|
||||
|
||||
def test_probe_cli_fails_while_held_and_passes_after_release(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
root = Path(d)
|
||||
proc, handle, _ = self.start(root)
|
||||
command = [sys.executable, str(BROKER), "probe", "--root", str(root),
|
||||
"--owner", "alice", "--repository", "ledger", "--issue", "7"]
|
||||
try:
|
||||
held = subprocess.run(command, capture_output=True, text=True)
|
||||
self.assertEqual(held.returncode, 1)
|
||||
self.assertFalse(json.loads(held.stdout)["free"])
|
||||
self.assertEqual(self.request(handle, "release")["status"], "released")
|
||||
proc.wait(timeout=2)
|
||||
free = subprocess.run(command, capture_output=True, text=True)
|
||||
self.assertEqual(free.returncode, 0)
|
||||
self.assertTrue(json.loads(free.stdout)["free"])
|
||||
finally:
|
||||
if proc.poll() is None:
|
||||
proc.kill(); proc.wait()
|
||||
proc.stdout.close(); proc.stderr.close()
|
||||
|
||||
def test_issue_probe_passes_when_another_workflow_holds_repository(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
root = Path(d)
|
||||
proc, handle, _ = self.start(root)
|
||||
base = [sys.executable, str(BROKER), "probe", "--root", str(root),
|
||||
"--owner", "alice", "--repository", "ledger", "--issue", "7"]
|
||||
repository = root / "repositories/alice/ledger/locks/repository.lock"
|
||||
try:
|
||||
self.assertEqual(self.request(handle, "release")["status"], "released")
|
||||
proc.wait(timeout=2)
|
||||
with ManagedLock(repository):
|
||||
issue = subprocess.run([*base, "--resource", "issue"], capture_output=True, text=True)
|
||||
all_resources = subprocess.run([*base, "--resource", "all"], capture_output=True, text=True)
|
||||
self.assertEqual(issue.returncode, 0)
|
||||
self.assertTrue(json.loads(issue.stdout)["issue_free"])
|
||||
self.assertEqual(all_resources.returncode, 1)
|
||||
self.assertFalse(json.loads(all_resources.stdout)["repository_free"])
|
||||
finally:
|
||||
if proc.poll() is None:
|
||||
proc.kill(); proc.wait()
|
||||
proc.stdout.close(); proc.stderr.close()
|
||||
|
||||
def test_new_handle_mode_is_exact_under_restrictive_umask(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
path = Path(d) / "handle.json"
|
||||
old = os.umask(0o777)
|
||||
try:
|
||||
broker._write_handle(path, {"token": "a" * 64, "socket": "/tmp/x"}, exclusive=True)
|
||||
finally:
|
||||
os.umask(old)
|
||||
self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o600)
|
||||
|
||||
def test_repository_scope_release_keeps_issue(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
root = Path(d); proc, handle, _ = self.start(root)
|
||||
try:
|
||||
self.assertEqual(self.request(handle, "acquire_repository")["status"], "repository-held")
|
||||
self.assertFalse(probe_free(root / "repositories/alice/ledger/locks/repository.lock"))
|
||||
self.assertEqual(self.request(handle, "release_repository")["status"], "held")
|
||||
self.assertTrue(probe_free(root / "repositories/alice/ledger/locks/repository.lock"))
|
||||
self.assertFalse(probe_free(root / "repositories/alice/ledger/locks/issue-7.lock"))
|
||||
finally:
|
||||
self.request(handle, "release"); proc.wait(timeout=2); proc.stdout.close(); proc.stderr.close()
|
||||
|
||||
def test_ttl_expiry_updates_handle_and_releases(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
root = Path(d); proc, handle, _ = self.start(root, ttl=1)
|
||||
proc.wait(timeout=4)
|
||||
state = json.loads(handle.read_text())
|
||||
self.assertEqual(state["status"], "expired")
|
||||
self.assertNotIn("token", state)
|
||||
self.assertTrue(probe_free(root / "repositories/alice/ledger/locks/issue-7.lock"))
|
||||
proc.stdout.close(); proc.stderr.close()
|
||||
|
||||
def test_expired_lease_rejects_renew_and_release_before_action(self):
|
||||
for action in ("renew", "release"):
|
||||
with self.subTest(action=action), tempfile.TemporaryDirectory() as d:
|
||||
root = Path(d)
|
||||
locks = root / "repositories/alice/ledger/locks"
|
||||
locks.mkdir(parents=True); os.chmod(locks, 0o700)
|
||||
(locks / "repository.lock").touch(mode=0o600)
|
||||
runtime = root / "runtime"; runtime.mkdir(mode=0o700)
|
||||
args = Namespace(root=str(root), owner="alice", repository="ledger", issue=7,
|
||||
role="gitea-impl-agent", run_id="expiry-test",
|
||||
handle=str(runtime / "handle.json"), socket=str(runtime / "broker.sock"),
|
||||
ttl_seconds=5.0, acquire_timeout=0.0)
|
||||
instance = broker.Broker(args)
|
||||
instance.set.issue.acquire()
|
||||
broker._write_handle(instance.handle, instance.state, exclusive=True)
|
||||
instance.handle_created = True; instance.started = True
|
||||
instance.deadline = time.monotonic() - 1
|
||||
request: dict[str, object] = {"action": action, "token": instance.token}
|
||||
if action == "renew": request["ttl_seconds"] = 30
|
||||
self.assertEqual(instance.handle_request(request), {"error": "expired"})
|
||||
self.assertTrue(probe_free(locks / "issue-7.lock"))
|
||||
terminal = json.loads(instance.handle.read_text())
|
||||
self.assertEqual(terminal["status"], "expired")
|
||||
self.assertNotIn("token", terminal)
|
||||
|
||||
def test_startup_failure_releases_issue_and_terminalizes_partial_handle(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
root = Path(d)
|
||||
locks = root / "repositories/alice/ledger/locks"
|
||||
locks.mkdir(parents=True); os.chmod(locks, 0o700)
|
||||
(locks / "repository.lock").touch(mode=0o600)
|
||||
runtime = root / "runtime"; runtime.mkdir(mode=0o700)
|
||||
handle = runtime / "handle.json"
|
||||
args = Namespace(root=str(root), owner="alice", repository="ledger", issue=7,
|
||||
role="gitea-impl-agent", run_id="startup-failure",
|
||||
handle=str(handle), socket=str(runtime / "broker.sock"),
|
||||
ttl_seconds=5.0, acquire_timeout=0.0)
|
||||
instance = broker.Broker(args)
|
||||
real_write = broker._write_handle
|
||||
|
||||
def partial_then_fail(path, state, *, exclusive=False, on_created=None):
|
||||
real_write(path, state, exclusive=exclusive, on_created=on_created)
|
||||
raise OSError("synthetic startup failure")
|
||||
|
||||
with patch.object(broker, "_write_handle", side_effect=partial_then_fail):
|
||||
with self.assertRaises(OSError):
|
||||
instance.run()
|
||||
self.assertTrue(probe_free(locks / "issue-7.lock"))
|
||||
terminal = json.loads(handle.read_text())
|
||||
self.assertIn(terminal["status"], ("released", "expired"))
|
||||
self.assertNotIn("token", terminal)
|
||||
|
||||
def test_startup_existing_handle_is_never_claimed_or_modified(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
root = Path(d)
|
||||
locks = root / "repositories/alice/ledger/locks"
|
||||
locks.mkdir(parents=True); os.chmod(locks, 0o700)
|
||||
(locks / "repository.lock").touch(mode=0o600)
|
||||
runtime = root / "runtime"; runtime.mkdir(mode=0o700)
|
||||
handle = runtime / "handle.json"
|
||||
original = b'{"foreign":"do-not-touch"}\n'
|
||||
handle.write_bytes(original); handle.chmod(0o600)
|
||||
args = Namespace(root=str(root), owner="alice", repository="ledger", issue=7,
|
||||
role="gitea-impl-agent", run_id="existing-handle",
|
||||
handle=str(handle), socket=str(runtime / "broker.sock"),
|
||||
ttl_seconds=5.0, acquire_timeout=0.0)
|
||||
instance = broker.Broker(args)
|
||||
with self.assertRaises(FileExistsError):
|
||||
instance.run()
|
||||
self.assertEqual(handle.read_bytes(), original)
|
||||
self.assertTrue(probe_free(locks / "issue-7.lock"))
|
||||
|
||||
def test_cli_client_compatibility_and_nonprinting_function(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
root = Path(d); proc, handle, _ = self.start(root)
|
||||
try:
|
||||
result = broker.request_handle(handle, "status")
|
||||
self.assertEqual(result["status"], "held")
|
||||
self.assertNotIn("token", json.dumps(result))
|
||||
cli = subprocess.run([sys.executable, str(BROKER), "request", "--handle", str(handle), "--action", "status"],
|
||||
capture_output=True, text=True)
|
||||
self.assertEqual(cli.returncode, 0)
|
||||
self.assertEqual(json.loads(cli.stdout)["status"], "held")
|
||||
finally:
|
||||
broker.request_handle(handle, "release"); proc.wait(timeout=2); proc.stdout.close(); proc.stderr.close()
|
||||
|
||||
def test_token_not_in_ready_or_errors_and_bad_handle_rejected(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
root = Path(d); proc, handle, _ = self.start(root)
|
||||
try:
|
||||
raw = handle.read_text()
|
||||
self.assertIn("token", raw)
|
||||
bad = root / "bad"; bad.symlink_to(handle)
|
||||
result = subprocess.run([sys.executable, str(BROKER), "request", "--handle", str(bad), "--action", "status"], capture_output=True, text=True)
|
||||
self.assertNotEqual(result.returncode, 0)
|
||||
self.assertNotIn(json.loads(raw)["token"], result.stderr)
|
||||
self.assertEqual(self.request(handle, "status", token=secrets.token_hex(32))["error"], "unauthorized")
|
||||
finally:
|
||||
self.request(handle, "release"); proc.wait(timeout=2); proc.stdout.close(); proc.stderr.close()
|
||||
|
||||
def test_unsafe_handle_and_identifiers_rejected(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
root = Path(d)
|
||||
for role in ("bad", "gitea-impl-agent/evil"):
|
||||
p = subprocess.run([sys.executable, str(BROKER), "serve", "--root", str(root), "--owner", "a", "--repository", "r", "--issue", "1", "--role", role, "--run-id", "x", "--handle", str(root/"h"), "--socket", str(root/"s"), "--ttl-seconds", "1"], capture_output=True, text=True)
|
||||
self.assertNotEqual(p.returncode, 0)
|
||||
self.assertNotIn("token", p.stderr.lower())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user