369 lines
18 KiB
Python
369 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""DPL append-only batch coordinator. Adapters must probe before and after writes."""
|
|
import argparse
|
|
import concurrent.futures
|
|
import contextlib
|
|
import fcntl
|
|
import hashlib
|
|
import ipaddress
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import signal
|
|
import sqlite3
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
PHASES = ('stage', 'dns', 'provision', 'relay', 'verify', 'publish', 'audit')
|
|
SHARED = {'stage': 'panel', 'dns': 'dns', 'relay': 'relay', 'publish': 'panel'}
|
|
STOP = threading.Event()
|
|
|
|
|
|
def digest(value):
|
|
return hashlib.sha256(json.dumps(value, sort_keys=True, ensure_ascii=False,
|
|
separators=(',', ':')).encode()).hexdigest()
|
|
|
|
|
|
def read_json(path):
|
|
return json.loads(Path(path).read_text())
|
|
|
|
|
|
def atomic_json(path, value):
|
|
path = Path(path)
|
|
tmp = path.with_name(path.name + '.tmp-' + str(os.getpid()) + '-' + str(threading.get_ident()))
|
|
with tmp.open('x') as f:
|
|
os.chmod(tmp, 0o600)
|
|
json.dump(value, f, ensure_ascii=False, sort_keys=True)
|
|
f.flush()
|
|
os.fsync(f.fileno())
|
|
os.replace(tmp, path)
|
|
|
|
|
|
def validate(plan):
|
|
if plan.get('version') != 1 or plan.get('mode') != 'append':
|
|
raise ValueError('Only version=1, mode=append is supported')
|
|
if not re.fullmatch(r'[a-z0-9-]{1,64}', plan.get('batch_id', '')):
|
|
raise ValueError('Invalid batch_id')
|
|
context = plan.get('context', {})
|
|
required = ('panel_vps', 'gateway_vps', 'jump_vps', 'app_root', 'panel_url',
|
|
'subscription_url', 'provision_commit', 'bundle', 'client_binary')
|
|
if any(not context.get(k) for k in required):
|
|
raise ValueError('Missing explicit context fields')
|
|
if not re.fullmatch(r'[0-9a-f]{40}', context['provision_commit']):
|
|
raise ValueError('A full reviewed, pushed Git commit is required')
|
|
for k in ('panel_url', 'subscription_url'):
|
|
if not re.fullmatch(r'https://[a-z0-9.-]+', context[k]):
|
|
raise ValueError('HTTPS origin without token/path required: ' + k)
|
|
adapter = plan.get('adapter')
|
|
if not isinstance(adapter, list) or not adapter or not all(isinstance(x, str) for x in adapter):
|
|
raise ValueError('adapter must be an argv array, never a shell string')
|
|
if not plan.get('nodes') or not isinstance(plan['nodes'], list):
|
|
raise ValueError('A nonempty node list is required')
|
|
seen = {k: set() for k in ('key', 'name', 'vps_id', 'node_address', 'tls_domain', 'relay', 'socks_port', 'number')}
|
|
for n in plan['nodes']:
|
|
if any(k not in n for k in ('key', 'country', 'sequence', 'flag', 'country_name', 'name',
|
|
'vps_id', 'node_address', 'listen_port', 'relay_host',
|
|
'relay_port', 'tls_domain', 'formal_group_id',
|
|
'expected_exit', 'socks_port')):
|
|
raise ValueError('Incomplete node specification')
|
|
if not re.fullmatch(r'[a-z0-9-]{1,48}', n['key']):
|
|
raise ValueError('Invalid deployment key')
|
|
if not re.fullmatch(r'[A-Z]{2}', n['country']):
|
|
raise ValueError('Country must be a two-letter code')
|
|
if type(n['sequence']) is not int or not 1 <= n['sequence'] <= 99:
|
|
raise ValueError('Display sequence must be 01..99')
|
|
expected = f"{n['flag']} {n['country_name']} {n['country']} - {n['sequence']:02d}"
|
|
flag = ''.join(chr(0x1F1E6 + ord(x) - ord('A')) for x in n['country'])
|
|
if n['name'] != expected or n['flag'] != flag:
|
|
raise ValueError('Display name/flag does not match country and sequence')
|
|
for k in ('listen_port', 'relay_port', 'socks_port'):
|
|
if type(n[k]) is not int or not 1024 <= n[k] <= 65535:
|
|
raise ValueError('Invalid port: ' + k)
|
|
if type(n['formal_group_id']) is not int or n['formal_group_id'] < 1:
|
|
raise ValueError('An existing formal group is required')
|
|
ipaddress.ip_address(n['node_address'])
|
|
ipaddress.ip_address(n['expected_exit'])
|
|
for k in ('relay_host', 'tls_domain'):
|
|
if not re.fullmatch(r'[a-z0-9](?:[a-z0-9.-]*[a-z0-9])?\.[a-z]{2,}', n[k]):
|
|
raise ValueError('Invalid DNS name: ' + k)
|
|
for k in ('vps_id',):
|
|
if not re.fullmatch(r'VPS-[0-9]{3}', n[k]):
|
|
raise ValueError('Use the exact CoDocs VPS ID')
|
|
identities = {k: n[k] for k in ('key', 'name', 'vps_id', 'node_address', 'tls_domain', 'socks_port')}
|
|
identities.update(relay=(n['relay_host'], n['relay_port']), number=(n['country'], n['sequence']))
|
|
for k, value in identities.items():
|
|
if value in seen[k]:
|
|
raise ValueError('Duplicate batch identity: ' + k)
|
|
seen[k].add(value)
|
|
# The plan carries references, never secret values. Adapters load credentials themselves.
|
|
def check(value):
|
|
if isinstance(value, dict):
|
|
for k, v in value.items():
|
|
if re.search(r'password|secret|token|private_key', k, re.I):
|
|
raise ValueError('Secret fields are not allowed in a batch plan')
|
|
check(v)
|
|
elif isinstance(value, list):
|
|
for v in value:
|
|
check(v)
|
|
check(plan)
|
|
return plan
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def lock(root, name):
|
|
path = root / ('lock-' + hashlib.sha256(name.encode()).hexdigest())
|
|
with path.open('a+') as f:
|
|
os.chmod(path, 0o600)
|
|
while True:
|
|
try:
|
|
fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
break
|
|
except BlockingIOError:
|
|
if STOP.wait(.1):
|
|
raise RuntimeError('Stopped while waiting for resource lock')
|
|
try:
|
|
yield
|
|
finally:
|
|
fcntl.flock(f, fcntl.LOCK_UN)
|
|
|
|
|
|
class Ledger:
|
|
def __init__(self, root):
|
|
self.root = Path(root).resolve()
|
|
if self.root == Path('/'):
|
|
raise ValueError('Invalid state root')
|
|
self.root.mkdir(parents=True, mode=0o700, exist_ok=True)
|
|
os.chmod(self.root, 0o700)
|
|
self.path = self.root / 'registry.sqlite3'
|
|
with self.db() as c:
|
|
c.executescript('''
|
|
CREATE TABLE IF NOT EXISTS batches(id TEXT PRIMARY KEY, digest TEXT NOT NULL);
|
|
CREATE TABLE IF NOT EXISTS claims(namespace TEXT, kind TEXT, value TEXT,
|
|
owner TEXT NOT NULL, PRIMARY KEY(namespace,kind,value));
|
|
CREATE TABLE IF NOT EXISTS phases(batch TEXT, node TEXT, phase TEXT,
|
|
status TEXT NOT NULL, attempts INTEGER NOT NULL DEFAULT 0,
|
|
updated REAL NOT NULL, PRIMARY KEY(batch,node,phase));
|
|
''')
|
|
os.chmod(self.path, 0o600)
|
|
|
|
def db(self):
|
|
c = sqlite3.connect(self.path, timeout=30)
|
|
c.execute('PRAGMA synchronous=FULL')
|
|
return c
|
|
|
|
def reserve(self, plan):
|
|
namespace = plan['context']['panel_vps'] + ':' + plan['context']['app_root']
|
|
with self.db() as c:
|
|
c.execute('BEGIN IMMEDIATE')
|
|
prior = c.execute('SELECT digest FROM batches WHERE id=?', (plan['batch_id'],)).fetchone()
|
|
if prior and prior[0] != digest(plan):
|
|
raise ValueError('Batch plan changed; resume requires the original exact plan')
|
|
for n in plan['nodes']:
|
|
owner = digest({'context': plan['context'], 'node': n})
|
|
claims = {k: n[k] for k in ('key', 'name', 'vps_id', 'node_address', 'tls_domain')}
|
|
claims['relay'] = [n['relay_host'], n['relay_port']]
|
|
claims['number'] = [n['country'], n['sequence']]
|
|
for kind, value in claims.items():
|
|
scope = namespace
|
|
if kind == 'relay':
|
|
scope = 'relay:' + n['relay_host']
|
|
elif kind in ('vps_id', 'node_address', 'tls_domain'):
|
|
scope = 'infrastructure'
|
|
encoded = json.dumps(value, ensure_ascii=False, sort_keys=True)
|
|
existing = c.execute('SELECT owner FROM claims WHERE namespace=? AND kind=? AND value=?',
|
|
(scope, kind, encoded)).fetchone()
|
|
if existing and existing[0] != owner:
|
|
raise ValueError('Another batch reserved identity: ' + kind)
|
|
c.execute('INSERT OR IGNORE INTO claims VALUES(?,?,?,?)', (scope, kind, encoded, owner))
|
|
# SOCKS listeners are shared by all batches using this gateway.
|
|
scope, value = 'gateway:' + plan['context']['gateway_vps'], str(n['socks_port'])
|
|
old = c.execute('SELECT owner FROM claims WHERE namespace=? AND kind=? AND value=?',
|
|
(scope, 'socks_port', value)).fetchone()
|
|
if old and old[0] != owner:
|
|
raise ValueError('Another batch reserved the gateway SOCKS port')
|
|
c.execute('INSERT OR IGNORE INTO claims VALUES(?,?,?,?)', (scope, 'socks_port', value, owner))
|
|
c.execute('INSERT OR IGNORE INTO batches VALUES(?,?)', (plan['batch_id'], digest(plan)))
|
|
|
|
def state(self, batch, key, phase):
|
|
with self.db() as c:
|
|
row = c.execute('SELECT status FROM phases WHERE batch=? AND node=? AND phase=?',
|
|
(batch, key, phase)).fetchone()
|
|
return row[0] if row else 'pending'
|
|
|
|
def mark(self, batch, key, phase, status, attempt=False):
|
|
with self.db() as c:
|
|
c.execute('''INSERT INTO phases VALUES(?,?,?,?,?,?) ON CONFLICT(batch,node,phase)
|
|
DO UPDATE SET status=excluded.status,attempts=phases.attempts+?,updated=excluded.updated''',
|
|
(batch, key, phase, status, int(attempt), time.time(), int(attempt)))
|
|
|
|
|
|
class Coordinator:
|
|
def __init__(self, plan, root, jobs=4, timeout=900):
|
|
self.plan = validate(plan)
|
|
self.ledger = Ledger(root)
|
|
self.jobs = jobs
|
|
self.timeout = timeout
|
|
self.ledger.reserve(plan)
|
|
self.dir = self.ledger.root / plan['batch_id']
|
|
self.dir.mkdir(mode=0o700, exist_ok=True)
|
|
self.plan_file = self.dir / 'plan.json'
|
|
atomic_json(self.plan_file, plan)
|
|
|
|
def call(self, n, phase, action):
|
|
fingerprint = digest({'context': self.plan['context'], 'node': n})
|
|
path = self.dir / (n['key'] + '.json')
|
|
atomic_json(path, dict(n, request_sha256=fingerprint))
|
|
argv = self.plan['adapter'] + ['--plan', str(self.plan_file), '--node', str(path),
|
|
'--phase', phase, '--action', action, '--state-dir', str(self.dir)]
|
|
with subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
|
|
start_new_session=True) as p:
|
|
deadline = time.monotonic() + self.timeout
|
|
while True:
|
|
try:
|
|
output, _ = p.communicate(timeout=.2)
|
|
break
|
|
except subprocess.TimeoutExpired:
|
|
if STOP.is_set() or time.monotonic() >= deadline:
|
|
os.killpg(p.pid, signal.SIGTERM)
|
|
try:
|
|
p.communicate(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
os.killpg(p.pid, signal.SIGKILL)
|
|
p.communicate()
|
|
raise RuntimeError('Adapter interrupted; remote outcome must be probed')
|
|
if p.returncode or len(output) > 65536:
|
|
raise RuntimeError('Adapter failed; no stdout/stderr persisted')
|
|
r = json.loads(output)
|
|
if (r.get('key'), r.get('phase'), r.get('request_sha256')) != (n['key'], phase, fingerprint):
|
|
raise RuntimeError('Adapter receipt identity mismatch')
|
|
if r.get('status') not in ('ready', 'absent', 'waiting', 'unknown', 'conflict'):
|
|
raise RuntimeError('Invalid adapter status')
|
|
return r
|
|
|
|
def phase(self, n, phase):
|
|
try:
|
|
return self._phase(n, phase)
|
|
except Exception:
|
|
self.ledger.mark(self.plan['batch_id'], n['key'], phase, 'unknown')
|
|
return False
|
|
|
|
def _phase(self, n, phase):
|
|
batch, key = self.plan['batch_id'], n['key']
|
|
resource = SHARED.get(phase, n['vps_id'])
|
|
# Common panel/dns/relay lock names across batches and processes.
|
|
namespace = self.plan['context']['panel_vps']
|
|
if resource in ('dns', 'relay'):
|
|
namespace = 'dpl-external'
|
|
with lock(self.ledger.root, namespace + ':' + resource):
|
|
r = self.call(n, phase, 'probe')
|
|
if r['status'] == 'ready':
|
|
self.ledger.mark(batch, key, phase, 'done')
|
|
return True
|
|
if r['status'] == 'conflict':
|
|
STOP.set()
|
|
if phase in ('preflight', 'verify', 'audit') or r['status'] != 'absent' or r.get('retry_safe') is not True:
|
|
self.ledger.mark(batch, key, phase, r['status'])
|
|
return False
|
|
if STOP.is_set():
|
|
return False
|
|
# Persist intent before a write; never infer success from the command exit code.
|
|
self.ledger.mark(batch, key, phase, 'running', attempt=True)
|
|
applied = self.call(n, phase, 'apply')
|
|
if applied['status'] == 'conflict':
|
|
self.ledger.mark(batch, key, phase, 'conflict')
|
|
STOP.set()
|
|
return False
|
|
r = self.call(n, phase, 'probe')
|
|
self.ledger.mark(batch, key, phase, 'done' if r['status'] == 'ready' else r['status'])
|
|
if r['status'] == 'conflict':
|
|
STOP.set()
|
|
return r['status'] == 'ready'
|
|
|
|
def worker(self, n):
|
|
try:
|
|
with lock(self.ledger.root, 'node:' + self.plan['context']['panel_vps'] + ':' + n['key']):
|
|
self.ledger.mark(self.plan['batch_id'], n['key'], 'coordinator', 'running')
|
|
# Already finished batches are re-audited; disabled QA need not be re-enabled.
|
|
if self.ledger.state(self.plan['batch_id'], n['key'], 'audit') == 'done':
|
|
success = self.phase(n, 'audit')
|
|
else:
|
|
success = True
|
|
for phase in PHASES:
|
|
if STOP.is_set() or not self.phase(n, phase):
|
|
success = False
|
|
break
|
|
self.ledger.mark(self.plan['batch_id'], n['key'], 'coordinator', 'done' if success else 'waiting')
|
|
return success
|
|
except Exception:
|
|
# Never print adapter payloads, commands, credentials or exception contents.
|
|
self.ledger.mark(self.plan['batch_id'], n['key'], 'coordinator', 'unknown')
|
|
return False
|
|
|
|
def run(self):
|
|
# Whole-batch barrier: no staging until every candidate passed current preflight.
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=self.jobs) as pool:
|
|
def check(n):
|
|
try:
|
|
return self.phase(n, 'preflight')
|
|
except Exception:
|
|
self.ledger.mark(self.plan['batch_id'], n['key'], 'preflight', 'unknown')
|
|
return False
|
|
checks = list(pool.map(check, self.plan['nodes']))
|
|
if not all(checks):
|
|
return False
|
|
return all(list(pool.map(self.worker, self.plan['nodes'])))
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser(description=__doc__)
|
|
p.add_argument('command', choices=('plan', 'run', 'status'))
|
|
p.add_argument('--plan', required=True)
|
|
p.add_argument('--state-root', required=True)
|
|
p.add_argument('--jobs', type=int, default=4)
|
|
p.add_argument('--timeout', type=int, default=900)
|
|
p.add_argument('--execute', action='store_true', help='Required for a run that may write remotely')
|
|
a = p.parse_args()
|
|
os.umask(0o077)
|
|
if not 1 <= a.jobs <= 16 or not 1 <= a.timeout <= 3600:
|
|
p.error('jobs must be 1..16; timeout 1..3600')
|
|
try:
|
|
plan = validate(read_json(a.plan))
|
|
if a.command == 'plan':
|
|
print(json.dumps({'batch_id': plan['batch_id'], 'nodes': len(plan['nodes']),
|
|
'plan_sha256': digest(plan), 'writes': False}))
|
|
return 0
|
|
if a.command == 'run' and not a.execute:
|
|
p.error('run requires --execute; plan does not perform remote requests')
|
|
if a.command == 'status':
|
|
path = Path(a.state_root).resolve() / 'registry.sqlite3'
|
|
rows = {}
|
|
if path.exists():
|
|
with sqlite3.connect(path.as_uri() + '?mode=ro', uri=True) as c:
|
|
prior = c.execute('SELECT digest FROM batches WHERE id=?', (plan['batch_id'],)).fetchone()
|
|
if prior and prior[0] != digest(plan):
|
|
raise ValueError('Batch plan differs from its reserved plan')
|
|
rows = {(key, phase): status for key, phase, status in c.execute(
|
|
'SELECT node,phase,status FROM phases WHERE batch=?', (plan['batch_id'],))}
|
|
print(json.dumps({n['key']: {phase: rows.get((n['key'], phase), 'pending')
|
|
for phase in ('preflight',) + PHASES + ('coordinator',)}
|
|
for n in plan['nodes']}))
|
|
return 0
|
|
runner = Coordinator(plan, a.state_root, a.jobs, a.timeout)
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
signal.signal(sig, lambda *_: STOP.set())
|
|
success = runner.run()
|
|
print(json.dumps({'batch_id': plan['batch_id'], 'complete': success,
|
|
'next': None if success else 'Inspect status, resolve waiting/conflict, rerun the exact plan'}))
|
|
return 0 if success else 2
|
|
except Exception as exc:
|
|
# Validation errors have no adapter/credential material.
|
|
print(json.dumps({'error': str(exc) if isinstance(exc, ValueError) else 'Coordinator failed; inspect protected state'}))
|
|
return 2
|
|
|
|
|
|
if __name__ == '__main__':
|
|
sys.exit(main())
|