2185 lines
91 KiB
Python
2185 lines
91 KiB
Python
import hashlib
|
|
import asyncio
|
|
import io
|
|
import json
|
|
import os
|
|
import socket
|
|
import ssl
|
|
from contextlib import asynccontextmanager, redirect_stdout
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from unittest import mock
|
|
from types import SimpleNamespace
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
APP_DIR = ROOT / 'app'
|
|
sys.path.insert(0, str(APP_DIR))
|
|
|
|
from starlette.testclient import TestClient
|
|
|
|
import worker_api as worker_api_module
|
|
import scanner
|
|
from result_bundle import BundleReservation, ResultBundleWriter, bundle_ready_path
|
|
from runtime_security import ensure_private_directory
|
|
from remote_worker_client import (
|
|
CLIENT_PROCESS_FAILURE_DETAIL,
|
|
CLIENT_STORAGE_FAILURE_DETAIL,
|
|
RunnerContainmentPending,
|
|
WorkerAssignmentCompatibilityError, WorkerClientError, WorkerHTTPClient,
|
|
WorkerHTTPError, WorkerNetworkError, WorkerSlot, ProgressOutbox, persisted_slot_ids,
|
|
default_worker_paths, parse_args, run_client, safe_worker_error_summary,
|
|
)
|
|
from worker_api import (
|
|
WorkerAPIError, WorkerService, _bounded_json, create_worker_app,
|
|
logger as worker_api_logger,
|
|
)
|
|
from scan_execution import (
|
|
PROTOCOL_VERSION,
|
|
normalize_docker_direct_execution_target,
|
|
remote_execution_identity,
|
|
remote_execution_snapshot_sha256,
|
|
validate_protocol2_remote_assignment,
|
|
)
|
|
from scanner import StagedResult
|
|
from scanner_db import (
|
|
PipelineCapacityUnavailable,
|
|
WorkerObservabilityConflictError,
|
|
WorkerProgressInactiveError,
|
|
)
|
|
from worker_contracts import (
|
|
AssignmentOutcome,
|
|
DiagnosticCategory,
|
|
DiagnosticKind,
|
|
ScanOutcome,
|
|
WorkerContractError,
|
|
WorkerPhase,
|
|
build_diagnostic_envelope,
|
|
decode_diagnostic_envelope,
|
|
encode_diagnostic_envelope,
|
|
validate_phase_transition,
|
|
)
|
|
|
|
|
|
def build_compatibility():
|
|
return {
|
|
'protocol_version': PROTOCOL_VERSION,
|
|
'bundle_format_version': 2,
|
|
'platform_tag': 'windows-x86_64',
|
|
'code_manifest_sha256': 'a' * 64,
|
|
'detector_policy_sha256': 'c' * 64,
|
|
}
|
|
|
|
|
|
def scan_policy():
|
|
return {
|
|
'drop_detectors': [],
|
|
'strict_git_provider_token_filter': True,
|
|
'trufflehog_stdout_max_mb': 32,
|
|
'trufflehog_stderr_max_mb': 8,
|
|
'result_bundle_max_event_bytes': 1024 * 1024,
|
|
'trufflehog_max_findings_per_target': 20000,
|
|
'trufflehog_job_memory_limit_bytes': 0,
|
|
'trufflehog_windows_job_cpu_weight': 0,
|
|
'trufflehog_windows_memory_priority': 0,
|
|
'trufflehog_diagnostic_max_lines': 2000,
|
|
'trufflehog_diagnostic_max_line_chars': 8192,
|
|
'trufflehog_diagnostic_max_line_bytes': 8192,
|
|
'trufflehog_diagnostic_max_errors': 200,
|
|
'trufflehog_diagnostic_max_warnings': 200,
|
|
'trufflehog_diagnostic_max_unclassified': 20,
|
|
}
|
|
|
|
|
|
def compatibility():
|
|
effective, _ = remote_execution_identity(
|
|
'github',
|
|
{'timeout_sec': 60, 'trufflehog_config': '@package/detector_policy'},
|
|
{'timeout_sec': 60.0, 'trufflehog_config': '@package/detector_policy'}, {},
|
|
{'candidate_max_items': 20, 'candidate_max_bytes': 4096},
|
|
scan_policy(),
|
|
)
|
|
return {**build_compatibility(), 'effective_config_sha256': effective}
|
|
|
|
|
|
def reservation_mapping():
|
|
bundle_id = 'd' * 32
|
|
return {
|
|
'reservation_id': 17,
|
|
'reservation_token': 'f' * 32,
|
|
'bundle_id': bundle_id,
|
|
'scan_event_id': 'e' * 32,
|
|
'queue_id': 19,
|
|
'claim_lease_token': 'remote-lease-token',
|
|
'declared_bundle_bytes': 1024 * 1024,
|
|
'ready_relative_path': f'ready/{bundle_id[:2]}/{bundle_id}.trb',
|
|
'source': 'github',
|
|
'platform': 'github',
|
|
'query': 'fixture',
|
|
'target': 'https://example.invalid/repo',
|
|
'normalized_target': 'https://example.invalid/repo',
|
|
'producer_instance_id': 'remote-api',
|
|
'producer_pid': 9,
|
|
'producer_creation_time': 'remote-v1',
|
|
'producer_executable': 'remote-device-9',
|
|
'assignment_kind': 'remote',
|
|
'remote_device_id': 9,
|
|
'remote_issued_at': '2026-09-17T00:00:00+00:00',
|
|
'remote_expires_at': '2026-09-18T00:00:00+00:00',
|
|
'remote_effective_config_sha256': compatibility()['effective_config_sha256'],
|
|
'attempts': 1,
|
|
}
|
|
|
|
|
|
def assignment():
|
|
reservation = reservation_mapping()
|
|
required = compatibility()
|
|
bound_plan = {'schema': 1, 'target': reservation['target']}
|
|
event_scan_options = {
|
|
'timeout_sec': 60.0, 'trufflehog_config': '@package/detector_policy',
|
|
}
|
|
scan_kwargs = {
|
|
**event_scan_options, 'token': 'fixture-token', 'git_plan': bound_plan,
|
|
}
|
|
limits = {'candidate_max_items': 20, 'candidate_max_bytes': 4096}
|
|
_effective, execution = remote_execution_identity(
|
|
'github', scan_kwargs, event_scan_options, {}, limits, scan_policy(),
|
|
)
|
|
snapshot = {
|
|
'schema': 1,
|
|
'compatibility': required,
|
|
'execution': execution,
|
|
'planning': {
|
|
'kind': 'exact_git_v1',
|
|
'git_baseline_depth': 100,
|
|
'git_ref_resolution_attempts': 2,
|
|
'git_ref_resolution_timeout_sec': 10.0,
|
|
'git_ref_resolution_max_bytes': 1024 * 1024,
|
|
},
|
|
'credential_ref': {'source': 'github', 'auth_entry': 'selected'},
|
|
}
|
|
return {
|
|
'reservation': reservation,
|
|
'deadlines': {
|
|
'target_scan_timeout_seconds': 60,
|
|
'result_upload_body_timeout_seconds': 1800,
|
|
'assignment_ttl_seconds': 86400,
|
|
'assignment_issued_at': reservation['remote_issued_at'],
|
|
'assignment_deadline_at': reservation['remote_expires_at'],
|
|
},
|
|
'compatibility': required,
|
|
'scan_kwargs': scan_kwargs,
|
|
'event_scan_options': event_scan_options,
|
|
'queue_policy': {},
|
|
'limits': limits,
|
|
'scan_policy': scan_policy(),
|
|
'execution_snapshot': snapshot,
|
|
'execution_snapshot_sha256': remote_execution_snapshot_sha256(snapshot),
|
|
'execution_plan': {
|
|
'kind': 'exact_git_v1',
|
|
'execution_target': reservation['target'],
|
|
'bound_plan': bound_plan,
|
|
},
|
|
}
|
|
|
|
|
|
def protocol1_assignment():
|
|
value = assignment()
|
|
value['compatibility']['protocol_version'] = 1
|
|
value['execution_snapshot_sha256'] = remote_execution_snapshot_sha256(
|
|
value['execution_snapshot'],
|
|
)
|
|
return value
|
|
|
|
|
|
def direct_assignment(source):
|
|
if source == 'dockerhub':
|
|
platform = 'docker'
|
|
planning_kind = 'docker_direct_v1'
|
|
target = 'docker.io/library/alpine@sha256:' + ('1' * 64)
|
|
parsed = normalize_docker_direct_execution_target(target)
|
|
execution_target = parsed['image']
|
|
normalized_target = parsed['normalized_target']
|
|
elif source == 'huggingface':
|
|
platform = 'huggingface'
|
|
planning_kind = 'huggingface_space_v1'
|
|
target = 'Owner/Space'
|
|
execution_target = target
|
|
normalized_target = target.lower()
|
|
else:
|
|
raise ValueError('unsupported direct source')
|
|
event_scan_options = {
|
|
'timeout_sec': 60.0, 'trufflehog_config': '@package/detector_policy',
|
|
}
|
|
limits = {'candidate_max_items': 20, 'candidate_max_bytes': 4096}
|
|
effective, execution = remote_execution_identity(
|
|
platform, event_scan_options, event_scan_options, {}, limits, scan_policy(),
|
|
)
|
|
required = {**build_compatibility(), 'effective_config_sha256': effective}
|
|
reservation = reservation_mapping()
|
|
reservation.update({
|
|
'source': source, 'platform': platform, 'target': target,
|
|
'normalized_target': normalized_target,
|
|
'remote_effective_config_sha256': effective,
|
|
})
|
|
snapshot = {
|
|
'schema': 1,
|
|
'compatibility': required,
|
|
'execution': execution,
|
|
'planning': {'kind': planning_kind},
|
|
'credential_ref': {'source': source, 'auth_entry': ''},
|
|
}
|
|
return {
|
|
'reservation': reservation,
|
|
'deadlines': {
|
|
'target_scan_timeout_seconds': 60,
|
|
'result_upload_body_timeout_seconds': 1800,
|
|
'assignment_ttl_seconds': 86400,
|
|
'assignment_issued_at': reservation['remote_issued_at'],
|
|
'assignment_deadline_at': reservation['remote_expires_at'],
|
|
},
|
|
'compatibility': required,
|
|
'scan_kwargs': dict(event_scan_options),
|
|
'event_scan_options': event_scan_options,
|
|
'queue_policy': {},
|
|
'limits': limits,
|
|
'scan_policy': scan_policy(),
|
|
'execution_snapshot': snapshot,
|
|
'execution_snapshot_sha256': remote_execution_snapshot_sha256(snapshot),
|
|
'execution_plan': {
|
|
'kind': planning_kind,
|
|
'execution_target': execution_target,
|
|
'bound_plan': None,
|
|
},
|
|
}
|
|
|
|
|
|
class WorkerServiceTests(unittest.TestCase):
|
|
@staticmethod
|
|
def _terminal_diagnostic(*, assignment_outcome, scan_outcome):
|
|
envelope = build_diagnostic_envelope(
|
|
occurrence_id='terminal-outcome-fixture',
|
|
reservation_id=17,
|
|
scan_event_id='e' * 32,
|
|
slot_id=0,
|
|
source='github',
|
|
phase=WorkerPhase.PREPARING,
|
|
kind=DiagnosticKind.ASSIGNMENT,
|
|
category=DiagnosticCategory.INTERNAL,
|
|
code='worker.terminal_fixture',
|
|
summary='terminal fixture',
|
|
retryable=False,
|
|
attempt=1,
|
|
assignment_outcome=assignment_outcome,
|
|
scan_outcome=scan_outcome,
|
|
occurred_at='2026-09-24T00:00:00Z',
|
|
captured_at='2026-09-24T00:00:00Z',
|
|
)
|
|
return json.loads(encode_diagnostic_envelope(envelope))
|
|
|
|
def test_terminal_rejects_bundle_outcome_and_duplicate_uids_before_database(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'bundles'), reject_reparse=True,
|
|
)
|
|
service = WorkerService(
|
|
'postgresql://fixture', root, lambda *_args: None,
|
|
db_factory=lambda **_kwargs: self.fail(
|
|
'invalid terminal report must not reach PostgreSQL'
|
|
),
|
|
)
|
|
identity = {'device_id': 9, 'token_sha256': 'a' * 64}
|
|
invalid = self._terminal_diagnostic(
|
|
assignment_outcome=AssignmentOutcome.ACCEPTED,
|
|
scan_outcome=ScanOutcome.ERROR,
|
|
)
|
|
with self.assertRaisesRegex(WorkerAPIError, 'terminal diagnostics'):
|
|
service.report_terminal(identity, 17, {
|
|
'failure_code': 'client_process_failed',
|
|
'detail': 'fixture',
|
|
'diagnostics': [invalid],
|
|
})
|
|
valid = self._terminal_diagnostic(
|
|
assignment_outcome=AssignmentOutcome.PREBUNDLE_FAILED,
|
|
scan_outcome=ScanOutcome.UNAVAILABLE,
|
|
)
|
|
with self.assertRaisesRegex(WorkerAPIError, 'terminal diagnostics'):
|
|
service.report_terminal(identity, 17, {
|
|
'failure_code': 'client_process_failed',
|
|
'detail': 'fixture',
|
|
'diagnostics': [valid, valid],
|
|
})
|
|
|
|
def test_protocol1_claim_only_reconciles_existing_request(self):
|
|
legacy_build = build_compatibility()
|
|
legacy_build['protocol_version'] = 1
|
|
calls = []
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'bundles'), reject_reparse=True,
|
|
)
|
|
service = WorkerService(
|
|
'postgresql://fixture', root,
|
|
lambda identity, request_id, supplied: (
|
|
calls.append((request_id, supplied)) or None
|
|
),
|
|
db_factory=lambda **kwargs: None,
|
|
)
|
|
with self.assertRaises(WorkerAPIError) as rejected:
|
|
service.claim(
|
|
{'device_id': 9},
|
|
{'request_id': 'f' * 32, 'build': legacy_build},
|
|
)
|
|
self.assertEqual(rejected.exception.status_code, 409)
|
|
self.assertEqual(rejected.exception.code, 'incompatible_protocol')
|
|
|
|
service.assignment_builder = lambda *_args: protocol1_assignment()
|
|
reconciled = service.claim(
|
|
{'device_id': 9},
|
|
{'request_id': 'f' * 32, 'build': legacy_build},
|
|
)
|
|
self.assertEqual(len(calls), 1)
|
|
self.assertEqual(reconciled['execution_snapshot']['compatibility'][
|
|
'protocol_version'
|
|
], 1)
|
|
self.assertEqual(
|
|
reconciled['execution_snapshot_sha256'],
|
|
remote_execution_snapshot_sha256(reconciled['execution_snapshot']),
|
|
)
|
|
|
|
def test_claim_requires_builder_to_preserve_remote_identity(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'bundles'), reject_reparse=True,
|
|
)
|
|
service = WorkerService(
|
|
'postgresql://fixture', root,
|
|
lambda identity, request_id, supplied: assignment(),
|
|
db_factory=lambda **kwargs: None,
|
|
)
|
|
result = service.claim(
|
|
{'device_id': 9},
|
|
{'request_id': 'f' * 32, 'build': build_compatibility()},
|
|
)
|
|
self.assertEqual(result['reservation']['reservation_id'], 17)
|
|
|
|
missing_deadlines = assignment()
|
|
del missing_deadlines['deadlines']
|
|
service.assignment_builder = lambda *args: missing_deadlines
|
|
with self.assertRaisesRegex(RuntimeError, 'payload shape'):
|
|
service.claim(
|
|
{'device_id': 9},
|
|
{'request_id': 'f' * 32, 'build': build_compatibility()},
|
|
)
|
|
|
|
conflicting = assignment()
|
|
conflicting['reservation']['remote_device_id'] = 10
|
|
service.assignment_builder = lambda *args: conflicting
|
|
with self.assertRaises(RuntimeError):
|
|
service.claim(
|
|
{'device_id': 9},
|
|
{'request_id': 'f' * 32, 'build': build_compatibility()},
|
|
)
|
|
|
|
def test_claim_preserves_terminal_request_reconciliation(self):
|
|
receipt = {
|
|
'receipt_id': '1' * 64, 'resolution': 'expired',
|
|
'reservation_id': 17, 'bundle_id': 'd' * 32,
|
|
'scan_event_id': 'e' * 32,
|
|
'resolved_at': '2026-09-18T00:00:00+00:00',
|
|
}
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'bundles'), reject_reparse=True,
|
|
)
|
|
service = WorkerService(
|
|
'postgresql://fixture', root,
|
|
lambda *_args: {'resolution': receipt},
|
|
db_factory=lambda **kwargs: None,
|
|
)
|
|
result = service.claim(
|
|
{'device_id': 9},
|
|
{'request_id': 'f' * 32, 'build': build_compatibility()},
|
|
)
|
|
self.assertEqual(result, {'resolution': receipt})
|
|
|
|
def test_presented_token_digest_reaches_each_post_auth_db_fence(self):
|
|
calls = []
|
|
|
|
class FenceDB:
|
|
enabled = True
|
|
|
|
def authenticate_remote_worker(self, token_sha256):
|
|
calls.append(('authenticate', token_sha256))
|
|
return {'device_id': 9, 'user_id': 3}
|
|
|
|
def remote_assignment_status(self, reservation_id, device_id, token_sha256):
|
|
calls.append(('status', token_sha256))
|
|
return {
|
|
'reservation_id': reservation_id,
|
|
'deadlines': {
|
|
'result_upload_body_timeout_seconds': 900,
|
|
'result_upload_body_timeout_availability': 'persisted',
|
|
},
|
|
}
|
|
|
|
def remote_assignment_transport(self, reservation_id, device_id, token_sha256):
|
|
calls.append(('transport', token_sha256))
|
|
return {'reservation_id': reservation_id}
|
|
|
|
def record_remote_worker_progress_event(
|
|
self, reservation_id, device_id, token_sha256, event,
|
|
):
|
|
calls.append(('progress', token_sha256))
|
|
return {
|
|
'sequence': event['sequence'],
|
|
'received_at': '2026-09-24T00:00:01+00:00',
|
|
'replayed': False,
|
|
}
|
|
|
|
def mark_result_bundle_ready(
|
|
self, reservation_id, metadata, remote_acceptance=None,
|
|
bundle_capacity_bytes=None,
|
|
):
|
|
calls.append(('accept', remote_acceptance['token_sha256']))
|
|
return {'reservation_id': reservation_id}
|
|
|
|
def report_remote_prebundle_failure(
|
|
self, reservation_id, device_id, token_sha256, payload,
|
|
):
|
|
calls.append(('terminal', token_sha256))
|
|
return {'reservation_id': reservation_id}
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
token = 'fixture-token-value'
|
|
token_sha256 = hashlib.sha256(token.encode('utf-8')).hexdigest()
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'bundles'), reject_reparse=True,
|
|
)
|
|
service = WorkerService(
|
|
'postgresql://fixture', root, lambda *_args: None,
|
|
db_factory=lambda **_kwargs: FenceDB(),
|
|
)
|
|
identity = service.authenticate(f'Bearer {token}')
|
|
status = service.status(identity, 17)
|
|
service.transport(identity, 17)
|
|
service.progress(identity, 17, {'sequence': 7})
|
|
service.accept_ready(
|
|
identity, {'reservation_id': 17, 'ready_relative_path': 'ready/a.trb'},
|
|
(),
|
|
SimpleNamespace(as_dict=lambda: {}), 'b' * 64,
|
|
)
|
|
service.report_terminal(
|
|
identity, 17,
|
|
{'failure_code': 'client_cancelled', 'detail': 'cancelled'},
|
|
)
|
|
self.assertEqual(
|
|
status['deadlines']['result_upload_body_timeout_seconds'], 900,
|
|
)
|
|
|
|
self.assertEqual(identity['token_sha256'], token_sha256)
|
|
self.assertEqual(
|
|
calls,
|
|
[(name, token_sha256) for name in (
|
|
'authenticate', 'status', 'transport', 'progress', 'accept', 'terminal',
|
|
)],
|
|
)
|
|
|
|
def test_bundle_capacity_backpressure_is_retryable_and_uses_aggregate_limit(self):
|
|
calls = []
|
|
|
|
class CapacityDB:
|
|
enabled = True
|
|
|
|
def mark_result_bundle_ready(
|
|
self, reservation_id, metadata, remote_acceptance=None,
|
|
bundle_capacity_bytes=None,
|
|
):
|
|
calls.append(bundle_capacity_bytes)
|
|
raise PipelineCapacityUnavailable('fixture full')
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'bundles'), reject_reparse=True,
|
|
)
|
|
service = WorkerService(
|
|
'postgresql://fixture', root, lambda *_args: None,
|
|
db_factory=lambda **_kwargs: CapacityDB(),
|
|
bundle_capacity_bytes=512 * 1024 * 1024,
|
|
)
|
|
with self.assertRaises(WorkerAPIError) as raised:
|
|
service.accept_ready(
|
|
{'device_id': 9, 'token_sha256': 'a' * 64},
|
|
{'reservation_id': 17, 'ready_relative_path': 'ready/a.trb'},
|
|
(), SimpleNamespace(as_dict=lambda: {}), 'b' * 64,
|
|
)
|
|
self.assertEqual(raised.exception.status_code, 503)
|
|
self.assertEqual(raised.exception.code, 'capacity_backpressure')
|
|
self.assertEqual(calls, [512 * 1024 * 1024])
|
|
|
|
def test_reap_expires_assignments_then_reconciles_drain_and_closes(self):
|
|
calls = []
|
|
|
|
class ReaperDB:
|
|
enabled = True
|
|
|
|
def reap_expired_remote_assignments(self, limit):
|
|
calls.append(('expire', limit))
|
|
return [{'receipt_id': 'a' * 64}]
|
|
|
|
def reconcile_runtime_drain(self):
|
|
calls.append(('reconcile',))
|
|
return None
|
|
|
|
def close(self):
|
|
calls.append(('close',))
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'bundles'), reject_reparse=True,
|
|
)
|
|
service = WorkerService(
|
|
'postgresql://fixture', root, lambda *_args: None,
|
|
db_factory=lambda **_kwargs: ReaperDB(), reaper_batch_size=7,
|
|
)
|
|
with mock.patch.object(
|
|
worker_api_module,
|
|
'reconcile_pending_host_results',
|
|
side_effect=lambda database: calls.append(('host-results', database)),
|
|
):
|
|
receipts = service.reap()
|
|
|
|
self.assertEqual(receipts, [{'receipt_id': 'a' * 64}])
|
|
self.assertEqual(calls, [
|
|
('expire', 7), ('reconcile',), ('host-results', mock.ANY), ('close',),
|
|
])
|
|
|
|
|
|
class FakeService:
|
|
def __init__(self, root):
|
|
self.bundle_root = root
|
|
self.max_bundle_bytes = 1024 * 1024
|
|
self.accepted = []
|
|
self.receipt = None
|
|
self.claim_retry_after_seconds = 5
|
|
|
|
def authenticate(self, authorization):
|
|
if authorization != 'Bearer fixture-token-value':
|
|
raise WorkerAPIError(401, 'unauthorized', 'worker credentials are invalid')
|
|
return {'device_id': 9, 'user_id': 3}
|
|
|
|
def claim(self, identity, payload):
|
|
return assignment()
|
|
|
|
def status(self, identity, reservation_id):
|
|
return {'reservation_id': reservation_id, 'state': 'scanning', 'expires_at': '2999-01-01T00:00:00+00:00'}
|
|
|
|
def transport(self, identity, reservation_id):
|
|
return {
|
|
**reservation_mapping(),
|
|
'id': reservation_id,
|
|
'state': 'scanning',
|
|
'remote_expires_at': '2999-01-01T00:00:00+00:00',
|
|
'remote_resolution_kind': (
|
|
'bundle_accepted' if self.receipt is not None else None
|
|
),
|
|
'remote_payload_sha256': (
|
|
self.receipt['payload_sha256'] if self.receipt is not None else None
|
|
),
|
|
'receipt': self.receipt,
|
|
}
|
|
|
|
def progress(self, identity, reservation_id, payload):
|
|
return {
|
|
'accepted': True,
|
|
'reservation_id': reservation_id,
|
|
'sequence': payload['sequence'],
|
|
'received_at': '2026-09-24T00:00:01+00:00',
|
|
'replayed': False,
|
|
}
|
|
|
|
def accept_ready(
|
|
self, identity, transport, effective_diagnostics, metadata, payload_sha256,
|
|
):
|
|
self.accepted.append((
|
|
identity, transport, effective_diagnostics, metadata, payload_sha256,
|
|
))
|
|
self.receipt = {
|
|
'receipt_id': '1' * 64,
|
|
'resolution': 'bundle_accepted',
|
|
'payload_sha256': payload_sha256,
|
|
'reservation_id': 17,
|
|
}
|
|
return self.receipt
|
|
|
|
def report_terminal(self, identity, reservation_id, payload):
|
|
return {'receipt_id': '2' * 64, 'resolution': 'prebundle_report'}
|
|
|
|
def reap(self):
|
|
return []
|
|
|
|
@asynccontextmanager
|
|
async def upload_scope(self, reservation_id):
|
|
yield
|
|
|
|
|
|
class WorkerAPIRouteTests(unittest.TestCase):
|
|
def _bundle(self, temp_dir):
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'client-bundles'), reject_reparse=True,
|
|
)
|
|
reservation = BundleReservation.from_mapping(reservation_mapping())
|
|
with ResultBundleWriter.open(root, reservation) as writer:
|
|
commit = writer.finish({'queue_status': 'done'})
|
|
path = os.path.join(root, *commit.relative_path.split('/'))
|
|
with open(path, 'rb') as handle:
|
|
return handle.read()
|
|
|
|
def _dless_error_bundle(self, temp_dir):
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'client-error-bundles'), reject_reparse=True,
|
|
)
|
|
reservation = BundleReservation.from_mapping(reservation_mapping())
|
|
with ResultBundleWriter.open(root, reservation) as writer:
|
|
writer.write_error('legacy first-upload error')
|
|
commit = writer.finish({'status': 'error', 'queue_status': 'failed'})
|
|
path = os.path.join(root, *commit.relative_path.split('/'))
|
|
with open(path, 'rb') as handle:
|
|
return handle.read()
|
|
|
|
def test_managed_file_traversal_lifespan_is_singleton_and_closed(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'server-bundles'), reject_reparse=True,
|
|
)
|
|
traversal = mock.Mock()
|
|
registry = object()
|
|
admin_service = SimpleNamespace(managed_file_roots=registry)
|
|
with mock.patch(
|
|
'worker_api.ManagedFileTraversal', return_value=traversal,
|
|
) as constructor:
|
|
app = create_worker_app(
|
|
FakeService(root), reaper_interval_seconds=3600,
|
|
admin_service=admin_service,
|
|
)
|
|
with TestClient(app):
|
|
self.assertIs(app.state.managed_file_traversal, traversal)
|
|
self.assertIsNone(app.state.managed_file_traversal)
|
|
constructor.assert_called_once_with(registry)
|
|
traversal.close.assert_called_once_with()
|
|
|
|
def test_managed_file_root_failure_does_not_disable_worker_routes(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'server-bundles'), reject_reparse=True,
|
|
)
|
|
service = FakeService(root)
|
|
admin_service = SimpleNamespace(managed_file_roots=object())
|
|
unavailable = RuntimeError('host path must not be logged')
|
|
unavailable.category = 'root_unavailable'
|
|
with mock.patch(
|
|
'worker_api.ManagedFileTraversal', side_effect=unavailable,
|
|
) as constructor, self.assertLogs(
|
|
worker_api_logger, level='ERROR',
|
|
) as logs:
|
|
app = create_worker_app(
|
|
service, reaper_interval_seconds=3600,
|
|
admin_service=admin_service,
|
|
)
|
|
with TestClient(app) as client:
|
|
self.assertIsNone(app.state.managed_file_traversal)
|
|
response = client.post(
|
|
'/api/v1/worker/claim',
|
|
json={
|
|
'request_id': 'f' * 32,
|
|
'build': build_compatibility(),
|
|
},
|
|
headers={'Authorization': 'Bearer fixture-token-value'},
|
|
)
|
|
self.assertEqual(response.status_code, 201)
|
|
constructor.assert_called_once_with(admin_service.managed_file_roots)
|
|
rendered = '\n'.join(logs.output)
|
|
self.assertIn('root_unavailable', rendered)
|
|
self.assertNotIn('host path', rendered)
|
|
|
|
def test_admin_disabled_does_not_construct_managed_file_traversal(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'server-bundles'), reject_reparse=True,
|
|
)
|
|
with mock.patch('worker_api.ManagedFileTraversal') as constructor:
|
|
app = create_worker_app(
|
|
FakeService(root), reaper_interval_seconds=3600,
|
|
)
|
|
with TestClient(app):
|
|
self.assertIsNone(app.state.managed_file_traversal)
|
|
constructor.assert_not_called()
|
|
|
|
def test_raw_bundle_upload_is_validated_published_and_receipted(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
server_root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'server-bundles'), reject_reparse=True,
|
|
)
|
|
service = FakeService(server_root)
|
|
payload = self._bundle(temp_dir)
|
|
digest = hashlib.sha256(payload).hexdigest()
|
|
with TestClient(create_worker_app(service, reaper_interval_seconds=3600)) as client:
|
|
response = client.put(
|
|
'/api/v1/worker/assignments/17/bundle', content=payload,
|
|
headers={
|
|
'Authorization': 'Bearer fixture-token-value',
|
|
'Content-Type': 'application/octet-stream',
|
|
'X-Truf-Payload-SHA256': digest,
|
|
},
|
|
)
|
|
self.assertEqual(response.status_code, 201, response.text)
|
|
self.assertEqual(response.json()['payload_sha256'], digest)
|
|
self.assertEqual(len(service.accepted), 1)
|
|
ready = os.path.join(server_root, 'ready', 'dd', f'{"d" * 32}.trb')
|
|
self.assertTrue(os.path.isfile(ready))
|
|
|
|
conflicting = bytearray(payload)
|
|
conflicting[-1] ^= 1
|
|
with TestClient(create_worker_app(service, reaper_interval_seconds=3600)) as client:
|
|
replay = client.put(
|
|
'/api/v1/worker/assignments/17/bundle', content=bytes(conflicting),
|
|
headers={
|
|
'Authorization': 'Bearer fixture-token-value',
|
|
'Content-Type': 'application/octet-stream',
|
|
'X-Truf-Payload-SHA256': digest,
|
|
},
|
|
)
|
|
self.assertEqual(replay.status_code, 400)
|
|
self.assertEqual(replay.json()['error']['code'], 'digest_mismatch')
|
|
|
|
def test_first_dless_error_upload_precomputes_before_partial_is_moved(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
server_root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'server-bundles'), reject_reparse=True,
|
|
)
|
|
service = FakeService(server_root)
|
|
payload = self._dless_error_bundle(temp_dir)
|
|
digest = hashlib.sha256(payload).hexdigest()
|
|
with TestClient(create_worker_app(
|
|
service, reaper_interval_seconds=3600,
|
|
)) as client:
|
|
response = client.put(
|
|
'/api/v1/worker/assignments/17/bundle', content=payload,
|
|
headers={
|
|
'Authorization': 'Bearer fixture-token-value',
|
|
'Content-Type': 'application/octet-stream',
|
|
'X-Truf-Payload-SHA256': digest,
|
|
},
|
|
)
|
|
self.assertEqual(response.status_code, 201, response.text)
|
|
self.assertEqual(len(service.accepted), 1)
|
|
self.assertEqual(len(service.accepted[0][2]), 1)
|
|
self.assertEqual(
|
|
service.accepted[0][2][0]['code'], 'legacy.error_frame',
|
|
)
|
|
|
|
def test_duplicate_and_stale_bundle_events_are_correlated_and_safe(self):
|
|
private = 'raw-finding-and-credential-sentinel'
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
server_root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'server-bundles'), reject_reparse=True,
|
|
)
|
|
service = FakeService(server_root)
|
|
payload = self._bundle(temp_dir)
|
|
digest = hashlib.sha256(payload).hexdigest()
|
|
headers = {
|
|
'Authorization': 'Bearer fixture-token-value',
|
|
'Content-Type': 'application/octet-stream',
|
|
'X-Truf-Payload-SHA256': digest,
|
|
}
|
|
with TestClient(create_worker_app(service, reaper_interval_seconds=3600)) as client:
|
|
accepted = client.put(
|
|
'/api/v1/worker/assignments/17/bundle', content=payload,
|
|
headers=headers,
|
|
)
|
|
original_transport = service.transport
|
|
|
|
def secret_transport(*args):
|
|
return {
|
|
**original_transport(*args),
|
|
'target': f'https://fixture.invalid/?token={private}',
|
|
'normalized_target': f'https://fixture.invalid/{private}',
|
|
}
|
|
|
|
service.transport = secret_transport
|
|
with self.assertLogs(worker_api_logger, level='INFO') as replay_logs:
|
|
replay = client.put(
|
|
'/api/v1/worker/assignments/17/bundle', content=payload,
|
|
headers=headers,
|
|
)
|
|
service.receipt = None
|
|
|
|
def stale_transport(*args):
|
|
return {
|
|
**original_transport(*args),
|
|
'state': 'expired',
|
|
'remote_expires_at': '2000-01-01T00:00:00+00:00',
|
|
'target': f'https://fixture.invalid/?token={private}',
|
|
'normalized_target': f'https://fixture.invalid/{private}',
|
|
}
|
|
|
|
service.transport = stale_transport
|
|
with self.assertLogs(worker_api_logger, level='INFO') as stale_logs:
|
|
stale = client.put(
|
|
'/api/v1/worker/assignments/17/bundle', content=payload,
|
|
headers=headers,
|
|
)
|
|
|
|
self.assertEqual(accepted.status_code, 201)
|
|
self.assertEqual(replay.status_code, 200)
|
|
self.assertEqual(stale.status_code, 410)
|
|
logged = '\n'.join((*replay_logs.output, *stale_logs.output))
|
|
self.assertIn('remote_result_replayed', logged)
|
|
self.assertIn('original_acceptance_returned', logged)
|
|
self.assertIn('remote_result_rejected', logged)
|
|
self.assertIn('stale_assignment', logged)
|
|
self.assertIn('target_sha256', logged)
|
|
self.assertNotIn(private, logged)
|
|
self.assertNotIn('fixture-token-value', logged)
|
|
|
|
def test_claim_and_terminal_routes_remain_json_bounded(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'server-bundles'), reject_reparse=True,
|
|
)
|
|
with TestClient(create_worker_app(FakeService(root), reaper_interval_seconds=3600)) as client:
|
|
claimed = client.post(
|
|
'/api/v1/worker/claim',
|
|
json={'request_id': 'f' * 32, 'build': build_compatibility()},
|
|
headers={'Authorization': 'Bearer fixture-token-value'},
|
|
)
|
|
terminal = client.post(
|
|
'/api/v1/worker/assignments/17/terminal',
|
|
json={'failure_code': 'client_cancelled', 'detail': 'fixture'},
|
|
headers={'Authorization': 'Bearer fixture-token-value'},
|
|
)
|
|
self.assertEqual(claimed.status_code, 201)
|
|
self.assertEqual(terminal.status_code, 200)
|
|
|
|
def test_progress_route_is_authenticated_canonical_and_non_renewing(self):
|
|
event = {
|
|
'schema': 1,
|
|
'sequence': 41,
|
|
'timestamp': '2026-09-24T00:00:00Z',
|
|
'instance_id': 'worker-instance',
|
|
'slot_id': 0,
|
|
'reservation_id': 17,
|
|
'source': 'github',
|
|
'type': 'slot.phase',
|
|
'phase': 'scanning',
|
|
'phase_started_at': '2026-09-24T00:00:00Z',
|
|
'scan_deadline_at': '2026-09-24T00:10:00Z',
|
|
'assignment_deadline_at': '2026-09-24T01:00:00Z',
|
|
'progress': {'items_seen': 3},
|
|
}
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'server-bundles'), reject_reparse=True,
|
|
)
|
|
with TestClient(create_worker_app(
|
|
FakeService(root), reaper_interval_seconds=3600,
|
|
)) as client:
|
|
response = client.post(
|
|
'/api/v1/worker/assignments/17/progress',
|
|
json=event,
|
|
headers={'Authorization': 'Bearer fixture-token-value'},
|
|
)
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertEqual(response.json(), {
|
|
'accepted': True,
|
|
'reservation_id': 17,
|
|
'sequence': 41,
|
|
'received_at': '2026-09-24T00:00:01+00:00',
|
|
'replayed': False,
|
|
})
|
|
|
|
def test_progress_route_rejects_missing_auth_inactive_and_conflicting_authority(self):
|
|
event = {
|
|
'schema': 1, 'sequence': 41,
|
|
'timestamp': '2026-09-24T00:00:00Z',
|
|
'instance_id': 'worker-instance', 'slot_id': 0,
|
|
'reservation_id': 17, 'source': 'github', 'type': 'slot.phase',
|
|
'phase': 'scanning', 'phase_started_at': '2026-09-24T00:00:00Z',
|
|
'scan_deadline_at': None, 'assignment_deadline_at': None,
|
|
'progress': {},
|
|
}
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'server-bundles'), reject_reparse=True,
|
|
)
|
|
service = FakeService(root)
|
|
with TestClient(create_worker_app(
|
|
service, reaper_interval_seconds=3600,
|
|
)) as client:
|
|
unauthorized = client.post(
|
|
'/api/v1/worker/assignments/17/progress', json=event,
|
|
)
|
|
wrong_token = client.post(
|
|
'/api/v1/worker/assignments/17/progress', json=event,
|
|
headers={'Authorization': 'Bearer wrong-token-value'},
|
|
)
|
|
service.progress = mock.Mock(side_effect=WorkerAPIError(
|
|
410, 'progress_stale',
|
|
'owned assignment is inactive for new progress',
|
|
))
|
|
inactive = client.post(
|
|
'/api/v1/worker/assignments/17/progress', json=event,
|
|
headers={'Authorization': 'Bearer fixture-token-value'},
|
|
)
|
|
service.progress = mock.Mock(side_effect=
|
|
WorkerObservabilityConflictError('exact sequence conflict')
|
|
)
|
|
conflict = client.post(
|
|
'/api/v1/worker/assignments/17/progress', json=event,
|
|
headers={'Authorization': 'Bearer fixture-token-value'},
|
|
)
|
|
self.assertEqual(unauthorized.status_code, 401)
|
|
self.assertEqual(wrong_token.status_code, 401)
|
|
self.assertEqual(inactive.status_code, 410)
|
|
self.assertEqual(inactive.json()['error']['code'], 'progress_stale')
|
|
self.assertEqual(conflict.status_code, 409)
|
|
|
|
def test_worker_service_maps_only_owned_inactive_progress_to_stale(self):
|
|
class DB:
|
|
enabled = True
|
|
|
|
def __init__(self, error):
|
|
self.error = error
|
|
|
|
def record_remote_worker_progress_event(self, *_args):
|
|
raise self.error
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'server-bundles'), reject_reparse=True,
|
|
)
|
|
identity = {'device_id': 9, 'token_sha256': 'a' * 64}
|
|
service = WorkerService(
|
|
'postgresql://fixture', root, lambda *_args: None,
|
|
db_factory=lambda **_kwargs: DB(WorkerProgressInactiveError('inactive')),
|
|
)
|
|
with self.assertRaises(WorkerAPIError) as stale:
|
|
service.progress(identity, 17, {'sequence': 1})
|
|
self.assertEqual((stale.exception.status_code, stale.exception.code), (
|
|
410, 'progress_stale',
|
|
))
|
|
service.db_factory = lambda **_kwargs: DB(
|
|
WorkerObservabilityConflictError('wrong owner'),
|
|
)
|
|
with self.assertRaises(WorkerObservabilityConflictError):
|
|
service.progress(identity, 17, {'sequence': 1})
|
|
|
|
def test_claim_empty_delay_and_terminal_reconciliation_are_explicit(self):
|
|
receipt = {
|
|
'receipt_id': '1' * 64, 'resolution': 'expired',
|
|
'reservation_id': 17, 'bundle_id': 'd' * 32,
|
|
'scan_event_id': 'e' * 32,
|
|
'resolved_at': '2026-09-18T00:00:00+00:00',
|
|
}
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
root = ensure_private_directory(
|
|
os.path.join(temp_dir, 'server-bundles'), reject_reparse=True,
|
|
)
|
|
service = FakeService(root)
|
|
service.claim = lambda *_args: {
|
|
'no_assignment': {'reason': 'capacity'},
|
|
}
|
|
with TestClient(create_worker_app(service, reaper_interval_seconds=3600)) as client:
|
|
empty = client.post(
|
|
'/api/v1/worker/claim',
|
|
json={'request_id': 'f' * 32, 'build': build_compatibility()},
|
|
headers={'Authorization': 'Bearer fixture-token-value'},
|
|
)
|
|
service.claim = lambda *_args: {'resolution': receipt}
|
|
with TestClient(create_worker_app(service, reaper_interval_seconds=3600)) as client:
|
|
resolved = client.post(
|
|
'/api/v1/worker/claim',
|
|
json={'request_id': 'f' * 32, 'build': build_compatibility()},
|
|
headers={'Authorization': 'Bearer fixture-token-value'},
|
|
)
|
|
|
|
self.assertEqual(empty.status_code, 204)
|
|
self.assertEqual(empty.headers['retry-after'], '5')
|
|
self.assertEqual(empty.headers['x-truf-no-work-reason'], 'capacity')
|
|
self.assertEqual(resolved.status_code, 200)
|
|
self.assertEqual(resolved.json(), {'resolution': receipt})
|
|
|
|
def test_json_body_idle_timeout_is_enforced(self):
|
|
class SlowRequest:
|
|
headers = {'content-type': 'application/json'}
|
|
|
|
async def stream(self):
|
|
await asyncio.sleep(0.05)
|
|
yield b'{}'
|
|
|
|
with self.assertRaises(WorkerAPIError) as raised:
|
|
asyncio.run(_bounded_json(
|
|
SlowRequest(), absolute_timeout=0.1, idle_timeout=0.01,
|
|
))
|
|
self.assertEqual(raised.exception.status_code, 408)
|
|
self.assertEqual(raised.exception.code, 'request_timeout')
|
|
|
|
|
|
class FakeClientAPI:
|
|
def __init__(self):
|
|
self.terminal_calls = 0
|
|
self.fail_terminal_once = False
|
|
self.terminal_reports = []
|
|
|
|
def status(self, reservation_id):
|
|
return {
|
|
'reservation_id': reservation_id, 'state': 'scanning', 'resolution': None,
|
|
'expires_at': '2999-01-01T00:00:00+00:00',
|
|
}
|
|
|
|
def upload(self, reservation_id, path):
|
|
digest = hashlib.sha256()
|
|
with open(path, 'rb') as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b''):
|
|
digest.update(chunk)
|
|
return {
|
|
'reservation_id': reservation_id, 'receipt_id': '1' * 64,
|
|
'bundle_id': 'd' * 32, 'scan_event_id': 'e' * 32,
|
|
'resolution': 'bundle_accepted', 'payload_sha256': digest.hexdigest(),
|
|
}
|
|
|
|
def terminal(self, reservation_id, report):
|
|
self.terminal_calls += 1
|
|
self.terminal_reports.append(json.loads(json.dumps(report)))
|
|
if self.fail_terminal_once and self.terminal_calls == 1:
|
|
raise WorkerClientError('temporary network failure')
|
|
encoded = json.dumps(
|
|
report, ensure_ascii=True, sort_keys=True, separators=(',', ':'),
|
|
).encode('utf-8')
|
|
return {
|
|
'reservation_id': reservation_id, 'receipt_id': '3' * 64,
|
|
'bundle_id': 'd' * 32, 'scan_event_id': 'e' * 32,
|
|
'resolution': 'prebundle_report',
|
|
'failure_code': report['failure_code'],
|
|
'payload_sha256': hashlib.sha256(encoded).hexdigest(),
|
|
}
|
|
|
|
|
|
class WorkerSlotRecoveryTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self._scanner_paths = (
|
|
scanner.scan_config.trufflehog_path,
|
|
scanner.scan_config.trufflehog_config,
|
|
)
|
|
|
|
def tearDown(self):
|
|
(
|
|
scanner.scan_config.trufflehog_path,
|
|
scanner.scan_config.trufflehog_config,
|
|
) = self._scanner_paths
|
|
|
|
def _roots(self, temp_dir):
|
|
state = ensure_private_directory(os.path.join(temp_dir, 'state'), reject_reparse=True)
|
|
bundles = ensure_private_directory(os.path.join(temp_dir, 'bundles'), reject_reparse=True)
|
|
for name in ('tmp', 'ready', 'quarantine'):
|
|
ensure_private_directory(os.path.join(bundles, name), reject_reparse=True)
|
|
return state, bundles
|
|
|
|
def _slot(self, temp_dir, api=None):
|
|
state, bundles = self._roots(temp_dir)
|
|
return WorkerSlot(
|
|
0, api or FakeClientAPI(), build_compatibility(), state, bundles,
|
|
package_runtime={
|
|
'code_manifest': {}, 'code_manifest_sha256': '9' * 64,
|
|
'trufflehog_path': 'trufflehog', 'git_path': 'git',
|
|
'detector_policy_path': 'detectors.yaml',
|
|
'bootstrap_path': str(APP_DIR / 'remote_worker_bootstrap.py'),
|
|
'capabilities': (
|
|
('github', 'github', 'exact_git_v1'),
|
|
('gitlab', 'gitlab', 'exact_git_v1'),
|
|
),
|
|
},
|
|
)
|
|
|
|
def _publish_bundle(self, slot):
|
|
reservation = BundleReservation.from_mapping(reservation_mapping())
|
|
with ResultBundleWriter.open(slot.bundle_root, reservation) as writer:
|
|
return writer.finish({'queue_status': 'done'})
|
|
|
|
def test_new_claim_uses_normalized_state_before_same_cycle_execution(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
api = FakeClientAPI()
|
|
api.claim = mock.Mock(return_value=assignment())
|
|
slot = self._slot(temp_dir, api)
|
|
captured = {}
|
|
|
|
def reject_execution(state):
|
|
captured.update(state)
|
|
raise WorkerAssignmentCompatibilityError('stop after state capture')
|
|
|
|
with mock.patch.object(slot, '_execute', side_effect=reject_execution):
|
|
with self.assertRaises(WorkerAssignmentCompatibilityError):
|
|
slot.step()
|
|
|
|
self.assertEqual(captured['schema'], 2)
|
|
self.assertIsNone(captured['runner'])
|
|
self.assertEqual(captured['retained_work'], [])
|
|
self.assertEqual(slot._load(), captured)
|
|
|
|
def test_restart_adopts_published_bundle_without_rescanning(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
slot = self._slot(temp_dir)
|
|
phases = []
|
|
|
|
def record_event(value):
|
|
phase = WorkerPhase(value['phase'])
|
|
if phases:
|
|
validate_phase_transition(phases[-1], phase)
|
|
phases.append(phase)
|
|
|
|
slot.event_callback = record_event
|
|
slot._save({'phase': 'assigned', 'assignment': assignment()})
|
|
self._publish_bundle(slot)
|
|
with mock.patch.object(
|
|
slot, '_launch_runner',
|
|
side_effect=AssertionError('scanner must not run'),
|
|
):
|
|
self.assertTrue(slot.step())
|
|
self.assertFalse(os.path.exists(slot.state_path))
|
|
self.assertEqual(phases[:3], [
|
|
WorkerPhase.ASSIGNED,
|
|
WorkerPhase.UPLOADING,
|
|
WorkerPhase.AWAITING_RECEIPT,
|
|
])
|
|
|
|
def test_failed_terminal_transport_replays_report_without_rescanning(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
api = FakeClientAPI()
|
|
api.fail_terminal_once = True
|
|
slot = self._slot(temp_dir, api)
|
|
slot._save({'phase': 'assigned', 'assignment': assignment()})
|
|
with mock.patch.object(
|
|
slot, '_execute', side_effect=RuntimeError('scanner failed'),
|
|
) as execute:
|
|
with self.assertRaises(WorkerClientError):
|
|
slot.step()
|
|
pending = slot._load()
|
|
self.assertEqual(pending['phase'], 'terminal_pending')
|
|
self.assertEqual(
|
|
pending['terminal']['detail'], CLIENT_PROCESS_FAILURE_DETAIL,
|
|
)
|
|
self.assertIn('scanner failed', json.dumps(pending))
|
|
self.assertTrue(slot.step())
|
|
self.assertEqual(execute.call_count, 1)
|
|
self.assertEqual(api.terminal_calls, 2)
|
|
self.assertEqual(api.terminal_reports[0], api.terminal_reports[1])
|
|
self.assertEqual(
|
|
api.terminal_reports[0]['failure_code'],
|
|
'client_process_failed',
|
|
)
|
|
self.assertEqual(
|
|
api.terminal_reports[0]['detail'], CLIENT_PROCESS_FAILURE_DETAIL,
|
|
)
|
|
self.assertEqual(len(api.terminal_reports[0]['diagnostics']), 1)
|
|
|
|
def test_old_two_field_terminal_report_replays_without_hash_change(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
api = FakeClientAPI()
|
|
slot = self._slot(temp_dir, api)
|
|
old_report = {
|
|
'failure_code': 'client_process_failed',
|
|
'detail': CLIENT_PROCESS_FAILURE_DETAIL,
|
|
}
|
|
slot._save({
|
|
'phase': 'terminal_pending',
|
|
'assignment': assignment(),
|
|
'terminal': old_report,
|
|
})
|
|
self.assertTrue(slot.step())
|
|
self.assertEqual(api.terminal_reports, [old_report])
|
|
encoded = json.dumps(
|
|
old_report, ensure_ascii=True, sort_keys=True,
|
|
separators=(',', ':'),
|
|
).encode('ascii')
|
|
self.assertEqual(
|
|
hashlib.sha256(encoded).hexdigest(),
|
|
hashlib.sha256(json.dumps(
|
|
api.terminal_reports[0], ensure_ascii=True, sort_keys=True,
|
|
separators=(',', ':'),
|
|
).encode('ascii')).hexdigest(),
|
|
)
|
|
|
|
def test_prebundle_diagnostic_report_is_bounded_persisted_and_replayed_exactly(self):
|
|
class CapturedFailure(RuntimeError):
|
|
body = b'b' * 20000
|
|
stdout = b'o' * 20000
|
|
stderr = b'e' * 20000
|
|
returncode = 1
|
|
process_name = 'fixture-scanner'
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
api = FakeClientAPI()
|
|
api.fail_terminal_once = True
|
|
slot = self._slot(temp_dir, api)
|
|
archived = []
|
|
|
|
def archive(envelope, materials):
|
|
archived.append((envelope, materials))
|
|
return {
|
|
'diagnostic_uid': envelope.diagnostic_uid,
|
|
'record': 'diagnostics/fixture.json',
|
|
'artifacts': {},
|
|
}
|
|
|
|
slot.diagnostic_callback = archive
|
|
state = {'phase': 'assigned', 'assignment': assignment()}
|
|
slot._save(state)
|
|
slot._event_phase = WorkerPhase.SCANNING
|
|
slot._capture_failure_diagnostic(
|
|
state, CapturedFailure('scanner failed with exact evidence'),
|
|
'client_process_failed', DiagnosticCategory.SCANNER,
|
|
)
|
|
with self.assertRaises(WorkerClientError):
|
|
slot._terminal(
|
|
state, 17, 'client_process_failed', CLIENT_PROCESS_FAILURE_DETAIL,
|
|
)
|
|
pending = slot._load()
|
|
encoded = json.dumps(
|
|
pending['terminal'], ensure_ascii=True, sort_keys=True,
|
|
separators=(',', ':'),
|
|
).encode('ascii')
|
|
self.assertLessEqual(len(encoded), 16 * 1024)
|
|
self.assertEqual(len(pending['terminal']['diagnostics']), 1)
|
|
self.assertEqual(len(archived), 1)
|
|
self.assertEqual(
|
|
archived[0][0].diagnostic_uid,
|
|
pending['terminal']['diagnostics'][0]['diagnostic_uid'],
|
|
)
|
|
self.assertEqual(archived[0][1]['body'], CapturedFailure.body)
|
|
self.assertEqual(archived[0][1]['stdout'], CapturedFailure.stdout)
|
|
self.assertEqual(archived[0][1]['stderr'], CapturedFailure.stderr)
|
|
|
|
restarted = self._slot(temp_dir, api)
|
|
self.assertTrue(restarted.step())
|
|
self.assertEqual(api.terminal_calls, 2)
|
|
self.assertEqual(api.terminal_reports[0], api.terminal_reports[1])
|
|
|
|
def test_storage_diagnostic_preserves_operation_path_and_traceback(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
slot = self._slot(temp_dir)
|
|
state = {'phase': 'assigned', 'assignment': assignment()}
|
|
state = slot._slot_state(state, slot_id=slot.slot_id)
|
|
slot._event_phase = WorkerPhase.SCANNING
|
|
try:
|
|
error = PermissionError(13, 'access denied', r'C:\fixture\work')
|
|
error.winerror = 5
|
|
raise error
|
|
except PermissionError as error:
|
|
slot._capture_failure_diagnostic(
|
|
state, error, 'client_storage_failed',
|
|
DiagnosticCategory.STORAGE,
|
|
operation='runner_protocol_create',
|
|
)
|
|
report = slot._fitted_terminal_report(
|
|
'client_storage_failed', CLIENT_STORAGE_FAILURE_DETAIL,
|
|
)
|
|
diagnostic = report['diagnostics'][0]
|
|
detail = json.loads(diagnostic['exception']['message'])
|
|
self.assertEqual(detail['operation'], 'runner_protocol_create')
|
|
self.assertEqual(detail['filename'], r'C:\fixture\work')
|
|
self.assertEqual(detail['winerror'], 5)
|
|
self.assertIn('PermissionError', detail['traceback'])
|
|
|
|
def test_contract_failure_diagnostic_preserves_invalid_field(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
slot = self._slot(temp_dir)
|
|
state = {'phase': 'assigned', 'assignment': assignment()}
|
|
slot._event_phase = WorkerPhase.SCANNING
|
|
slot._capture_failure_diagnostic(
|
|
state, WorkerContractError('phase_started_at'),
|
|
'client_process_failed', DiagnosticCategory.INTERNAL,
|
|
)
|
|
|
|
report = slot._fitted_terminal_report(
|
|
'client_process_failed', CLIENT_PROCESS_FAILURE_DETAIL,
|
|
)
|
|
diagnostic = report['diagnostics'][0]
|
|
self.assertEqual(
|
|
diagnostic['code'],
|
|
'worker_contract_invalid.phase_started_at',
|
|
)
|
|
self.assertEqual(diagnostic['category'], 'protocol')
|
|
self.assertIn('phase_started_at', diagnostic['summary'])
|
|
|
|
def test_reused_slot_clears_diagnostics_after_terminal_resolution(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
api = FakeClientAPI()
|
|
slot = self._slot(temp_dir, api)
|
|
first = {'phase': 'assigned', 'assignment': assignment()}
|
|
slot._event_phase = WorkerPhase.SCANNING
|
|
slot._capture_failure_diagnostic(
|
|
first, RuntimeError('first reservation failure'),
|
|
'client_process_failed', DiagnosticCategory.INTERNAL,
|
|
)
|
|
self.assertTrue(slot._terminal(
|
|
first, 17, 'client_process_failed', CLIENT_PROCESS_FAILURE_DETAIL,
|
|
))
|
|
self.assertEqual(slot._diagnostics, [])
|
|
self.assertEqual(slot._transport_diagnostics, [])
|
|
|
|
second_assignment = assignment()
|
|
second_assignment['reservation'] = dict(
|
|
second_assignment['reservation'],
|
|
reservation_id=18,
|
|
bundle_id='1' * 32,
|
|
scan_event_id='2' * 32,
|
|
)
|
|
second = {'phase': 'assigned', 'assignment': second_assignment}
|
|
slot._event_phase = WorkerPhase.SCANNING
|
|
slot._capture_failure_diagnostic(
|
|
second, RuntimeError('second reservation failure'),
|
|
'client_process_failed', DiagnosticCategory.INTERNAL,
|
|
)
|
|
report = slot._fitted_terminal_report(
|
|
'client_process_failed', CLIENT_PROCESS_FAILURE_DETAIL,
|
|
)
|
|
self.assertEqual(len(report['diagnostics']), 1)
|
|
self.assertEqual(report['diagnostics'][0]['reservation_id'], 18)
|
|
|
|
def test_unicode_adversarial_terminal_fit_uses_original_bytes_without_archive(self):
|
|
class AdversarialFailure(RuntimeError):
|
|
body = bytes(range(256)) * 100
|
|
stdout = ('雪' * 12000).encode('utf-8')
|
|
stderr = b'raw-stderr\x00bytes' * 2000
|
|
returncode = 23
|
|
status_code = 503
|
|
content_type = 'application/octet-stream'
|
|
|
|
message = '雪' * 10000
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
api = FakeClientAPI()
|
|
api.fail_terminal_once = True
|
|
slot = self._slot(temp_dir, api)
|
|
slot.diagnostic_callback = mock.Mock(
|
|
side_effect=OSError('optional archive unavailable'),
|
|
)
|
|
state = {'phase': 'assigned', 'assignment': assignment()}
|
|
slot._event_phase = WorkerPhase.SCANNING
|
|
self.assertIsNone(slot._capture_failure_diagnostic(
|
|
state, AdversarialFailure(message),
|
|
'client_process_failed', DiagnosticCategory.INTERNAL,
|
|
))
|
|
self.assertEqual(len(slot._transport_diagnostics), 1)
|
|
with self.assertRaises(WorkerClientError):
|
|
slot._terminal(
|
|
state, 17, 'client_process_failed', CLIENT_PROCESS_FAILURE_DETAIL,
|
|
)
|
|
pending = slot._load()['terminal']
|
|
encoded = json.dumps(
|
|
pending, ensure_ascii=True, sort_keys=True, separators=(',', ':'),
|
|
).encode('ascii')
|
|
self.assertLessEqual(len(encoded), 16 * 1024)
|
|
envelope = decode_diagnostic_envelope(json.dumps(
|
|
pending['diagnostics'][0], ensure_ascii=True, sort_keys=True,
|
|
separators=(',', ':'),
|
|
).encode('ascii'))
|
|
self.assertEqual(
|
|
envelope.process.stdout.original_size,
|
|
len(AdversarialFailure.stdout),
|
|
)
|
|
self.assertEqual(
|
|
envelope.process.stdout.sha256,
|
|
hashlib.sha256(AdversarialFailure.stdout).hexdigest(),
|
|
)
|
|
self.assertEqual(
|
|
envelope.http.body.sha256,
|
|
hashlib.sha256(AdversarialFailure.body).hexdigest(),
|
|
)
|
|
restarted = self._slot(temp_dir, api)
|
|
self.assertTrue(restarted.step())
|
|
self.assertEqual(api.terminal_reports[0], api.terminal_reports[1])
|
|
|
|
def test_progress_outbox_retries_global_sequence_gaps_and_recovers_cursor(self):
|
|
events = [
|
|
{'sequence': 4, 'reservation_id': 17},
|
|
{'sequence': 5, 'reservation_id': None},
|
|
{'sequence': 9, 'reservation_id': 18},
|
|
]
|
|
|
|
class API:
|
|
def __init__(self):
|
|
self.calls = []
|
|
self.failures = [
|
|
WorkerClientError('temporary transport failure'),
|
|
WorkerHTTPError(409, 'reservation_conflict', 'sequence conflict'),
|
|
WorkerHTTPError(410, 'assignment_expired', 'not explicit'),
|
|
WorkerHTTPError(410, 'progress_stale', 'owned stale'),
|
|
]
|
|
|
|
def progress(self, reservation_id, event, *, timeout_seconds=None):
|
|
self.calls.append((reservation_id, event['sequence']))
|
|
if self.failures:
|
|
raise self.failures.pop(0)
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
state = ensure_private_directory(
|
|
os.path.join(temp_dir, 'state'), reject_reparse=True,
|
|
)
|
|
|
|
def read(after, limit):
|
|
return [event for event in events if event['sequence'] > after][:limit]
|
|
|
|
api = API()
|
|
outbox = ProgressOutbox(api, state, read)
|
|
with self.assertRaises(WorkerClientError):
|
|
outbox.publish_once()
|
|
self.assertEqual(outbox.sequence, 0)
|
|
with self.assertRaises(WorkerHTTPError):
|
|
outbox.publish_once()
|
|
self.assertEqual(outbox.sequence, 0)
|
|
with self.assertRaises(WorkerHTTPError):
|
|
outbox.publish_once()
|
|
self.assertEqual(outbox.sequence, 0)
|
|
self.assertTrue(outbox.publish_once())
|
|
self.assertEqual(outbox.sequence, 9)
|
|
self.assertEqual(
|
|
api.calls, [
|
|
(17, 4), (17, 4), (17, 4), (17, 4), (18, 9),
|
|
],
|
|
)
|
|
recovered = ProgressOutbox(api, state, read)
|
|
self.assertEqual(recovered.sequence, 9)
|
|
self.assertFalse(recovered.publish_once())
|
|
|
|
def test_progress_outbox_performs_bounded_final_shutdown_drain(self):
|
|
events = [{'sequence': 7, 'reservation_id': 17}]
|
|
|
|
class API:
|
|
calls = []
|
|
|
|
def progress(self, reservation_id, event, *, timeout_seconds=None):
|
|
self.calls.append((reservation_id, event['sequence']))
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
state = ensure_private_directory(
|
|
os.path.join(temp_dir, 'state'), reject_reparse=True,
|
|
)
|
|
outbox = ProgressOutbox(
|
|
API(), state,
|
|
lambda after, _limit: [
|
|
event for event in events if event['sequence'] > after
|
|
],
|
|
)
|
|
stopping = threading.Event()
|
|
started = time.monotonic()
|
|
self.assertTrue(outbox.drain(1.0))
|
|
self.assertLess(time.monotonic() - started, 1.0)
|
|
self.assertEqual(outbox.sequence, 7)
|
|
self.assertEqual(API.calls, [(17, 7)])
|
|
|
|
def test_progress_outbox_deadline_bounds_blocking_api_and_forced_stop_skips_drain(self):
|
|
events = [{'sequence': 7, 'reservation_id': 17}]
|
|
|
|
class BlockingAPI:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def progress(self, reservation_id, event, *, timeout_seconds=None):
|
|
self.calls.append((reservation_id, event['sequence'], timeout_seconds))
|
|
time.sleep(float(timeout_seconds) + 0.01)
|
|
raise TimeoutError('bounded blocking API fixture')
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
state = ensure_private_directory(
|
|
os.path.join(temp_dir, 'state'), reject_reparse=True,
|
|
)
|
|
api = BlockingAPI()
|
|
outbox = ProgressOutbox(
|
|
api, state,
|
|
lambda after, _limit: [
|
|
event for event in events if event['sequence'] > after
|
|
],
|
|
)
|
|
started = time.monotonic()
|
|
self.assertFalse(outbox.drain(0.15))
|
|
self.assertLess(time.monotonic() - started, 0.35)
|
|
self.assertEqual(outbox.sequence, 0)
|
|
self.assertLessEqual(api.calls[0][2], 0.151)
|
|
|
|
stopping = threading.Event()
|
|
stopping.set()
|
|
outbox.run(stopping)
|
|
self.assertEqual(len(api.calls), 1)
|
|
|
|
def test_progress_http_absolute_deadline_closes_blocked_transport(self):
|
|
closed = threading.Event()
|
|
|
|
class Connection:
|
|
def request(self, *_args, **_kwargs):
|
|
closed.wait(5)
|
|
raise OSError('transport closed at absolute deadline')
|
|
|
|
def close(self):
|
|
closed.set()
|
|
|
|
client = object.__new__(WorkerHTTPClient)
|
|
client.token = 'x' * 32
|
|
client._progress_connections = set()
|
|
client._progress_connections_lock = threading.Lock()
|
|
client._connection = lambda _timeout=None: Connection()
|
|
started = time.monotonic()
|
|
with self.assertRaises(WorkerNetworkError) as raised:
|
|
client.progress(
|
|
17, {'sequence': 1}, timeout_seconds=0.1,
|
|
)
|
|
self.assertIsInstance(raised.exception.__cause__, OSError)
|
|
self.assertNotIn('absolute deadline', str(raised.exception))
|
|
self.assertLess(time.monotonic() - started, 0.3)
|
|
self.assertTrue(closed.is_set())
|
|
self.assertEqual(client._progress_connections, set())
|
|
|
|
def test_error_summary_never_includes_exception_text(self):
|
|
secret = 'token-and-private-path-must-not-be-logged'
|
|
for error, expected in (
|
|
(RuntimeError(secret), 'worker operation failed'),
|
|
(OSError(secret), 'local I/O operation failed'),
|
|
(WorkerNetworkError(secret), 'worker network operation failed'),
|
|
(WorkerClientError(secret), 'worker protocol or state validation failed'),
|
|
(WorkerHTTPError(409, secret, secret), 'worker API request failed (HTTP 409)'),
|
|
):
|
|
summary = safe_worker_error_summary(error)
|
|
self.assertEqual(summary, expected)
|
|
self.assertNotIn(secret, summary)
|
|
|
|
def test_status_transport_oserror_is_network_error(self):
|
|
secret = 'network-detail-must-not-be-logged'
|
|
|
|
class Connection:
|
|
def request(self, *_args, **_kwargs):
|
|
raise OSError(secret)
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
client = object.__new__(WorkerHTTPClient)
|
|
client.token = 'x' * 32
|
|
client._connection = Connection
|
|
with self.assertRaises(WorkerNetworkError) as raised:
|
|
client.status(17)
|
|
self.assertIsInstance(raised.exception.__cause__, OSError)
|
|
self.assertEqual(
|
|
safe_worker_error_summary(raised.exception),
|
|
'worker network operation failed',
|
|
)
|
|
self.assertNotIn(secret, str(raised.exception))
|
|
|
|
def test_client_loop_logs_only_redacted_error_summary(self):
|
|
sentinels = (
|
|
'device-token-must-not-be-logged',
|
|
'task-credential-must-not-be-logged',
|
|
'raw-finding-must-not-be-logged',
|
|
)
|
|
stepped = threading.Event()
|
|
observed_compatibility = []
|
|
observed_package_runtime = []
|
|
|
|
class FailingSlot:
|
|
def __init__(self, _slot_id, _api, compatibility, *_args, **_kwargs):
|
|
observed_compatibility.append(compatibility)
|
|
observed_package_runtime.append(_kwargs.get('package_runtime'))
|
|
self.claim_enabled = True
|
|
self.state_path = 'unused'
|
|
self.retry_after_seconds = None
|
|
|
|
def step(self):
|
|
stepped.set()
|
|
raise RuntimeError(' '.join(sentinels))
|
|
|
|
def interrupt_after_step(_seconds):
|
|
if not stepped.wait(3):
|
|
raise AssertionError('worker slot did not execute')
|
|
raise KeyboardInterrupt
|
|
|
|
package = {
|
|
'manifest': {'capabilities': [
|
|
{
|
|
'source': 'github', 'platform': 'github',
|
|
'planning_kind': 'exact_git_v1',
|
|
},
|
|
{
|
|
'source': 'gitlab', 'platform': 'gitlab',
|
|
'planning_kind': 'exact_git_v1',
|
|
},
|
|
]},
|
|
'build_compatibility': build_compatibility(),
|
|
'code_manifest': {},
|
|
'code_manifest_sha256': '9' * 64,
|
|
'trufflehog_path': 'trufflehog',
|
|
'git_path': 'git',
|
|
'detector_policy_path': 'detectors.yaml',
|
|
'runtime_trees': {},
|
|
}
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
args = SimpleNamespace(
|
|
package_manifest=os.path.join(temp_dir, 'worker-package.json'),
|
|
state_dir=os.path.join(temp_dir, 'state'),
|
|
bundle_dir=os.path.join(temp_dir, 'bundles'),
|
|
work_dir=os.path.join(temp_dir, 'work'),
|
|
server='https://worker.example', token=sentinels[0],
|
|
http_timeout=3, parallelism=1, poll_seconds=0.1,
|
|
error_delay_seconds=0.1,
|
|
)
|
|
output = io.StringIO()
|
|
with mock.patch(
|
|
'remote_worker_client.verify_worker_package', return_value=package,
|
|
), mock.patch(
|
|
'remote_worker_client.WorkerHTTPClient', return_value=object(),
|
|
), mock.patch(
|
|
'remote_worker_client.WorkerSlot', FailingSlot,
|
|
), mock.patch(
|
|
'remote_worker_client.time.sleep', side_effect=interrupt_after_step,
|
|
), redirect_stdout(output):
|
|
run_client(args)
|
|
|
|
rendered = output.getvalue()
|
|
self.assertIn('worker slot 0: worker operation failed', rendered)
|
|
self.assertEqual(observed_compatibility, [build_compatibility()])
|
|
self.assertEqual(set(observed_package_runtime[0]), {
|
|
'code_manifest', 'code_manifest_sha256', 'trufflehog_path',
|
|
'git_path', 'detector_policy_path', 'capabilities', 'bootstrap_path',
|
|
})
|
|
for sentinel in sentinels:
|
|
self.assertNotIn(sentinel, rendered)
|
|
|
|
def test_containment_pending_exits_engine_for_supervisor_escalation(self):
|
|
class PendingSlot:
|
|
def __init__(self, slot_id, *_args, **_kwargs):
|
|
self.slot_id = slot_id
|
|
self.claim_enabled = True
|
|
self.state_path = 'pending-authority'
|
|
self.retry_after_seconds = None
|
|
|
|
def step(self):
|
|
raise RunnerContainmentPending(
|
|
'fixture containment pending',
|
|
)
|
|
|
|
package = {
|
|
'manifest': {'capabilities': [{
|
|
'source': 'gitlab', 'platform': 'gitlab',
|
|
'planning_kind': 'exact_git_v1',
|
|
}]},
|
|
'build_compatibility': build_compatibility(),
|
|
'code_manifest': {}, 'code_manifest_sha256': '9' * 64,
|
|
'trufflehog_path': 'trufflehog', 'git_path': 'git',
|
|
'detector_policy_path': 'detectors.yaml', 'runtime_trees': {},
|
|
}
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
args = SimpleNamespace(
|
|
package_manifest=os.path.join(temp_dir, 'worker-package.json'),
|
|
state_dir=os.path.join(temp_dir, 'state'),
|
|
bundle_dir=os.path.join(temp_dir, 'bundles'),
|
|
work_dir=os.path.join(temp_dir, 'work'),
|
|
server='https://worker.example', token='x' * 32,
|
|
http_timeout=3, parallelism=1, poll_seconds=0.01,
|
|
error_delay_seconds=0.01,
|
|
)
|
|
with mock.patch(
|
|
'remote_worker_client.verify_worker_package', return_value=package,
|
|
), mock.patch(
|
|
'remote_worker_client.WorkerHTTPClient', return_value=object(),
|
|
), mock.patch(
|
|
'remote_worker_client.WorkerSlot', PendingSlot,
|
|
):
|
|
self.assertEqual(run_client(args), 1)
|
|
|
|
def test_https_client_validates_ca_and_hostname(self):
|
|
certificate = ROOT / 'tests' / 'fixtures' / 'worker_tls_cert.pem'
|
|
private_key = ROOT / 'tests' / 'fixtures' / 'worker_tls_key.pem'
|
|
server_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
|
server_context.load_cert_chain(certificate, private_key)
|
|
listener = socket.create_server(('127.0.0.1', 0))
|
|
listener.settimeout(0.1)
|
|
port = listener.getsockname()[1]
|
|
stopping = threading.Event()
|
|
|
|
def serve():
|
|
while not stopping.is_set():
|
|
try:
|
|
connection, _address = listener.accept()
|
|
except TimeoutError:
|
|
continue
|
|
except OSError:
|
|
return
|
|
try:
|
|
with server_context.wrap_socket(connection, server_side=True) as secured:
|
|
secured.recv(1)
|
|
except OSError:
|
|
connection.close()
|
|
|
|
thread = threading.Thread(target=serve, daemon=True)
|
|
thread.start()
|
|
token = 'fixture-worker-token-not-a-secret'
|
|
try:
|
|
with mock.patch.dict(os.environ, {'SSL_CERT_FILE': str(certificate)}):
|
|
trusted = WorkerHTTPClient(f'https://localhost:{port}', token, 3)
|
|
self.assertTrue(trusted.ssl_context.check_hostname)
|
|
self.assertEqual(trusted.ssl_context.verify_mode, ssl.CERT_REQUIRED)
|
|
connection = trusted._connection()
|
|
connection.connect()
|
|
connection.close()
|
|
|
|
wrong_host = WorkerHTTPClient(f'https://127.0.0.1:{port}', token, 3)
|
|
with self.assertRaises(ssl.SSLCertVerificationError):
|
|
wrong_host._connection().connect()
|
|
|
|
environment = dict(os.environ)
|
|
environment.pop('SSL_CERT_FILE', None)
|
|
environment.pop('SSL_CERT_DIR', None)
|
|
with mock.patch.dict(os.environ, environment, clear=True):
|
|
untrusted = WorkerHTTPClient(f'https://localhost:{port}', token, 3)
|
|
with self.assertRaises(ssl.SSLCertVerificationError):
|
|
untrusted._connection().connect()
|
|
finally:
|
|
stopping.set()
|
|
listener.close()
|
|
thread.join(timeout=3)
|
|
|
|
def test_bootstrap_paths_are_derived_from_the_installation(self):
|
|
windows = default_worker_paths(
|
|
r'C:\worker\app\remote_worker_client.py', platform_name='nt',
|
|
environ={'LOCALAPPDATA': r'C:\Users\worker\AppData\Local'},
|
|
home=r'C:\Users\worker',
|
|
)
|
|
self.assertEqual(windows['package_manifest'], r'C:\worker\worker-package.json')
|
|
self.assertEqual(
|
|
windows['state_dir'], r'C:\Users\worker\AppData\Local\TRUF\RemoteWorker',
|
|
)
|
|
self.assertEqual(
|
|
windows['work_dir'], r'C:\Users\worker\AppData\Local\TRUF\RemoteWorker\work',
|
|
)
|
|
linux = default_worker_paths(
|
|
'/opt/truf-worker/app/remote_worker_client.py', platform_name='posix',
|
|
environ={}, home='/home/worker',
|
|
)
|
|
self.assertEqual(linux['package_manifest'], '/opt/truf-worker/worker-package.json')
|
|
self.assertEqual(linux['state_dir'], '/home/worker/.local/state/truf/remote-worker')
|
|
self.assertEqual(
|
|
linux['bundle_dir'], '/home/worker/.local/share/truf/remote-worker/bundles',
|
|
)
|
|
self.assertEqual(
|
|
linux['work_dir'], '/home/worker/.local/share/truf/remote-worker/work',
|
|
)
|
|
|
|
def test_bootstrap_requires_only_server_token_and_optional_parallelism(self):
|
|
args = parse_args([
|
|
'--server', 'https://worker.example', '--token', 'x' * 32,
|
|
'--parallelism', '3',
|
|
])
|
|
self.assertEqual(args.parallelism, 3)
|
|
self.assertTrue(args.package_manifest.endswith('worker-package.json'))
|
|
self.assertTrue(args.state_dir)
|
|
self.assertTrue(args.bundle_dir)
|
|
self.assertTrue(args.work_dir)
|
|
|
|
def test_bootstrap_rejects_client_authored_operational_overrides(self):
|
|
common = ['--server', 'https://worker.example', '--token', 'x' * 32]
|
|
for option, value in (
|
|
('--package-manifest', 'foreign.json'),
|
|
('--state-dir', 'foreign-state'),
|
|
('--bundle-dir', 'foreign-bundles'),
|
|
('--poll-seconds', '1'),
|
|
('--error-delay-seconds', '1'),
|
|
('--http-timeout', '1'),
|
|
):
|
|
with self.subTest(option=option, value=value), self.assertRaises(SystemExit):
|
|
parse_args(common + [option, value])
|
|
|
|
def test_http_claim_parses_server_wait_and_terminal_resolution(self):
|
|
client = object.__new__(WorkerHTTPClient)
|
|
client._json_request = lambda *_args, **_kwargs: (
|
|
204, None, '17', 'empty_queue',
|
|
)
|
|
self.assertEqual(
|
|
client.claim('f' * 32, build_compatibility()),
|
|
{'retry_after_seconds': 17, 'reason': 'empty_queue'},
|
|
)
|
|
|
|
receipt = {
|
|
'receipt_id': '1' * 64, 'resolution': 'expired',
|
|
'reservation_id': 17, 'bundle_id': 'd' * 32,
|
|
'scan_event_id': 'e' * 32,
|
|
}
|
|
client._json_request = lambda *_args, **_kwargs: (
|
|
200, {'resolution': receipt}, None, None,
|
|
)
|
|
self.assertEqual(
|
|
client.claim('f' * 32, build_compatibility()),
|
|
{'claim_resolution': receipt},
|
|
)
|
|
|
|
client._json_request = lambda *_args, **_kwargs: (
|
|
204, None, '301', 'capacity',
|
|
)
|
|
with self.assertRaisesRegex(WorkerClientError, 'retry delay'):
|
|
client.claim('f' * 32, build_compatibility())
|
|
|
|
def test_terminal_awaiting_receipt_callback_runs_after_request_send(self):
|
|
order = []
|
|
|
|
class Response:
|
|
status = 200
|
|
|
|
@staticmethod
|
|
def read(_maximum):
|
|
return b'{}'
|
|
|
|
@staticmethod
|
|
def getheader(_name):
|
|
return None
|
|
|
|
class Connection:
|
|
def request(self, *_args, **_kwargs):
|
|
order.append('request_sent')
|
|
|
|
def getresponse(self):
|
|
order.append('response_wait')
|
|
return Response()
|
|
|
|
def close(self):
|
|
order.append('closed')
|
|
|
|
client = object.__new__(WorkerHTTPClient)
|
|
client.token = 'x' * 32
|
|
client._connection = Connection
|
|
client.terminal(
|
|
17, {'failure_code': 'client_process_failed', 'detail': 'detail'},
|
|
response_wait_callback=lambda: order.append('awaiting_receipt'),
|
|
)
|
|
self.assertEqual(order, [
|
|
'request_sent', 'awaiting_receipt', 'response_wait', 'closed',
|
|
])
|
|
|
|
def test_claim_retry_guidance_and_terminal_resolution_clear_pending_request(self):
|
|
class ClaimAPI(FakeClientAPI):
|
|
def __init__(self, result):
|
|
super().__init__()
|
|
self.result = result
|
|
|
|
def claim(self, request_id, compatibility):
|
|
return self.result
|
|
|
|
def status(self, reservation_id):
|
|
raise AssertionError('resolved or empty claim must not query assignment status')
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
waiting = self._slot(
|
|
temp_dir, ClaimAPI({
|
|
'retry_after_seconds': 17, 'reason': 'capacity',
|
|
}),
|
|
)
|
|
events = []
|
|
waiting.event_callback = events.append
|
|
self.assertFalse(waiting.step())
|
|
self.assertEqual(waiting.retry_after_seconds, 17)
|
|
self.assertIsNone(waiting._load())
|
|
self.assertEqual(events[-1]['progress']['reason'], 'capacity')
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
unknown = self._slot(temp_dir, ClaimAPI(None))
|
|
events = []
|
|
unknown.event_callback = events.append
|
|
self.assertFalse(unknown.step())
|
|
self.assertNotIn('reason', events[-1]['progress'])
|
|
|
|
receipt = {
|
|
'receipt_id': '1' * 64, 'resolution': 'expired',
|
|
'reservation_id': 17, 'bundle_id': 'd' * 32,
|
|
'scan_event_id': 'e' * 32,
|
|
}
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
resolved = self._slot(
|
|
temp_dir, ClaimAPI({'claim_resolution': receipt}),
|
|
)
|
|
self.assertTrue(resolved.step())
|
|
self.assertIsNone(resolved._load())
|
|
|
|
def test_package_source_allowlist_rejects_assignment_before_scanning(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
slot = self._slot(temp_dir)
|
|
slot.package_runtime['capabilities'] = (
|
|
('gitlab', 'gitlab', 'exact_git_v1'),
|
|
)
|
|
state = {'phase': 'assigned', 'assignment': assignment()}
|
|
with mock.patch.object(
|
|
slot, '_launch_runner',
|
|
side_effect=AssertionError('scanner must not run'),
|
|
) as execute:
|
|
with self.assertRaisesRegex(WorkerClientError, 'incompatible'):
|
|
slot._execute(state)
|
|
execute.assert_not_called()
|
|
|
|
def test_direct_capability_assignment_dispatches_validated_target(self):
|
|
for source, platform, planning_kind in (
|
|
('dockerhub', 'docker', 'docker_direct_v1'),
|
|
('huggingface', 'huggingface', 'huggingface_space_v1'),
|
|
):
|
|
with self.subTest(source=source), tempfile.TemporaryDirectory() as temp_dir:
|
|
api = FakeClientAPI()
|
|
slot = self._slot(temp_dir, api)
|
|
slot.package_runtime['capabilities'] = ((source, platform, planning_kind),)
|
|
pending = direct_assignment(source)
|
|
state = {'phase': 'assigned', 'assignment': pending}
|
|
with mock.patch(
|
|
'remote_worker_client.validate_protocol2_remote_assignment',
|
|
wraps=validate_protocol2_remote_assignment,
|
|
) as validate, mock.patch.object(
|
|
slot, '_new_runner', return_value=object(),
|
|
), mock.patch.object(
|
|
slot, '_launch_runner', side_effect=lambda value: {
|
|
**value, 'phase': 'bundle_ready', 'bundle': {'source': source},
|
|
},
|
|
):
|
|
completed = slot._execute(state)
|
|
self.assertEqual(completed['phase'], 'bundle_ready')
|
|
validated = validate_protocol2_remote_assignment(
|
|
*validate.call_args.args, **validate.call_args.kwargs,
|
|
)
|
|
self.assertEqual(validated['planning_kind'], planning_kind)
|
|
self.assertEqual(
|
|
validated['execution_target'],
|
|
pending['execution_plan']['execution_target'],
|
|
)
|
|
self.assertEqual(api.terminal_calls, 0)
|
|
|
|
def test_direct_assignment_tampering_fails_before_scanner_start(self):
|
|
mutations = (
|
|
('snapshot-hash', lambda item: item.update(
|
|
execution_snapshot_sha256='0' * 64,
|
|
)),
|
|
('credential-source', lambda item: item[
|
|
'execution_snapshot'
|
|
]['credential_ref'].update(source='gitlab')),
|
|
('reservation-platform', lambda item: item['reservation'].update(
|
|
platform='huggingface',
|
|
)),
|
|
('plan-kind', lambda item: item['execution_plan'].update(
|
|
kind='huggingface_space_v1',
|
|
)),
|
|
('plan-target', lambda item: item['execution_plan'].update(
|
|
execution_target='docker.io/library/busybox@sha256:' + ('2' * 64),
|
|
)),
|
|
('private-token', lambda item: item['scan_kwargs'].update(
|
|
token='must-not-cross-assignment-boundary',
|
|
)),
|
|
)
|
|
for label, mutate in mutations:
|
|
with self.subTest(label=label), tempfile.TemporaryDirectory() as temp_dir:
|
|
slot = self._slot(temp_dir)
|
|
slot.package_runtime['capabilities'] = (
|
|
('dockerhub', 'docker', 'docker_direct_v1'),
|
|
)
|
|
pending = direct_assignment('dockerhub')
|
|
mutate(pending)
|
|
with mock.patch.object(
|
|
slot, '_launch_runner',
|
|
side_effect=AssertionError('scanner must not run'),
|
|
) as execute:
|
|
with self.assertRaises(WorkerAssignmentCompatibilityError):
|
|
slot._execute({'phase': 'assigned', 'assignment': pending})
|
|
execute.assert_not_called()
|
|
|
|
def test_incompatible_persisted_ready_bundle_still_uploads(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
api = FakeClientAPI()
|
|
slot = self._slot(temp_dir, api)
|
|
slot.package_runtime['capabilities'] = (
|
|
('dockerhub', 'docker', 'docker_direct_v1'),
|
|
)
|
|
pending = assignment()
|
|
slot._save({'phase': 'bundle_ready', 'assignment': pending})
|
|
self._publish_bundle(slot)
|
|
with mock.patch.object(
|
|
slot, '_launch_runner',
|
|
side_effect=AssertionError('scanner must not run'),
|
|
) as execute:
|
|
self.assertTrue(slot.step())
|
|
execute.assert_not_called()
|
|
self.assertEqual(api.terminal_calls, 0)
|
|
self.assertIsNone(slot._load())
|
|
|
|
def test_effective_scan_policy_mismatch_rejects_before_scanning(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
slot = self._slot(temp_dir)
|
|
changed = assignment()
|
|
changed['scan_policy']['drop_detectors'] = ['unexpected']
|
|
state = {'phase': 'assigned', 'assignment': changed}
|
|
with mock.patch.object(
|
|
slot, '_launch_runner',
|
|
side_effect=AssertionError('scanner must not run'),
|
|
) as execute:
|
|
with self.assertRaises(WorkerAssignmentCompatibilityError):
|
|
slot._execute(state)
|
|
execute.assert_not_called()
|
|
|
|
def test_state_write_failure_after_publish_uploads_without_terminal_report(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
api = FakeClientAPI()
|
|
slot = self._slot(temp_dir, api)
|
|
slot._save({'phase': 'assigned', 'assignment': assignment()})
|
|
state = slot._load()
|
|
self._publish_bundle(slot)
|
|
with mock.patch.object(
|
|
slot, '_save', side_effect=OSError('state disk unavailable'),
|
|
):
|
|
with self.assertRaises(OSError):
|
|
slot._adopt_ready_bundle(state)
|
|
self.assertTrue(slot.step())
|
|
self.assertEqual(api.terminal_calls, 0)
|
|
|
|
def test_execute_serializes_the_real_staged_result_contract(self):
|
|
staged = StagedResult(
|
|
target='https://github.com/example/repo.git',
|
|
scan_event_id='e' * 32,
|
|
bundle_id='d' * 32,
|
|
reservation_id=17,
|
|
scan_event_hash='4' * 64,
|
|
actual_bytes=4096,
|
|
relative_path='ready/dd/' + ('d' * 32) + '.trb',
|
|
frame_count=4,
|
|
finding_count=1,
|
|
error_count=0,
|
|
candidate_count=1,
|
|
queue_status='done',
|
|
source_failure=False,
|
|
source_failure_category='',
|
|
source_failure_auth_related=False,
|
|
first_error='',
|
|
)
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
slot = self._slot(temp_dir)
|
|
state = {'phase': 'assigned', 'assignment': assignment()}
|
|
state['runner'] = {
|
|
'root_name': 'worker-assignment-0-17-' + ('a' * 32),
|
|
}
|
|
outcome = {
|
|
'status': 'succeeded',
|
|
'bundle': {'commit': staged.as_dict()},
|
|
}
|
|
with mock.patch(
|
|
'remote_worker_client.adopt_runner_bundle',
|
|
return_value=SimpleNamespace(scan_event_hash=staged.scan_event_hash),
|
|
), mock.patch.object(
|
|
slot, '_retain_runner_work', return_value='abandoned/fixture',
|
|
), mock.patch.object(slot, '_save'):
|
|
result = slot._adopt_runner_outcome(state, outcome)
|
|
self.assertEqual(result['phase'], 'bundle_ready')
|
|
self.assertEqual(result['bundle']['scan_event_hash'], staged.scan_event_hash)
|
|
self.assertEqual(
|
|
result['bundle']['canonical_scan_event_hash'], staged.scan_event_hash,
|
|
)
|
|
|
|
def test_expired_assignment_is_not_executed(self):
|
|
class ExpiredAPI(FakeClientAPI):
|
|
def status(self, reservation_id):
|
|
return {
|
|
'reservation_id': reservation_id, 'state': 'scanning',
|
|
'expires_at': '2000-01-01T00:00:00+00:00', 'resolution': None,
|
|
}
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
slot = self._slot(temp_dir, ExpiredAPI())
|
|
slot._save({'phase': 'assigned', 'assignment': assignment()})
|
|
with mock.patch.object(
|
|
slot, '_execute',
|
|
side_effect=AssertionError('expired assignment must not execute'),
|
|
) as execute:
|
|
with self.assertRaisesRegex(WorkerClientError, 'deadline has passed'):
|
|
slot.step()
|
|
execute.assert_not_called()
|
|
|
|
def test_conflicting_upload_retries_once_then_reports_terminal_failure(self):
|
|
class ConflictAPI(FakeClientAPI):
|
|
def upload(self, reservation_id, path):
|
|
raise WorkerHTTPError(409, 'reservation_conflict', 'conflict')
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
slot = self._slot(temp_dir, ConflictAPI())
|
|
slot._save({'phase': 'bundle_ready', 'assignment': assignment()})
|
|
self._publish_bundle(slot)
|
|
self.assertFalse(slot.step())
|
|
pending = slot._load()
|
|
self.assertEqual(pending['phase'], 'bundle_ready')
|
|
self.assertEqual(pending['transport_conflict']['attempts'], 1)
|
|
self.assertTrue(os.path.exists(bundle_ready_path(slot.bundle_root, 'd' * 32)))
|
|
self.assertTrue(slot.step())
|
|
self.assertIsNone(slot._load())
|
|
self.assertEqual(slot.api.terminal_calls, 1)
|
|
self.assertEqual(
|
|
slot.api.terminal_reports[0]['failure_code'],
|
|
'client_process_failed',
|
|
)
|
|
self.assertFalse(os.path.exists(bundle_ready_path(slot.bundle_root, 'd' * 32)))
|
|
self.assertFalse(os.path.exists(slot.stale_path))
|
|
|
|
def test_restart_finishes_cleanup_after_stale_outcome_was_recorded(self):
|
|
class MissingAPI(FakeClientAPI):
|
|
def status(self, reservation_id):
|
|
raise WorkerHTTPError(404, 'assignment_not_found', 'private detail')
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
slot = self._slot(temp_dir, MissingAPI())
|
|
slot._save({'phase': 'assigned', 'assignment': assignment()})
|
|
original_remove = slot._remove_state
|
|
with mock.patch.object(slot, '_remove_state', side_effect=OSError('interrupted')):
|
|
with self.assertRaises(OSError):
|
|
slot.step()
|
|
self.assertTrue(os.path.isfile(slot.state_path))
|
|
self.assertTrue(os.path.isfile(slot.stale_path))
|
|
with mock.patch.object(slot, '_remove_state', side_effect=original_remove):
|
|
self.assertTrue(slot.step())
|
|
self.assertFalse(os.path.exists(slot.state_path))
|
|
|
|
def test_terminal_receipt_must_match_pending_report(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
slot = self._slot(temp_dir)
|
|
state = {'phase': 'assigned', 'assignment': assignment()}
|
|
receipt = {
|
|
'reservation_id': 17, 'bundle_id': 'd' * 32,
|
|
'scan_event_id': 'e' * 32, 'receipt_id': '3' * 64,
|
|
'resolution': 'prebundle_report', 'failure_code': 'client_process_failed',
|
|
'payload_sha256': '4' * 64,
|
|
}
|
|
with self.assertRaisesRegex(WorkerClientError, 'pending state'):
|
|
slot._resolved(state, receipt)
|
|
|
|
def test_persisted_high_slot_is_discovered_but_cannot_claim_new_work(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
state, bundles = self._roots(temp_dir)
|
|
high = WorkerSlot(
|
|
7, FakeClientAPI(), build_compatibility(), state, bundles,
|
|
claim_enabled=False,
|
|
)
|
|
high._save({'phase': 'assigned', 'assignment': assignment()})
|
|
self.assertEqual(persisted_slot_ids(state), {7})
|
|
high._remove_state()
|
|
self.assertFalse(high.step())
|
|
self.assertFalse(os.path.exists(high.state_path))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|