"""Isolated HTTPS/PostgreSQL harness for real packaged worker clients.""" import sys sys.dont_write_bytecode = True import hashlib import json import os from pathlib import Path import runpy import signal import stat import subprocess import threading import time from types import SimpleNamespace from urllib.parse import urlsplit APP = Path('/opt/truf/app') DATA = Path('/data') CONTROL = DATA / 'control' POSTGRES = DATA / 'postgres' SOCKET = DATA / 'postgres-socket' BUNDLES = DATA / 'bundles' FIXTURE = Path('/fixture') PORT = 8443 DB_PORT = 55432 DB_URL = f'postgresql://truf@127.0.0.1:{DB_PORT}/packaged_worker_e2e' MAX_CONTROL_BYTES = 1024 * 1024 COMPLETION_REQUIREMENTS = frozenset(( 'normalized row counts', 'queue completion', 'remote reservation completion', 'bundle completion', 'scan completion', 'native findings', 'candidate routing', 'exact fixture commits', 'server bundle spool cleanup', 'bundle capacity release', 'source coverage', 'direct assignment planning', 'exact fixture targets', )) def fail(message): raise RuntimeError('packaged worker E2E: ' + message) def require(condition, message): if not condition: fail(message) def private_directory(path, create=False): path = Path(path) if create: path.mkdir(mode=0o700, parents=True, exist_ok=True) os.chmod(path, 0o700) details = path.stat(follow_symlinks=False) require(stat.S_ISDIR(details.st_mode) and details.st_uid == os.getuid(), 'private directory') require(stat.S_IMODE(details.st_mode) == 0o700, 'private directory mode') return path def write_json(path, value): payload = json.dumps( value, ensure_ascii=True, sort_keys=True, separators=(',', ':'), ).encode('ascii') require(len(payload) <= MAX_CONTROL_BYTES, 'control payload bound') temporary = Path(str(path) + '.tmp') try: temporary.unlink() except FileNotFoundError: pass descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) with os.fdopen(descriptor, 'wb') as handle: handle.write(payload) handle.flush() os.fsync(handle.fileno()) os.replace(temporary, path) def read_json(path): details = Path(path).stat(follow_symlinks=False) require(stat.S_ISREG(details.st_mode) and details.st_size <= MAX_CONTROL_BYTES, 'fixture JSON bound') with open(path, 'rb') as handle: payload = handle.read(MAX_CONTROL_BYTES + 1) require(len(payload) <= MAX_CONTROL_BYTES, 'fixture JSON bound') value = json.loads(payload.decode('ascii')) require(isinstance(value, dict), 'fixture JSON shape') return value def run(command, timeout=60): completed = subprocess.run( command, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=timeout, check=False, env={ 'PATH': '/usr/lib/postgresql/16/bin:/usr/local/bin:/usr/bin:/bin', 'HOME': str(DATA / 'home'), 'LANG': 'C.UTF-8', 'LC_ALL': 'C.UTF-8', }, ) if completed.returncode: fail('native command failed: ' + Path(command[0]).name) require(len(completed.stdout) <= MAX_CONTROL_BYTES and len(completed.stderr) <= MAX_CONTROL_BYTES, 'native command output bound') return completed def start_postgres(): private_directory(DATA) for path in (CONTROL, SOCKET, BUNDLES, DATA / 'home'): private_directory(path, create=True) first = not (POSTGRES / 'PG_VERSION').exists() if first: private_directory(POSTGRES, create=True) run([ '/usr/lib/postgresql/16/bin/initdb', '--pgdata', str(POSTGRES), '--username=truf', '--auth=trust', '--encoding=UTF8', '--no-locale', ]) run([ '/usr/lib/postgresql/16/bin/pg_ctl', '-D', str(POSTGRES), '-w', 'start', '-l', str(DATA / 'postgres.log'), '-o', f'-k {SOCKET} -h 127.0.0.1 -p {DB_PORT}', ]) if first: run([ '/usr/lib/postgresql/16/bin/createdb', '-h', str(SOCKET), '-p', str(DB_PORT), '-U', 'truf', 'packaged_worker_e2e', ]) return first def stop_postgres(): if (POSTGRES / 'postmaster.pid').exists(): try: run([ '/usr/lib/postgresql/16/bin/pg_ctl', '-D', str(POSTGRES), '-w', '-m', 'fast', 'stop', ], timeout=30) except Exception: pass def source_args(platform): return SimpleNamespace( platform=platform, exact_git_planning_enabled=platform == 'gitlab', workers=2, timeout=240, save_dir=str(DATA), detectors='OpenAI', exclude_detectors='', drop_detectors='', no_verification=True, trufflehog_config=str(APP / 'trufflehog-custom-detectors.yaml'), token='', scan_full_history=False, max_depth=25, git_baseline_depth=25, max_commit_age_days=0, commit_lookup_pages=1, skip_if_commit_lookup_fails=True, result_bundle_max_event_bytes=1 << 20, result_bundle_max_items=20, result_bundle_max_total_bytes=32 << 20, projection_backlog_max_items=20, projection_backlog_max_bytes=32 << 20, projection_backlog_headroom_bytes=2 << 20, keycheck_queue_max_items=200, keycheck_queue_max_bytes=8 << 20, pipeline_quarantine_max_items=20, pipeline_quarantine_max_bytes=8 << 20, keycheck_candidates_per_event=50, keycheck_candidate_bytes_per_event=1 << 20, target_retry_max_attempts=3, target_retry_base_delay_sec=60, target_retry_max_delay_sec=600, target_timeout_retry_delay_sec=300, max_active_scans=2, admission_resolution_attempts=2, admission_resolution_seconds=1, admission_resolution_retry_delay_sec=0.01, target_claim_order='oldest', git_ref_resolution_attempts=1, git_ref_resolution_timeout_sec=1, git_ref_resolution_max_bytes=1 << 20, ) class Harness: def __init__(self, fixture): self.fixture = fixture self.targets = tuple(fixture.get('targets') or ()) self.repositories = dict(fixture.get('repositories') or {}) require(len(self.targets) == 2 and set(self.targets) == set(self.repositories), 'exactly two fixture targets required') self.direct_targets = dict(fixture.get('direct_targets') or {}) require( set(self.direct_targets) == {'dockerhub', 'huggingface'} and all(isinstance(value, str) and value for value in self.direct_targets.values()), 'direct fixture targets required', ) self.token = str(os.environ.get('TRUF_WORKER_E2E_TOKEN') or '') self.phase = str(os.environ.get('TRUF_WORKER_E2E_PHASE') or '') require(16 <= len(self.token) <= 512 and self.phase in ('windows', 'linux'), 'phase credentials') self.stop = threading.Event() self.ingester_ready = threading.Event() self.ingester_error = [] self.claims = 0 self.status_checks = 0 self.git_reservation_ids = frozenset() self.claim_lock = threading.Lock() self.server = None def planner(self, args, db_url, source, claim, scan_kwargs, remote_credential=None): from scanner_db import ScannerDB target = str(claim['target']) repository = dict(self.repositories.get(target) or {}) parsed = urlsplit(target) repo_path = parsed.path.removeprefix('/').removesuffix('.git') resolution = { 'provider': 'gitlab', 'repo_url': target, 'repo_path': repo_path, 'branch': 'main', 'ref': 'refs/heads/main', 'head_sha': str(repository.get('head_sha') or ''), 'ref_source': 'provider_default', } db = ScannerDB(db_url=db_url, initialize=False) try: return db.bind_git_scan_plan( claim['reservation_id'], claim['claim_lease_token'], resolution, 25, remote_credential=remote_credential, ) finally: db.close() def initialize_database(self, first): from scanner_db import ScannerDB, migrate_runtime_safety_schema db = ScannerDB(db_url=DB_URL, initialize=False) require(db.enabled and db.conn.is_postgres, 'PostgreSQL connection') try: if first: migrate_runtime_safety_schema(db, initialize_base=True) db.record_final_cutover({'fixture': 'packaged-worker-e2e-v1'}) provisioned = db.provision_remote_worker_device( 'packaged-worker-e2e-' + self.phase, 'packaged-worker-e2e-' + self.phase, hashlib.sha256(self.token.encode('utf-8')).hexdigest(), 2, ) require(provisioned['active_assignment_cap'] == 2, 'device capacity') require(db.enqueue_targets( 'gitlab', 'gitlab', 'packaged-worker-e2e', self.targets, ) == 2, 'target enqueue') write_json(CONTROL / 'prepared.json', { 'schema': 1, 'phase': self.phase, 'target_count': 2, }) else: db.require_runtime_safety_schema() db.require_final_cutover() marker = read_json(CONTROL / 'prepared.json') require(marker == {'schema': 1, 'phase': self.phase, 'target_count': 2}, 'prepared marker') finally: db.close() def assignment_builder(self): from worker_assignment import RemoteGitAssignmentBuilder return RemoteGitAssignmentBuilder( DB_URL, str(BUNDLES), { 'gitlab': source_args('gitlab'), 'dockerhub': source_args('docker'), 'huggingface': source_args('huggingface'), }, { 'windows': { 'package_manifest': FIXTURE / 'windows-manifest.json', 'sources': ['gitlab', 'dockerhub', 'huggingface'], }, 'linux': { 'package_manifest': FIXTURE / 'linux-manifest.json', 'sources': ['gitlab', 'dockerhub', 'huggingface'], }, }, 'packaged-worker-e2e-' + self.phase, assignment_ttl_seconds=600, planner=self.planner, ) def enqueue_direct_targets(self): from scanner_db import ScannerDB db = ScannerDB(db_url=DB_URL, initialize=False) try: rows = db.conn.execute( 'SELECT id FROM result_reservations ORDER BY id' ).fetchall() self.git_reservation_ids = frozenset(int(row['id']) for row in rows) require(len(self.git_reservation_ids) == 2, 'initial Git reservations') require(db.enqueue_targets( 'dockerhub', 'docker', 'packaged-worker-e2e', [self.direct_targets['dockerhub']], ) == 1, 'Docker target enqueue') require(db.enqueue_targets( 'huggingface', 'huggingface', 'packaged-worker-e2e', [self.direct_targets['huggingface']], ) == 1, 'HuggingFace target enqueue') finally: db.close() def ingester_loop(self): from result_ingester import ResultIngester from scanner_db import ScannerDB db = ScannerDB(db_url=DB_URL, initialize=False) ingester = None try: ingester = ResultIngester( db, str(BUNDLES), 'packaged-worker-e2e-' + self.phase, lease_seconds=30, ).start() self.ingester_ready.set() heartbeat = time.monotonic() while not self.stop.is_set(): progressed = ingester.process_one() if time.monotonic() - heartbeat >= 5: require(ingester.heartbeat(), 'ingester heartbeat') heartbeat = time.monotonic() if not progressed: self.stop.wait(0.05) except Exception as exc: self.ingester_error.append(type(exc).__name__) self.ingester_ready.set() self.stop.set() finally: if ingester is not None: ingester.stop('harness stopping' if self.ingester_error else '') db.close() def claim_complete(self): with self.claim_lock: self.claims += 1 def status_complete(self): with self.claim_lock: self.status_checks += 1 if self.status_checks != 2: return require(self.claims == 2, 'status fencing before both claims') write_json(CONTROL / 'outage.json', { 'schema': 1, 'claim_count': 2, 'phase': self.phase, }) if self.server is not None: self.server.should_exit = True def app(self, stop_after_claims): from worker_api import WorkerService, create_worker_app service = WorkerService( DB_URL, str(BUNDLES), self.assignment_builder(), max_bundle_bytes=32 << 20, claim_retry_after_seconds=1, ) app = create_worker_app(service, reaper_interval_seconds=10) if not stop_after_claims: harness = self class DirectStatusFence: async def __call__(self, scope, receive, send): path = str(scope.get('path') or '') parts = path.split('/') reservation_id = 0 if ( scope.get('type') == 'http' and scope.get('method') == 'GET' and len(parts) == 6 and parts[1:5] == ['api', 'v1', 'worker', 'assignments'] ): try: reservation_id = int(parts[5]) except ValueError: reservation_id = 0 if ( reservation_id > 0 and reservation_id not in harness.git_reservation_ids and not (CONTROL / 'direct-ready').is_file() ): body = b'{"code":"fixture_not_ready"}' await send({ 'type': 'http.response.start', 'status': 503, 'headers': [ (b'content-type', b'application/json'), (b'content-length', str(len(body)).encode('ascii')), ], }) await send({'type': 'http.response.body', 'body': body}) return await app(scope, receive, send) return DirectStatusFence() harness = self class StopAfterClaims: async def __call__(self, scope, receive, send): status = None async def wrapped(message): nonlocal status if message['type'] == 'http.response.start': status = int(message['status']) await send(message) if message['type'] != 'http.response.body' or message.get('more_body'): return path = str(scope.get('path') or '') if path == '/api/v1/worker/claim' and status == 201: harness.claim_complete() elif ( scope.get('method') == 'GET' and status == 200 and path.startswith('/api/v1/worker/assignments/') and path.count('/') == 5 ): harness.status_complete() await app(scope, receive, wrapped) return StopAfterClaims() def evidence(self): from scanner_db import ScannerDB db = ScannerDB(db_url=DB_URL, initialize=False) try: def rows(statement, values=()): return [dict(row) for row in db.conn.execute(statement, values).fetchall()] counts = { table: int(rows(f'SELECT COUNT(*) AS count FROM {table}')[0]['count']) for table in ( 'target_queue', 'result_reservations', 'result_bundles', 'target_scans', 'scan_result_compat', 'findings', 'finding_compat_payloads', 'finding_uid_map', 'keycheck_candidates', 'keycheck_credentials', 'errors', 'pipeline_quarantine', ) } expected = { name: (4 if name in { 'target_queue', 'result_reservations', 'result_bundles', 'target_scans', 'scan_result_compat', } else 2) for name in counts } expected.update({'errors': 0, 'pipeline_quarantine': 0}) require(counts == expected, 'normalized row counts') queue = rows('SELECT * FROM target_queue ORDER BY target') reservations = rows('SELECT * FROM result_reservations ORDER BY id') bundles = rows('SELECT * FROM result_bundles ORDER BY reservation_id') scans = rows('SELECT * FROM target_scans ORDER BY target') findings = rows('SELECT * FROM findings ORDER BY target') candidates = rows('SELECT * FROM keycheck_candidates ORDER BY target') require( all(row['status'] == 'done' and row['attempts'] == 1 and row['lease_token'] is None and row['current_result_reservation_id'] is None for row in queue), 'queue completion', ) require( {(row['source'], row['target']) for row in queue} == { *(('gitlab', target) for target in self.targets), ('dockerhub', self.direct_targets['dockerhub']), ('huggingface', self.direct_targets['huggingface']), }, 'exact fixture targets', ) require( all(row['state'] == 'acknowledged' and row['assignment_kind'] == 'remote' and row['remote_resolution_kind'] == 'bundle_accepted' and row['bundle_credit_released'] == 1 for row in reservations), 'remote reservation completion', ) reservation_by_id = {int(row['id']): row for row in reservations} for row in bundles: source = reservation_by_id[int(row['reservation_id'])]['source'] expected_findings = 1 if source == 'gitlab' else 0 require( row['state'] == 'acknowledged' and row['actual_bytes'] > 0 and row['finding_count'] == row['candidate_count'] == expected_findings and row['error_count'] == 0, 'bundle completion', ) scan_contract = { 'gitlab': ('gitlab', 'found', 1), 'dockerhub': ('docker', 'clean', 0), 'huggingface': ('huggingface', 'clean', 0), } for row in scans: scan_type, status, finding_count = scan_contract.get(row['source'], (None, None, None)) require( row['scan_type'] == scan_type and row['status'] == status and row['findings_count'] == finding_count and row['error_count'] == 0 and row['queue_completion_applied'] == 1, 'scan completion', ) require(all(row['detector_name'] == 'OpenAI' and not row['verified'] and row['file_path'] == 'synthetic.env' for row in findings), 'native findings') require(all(row['service'] == row['routed_service'] == 'openai' and row['state'] == 'pending' for row in candidates), 'candidate routing') expected_heads = { target: str(self.repositories[target]['head_sha']) for target in self.targets } require({row['target']: row['commit_hash'] for row in findings} == expected_heads, 'exact fixture commits') source_counts = { source: sum(row['source'] == source for row in reservations) for source in ('gitlab', 'dockerhub', 'huggingface') } require( source_counts == {'gitlab': 2, 'dockerhub': 1, 'huggingface': 1}, 'source coverage', ) planning_counts = {} expected_planning = { 'gitlab': 'exact_git_v1', 'dockerhub': 'docker_direct_v1', 'huggingface': 'huggingface_space_v1', } for row in reservations: snapshot = json.loads(row['remote_execution_snapshot_json']) source = row['source'] planning = snapshot['planning']['kind'] require(planning == expected_planning[source], 'direct assignment planning') if source != 'gitlab': require( snapshot['credential_ref'] == { 'source': source, 'auth_entry': '', }, 'direct assignment planning', ) planning_counts[planning] = planning_counts.get(planning, 0) + 1 require(not any(BUNDLES.rglob('*.trb')), 'server bundle spool cleanup') capacity = db.pipeline_capacity_snapshot() require(capacity['bundle_items'] == capacity['bundle_bytes'] == 0, 'bundle capacity release') return { 'schema': 1, 'phase': self.phase, 'counts': counts, 'claim_count': len(reservations), 'detectors': sorted({row['detector_name'] for row in findings}), 'candidate_services': sorted({row['service'] for row in candidates}), 'secret_hashes': sorted(row['secret_hash'] for row in findings), 'commit_hashes': sorted(row['commit_hash'] for row in findings), 'receipt_count': len({row['remote_receipt_id'] for row in reservations}), 'source_counts': source_counts, 'planning_counts': planning_counts, 'capacity': { 'bundle_items': capacity['bundle_items'], 'bundle_bytes': capacity['bundle_bytes'], 'keycheck_items': capacity['keycheck_items'], }, } finally: db.close() def monitor_completion(self): while not self.stop.wait(0.1): if self.ingester_error: return try: evidence = self.evidence() except Exception as exc: reason = type(exc).__name__ message = str(exc) prefix = 'packaged worker E2E: ' if ( isinstance(exc, RuntimeError) and message.startswith(prefix) and message[len(prefix):] in COMPLETION_REQUIREMENTS ): reason = message[len(prefix):] write_json(CONTROL / 'completion-wait.json', { 'schema': 1, 'phase': self.phase, 'reason': reason, }) continue try: (CONTROL / 'completion-wait.json').unlink() except FileNotFoundError: pass write_json(CONTROL / 'completed.json', evidence) return def serve(self, stop_after_claims): import uvicorn server = uvicorn.Server(uvicorn.Config( self.app(stop_after_claims), host='0.0.0.0', port=PORT, ssl_certfile=str(FIXTURE / 'worker_tls_cert.pem'), ssl_keyfile=str(FIXTURE / 'worker_tls_key.pem'), access_log=False, log_level='warning', server_header=False, )) self.server = server server.run() self.server = None def run(self, first): self.initialize_database(first) thread = threading.Thread(target=self.ingester_loop, name='result-ingester', daemon=True) thread.start() require(self.ingester_ready.wait(30) and not self.ingester_error, 'ingester startup') self.serve(stop_after_claims=first) require(first and (CONTROL / 'outage.json').is_file(), 'planned API outage') while not (CONTROL / 'restore').is_file(): require(not self.stop.wait(0.1), 'harness stopped before restore') self.enqueue_direct_targets() monitor = threading.Thread(target=self.monitor_completion, name='completion-monitor', daemon=True) monitor.start() self.serve(stop_after_claims=False) self.stop.set() thread.join(10) monitor.join(2) require(not self.ingester_error, 'result ingester failure') def main(): require(sys.platform == 'linux' and os.getuid() == os.getgid() == 10001, 'Linux UID 10001 required') require(sys.flags.isolated and sys.flags.no_site and sys.flags.dont_write_bytecode, 'isolated Python required') os.umask(0o077) sys.path.insert(0, str(APP)) bootstrap = runpy.run_path(str(APP / 'child_bootstrap.py')) bootstrap['_enable_dependency_paths']('supervisor') fixture = read_json(FIXTURE / 'fixture.json') first = start_postgres() harness = Harness(fixture) def terminate(_signum, _frame): if harness.server is not None: harness.server.should_exit = True harness.stop.set() signal.signal(signal.SIGTERM, terminate) signal.signal(signal.SIGINT, terminate) try: harness.run(first) finally: harness.stop.set() stop_postgres() if __name__ == '__main__': main()