Files
truf-server/tests/edge_e2e_backend.py
T
2026-09-30 20:30:56 +03:00

434 lines
15 KiB
Python

"""Private PostgreSQL and loopback backend for the standalone edge E2E."""
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
APP = Path('/opt/truf/app')
DATA = Path('/data')
CONTROL = DATA / 'control'
POSTGRES = DATA / 'postgres'
SOCKET = DATA / 'postgres-socket'
BUNDLES = DATA / 'bundles'
DB_PORT = 55433
DB_URL = f'postgresql://truf@127.0.0.1:{DB_PORT}/edge_e2e'
BACKEND_PORT = 8766
SUPERVISOR_ID = 'standalone-edge-e2e'
TARGET = 'https://gitlab.com/truf-edge-e2e/repository.git'
MAX_OUTPUT = 1024 * 1024
def require(condition, label):
if not condition:
raise RuntimeError('standalone edge E2E backend: ' + label)
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()
and stat.S_IMODE(details.st_mode) == 0o700,
'private directory',
)
return path
def write_json(path, value):
payload = json.dumps(
value, ensure_ascii=True, sort_keys=True, separators=(',', ':'),
).encode('ascii') + b'\n'
require(len(payload) <= MAX_OUTPUT, '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 run(command, timeout=120):
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',
},
)
require(
len(completed.stdout) <= MAX_OUTPUT and len(completed.stderr) <= MAX_OUTPUT,
'native command output bound',
)
require(completed.returncode == 0, 'native command failed: ' + Path(command[0]).name)
return completed
def start_postgres():
private_directory(DATA)
for path in (CONTROL, SOCKET, BUNDLES, DATA / 'home'):
private_directory(path, create=True)
require(not (POSTGRES / 'PG_VERSION').exists(), 'PostgreSQL volume is not fresh')
private_directory(POSTGRES, create=True)
run([
'/usr/lib/postgresql/16/bin/initdb', '--pgdata', str(POSTGRES),
'--username=truf', '--auth=trust', '--encoding=UTF8', '--no-locale',
])
with open(POSTGRES / 'pg_hba.conf', 'a', encoding='ascii') as handle:
handle.write('\n# Disposable internal E2E network only.\nhost edge_e2e truf 0.0.0.0/0 trust\n')
run([
'/usr/lib/postgresql/16/bin/pg_ctl', '-D', str(POSTGRES), '-w', 'start',
'-l', str(DATA / 'postgres.log'),
'-o', f'-k {SOCKET} -h 0.0.0.0 -p {DB_PORT}',
])
run([
'/usr/lib/postgresql/16/bin/createdb', '-h', str(SOCKET),
'-p', str(DB_PORT), '-U', 'truf', 'edge_e2e',
])
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=1, timeout=60, save_dir=str(DATA), detectors='',
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=100,
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=1,
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,
)
def package_manifest():
from lifecycle_authority import (
GIT_MANIFEST_NAME, REMOTE_WORKER_CODE_AUTHORITY_FILES,
TRUFFLEHOG_MANIFEST_NAME,
)
from result_bundle import FORMAT_VERSION
from scan_execution import PROTOCOL_VERSION
policy_digest = hashlib.sha256(
(APP / 'trufflehog-custom-detectors.yaml').read_bytes(),
).hexdigest()
files = {
name: {'path': f'app/{name}', 'sha256': '1' * 64}
for name in REMOTE_WORKER_CODE_AUTHORITY_FILES
}
return {
'schema': 3,
'protocol_version': PROTOCOL_VERSION,
'bundle_format_version': FORMAT_VERSION,
'platform_tag': 'linux-x86_64',
'capabilities': [
{
'source': 'gitlab', 'platform': 'gitlab',
'planning_kind': 'exact_git_v1',
},
{
'source': 'dockerhub', 'platform': 'docker',
'planning_kind': 'docker_direct_v1',
},
{
'source': 'huggingface', 'platform': 'huggingface',
'planning_kind': 'huggingface_space_v1',
},
],
'app_root': 'app',
'files': files,
'executables': {
TRUFFLEHOG_MANIFEST_NAME: {
'path': 'bin/trufflehog', 'sha256': '2' * 64,
},
GIT_MANIFEST_NAME: {
'path': 'runtime/git/bin/git', 'sha256': '3' * 64,
},
},
'assets': {
'detector_policy': {
'path': 'app/trufflehog-custom-detectors.yaml',
'sha256': policy_digest,
},
},
'runtime_trees': {
'git': {'path': 'runtime/git', 'sha256': '4' * 64, 'file_count': 1},
},
}
def worker_build():
from worker_package import worker_package_build_compatibility
return worker_package_build_compatibility(package_manifest())
class Harness:
def __init__(self):
self.stop = threading.Event()
self.ingester_ready = threading.Event()
self.ingester_error = []
self.server = None
self.tokens = {
name: str(os.environ.get(environment) or '')
for name, environment in {
'good': 'TRUF_EDGE_E2E_GOOD_TOKEN',
'wrong': 'TRUF_EDGE_E2E_WRONG_TOKEN',
'revoked': 'TRUF_EDGE_E2E_REVOKED_TOKEN',
}.items()
}
self.edge_marker = str(os.environ.get('TRUF_ADMIN_EDGE_MARKER') or '')
require(
all(32 <= len(value) <= 512 for value in self.tokens.values())
and len(set(self.tokens.values())) == 3,
'test token configuration',
)
require(
len(self.edge_marker) == 64
and all(character in '0123456789abcdef' for character in self.edge_marker),
'edge marker configuration',
)
def initialize_database(self):
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:
migrate_runtime_safety_schema(db, initialize_base=True)
db.record_final_cutover({'fixture': 'standalone-edge-e2e-v1'})
for name, token in self.tokens.items():
result = db.provision_remote_worker_device(
'edge-user-' + name, 'edge-device-' + name,
hashlib.sha256(token.encode('utf-8')).hexdigest(), 1,
)
require(result['device_key'] == 'edge-device-' + name, 'device provisioning')
require(
db.set_remote_worker_device_revoked('edge-device-revoked', True),
'revoked fixture',
)
require(
db.enqueue_targets('gitlab', 'gitlab', 'standalone-edge-e2e', [TARGET]) == 1,
'target enqueue',
)
finally:
db.close()
@staticmethod
def planner(args, db_url, source, claim, scan_kwargs, remote_credential=None):
from scanner_db import ScannerDB
resolution = {
'provider': 'gitlab', 'repo_url': TARGET,
'repo_path': 'truf-edge-e2e/repository', 'branch': 'main',
'ref': 'refs/heads/main', 'head_sha': 'a' * 40,
'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 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'),
},
{
'linux': {
'package_manifest': package_manifest(),
'sources': ['gitlab', 'dockerhub', 'huggingface'],
},
},
SUPERVISOR_ID, assignment_ttl_seconds=600, planner=self.planner,
)
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), SUPERVISOR_ID, 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('edge E2E stopping' if self.ingester_error else '')
db.close()
def app(self):
from admin_api import AdminService
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,
)
admin = AdminService(
DB_URL, 'https://localhost', self.edge_marker,
db_factory=service.db_factory,
)
app = create_worker_app(
service, reaper_interval_seconds=30, admin_service=admin,
)
marker = self.edge_marker.encode('ascii')
class BackendEvidence:
async def __call__(self, scope, receive, send):
if scope.get('type') == 'http':
path = str(scope.get('path') or '')
marker_headers = [
value for name, value in scope.get('headers', ())
if name.lower() == b'x-truf-admin-edge'
]
operator_headers = [
value for name, value in scope.get('headers', ())
if name.lower() == b'x-truf-admin-operator'
]
if path.startswith('/admin-internal'):
write_json(CONTROL / 'last-admin-request.json', {
'schema': 1, 'path': path,
'marker_authorized': marker_headers == [marker],
'operators': [
value.decode('ascii', errors='strict')
for value in operator_headers
],
})
elif path.startswith('/api/v1/worker/'):
write_json(CONTROL / 'last-worker-request.json', {
'schema': 1, 'admin_headers_absent': (
not marker_headers and not operator_headers
),
})
await app(scope, receive, send)
return BackendEvidence()
def run(self):
import uvicorn
self.initialize_database()
ingester = threading.Thread(
target=self.ingester_loop, name='result-ingester', daemon=True,
)
ingester.start()
require(
self.ingester_ready.wait(30) and not self.ingester_error,
'ingester startup',
)
write_json(CONTROL / 'ready.json', {
'schema': 1, 'backend': 'private-network', 'postgres': 'fresh',
'sources': ['gitlab', 'dockerhub', 'huggingface'],
})
self.server = uvicorn.Server(uvicorn.Config(
self.app(), host='0.0.0.0', port=BACKEND_PORT,
access_log=False, log_level='warning', server_header=False,
proxy_headers=False,
))
try:
self.server.run()
finally:
self.server = None
self.stop.set()
ingester.join(15)
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')
start_postgres()
harness = Harness()
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()
finally:
harness.stop.set()
stop_postgres()
if __name__ == '__main__':
main()