import hashlib import os import re from dataclasses import asdict, dataclass from datetime import datetime, timedelta from types import MappingProxyType, SimpleNamespace from console_runner import ( _GitResolutionFailure, git_resolution_failure_result, prepare_scan_options, reserve_v2_admission_with_recovery, resolve_and_bind_git_claim, validate_v2_capacity_model, ) from process_identity import current_process_identity from result_bundle import BundleReservation, ResultBundleReader, bundle_ready_path from lifecycle_authority import DISCOVERY_PRODUCER_SOURCES from scan_execution import ( PROTOCOL_VERSION, QueueDispositionPolicy, ScanCompatibility, ScanExecutionError, WorkerBuildCompatibility, canonical_json_sha256, normalize_docker_direct_execution_snapshot, normalize_docker_direct_execution_target, normalize_exact_git_execution_snapshot, normalize_huggingface_space_execution_snapshot, normalize_huggingface_space_execution_target, normalize_remote_scan_policy, remote_execution_identity, remote_execution_snapshot_sha256, stage_scan_result_in_scope, validate_scan_kwargs, validate_worker_build_compatibility, ) from scanner_db import ScannerDB, ScanEventConflictError from scanner import scan_config from runtime_security import sha256_file from worker_contracts import ( DIAGNOSTIC_PROJECTION_VERSION, ordered_diagnostic_uid_set_sha256, ) from worker_package import ( KNOWN_WORKER_PACKAGE_CAPABILITIES, PACKAGE_DETECTOR_POLICY, load_worker_package_manifest, normalize_worker_package_manifest, worker_package_build_compatibility, ) DEFAULT_ASSIGNMENT_TTL_SECONDS = 24 * 60 * 60 DEFAULT_RESULT_UPLOAD_BODY_TIMEOUT_SECONDS = 30 * 60 PROTOCOL1_NEW_CLAIM_SOURCES = frozenset(('github', 'gitlab')) SUPPORTED_REMOTE_GIT_SOURCES = PROTOCOL1_NEW_CLAIM_SOURCES _EXACT_GIT_ASSIGNMENT_FLOW = object() _DIRECT_ASSIGNMENT_FLOW = object() NO_WORK_REASONS = frozenset(( 'empty_queue', 'assignment_cap', 'dispatch_paused', 'capacity', 'compatibility', )) _ADMISSION_NO_WORK_REASONS = { 'no_claimable_target': 'empty_queue', 'remote_user_quota_closed': 'assignment_cap', 'remote_global_quota_closed': 'assignment_cap', 'dispatch_gate_closed': 'dispatch_paused', 'pipeline_capacity_closed': 'capacity', 'quarantine_capacity_conflict': 'capacity', 'quarantine_admission_closed': 'capacity', 'ingester_not_ready': 'capacity', } _CANONICAL_CAPABILITY_NAME = re.compile(r'[a-z][a-z0-9_]{0,63}\Z') @dataclass(frozen=True) class AssignmentCapability: source: str platform: str planning_kind: str def as_dict(self): return { 'source': self.source, 'platform': self.platform, 'planning_kind': self.planning_kind, } @dataclass(frozen=True) class AssignmentSourceAdapter: queue_source: str worker_platform: str planning_kind: str package_capability: AssignmentCapability snapshot_validator: object assignment_flow: object = None def validate_source_args(self, args): if str(getattr(args, 'platform', '') or '').strip().lower() != self.worker_platform: raise ValueError('remote assignment source platform does not match its adapter') if self.assignment_flow is None: raise ValueError('remote assignment source is not available for new claims') if self.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW and ( self.planning_kind != 'exact_git_v1' or not bool(getattr(args, 'exact_git_planning_enabled', False)) ): raise ValueError('remote Git sources require enabled exact Git planning') if self.assignment_flow is _DIRECT_ASSIGNMENT_FLOW: if self.planning_kind not in { 'docker_direct_v1', 'huggingface_space_v1', } or any(str(getattr(args, name, '') or '') for name in ( 'token', 'docker_username', 'docker_token', 'auth_name', )): raise ValueError('remote direct sources must be credential-free') elif self.assignment_flow is not _EXACT_GIT_ASSIGNMENT_FLOW: raise ValueError('remote assignment source flow is invalid') def validate_snapshot(self, value): normalized = self.snapshot_validator(value) if ( normalized['execution']['source'] != self.worker_platform or normalized['planning']['kind'] != self.planning_kind or normalized['credential_ref']['source'] != self.queue_source ): raise ScanExecutionError('remote execution snapshot does not match its source adapter') return normalized def _exact_git_snapshot(value): return normalize_exact_git_execution_snapshot(value) def _adapter(source, platform, planning_kind, *, assignment_flow=None, validator=None): capability = AssignmentCapability(source, platform, planning_kind) return AssignmentSourceAdapter( source, platform, planning_kind, capability, validator, assignment_flow, ) LEGACY_GITHUB_ASSIGNMENT_ADAPTER = _adapter( 'github', 'github', 'exact_git_v1', assignment_flow=_EXACT_GIT_ASSIGNMENT_FLOW, validator=_exact_git_snapshot, ) CORE_ASSIGNMENT_SOURCE_ADAPTERS = MappingProxyType({ 'gitlab': _adapter( 'gitlab', 'gitlab', 'exact_git_v1', assignment_flow=_EXACT_GIT_ASSIGNMENT_FLOW, validator=_exact_git_snapshot, ), 'dockerhub': _adapter( 'dockerhub', 'docker', 'docker_direct_v1', assignment_flow=_DIRECT_ASSIGNMENT_FLOW, validator=normalize_docker_direct_execution_snapshot, ), 'huggingface': _adapter( 'huggingface', 'huggingface', 'huggingface_space_v1', assignment_flow=_DIRECT_ASSIGNMENT_FLOW, validator=normalize_huggingface_space_execution_snapshot, ), }) ASSIGNMENT_SOURCE_ADAPTERS = MappingProxyType({ 'github': LEGACY_GITHUB_ASSIGNMENT_ADAPTER, **CORE_ASSIGNMENT_SOURCE_ADAPTERS, }) PROTOCOL2_NEW_CLAIM_SOURCES = frozenset(CORE_ASSIGNMENT_SOURCE_ADAPTERS) def _validate_assignment_source_adapters(): if tuple(CORE_ASSIGNMENT_SOURCE_ADAPTERS) != tuple(DISCOVERY_PRODUCER_SOURCES): raise RuntimeError('core assignment adapters do not match discovery producers') pairs = set() for key, adapter in ASSIGNMENT_SOURCE_ADAPTERS.items(): capability = adapter.package_capability identities = ( key, adapter.queue_source, adapter.worker_platform, adapter.planning_kind, ) if any(_CANONICAL_CAPABILITY_NAME.fullmatch(value) is None for value in identities): raise RuntimeError('remote assignment adapter identity is invalid') if key != adapter.queue_source or capability.as_dict() != { 'source': adapter.queue_source, 'platform': adapter.worker_platform, 'planning_kind': adapter.planning_kind, }: raise RuntimeError('remote assignment adapter capability is inconsistent') pair = (adapter.queue_source, adapter.worker_platform) if pair in pairs: raise RuntimeError('remote assignment adapter source/platform is duplicated') pairs.add(pair) if { source for source, adapter in ASSIGNMENT_SOURCE_ADAPTERS.items() if adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW } != set(PROTOCOL1_NEW_CLAIM_SOURCES): raise RuntimeError('protocol-1 assignment source registry is inconsistent') if { source for source, adapter in ASSIGNMENT_SOURCE_ADAPTERS.items() if adapter.assignment_flow is _DIRECT_ASSIGNMENT_FLOW } != {'dockerhub', 'huggingface'}: raise RuntimeError('direct assignment source registry is inconsistent') if { ( adapter.package_capability.source, adapter.package_capability.platform, adapter.package_capability.planning_kind, ) for adapter in ASSIGNMENT_SOURCE_ADAPTERS.values() } != set(KNOWN_WORKER_PACKAGE_CAPABILITIES): raise RuntimeError('worker package capabilities do not match source adapters') _validate_assignment_source_adapters() def assignment_source_adapter(source): source = str(source or '').strip().lower() adapter = ASSIGNMENT_SOURCE_ADAPTERS.get(source) if adapter is None: raise ValueError('unsupported remote assignment source') return adapter def _stable_id(device_id, request_id, label): payload = f'truf-worker-v1\0{int(device_id)}\0{request_id}\0{label}'.encode('ascii') return hashlib.sha256(payload).hexdigest()[:32] def _hash_file(path): digest = hashlib.sha256() with open(path, 'rb', buffering=0) as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b''): digest.update(chunk) return digest.hexdigest() class RemoteAssignmentBuilder: """Central remote admission builder routed through closed source adapters.""" def __init__( self, db_url, bundle_root, source_args, compatibility_profiles, supervisor_instance_id, *, assignment_ttl_seconds=DEFAULT_ASSIGNMENT_TTL_SECONDS, assignment_ttl_seconds_by_source=None, result_upload_body_timeout_seconds=DEFAULT_RESULT_UPLOAD_BODY_TIMEOUT_SECONDS, credential_refs=None, db_factory=ScannerDB, admission=reserve_v2_admission_with_recovery, planner=resolve_and_bind_git_claim, ): self.db_url = str(db_url or '') self.bundle_root = os.path.abspath(bundle_root) self.supervisor_instance_id = str(supervisor_instance_id or '') if not self.db_url or not self.supervisor_instance_id: raise ValueError('remote assignment builder requires database and supervisor identities') if ( type(assignment_ttl_seconds) is not int or not 60 <= assignment_ttl_seconds <= 7 * 24 * 60 * 60 ): raise ValueError('remote assignment lifetime must be between one minute and seven days') self.assignment_ttl_seconds = assignment_ttl_seconds raw_ttl_overrides = assignment_ttl_seconds_by_source if raw_ttl_overrides is None: raw_ttl_overrides = {} if ( type(raw_ttl_overrides) is not dict or set(raw_ttl_overrides) - set(CORE_ASSIGNMENT_SOURCE_ADAPTERS) ): raise ValueError('remote assignment lifetime overrides contain unsupported sources') self.assignment_ttl_seconds_by_source = {} for source, value in raw_ttl_overrides.items(): if type(value) is not int or not 60 <= value <= 7 * 24 * 60 * 60: raise ValueError('remote assignment lifetime override is outside its bounds') self.assignment_ttl_seconds_by_source[source] = value if ( type(result_upload_body_timeout_seconds) is not int or not 30 <= result_upload_body_timeout_seconds <= 24 * 60 * 60 ): raise ValueError('remote result upload body timeout is outside its bounds') self.result_upload_body_timeout_seconds = result_upload_body_timeout_seconds self.db_factory = db_factory self.admission = admission self.planner = planner self.source_args = {} self.source_adapters = {} for source, args in dict(source_args or {}).items(): source = str(source or '').strip().lower() adapter = assignment_source_adapter(source) adapter.validate_source_args(args) self.source_args[source] = args self.source_adapters[source] = adapter if not self.source_args: raise ValueError('remote assignment builder has no supported sources') raw_credential_refs = dict(credential_refs or {}) if set(raw_credential_refs) - set(self.source_args): raise ValueError('remote credential references contain unsupported sources') self.credential_refs = {} for source in self.source_args: reference = str(raw_credential_refs.get(source) or '') if len(reference) > 128 or '\x00' in reference: raise ValueError('remote credential reference is invalid') if ( self.source_adapters[source].assignment_flow is _DIRECT_ASSIGNMENT_FLOW and reference ): raise ValueError('remote direct source credential reference must be empty') self.credential_refs[source] = reference self.compatibility_profiles = {} for profile_name, value in dict(compatibility_profiles or {}).items(): profile_name = str(profile_name or '').strip() profile = dict(value or {}) if set(profile) - {'package_manifest', 'sources'}: raise ValueError('remote compatibility profile shape is invalid') package_value = profile.get('package_manifest') package = ( load_worker_package_manifest(package_value) if isinstance(package_value, (str, os.PathLike)) else normalize_worker_package_manifest(package_value) ) required_build = WorkerBuildCompatibility.from_mapping( worker_package_build_compatibility(package), ) package_capabilities = frozenset( ( item['source'], item['platform'], item['planning_kind'], ) for item in package['capabilities'] ) package_sources = frozenset( source for source, _platform, _planning in package_capabilities ) allowed_sources = frozenset( str(item or '').strip().lower() for item in (profile.get('sources') or package_sources) ) if ( not profile_name or not allowed_sources or not allowed_sources <= set(self.source_args) or any( ( self.source_adapters[source].package_capability.source, self.source_adapters[source].package_capability.platform, self.source_adapters[source].package_capability.planning_kind, ) not in package_capabilities for source in allowed_sources ) ): raise ValueError('remote compatibility profile has invalid sources') key = tuple(required_build.as_dict().values()) if key in self.compatibility_profiles: raise ValueError('remote compatibility profile identity is duplicated') self.compatibility_profiles[key] = { 'profile_name': profile_name, 'build': required_build, 'package_manifest': package, 'sources': allowed_sources, 'capabilities': package_capabilities, } if not self.compatibility_profiles: raise ValueError('remote assignment builder has no compatibility profiles') def compatibility_snapshot(self): profiles = [] for profile in sorted( self.compatibility_profiles.values(), key=lambda item: item['profile_name'], ): profiles.append({ 'profile_name': profile['profile_name'], **profile['build'].as_dict(), 'sources': sorted(profile['sources']), 'capabilities': [ { 'source': source, 'platform': platform, 'planning_kind': planning_kind, } for source, platform, planning_kind in sorted( profile['capabilities'], ) ], }) required_capabilities = [ self.source_adapters[source].package_capability.as_dict() for source in sorted(self.source_adapters) ] return { 'profiles': profiles, 'required_capabilities': required_capabilities, } @staticmethod def _queue_policy(args): return QueueDispositionPolicy( target_retry_max_attempts=int(getattr(args, 'target_retry_max_attempts', 3) or 3), target_retry_base_delay_sec=int(getattr(args, 'target_retry_base_delay_sec', 3600) or 3600), target_retry_max_delay_sec=int(getattr(args, 'target_retry_max_delay_sec', 86400) or 86400), target_timeout_retry_delay_sec=int(getattr(args, 'target_timeout_retry_delay_sec', 21600) or 21600), ) @staticmethod def _capacity(args): return { 'bundle_items': int(getattr(args, 'result_bundle_max_items', 10000)), 'bundle_bytes': int(getattr(args, 'result_bundle_max_total_bytes', 3 << 30)), 'projection_items': int(getattr(args, 'projection_backlog_max_items', 10000)), 'projection_bytes': int(getattr(args, 'projection_backlog_max_bytes', 2 << 30)), 'projection_headroom_bytes': int(getattr( args, 'projection_backlog_headroom_bytes', 0, )), 'keycheck_items': int(getattr(args, 'keycheck_queue_max_items', 100000)), 'keycheck_bytes': int(getattr(args, 'keycheck_queue_max_bytes', 512 << 20)), 'quarantine_items': int(getattr(args, 'pipeline_quarantine_max_items', 10000)), 'quarantine_bytes': int(getattr(args, 'pipeline_quarantine_max_bytes', 1 << 30)), } @staticmethod def _scan_policy(args): def setting(name, default): return getattr(args, name, getattr(scan_config, name, default)) return normalize_remote_scan_policy({ 'drop_detectors': setting('drop_detectors', ()), 'strict_git_provider_token_filter': bool(setting( 'strict_git_provider_token_filter', True, )), 'trufflehog_stdout_max_mb': int(setting('trufflehog_stdout_max_mb', 32)), 'trufflehog_stderr_max_mb': int(setting('trufflehog_stderr_max_mb', 8)), 'result_bundle_max_event_bytes': int(setting( 'result_bundle_max_event_bytes', 64 << 20, )), 'trufflehog_max_findings_per_target': int(setting( 'trufflehog_max_findings_per_target', 20000, )), 'trufflehog_job_memory_limit_bytes': int(setting( 'trufflehog_job_memory_limit_bytes', 0, )), 'trufflehog_windows_job_cpu_weight': int(setting( 'trufflehog_windows_job_cpu_weight', 0, )), 'trufflehog_windows_memory_priority': int(setting( 'trufflehog_windows_memory_priority', 0, )), 'trufflehog_diagnostic_max_lines': int(setting( 'trufflehog_diagnostic_max_lines', 2000, )), 'trufflehog_diagnostic_max_line_chars': int(setting( 'trufflehog_diagnostic_max_line_chars', 8192, )), 'trufflehog_diagnostic_max_line_bytes': int(setting( 'trufflehog_diagnostic_max_line_bytes', 8192, )), 'trufflehog_diagnostic_max_errors': int(setting( 'trufflehog_diagnostic_max_errors', 200, )), 'trufflehog_diagnostic_max_warnings': int(setting( 'trufflehog_diagnostic_max_warnings', 200, )), 'trufflehog_diagnostic_max_unclassified': int(setting( 'trufflehog_diagnostic_max_unclassified', 20, )), }) def _db(self, application_name): db = self.db_factory(db_url=self.db_url, initialize=False) if not db.enabled: db.close() raise RuntimeError('remote assignment PostgreSQL connection is unavailable') db.set_application_name(application_name) return db def _accept_staged(self, identity, claim, staged): ready_path = bundle_ready_path(self.bundle_root, claim['bundle_id']) relative_path = os.path.relpath(ready_path, self.bundle_root).replace(os.sep, '/') staged_relative_path = str(getattr(staged, 'relative_path', '') or '').replace( '\\', '/' ) if staged_relative_path and staged_relative_path != relative_path: raise RuntimeError('published remote bundle path conflicts with its reservation') digest = _hash_file(ready_path) reader = ResultBundleReader( ready_path, max_event_bytes=int(claim['declared_bundle_bytes']), ) metadata = reader.validate() effective_diagnostics = reader.effective_diagnostics() values = metadata.as_dict() values['effective_diagnostic_count'] = len(effective_diagnostics) values['effective_diagnostic_projection_version'] = ( DIAGNOSTIC_PROJECTION_VERSION ) values['effective_diagnostic_uids_sha256'] = ( ordered_diagnostic_uid_set_sha256(effective_diagnostics) ) values['relative_path'] = relative_path db = self._db('truf-worker-planning-result') try: return db.mark_result_bundle_ready( int(claim['reservation_id']), values, remote_acceptance={ 'device_id': int(identity['device_id']), 'payload_sha256': digest, 'token_sha256': str(identity['token_sha256']), }, bundle_capacity_bytes=self._capacity( self.source_args[claim['source']] )['bundle_bytes'], ) finally: db.close() def _recover_published(self, identity, claim): path = bundle_ready_path(self.bundle_root, claim['bundle_id']) if not os.path.lexists(path): return False metadata = ResultBundleReader( path, max_event_bytes=int(claim['declared_bundle_bytes']), ).validate() if metadata.header != BundleReservation.from_mapping(claim).header(): raise RuntimeError('published remote bundle conflicts with its reservation') self._accept_staged(identity, claim, metadata) return True def _bound_plan(self, identity, claim): db = self._db('truf-worker-plan-reconcile') try: return db.remote_bound_git_scan_plan( int(claim['reservation_id']), int(identity['device_id']), str(claim['claim_lease_token']), str(identity['token_sha256']), ) finally: db.close() def _reconcile_request(self, identity, request_id): db = self._db('truf-worker-request-reconcile') try: return db.reconcile_remote_assignment_request( request_id, int(identity['device_id']), str(identity['token_sha256']), ) finally: db.close() def _snapshot(self, adapter, args, compatibility, execution): planning = {'kind': adapter.planning_kind} if adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW: planning.update({ 'git_baseline_depth': max( 1, int(getattr(args, 'git_baseline_depth', 100) or 100), ), 'git_ref_resolution_attempts': max( 1, int(getattr(args, 'git_ref_resolution_attempts', 2) or 2), ), 'git_ref_resolution_timeout_sec': max( 0.1, float(getattr(args, 'git_ref_resolution_timeout_sec', 10) or 10), ), 'git_ref_resolution_max_bytes': max( 1024, int(getattr(args, 'git_ref_resolution_max_bytes', 1 << 20) or (1 << 20)), ), }) return adapter.validate_snapshot({ 'schema': 1, 'compatibility': compatibility.as_dict(), 'execution': execution, 'planning': planning, 'credential_ref': { 'source': adapter.queue_source, 'auth_entry': self.credential_refs[adapter.queue_source], }, }) def _rehydrate(self, client, claim, snapshot, snapshot_sha256=None): adapter = assignment_source_adapter(claim.get('source')) if ( self.source_adapters.get(adapter.queue_source) is not adapter or str(claim.get('platform') or '') != adapter.worker_platform or adapter.assignment_flow not in { _EXACT_GIT_ASSIGNMENT_FLOW, _DIRECT_ASSIGNMENT_FLOW, } ): raise RuntimeError('remote assignment source adapter is unavailable') snapshot = adapter.validate_snapshot(snapshot) if ( snapshot_sha256 is not None and str(snapshot_sha256) != remote_execution_snapshot_sha256(snapshot) ): raise ScanEventConflictError('remote execution snapshot changed after admission') required = ScanCompatibility.from_mapping(snapshot['compatibility']) required_build = WorkerBuildCompatibility.from_mapping({ 'protocol_version': required.protocol_version, 'bundle_format_version': required.bundle_format_version, 'platform_tag': required.platform_tag, 'code_manifest_sha256': required.code_manifest_sha256, 'detector_policy_sha256': required.detector_policy_sha256, }) if required.protocol_version not in {1, PROTOCOL_VERSION}: raise ScanEventConflictError( 'remote execution snapshot protocol is unsupported' ) validate_worker_build_compatibility( required_build, client, expected_protocol_version=required.protocol_version, ) source = adapter.queue_source args = self.source_args.get(source) if args is None or snapshot['credential_ref'] != { 'source': source, 'auth_entry': self.credential_refs.get(source), }: raise RuntimeError('remote assignment credential reference is unavailable') values = vars(args).copy() values.update(snapshot['planning']) values.pop('kind', None) planning_args = SimpleNamespace(**values) event_scan_options = dict(snapshot['execution']['scan_kwargs']) scan_kwargs = dict(event_scan_options) if adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW: scan_kwargs['token'] = str(getattr(args, 'token', '') or '') return ( adapter, planning_args, snapshot, scan_kwargs, event_scan_options, QueueDispositionPolicy(**snapshot['execution']['queue_policy']), dict(snapshot['execution']['limits']), dict(snapshot['execution']['scan_policy']), required, ) def _complete_claim( self, identity, client, claim, snapshot, *, snapshot_sha256=None, recovered_execution_plan=None, ): ( adapter, args, snapshot, scan_kwargs, event_scan_options, queue_policy, limits, scan_policy, required, ) = self._rehydrate(client, claim, snapshot, snapshot_sha256) device_id = int(identity['device_id']) status_db = self._db('truf-worker-claim-reconcile') try: status = status_db.remote_assignment_status( int(claim['reservation_id']), device_id, str(identity['token_sha256']), ) finally: status_db.close() if status and status.get('receipt_id'): return None if self._recover_published(identity, claim): return None if adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW: execution_target = str(claim.get('target') or '') elif adapter.planning_kind == 'docker_direct_v1': execution_target = normalize_docker_direct_execution_target( claim.get('target'), )['image'] elif adapter.planning_kind == 'huggingface_space_v1': execution_target = normalize_huggingface_space_execution_target( claim.get('target'), ) else: raise RuntimeError('remote assignment execution plan is unavailable') bound_plan = None if recovered_execution_plan is not None: if ( not isinstance(recovered_execution_plan, dict) or set(recovered_execution_plan) != { 'kind', 'execution_target', 'bound_plan', } or recovered_execution_plan['kind'] != adapter.planning_kind or str(recovered_execution_plan['execution_target']) != execution_target or ( adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW and not isinstance(recovered_execution_plan['bound_plan'], dict) ) or ( adapter.assignment_flow is _DIRECT_ASSIGNMENT_FLOW and recovered_execution_plan['bound_plan'] is not None ) ): raise ScanEventConflictError( 'recovered execution plan conflicts with its assignment' ) bound_plan = recovered_execution_plan['bound_plan'] if ( adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW and bound_plan is None ): bound_plan = self._bound_plan(identity, claim) if ( adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW and bound_plan is None ): try: bound_plan = self.planner( args, self.db_url, adapter.queue_source, claim, scan_kwargs, remote_credential={ 'device_id': device_id, 'token_sha256': str(identity['token_sha256']), }, ) except ScanEventConflictError: bound_plan = self._bound_plan(identity, claim) if bound_plan is None: raise except _GitResolutionFailure as failure: bound_plan = self._bound_plan(identity, claim) if bound_plan is None: result = git_resolution_failure_result(args, claim, scan_kwargs, failure) staged = stage_scan_result_in_scope( result, claim, self.bundle_root, event_scan_options, queue_policy, attempts=int(claim.get('attempts') or 0), candidate_max_items=limits['candidate_max_items'], candidate_max_bytes=limits['candidate_max_bytes'], require_s_drive=False, ) self._accept_staged(identity, claim, staged) return None assigned_scan_kwargs = dict(scan_kwargs) if adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW: assigned_scan_kwargs['git_plan'] = bound_plan execution_plan = { 'kind': adapter.planning_kind, 'execution_target': execution_target, 'bound_plan': bound_plan, } issued_at = str(claim.get('remote_issued_at') or '') deadline_at = str(claim.get('remote_expires_at') or '') try: issued = datetime.fromisoformat(issued_at) deadline = datetime.fromisoformat(deadline_at) except ValueError as exc: raise RuntimeError('remote assignment reservation deadline is invalid') from exc if ( issued.tzinfo is None or deadline.tzinfo is None or issued.utcoffset() != timedelta(0) or deadline.utcoffset() != timedelta(0) or issued.isoformat(timespec='seconds') != issued_at or deadline.isoformat(timespec='seconds') != deadline_at ): raise RuntimeError('remote assignment reservation deadline is invalid') assignment_ttl = (deadline - issued).total_seconds() if ( not assignment_ttl.is_integer() or not 60 <= assignment_ttl <= 7 * 24 * 60 * 60 ): raise RuntimeError('remote assignment reservation lifetime is invalid') target_scan_timeout = scan_kwargs.get('timeout_sec') if ( isinstance(target_scan_timeout, bool) or not isinstance(target_scan_timeout, (int, float)) or not float(target_scan_timeout).is_integer() or target_scan_timeout <= 0 ): raise RuntimeError('remote assignment target scan timeout is invalid') persisted_upload_timeout = claim.get( 'remote_result_upload_body_timeout_seconds' ) if persisted_upload_timeout is None: persisted_upload_timeout = self.result_upload_body_timeout_seconds if ( isinstance(persisted_upload_timeout, bool) or not isinstance(persisted_upload_timeout, int) or not 30 <= persisted_upload_timeout <= 24 * 60 * 60 ): raise RuntimeError('remote assignment result upload timeout is invalid') reservation = dict(claim) reservation.pop('remote_result_upload_body_timeout_seconds', None) return { 'reservation': reservation, 'deadlines': { 'target_scan_timeout_seconds': int(target_scan_timeout), 'result_upload_body_timeout_seconds': persisted_upload_timeout, 'assignment_ttl_seconds': int(assignment_ttl), 'assignment_issued_at': issued_at, 'assignment_deadline_at': deadline_at, }, 'compatibility': required.as_dict(), 'scan_kwargs': assigned_scan_kwargs, 'event_scan_options': event_scan_options, 'queue_policy': asdict(queue_policy), 'limits': limits, 'scan_policy': scan_policy, 'execution_snapshot': snapshot, 'execution_snapshot_sha256': remote_execution_snapshot_sha256(snapshot), 'execution_plan': execution_plan, } def __call__(self, identity, request_id, client_compatibility): client = WorkerBuildCompatibility.from_mapping(client_compatibility) reconciled = self._reconcile_request(identity, request_id) if reconciled is not None: if reconciled.get('state') != 'committed': return None if reconciled.get('receipt'): return {'resolution': dict(reconciled['receipt'])} execution_plan = reconciled.get('execution_plan') if execution_plan is None: if assignment_source_adapter( reconciled['claim'].get('source') ).assignment_flow is not _EXACT_GIT_ASSIGNMENT_FLOW: raise ScanEventConflictError( 'direct assignment execution plan is unavailable' ) execution_plan = { 'kind': 'exact_git_v1', 'execution_target': reconciled['claim']['target'], 'bound_plan': reconciled.get('git_plan'), } return self._complete_claim( identity, client, reconciled['claim'], reconciled['execution_snapshot'], snapshot_sha256=reconciled['execution_snapshot_sha256'], recovered_execution_plan=execution_plan, ) profile = self.compatibility_profiles.get(tuple(client.as_dict().values())) if profile is None: return {'no_assignment': {'reason': 'compatibility'}} required_build = validate_worker_build_compatibility(profile['build'], client) candidates = [] for source in sorted(profile['sources']): adapter = self.source_adapters[source] capability = adapter.package_capability if ( capability.source, capability.platform, capability.planning_kind, ) not in profile['capabilities']: raise RuntimeError( 'remote package capability changed after profile validation' ) values = vars(self.source_args[source]).copy() central_policy = str(values.get('trufflehog_config') or '') if ( not central_policy or not os.path.isfile(central_policy) or sha256_file(central_policy) != required_build.detector_policy_sha256 ): raise RuntimeError('remote package detector policy does not match central authority') values['trufflehog_config'] = PACKAGE_DETECTOR_POLICY source_args = SimpleNamespace(**values) source_scan_kwargs = validate_scan_kwargs( adapter.worker_platform, prepare_scan_options(source_args, 1, quiet=True), ) source_event_options = { key: value for key, value in source_scan_kwargs.items() if key != 'token' } source_queue_policy = self._queue_policy(source_args) source_limits = { 'candidate_max_items': int(getattr(source_args, 'keycheck_candidates_per_event', 2000)), 'candidate_max_bytes': int(getattr(source_args, 'keycheck_candidate_bytes_per_event', 2 << 20)), } source_scan_policy = self._scan_policy(source_args) effective_identity, execution = remote_execution_identity( adapter.worker_platform, source_scan_kwargs, source_event_options, source_queue_policy, source_limits, source_scan_policy, ) required = ScanCompatibility( protocol_version=required_build.protocol_version, bundle_format_version=required_build.bundle_format_version, platform_tag=required_build.platform_tag, code_manifest_sha256=required_build.code_manifest_sha256, effective_config_sha256=effective_identity, detector_policy_sha256=required_build.detector_policy_sha256, ) snapshot = self._snapshot(adapter, source_args, required, execution) candidates.append(( adapter, source_args, source_scan_kwargs, source_event_options, source_queue_policy, source_limits, source_scan_policy, required, snapshot, )) if not candidates: raise RuntimeError('remote compatibility profile has no central source configuration') device_id = int(identity['device_id']) producer = current_process_identity().as_dict() start = int(request_id[:16], 16) % len(candidates) ordered_candidates = candidates[start:] + candidates[:start] multiple_sources = len(ordered_candidates) > 1 no_work_reasons = [] for candidate in ordered_candidates: ( adapter, args, scan_kwargs, event_scan_options, queue_policy, limits, scan_policy, required, snapshot, ) = candidate admission_token = ( _stable_id( device_id, request_id, f'admission:{adapter.queue_source}', ) if multiple_sources else request_id ) if multiple_sources: recovered = self._reconcile_request(identity, admission_token) if recovered is not None: if recovered.get('state') != 'committed': continue if recovered.get('receipt'): return {'resolution': dict(recovered['receipt'])} execution_plan = recovered.get('execution_plan') if execution_plan is None: if assignment_source_adapter( recovered['claim'].get('source') ).assignment_flow is not _EXACT_GIT_ASSIGNMENT_FLOW: raise ScanEventConflictError( 'direct assignment execution plan is unavailable' ) execution_plan = { 'kind': 'exact_git_v1', 'execution_target': recovered['claim']['target'], 'bound_plan': recovered.get('git_plan'), } return self._complete_claim( identity, client, recovered['claim'], recovered['execution_snapshot'], snapshot_sha256=recovered[ 'execution_snapshot_sha256' ], recovered_execution_plan=execution_plan, ) max_event_bytes = int(getattr( args, 'result_bundle_max_event_bytes', 64 << 20, )) remote_reserve_bytes = int(getattr( args, 'remote_assignment_reserve_bytes', 2 << 20, )) remote_max_active = int(getattr( args, 'remote_assignment_max_active', 50, )) validate_v2_capacity_model( int(getattr(args, 'max_active_scans', 1) or 1), max_event_bytes, int(getattr(args, 'projection_backlog_max_bytes', 2 << 30)), int(getattr( args, 'projection_backlog_headroom_bytes', max_event_bytes * 2, )), ) outcome = self.admission( self.db_url, adapter.queue_source, adapter.worker_platform, producer, self.supervisor_instance_id, max_event_bytes, remote_reserve_bytes, limits['candidate_max_items'], limits['candidate_max_bytes'], lease_seconds=self.assignment_ttl_seconds_by_source.get( adapter.queue_source, self.assignment_ttl_seconds, ), max_attempts=int(getattr( args, 'target_retry_max_attempts', 3, ) or 3), capacity_limits=self._capacity(args), run_id=None, cycle_id=None, reservation_token=admission_token, bundle_id=_stable_id(device_id, admission_token, 'bundle'), scan_event_id=_stable_id(device_id, admission_token, 'event'), resolution_attempts=int(getattr( args, 'admission_resolution_attempts', 8, ) or 8), resolution_seconds=float(getattr( args, 'admission_resolution_seconds', 30, ) or 30), retry_delay=float(getattr( args, 'admission_resolution_retry_delay_sec', 0.2, ) or 0.2), claim_order=str(getattr( args, 'target_claim_order', 'oldest', ) or 'oldest'), final_cutover=True, reserved_bundle_bytes=remote_reserve_bytes, remote_max_active=remote_max_active, db_factory=self.db_factory, remote_assignment={ 'user_id': int(identity['user_id']), 'device_id': device_id, 'effective_config_sha256': required.effective_config_sha256, 'client_compat_sha256': canonical_json_sha256( client.as_dict() ), 'token_sha256': str(identity['token_sha256']), 'result_upload_body_timeout_seconds': ( self.result_upload_body_timeout_seconds ), 'execution_snapshot': snapshot, }, ) if outcome.claim is not None: return self._complete_claim( identity, client, outcome.claim, snapshot, ) reason = _ADMISSION_NO_WORK_REASONS.get(str( getattr(outcome, 'reason', None) or '' )) if reason: no_work_reasons.append(reason) if no_work_reasons: priority = ( 'dispatch_paused', 'assignment_cap', 'capacity', 'empty_queue', ) return {'no_assignment': {'reason': next( reason for reason in priority if reason in no_work_reasons )}} return None RemoteGitAssignmentBuilder = RemoteAssignmentBuilder