diff --git a/compute_horde/compute_horde/miner_client/organic.py b/compute_horde/compute_horde/miner_client/organic.py index e9e67cb47..d4737ec15 100644 --- a/compute_horde/compute_horde/miner_client/organic.py +++ b/compute_horde/compute_horde/miner_client/organic.py @@ -56,6 +56,7 @@ JobAcceptedReceiptPayload, JobFinishedReceiptPayload, JobStartedReceiptPayload, + clamp_ttl, ) from compute_horde.transport import ( AbstractTransport, @@ -107,11 +108,11 @@ class OrganicMinerClient(AbstractMinerClient[MinerToValidatorMessage, ValidatorT client = OrganicMinerClient(...) async with client: try: - await client.send_model(V0InitialJobRequest(...)) + await client.send_model(V0InitialJobRequest(...).signed(client.my_keypair)) msg = await client.job_accepted_future # handle msg - await client.send_model(V0JobRequest(...)) + await client.send_model(V0JobRequest(...).signed(client.my_keypair)) msg = client.job_finished_future # handle msg except MinerRejectedJob: @@ -160,7 +161,9 @@ def __init__( self.miner_machine_specs: MachineSpecs | None = None name = f"{miner_hotkey}({miner_address}:{miner_port})" - transport = transport or WSTransport(name, self.miner_url()) + transport = transport or WSTransport( + name, self.miner_url(), reconnect_greeting=self._serialized_authentication_message + ) super().__init__(name, transport) @cached_property @@ -344,6 +347,9 @@ def generate_authentication_message(self) -> ValidatorAuthForMiner: msg.signature = sign_blob(self.my_keypair, msg.blob_for_signing()) return msg + def _serialized_authentication_message(self) -> str: + return self.generate_authentication_message().model_dump_json() + def generate_job_started_receipt_message( self, executor_class: ExecutorClass, @@ -373,7 +379,7 @@ def generate_job_accepted_receipt_message( validator_hotkey=self.my_hotkey, timestamp=datetime.datetime.now(datetime.UTC), time_accepted=time_accepted, - ttl=ttl, + ttl=clamp_ttl(ttl), ) return V0JobAcceptedReceiptRequest( payload=receipt_payload, @@ -415,7 +421,8 @@ def generate_job_finished_receipt_message( validator_hotkey=self.my_hotkey, timestamp=datetime.datetime.now(datetime.UTC), time_started=time_started, - time_took_us=int(time_took_seconds * 1_000_000), + # Wall-clock durations go negative after a backwards clock step; the payload rejects that. + time_took_us=max(0, int(time_took_seconds * 1_000_000)), block_numbers=allowance_blocks, score_str=f"{score:.6f}", ) @@ -471,7 +478,7 @@ async def send_initial_job_request( if job_details.job_timing is not None else None, streaming_details=job_details.streaming_details, - ), + ).signed(self.my_keypair), ) async def send_model( @@ -484,7 +491,8 @@ async def send_model( async def connect(self) -> None: await super().connect() - await self.transport.send(self.generate_authentication_message().model_dump_json()) + # The default WSTransport authenticates every reopened connection via reconnect_greeting. + await self.transport.send(self._serialized_authentication_message()) @dataclass @@ -619,7 +627,7 @@ async def execute_organic_job_on_miner( volume=None, # Was sent in the initial request output_upload=job_details.output, artifacts_dir=job_details.artifacts_dir, - ) + ).signed(client.my_keypair) ) ## STAGE: Volume download diff --git a/compute_horde/compute_horde/protocol_consts.py b/compute_horde/compute_horde/protocol_consts.py index 72efa7e18..3e5395f3c 100644 --- a/compute_horde/compute_horde/protocol_consts.py +++ b/compute_horde/compute_horde/protocol_consts.py @@ -6,6 +6,10 @@ else: from backports.strenum import StrEnum # noqa: UP035 +# Miner hotkey that a validator puts in the authentication frames and receipts it sends to its +# trusted miner, which it reaches by address rather than by a hotkey registered on the subnet. +TRUSTED_MINER_FAKE_KEY = "0" * 48 + class JobParticipantType(StrEnum): UNKNOWN = "unknown" @@ -151,6 +155,7 @@ class JobRejectionReason(StrEnum): UNKNOWN = "unknown" BUSY = "busy" INVALID_SIGNATURE = "invalid_signature" + INVALID_RECEIPT = "invalid_receipt" NO_MINER_FOR_JOB = "no_miner_for_job" VALIDATOR_BLACKLISTED = "validator_blacklisted" diff --git a/compute_horde/compute_horde/protocol_messages.py b/compute_horde/compute_horde/protocol_messages.py index 64a69bdcb..99948ac2b 100644 --- a/compute_horde/compute_horde/protocol_messages.py +++ b/compute_horde/compute_horde/protocol_messages.py @@ -1,11 +1,13 @@ import enum -from typing import Annotated, Literal, TypeAlias +import json +from typing import Annotated, Any, Literal, Self, TypeAlias +import bittensor from compute_horde_core.executor_class import ExecutorClass from compute_horde_core.output_upload import OutputUpload from compute_horde_core.streaming import StreamingDetails from compute_horde_core.volume import Volume -from pydantic import AliasChoices, BaseModel, Field, JsonValue +from pydantic import AliasChoices, BaseModel, Field, JsonValue, PrivateAttr, TypeAdapter from typing_extensions import deprecated from compute_horde.base.docker import DockerRunOptionsPreset @@ -22,7 +24,7 @@ JobStartedReceiptPayload, Receipt, ) -from compute_horde.utils import MachineSpecs +from compute_horde.utils import MachineSpecs, canonical_json, sign_blob, verify_blob_signature # NOTE: # miner.ec = Miner's executor consumer @@ -78,8 +80,43 @@ class V0MainHotkeyMessage(BaseModel): main_hotkey: str | None = None +# validator -> miner.vc +class ValidatorSignedMessage(BaseModel): + """ + Message whose whole content is signed with the validator hotkey, so the miner can tell a + request authored by the authenticated validator from one injected into its socket. + + The signature covers the message's JSON object without validator_signature, in canonical + form. A message from parse_validator_to_miner_message is verified against the object as it + arrived rather than against its own fields, so a field known to only one side keeps the + signature valid. Copies keep that object too, so verify a received message before changing + it. + """ + + validator_signature: str | None = None + _received_content: dict[str, Any] | None = PrivateAttr(default=None) + + def blob_for_signing(self) -> str: + content = self._received_content + if content is None: + content = json.loads(self.model_dump_json(exclude={"validator_signature"})) + return canonical_json(content) + + def signed(self, keypair: bittensor.Keypair) -> Self: + return self.model_copy( + update={"validator_signature": sign_blob(keypair, self.blob_for_signing())} + ) + + def is_signed_by(self, validator_hotkey: str) -> bool: + if self.validator_signature is None: + return False + return verify_blob_signature( + validator_hotkey, self.blob_for_signing(), self.validator_signature + ) + + # validator -> miner.vc -> miner.ec -> executor -class V0InitialJobRequest(BaseModel): +class V0InitialJobRequest(ValidatorSignedMessage): class ExecutorTimingDetails(BaseModel): allowed_leeway: int download_time_limit: int @@ -172,7 +209,7 @@ class V0ExecutionDoneRequest(BaseModel): # validator -> miner.vc -> miner.ec -> executor -class V0JobRequest(BaseModel): +class V0JobRequest(ValidatorSignedMessage): message_type: Literal["V0JobRequest"] = "V0JobRequest" job_uuid: str executor_class: ExecutorClass @@ -264,6 +301,32 @@ class V0HordeFailedRequest(BaseModel): Field(discriminator="message_type"), ] +_validator_to_miner_message: TypeAdapter[ValidatorToMinerMessage] = TypeAdapter( + ValidatorToMinerMessage +) + + +def parse_validator_to_miner_message(raw: str | bytes) -> ValidatorToMinerMessage: + """ + Parses a frame sent by a validator. A signed message is validated from the same JSON object + that its signature is checked against, so with duplicate keys, too, the miner acts only on + content the validator signed. + + Raises ValueError, such as json.JSONDecodeError or pydantic.ValidationError, for a malformed + frame. + """ + try: + content = json.loads(raw) + message = _validator_to_miner_message.validate_python(content) + except RecursionError as ex: + raise ValueError("JSON is nested too deeply") from ex + if isinstance(message, ValidatorSignedMessage): + message._received_content = { + key: value for key, value in content.items() if key != "validator_signature" + } + return message + + MinerToExecutorMessage = Annotated[ GenericError | V0InitialJobRequest | V0JobRequest, Field(discriminator="message_type"), diff --git a/compute_horde/compute_horde/receipts/models.py b/compute_horde/compute_horde/receipts/models.py index 501336b20..92a9aa9fb 100644 --- a/compute_horde/compute_horde/receipts/models.py +++ b/compute_horde/compute_horde/receipts/models.py @@ -8,10 +8,12 @@ from compute_horde.executor_class import DEFAULT_EXECUTOR_CLASS from compute_horde.receipts import ReceiptType from compute_horde.receipts.schemas import ( + SCORE_STR_PATTERN, JobAcceptedReceiptPayload, JobFinishedReceiptPayload, JobStartedReceiptPayload, Receipt, + parse_score_str, ) @@ -181,8 +183,18 @@ class JobFinishedReceipt(AbstractReceipt): def time_took(self): return timedelta(microseconds=self.time_took_us) - def score(self): - return float(self.score_str) + def score(self) -> float: + """ + Raises ValueError for a persisted row whose score_str predates wire-format validation. + """ + return parse_score_str(self.score_str) + + @classmethod + def with_malformed_score(cls) -> models.QuerySet["JobFinishedReceipt"]: + """ + Rows for which score() raises ValueError. + """ + return cls.objects.exclude(score_str__regex=SCORE_STR_PATTERN.pattern) def to_receipt(self) -> Receipt: if self.miner_signature is None: diff --git a/compute_horde/compute_horde/receipts/schemas.py b/compute_horde/compute_horde/receipts/schemas.py index 1d1f50bc4..1a7413d7f 100644 --- a/compute_horde/compute_horde/receipts/schemas.py +++ b/compute_horde/compute_horde/receipts/schemas.py @@ -1,11 +1,28 @@ import datetime import enum import json +import re +import uuid from typing import Annotated, Literal import bittensor from compute_horde_core.executor_class import ExecutorClass -from pydantic import BaseModel, Field +from pydantic import AfterValidator, AwareDatetime, BaseModel, Field + +# Receipts are signed by their issuers but transferred from miners, which can forge them for any +# hotkey they control. The bounds below keep every accepted receipt storable in the receipt tables +# and keep its arithmetic (timestamp + ttl, score sums) finite and in range. +MAX_SCORE_STR_LENGTH = 256 +MAX_TTL_SECONDS = 30 * 24 * 60 * 60 +MAX_TIME_TOOK_US = 365 * 24 * 60 * 60 * 1_000_000 +MAX_BLOCK_NUMBER = 2**32 +MAX_BLOCK_NUMBERS = 10_000 +MIN_TIMESTAMP_YEAR = 2000 +MAX_TIMESTAMP_YEAR = 2100 + +# Executor seconds as validators write them: an unsigned decimal, no exponent, no sign, no NaN. +# Usable both as a Python pattern and as a POSIX regex in database lookups. +SCORE_STR_PATTERN = re.compile(r"^[0-9]+(\.[0-9]+)?$") class ReceiptType(enum.StrEnum): @@ -14,11 +31,64 @@ class ReceiptType(enum.StrEnum): JobFinishedReceipt = "JobFinishedReceipt" +def parse_score_str(score_str: str) -> float: + """ + Parse the executor-seconds cost carried by a job finished receipt. + + Raises ValueError for anything but an unsigned decimal of at most MAX_SCORE_STR_LENGTH + characters, so the result is always finite and non-negative. + """ + if len(score_str) > MAX_SCORE_STR_LENGTH or not SCORE_STR_PATTERN.fullmatch(score_str): + raise ValueError(f"Malformed score_str: {score_str[:MAX_SCORE_STR_LENGTH]!r}") + return float(score_str) + + +def _validate_score_str(value: str) -> str: + parse_score_str(value) + return value + + +def _validate_job_uuid(value: str) -> str: + try: + canonical = str(uuid.UUID(value)) + except ValueError: + raise ValueError(f"job_uuid is not a UUID: {value[:64]!r}") from None + if canonical != value: + raise ValueError(f"job_uuid must be in canonical form ({canonical}), got {value!r}") + return value + + +def _validate_timestamp(value: datetime.datetime) -> datetime.datetime: + if not MIN_TIMESTAMP_YEAR <= value.year < MAX_TIMESTAMP_YEAR: + raise ValueError( + f"Timestamp year must be within [{MIN_TIMESTAMP_YEAR}, {MAX_TIMESTAMP_YEAR}), " + f"got {value.isoformat()}" + ) + return value + + +JobUuid = Annotated[str, AfterValidator(_validate_job_uuid)] +ScoreStr = Annotated[str, AfterValidator(_validate_score_str)] +ReceiptTimestamp = Annotated[AwareDatetime, AfterValidator(_validate_timestamp)] +TtlSeconds = Annotated[int, Field(ge=0, le=MAX_TTL_SECONDS)] +BlockNumber = Annotated[int, Field(ge=0, le=MAX_BLOCK_NUMBER)] + + +def clamp_ttl(ttl: int) -> int: + """ + Clamp a ttl summed from job time limits to MAX_TTL_SECONDS, the largest that payloads accept. + + Validators and miners evict receipts two days after their timestamp, long before + MAX_TTL_SECONDS runs out, so clamping does not change how long a receipt counts as valid. + """ + return min(ttl, MAX_TTL_SECONDS) + + class BaseReceiptPayload(BaseModel): - job_uuid: str + job_uuid: JobUuid miner_hotkey: str validator_hotkey: str - timestamp: datetime.datetime # when the receipt was generated + timestamp: ReceiptTimestamp # when the receipt was generated def blob_for_signing(self): # pydantic v2 does not support sort_keys anymore. @@ -29,22 +99,24 @@ class JobStartedReceiptPayload(BaseReceiptPayload): receipt_type: Literal[ReceiptType.JobStartedReceipt] = ReceiptType.JobStartedReceipt executor_class: ExecutorClass is_organic: bool - ttl: int + ttl: TtlSeconds class JobAcceptedReceiptPayload(BaseReceiptPayload): receipt_type: Literal[ReceiptType.JobAcceptedReceipt] = ReceiptType.JobAcceptedReceipt - time_accepted: datetime.datetime - ttl: int # unused + time_accepted: ReceiptTimestamp + ttl: TtlSeconds class JobFinishedReceiptPayload(BaseReceiptPayload): receipt_type: Literal[ReceiptType.JobFinishedReceipt] = ReceiptType.JobFinishedReceipt - time_started: datetime.datetime - time_took_us: int # micro-seconds - score_str: str - block_numbers: list[int] = Field( - default_factory=list, description="List of block numbers used to pay for this job" + time_started: ReceiptTimestamp + time_took_us: Annotated[int, Field(ge=0, le=MAX_TIME_TOOK_US)] # micro-seconds + score_str: ScoreStr + block_numbers: list[BlockNumber] = Field( + default_factory=list, + max_length=MAX_BLOCK_NUMBERS, + description="List of block numbers used to pay for this job", ) @property @@ -53,7 +125,7 @@ def time_took(self): @property def score(self): - return float(self.score_str) + return parse_score_str(self.score_str) ReceiptPayload = Annotated[ @@ -75,25 +147,31 @@ class BadValidatorReceiptSignature(BadReceiptSignature): pass +def _signature_is_valid(hotkey: str, blob: str, signature: str) -> bool: + try: + return bool(bittensor.Keypair(ss58_address=hotkey).verify(blob, signature)) + except Exception: + # A hotkey or signature that cannot even be decoded is as bad as a wrong one. + return False + + class Receipt(BaseModel): payload: ReceiptPayload validator_signature: str miner_signature: str def verify_miner_signature(self, throw=False): - miner_keypair = bittensor.Keypair(ss58_address=self.payload.miner_hotkey) - is_valid = miner_keypair.verify(self.payload.blob_for_signing(), self.miner_signature) + is_valid = _signature_is_valid( + self.payload.miner_hotkey, self.payload.blob_for_signing(), self.miner_signature + ) if throw and not is_valid: raise BadMinerReceiptSignature(self) return is_valid def verify_validator_signature(self, throw=False): - validator_keypair = bittensor.Keypair(ss58_address=self.payload.validator_hotkey) - blob = self.payload.blob_for_signing() - try: - is_valid = validator_keypair.verify(blob, self.validator_signature) - except Exception: - is_valid = False + is_valid = _signature_is_valid( + self.payload.validator_hotkey, self.payload.blob_for_signing(), self.validator_signature + ) if throw and not is_valid: raise BadValidatorReceiptSignature(self) return is_valid diff --git a/compute_horde/compute_horde/receipts/transfer.py b/compute_horde/compute_horde/receipts/transfer.py index 13c8d1087..af8fc6e8f 100644 --- a/compute_horde/compute_horde/receipts/transfer.py +++ b/compute_horde/compute_horde/receipts/transfer.py @@ -1,4 +1,5 @@ import asyncio +import datetime import logging from collections import defaultdict from collections.abc import Sequence @@ -11,6 +12,7 @@ from aiohttp import ClientTimeout from asgiref.sync import sync_to_async from django.core.exceptions import ImproperlyConfigured +from django.db import DataError, IntegrityError from compute_horde.receipts.models import ReceiptModel, receipt_to_django_model from compute_horde.receipts.schemas import ( @@ -24,6 +26,29 @@ MinerInfo: TypeAlias = tuple[str, str, int] +# Errors a single receipt row can cause while being written. Connection-level database errors are +# deliberately not here: retrying rows one by one would not help and the transfer must fail loudly. +_ROW_STORAGE_ERRORS = (DataError, IntegrityError, ValueError, TypeError, OverflowError) + +_STORAGE_BATCH_SIZE = 1000 + +# Receipts age out by their own timestamp, so a future-dated receipt outlives the retention period. +# The skew tolerates the issuing validator's clock running ahead of the local one. +MAX_RECEIPT_TIMESTAMP_SKEW = datetime.timedelta(hours=1) + + +class FutureDatedReceipt(Exception): + """ + The receipt's timestamp is more than MAX_RECEIPT_TIMESTAMP_SKEW ahead of the local clock. + """ + + def __init__(self, receipt: Receipt): + super().__init__( + f"Receipt {receipt.payload.job_uuid} is dated {receipt.payload.timestamp.isoformat()}, " + f"more than {MAX_RECEIPT_TIMESTAMP_SKEW} ahead of the local clock" + ) + self.receipt = receipt + class TransferException(Exception): """ @@ -38,13 +63,11 @@ def __init__(self, miner: MinerInfo, page: int): class LineException(Exception): """ - A line exception means that, while processing a correct response, one of the receipt lines could not be processed. + A line exception means that, while processing a correct response, one of the receipt lines + could not be processed: it failed validation, signature verification or storage. """ - def __init__( - self, - cause: pydantic.ValidationError | BadMinerReceiptSignature | BadValidatorReceiptSignature, - ): + def __init__(self, cause: Exception): super().__init__(f"Transferred line skipped: {type(cause).__name__}: {cause}") self.__cause__ = self.cause = cause @@ -112,6 +135,15 @@ async def rate_limited_transfer( transfer_errors: list[TransferException] = [] line_errors: list[LineException] = [] + async def store_bucket(model_type: type[ReceiptModel]) -> None: + nonlocal total_receipts + bucket = receipts_by_type[model_type] + n_stored, storage_errors = await _store_receipts(model_type, bucket) + total_receipts += n_stored + line_errors.extend(storage_errors) + logger.info(f"Stored {n_stored} {model_type.__name__} receipts") + bucket.clear() + # Wait for transfer tasks in parallel, handle a batch of receipts as soon as any is available for transfer_task in asyncio.as_completed(transfer_tasks): try: @@ -120,23 +152,16 @@ async def rate_limited_transfer( successful_transfers += 1 for receipt in transferred_batch: model = receipt_to_django_model(receipt) - model_type = type(model) - bucket = receipts_by_type[model_type] - bucket.append(model) - if len(bucket) >= 1000: - await model_type.objects.abulk_create(bucket, ignore_conflicts=True) # type: ignore - total_receipts += len(bucket) - logger.info(f"Stored {len(bucket)} {model_type.__name__} receipts") - bucket.clear() + receipts_by_type[type(model)].append(model) + if len(receipts_by_type[type(model)]) >= _STORAGE_BATCH_SIZE: + await store_bucket(type(model)) except TransferException as e: transfer_errors.append(e) continue # Insert the remainder of the receipts - for model_type, bucket in receipts_by_type.items(): - await model_type.objects.abulk_create(bucket, ignore_conflicts=True) # type: ignore - total_receipts += len(bucket) - logger.info(f"Stored {len(bucket)} {model_type.__name__} receipts") + for model_type in receipts_by_type: + await store_bucket(model_type) return TransferResult(total_receipts, successful_transfers, transfer_errors, line_errors) @@ -219,9 +244,12 @@ def _to_valid_receipts( """ receipts: list[Receipt] = [] line_errors: list[LineException] = [] + latest_accepted_timestamp = datetime.datetime.now(datetime.UTC) + MAX_RECEIPT_TIMESTAMP_SKEW for line in BytesIO(received_bytes): try: receipt = Receipt.model_validate_json(line) + if receipt.payload.timestamp > latest_accepted_timestamp: + raise FutureDatedReceipt(receipt) receipt.verify_miner_signature(throw=True) receipt.verify_validator_signature(throw=True) receipts.append(receipt) @@ -253,4 +281,52 @@ def _to_valid_receipts( ) line_errors.append(LineException(e)) continue + except FutureDatedReceipt as e: + logger.info( + "Skipping %s dated %s, beyond the accepted clock skew: %s", + e.receipt.payload.receipt_type, + e.receipt.payload.timestamp.isoformat(), + e.receipt.payload.job_uuid, + ) + line_errors.append(LineException(e)) + continue return receipts, line_errors + + +async def _store_receipts( + model_type: type[ReceiptModel], receipts: list[ReceiptModel] +) -> tuple[int, list[LineException]]: + """ + Insert receipts in one batch, skipping the ones already stored. + + Validated receipts fit the receipt tables, so the batch insert is expected to succeed. If the + database still rejects it, the rows are retried one by one so that only the offending receipt + is lost: receipts of all miners share a batch and one miner must not be able to drop the others. + Returns the number of rows handed to the database and the errors of the skipped rows. + """ + try: + await model_type.objects.abulk_create(receipts, ignore_conflicts=True) # type: ignore + return len(receipts), [] + except _ROW_STORAGE_ERRORS as e: + logger.warning( + "Batch insert of %d %s receipts rejected by the database, retrying one by one: %s", + len(receipts), + model_type.__name__, + e, + ) + + n_stored = 0 + line_errors: list[LineException] = [] + for receipt in receipts: + try: + await model_type.objects.abulk_create([receipt], ignore_conflicts=True) # type: ignore + n_stored += 1 + except _ROW_STORAGE_ERRORS as e: + logger.info( + "Skipping %s receipt %s rejected by the database: %s", + model_type.__name__, + receipt.job_uuid, + e, + ) + line_errors.append(LineException(e)) + return n_stored, line_errors diff --git a/compute_horde/compute_horde/transport/ws.py b/compute_horde/compute_horde/transport/ws.py index 75bdf0781..35b884249 100644 --- a/compute_horde/compute_horde/transport/ws.py +++ b/compute_horde/compute_horde/transport/ws.py @@ -1,6 +1,7 @@ import asyncio import logging import random +from collections.abc import Callable import websockets @@ -16,7 +17,7 @@ class WSTransport(AbstractTransport): An instance needs to be started before sending or receiving messages and stopped after use. That's where the actual connection is established and terminated. Not started instances will raise a runtime error when trying to send or receive messages. ```python - transport = WSTransport("ws://example.com") + transport = WSTransport("miner", "ws://example.com") await transport.start() await transport.send("Hello, world!") @@ -29,13 +30,19 @@ class WSTransport(AbstractTransport): Continuously receiving messages can be done using an async for loop. ```python - transport = WSTransport("ws://example.com") + transport = WSTransport("miner", "ws://example.com") await transport.start() async for message in transport: print(message) - ``` + await transport.stop() + ``` + + `send()` and `receive()` reopen a dropped connection, which the peer sees as a new one. + A protocol that starts with a greeting, such as an authentication message, passes + `reconnect_greeting`: it is called for every reopened connection and its result is sent + before anything else. `start()` does not send it; the caller sends the first greeting. """ def __init__( @@ -46,12 +53,14 @@ def __init__( max_retries: int = 5, base_retry_delay: float = 1, retry_jitter: float = 1, + reconnect_greeting: Callable[[], str] | None = None, ): super().__init__(name) self.url = url self.base_retry_delay = base_retry_delay self.retry_jitter = retry_jitter self.max_retries = max_retries + self.reconnect_greeting = reconnect_greeting self.connect_lock = asyncio.Lock() self._ws: websockets.ClientConnection | None = None @@ -73,7 +82,13 @@ async def stop(self) -> None: if self._ws and self._ws.state is websockets.State.OPEN: await self._ws.close() - async def connect(self): + async def connect(self, *, greeting: Callable[[], str] | None = None) -> None: + """ + Opens the connection unless it is already open; `greeting` produces the message sent + on it before any other. An `OSError` or websockets error while producing or sending the + greeting counts as a failed attempt; any other error propagates. Either way the new + connection is closed. + """ if self._ws and self._ws.state is websockets.State.OPEN: return @@ -83,9 +98,7 @@ async def connect(self): while self.max_retries == 0 or attempt < self.max_retries: try: - self._ws = await websockets.connect( - self.url, max_size=50 * (2**20), ping_timeout=120 - ) # 50MB - ping timeout 2min in case of blocking job sends + self._ws = await self._open_connection(greeting) logger.info(f"Connected to {self.name} after {attempt} attempts") return except (websockets.WebSocketException, OSError): @@ -99,6 +112,27 @@ async def connect(self): f"Could not connect to {self.name} after {attempt} attempts in {time_took:0.2f} seconds" ) + async def _open_connection( + self, greeting: Callable[[], str] | None + ) -> websockets.ClientConnection: + ws = await websockets.connect( + self.url, max_size=50 * (2**20), ping_timeout=120 + ) # 50MB - ping timeout 2min in case of blocking job sends + if greeting is not None: + # The connection is not self._ws yet, so stop() would not close it. + try: + await ws.send(greeting()) + except BaseException: + await ws.close() + raise + return ws + + async def _reconnect(self) -> None: + # send() and receive() notice a dropped connection independently; the lock makes the + # later one find the connection already reopened instead of opening a second one. + async with self.connect_lock: + await self.connect(greeting=self.reconnect_greeting) + async def send(self, data: str | bytes) -> None: while True: try: @@ -110,7 +144,7 @@ async def send(self, data: str | bytes) -> None: return except (websockets.WebSocketException, OSError): logger.info(f"Could not send msg to {self.name}. Reconnecting...") - await self.connect() + await self._reconnect() async def receive(self) -> str | bytes: while True: @@ -120,4 +154,4 @@ async def receive(self) -> str | bytes: return msg except (websockets.WebSocketException, OSError): logger.info(f"Could not receive msg from {self.name}. Reconnecting...") - await self.connect() + await self._reconnect() diff --git a/compute_horde/compute_horde/utils.py b/compute_horde/compute_horde/utils.py index 1126bac1e..76bd5a50f 100644 --- a/compute_horde/compute_horde/utils.py +++ b/compute_horde/compute_horde/utils.py @@ -1,6 +1,7 @@ import asyncio import dataclasses import datetime +import json from collections.abc import Awaitable, Callable from functools import wraps from typing import Any, ParamSpec, TypeVar @@ -150,6 +151,25 @@ def sign_blob(kp: bittensor.Keypair, blob: str) -> str: return f"0x{kp.sign(blob).hex()}" +def verify_blob_signature(hotkey: str, blob: str, signature: str) -> bool: + """ + Checks a sign_blob signature of the blob against the hotkey; a malformed signature or + hotkey counts as invalid. + """ + try: + return bool(bittensor.Keypair(ss58_address=hotkey).verify(blob, signature)) + except ValueError: + return False + + +def canonical_json(data: Any) -> str: + """ + Serialises JSON-compatible data to text that depends only on its content: sorted keys, + no whitespace. + """ + return json.dumps(data, sort_keys=True, separators=(",", ":")) + + P = ParamSpec("P") R = TypeVar("R") diff --git a/compute_horde/tests/test_organic_miner_client.py b/compute_horde/tests/test_organic_miner_client.py index 95e5d8acc..958b6d677 100644 --- a/compute_horde/tests/test_organic_miner_client.py +++ b/compute_horde/tests/test_organic_miner_client.py @@ -1,13 +1,21 @@ +import asyncio +import datetime import uuid +from unittest import mock import bittensor import pytest -from pydantic import BaseModel +from compute_horde_core.executor_class import ExecutorClass +from compute_horde_core.streaming import StreamingDetails +from compute_horde_core.volume import InlineVolume +from freezegun import freeze_time +from pydantic import BaseModel, TypeAdapter from compute_horde.miner_client.organic import ( MinerRejectedJob, MinerReportedHordeFailed, MinerReportedJobFailed, + OrganicJobDetails, OrganicMinerClient, ) from compute_horde.protocol_consts import HordeFailureReason, JobParticipantType @@ -16,10 +24,15 @@ V0DeclineJobRequest, V0ExecutorReadyRequest, V0HordeFailedRequest, + V0InitialJobRequest, V0JobFailedRequest, V0JobFinishedRequest, + ValidatorAuthForMiner, + ValidatorToMinerMessage, ) +from compute_horde.receipts.schemas import JobFinishedReceiptPayload from compute_horde.transport import StubTransport +from tests.ws_server import ReconnectingWSTestServer JOB_UUID = "b4793a02-33a2-4a49-b4e2-4a7b903847e7" @@ -176,3 +189,103 @@ async def test_organic_miner_client__skip_different_job__final_future(final_msg, await miner_client.handle_message(final_msg) assert not miner_client.job_finished_future.done() + + +@pytest.mark.asyncio +async def test_organic_miner_client__reauthenticates_after_reconnect(keypair, unused_tcp_port): + # The reconnect handshake takes a second, so the renewed authentication message carries a + # later timestamp than the first one. + async with ReconnectingWSTestServer( + port=unused_tcp_port, drop_first_connection_after=1, handshake_delay=1 + ) as miner: + async with OrganicMinerClient( + miner_hotkey="mock", + miner_address=miner.host, + miner_port=miner.port, + job_uuid=JOB_UUID, + my_keypair=keypair, + ): + received = [await asyncio.wait_for(miner.received.get(), 3) for _ in range(2)] + + messages = [ + (connection, TypeAdapter(ValidatorToMinerMessage).validate_json(raw)) + for connection, raw in received + ] + auth = ValidatorAuthForMiner.model_construct( + validator_hotkey=keypair.ss58_address, + miner_hotkey="mock", + timestamp=mock.ANY, + signature=mock.ANY, + ) + assert messages == [(1, auth), (2, auth)] + (_, first_auth), (_, second_auth) = messages + assert second_auth.timestamp > first_auth.timestamp + assert keypair.verify(second_auth.blob_for_signing(), second_auth.signature) + + +@pytest.mark.asyncio +async def test_organic_miner_client__signs_initial_job_request(keypair): + miner_client = get_miner_client(keypair) + volume = InlineVolume(contents="aW5wdXQ=", relative_path="input") + job_details = OrganicJobDetails( + job_uuid=JOB_UUID, + executor_class=ExecutorClass.always_on__llm__a6000, + docker_image="mock", + job_timing=OrganicJobDetails.TimingDetails(1, 2, 3, 4, 5), + volume=volume, + streaming_details=StreamingDetails(public_key="certificate"), + ) + receipt_payload, receipt_signature = miner_client.generate_job_started_receipt_message( + ExecutorClass.always_on__llm__a6000, ttl=5 + ) + + await miner_client.send_initial_job_request(job_details, receipt_payload, receipt_signature) + + [sent] = [ + TypeAdapter(ValidatorToMinerMessage).validate_json(message) + for message in miner_client.transport.sent_messages + ] + assert sent == V0InitialJobRequest( + job_uuid=JOB_UUID, + executor_class=ExecutorClass.always_on__llm__a6000, + docker_image="mock", + timeout_seconds=300, + volume=volume, + job_started_receipt_payload=receipt_payload, + job_started_receipt_signature=receipt_signature, + streaming_details=StreamingDetails(public_key="certificate"), + executor_timing=V0InitialJobRequest.ExecutorTimingDetails( + allowed_leeway=1, + download_time_limit=2, + execution_time_limit=3, + upload_time_limit=4, + streaming_start_time_limit=5, + ), + validator_signature=sent.validator_signature, + ) + assert sent.is_signed_by(keypair.ss58_address) + + +@pytest.mark.asyncio +async def test_organic_miner_client__job_finished_receipt_clamps_negative_duration(keypair): + miner_client = get_miner_client(keypair) + started_at = datetime.datetime(2025, 3, 4, 12, 0, tzinfo=datetime.UTC) + + with freeze_time(started_at): + message = miner_client.generate_job_finished_receipt_message( + started_timestamp=started_at.timestamp(), + time_took_seconds=-0.5, + score=1.5, + allowance_blocks=[7, 8], + ) + + assert message.payload == JobFinishedReceiptPayload( + job_uuid=JOB_UUID, + miner_hotkey="mock", + validator_hotkey=keypair.ss58_address, + timestamp=started_at, + time_started=started_at, + time_took_us=0, + block_numbers=[7, 8], + score_str="1.500000", + ) diff --git a/compute_horde/tests/test_receipts_schemas.py b/compute_horde/tests/test_receipts_schemas.py new file mode 100644 index 000000000..9231051a9 --- /dev/null +++ b/compute_horde/tests/test_receipts_schemas.py @@ -0,0 +1,256 @@ +import datetime + +import pytest +from compute_horde_core.executor_class import ExecutorClass +from pydantic import ValidationError + +from compute_horde.receipts.schemas import ( + MAX_BLOCK_NUMBER, + MAX_BLOCK_NUMBERS, + MAX_SCORE_STR_LENGTH, + MAX_TIME_TOOK_US, + MAX_TTL_SECONDS, + BadMinerReceiptSignature, + BadValidatorReceiptSignature, + JobAcceptedReceiptPayload, + JobFinishedReceiptPayload, + JobStartedReceiptPayload, + Receipt, + parse_score_str, +) +from compute_horde.utils import sign_blob +from tests.utils import random_keypair + +JOB_UUID = "3342460e-4a99-438b-8757-795f4cb348dd" +MINER_HOTKEY = "5FHneW46xGXgs5mUiveU4sbTyGBzmstUspZC92UhjJM694ty" +VALIDATOR_HOTKEY = "5HGjWAeFDfFCWPsjFQdVV2Msvz2XtMktvgocEZcCj68kUMaw" +TIMESTAMP = datetime.datetime(2025, 3, 4, 5, 6, 7, 123456, tzinfo=datetime.UTC) + +COMMON_FIELDS = { + "job_uuid": JOB_UUID, + "miner_hotkey": MINER_HOTKEY, + "validator_hotkey": VALIDATOR_HOTKEY, + "timestamp": TIMESTAMP, +} +STARTED_FIELDS = { + **COMMON_FIELDS, + "executor_class": ExecutorClass.always_on__gpu_24gb, + "is_organic": True, + "ttl": 300, +} +ACCEPTED_FIELDS = { + **COMMON_FIELDS, + "time_accepted": datetime.datetime(2025, 3, 4, 5, 6, 8, tzinfo=datetime.UTC), + "ttl": 123, +} +FINISHED_FIELDS = { + **COMMON_FIELDS, + "time_started": datetime.datetime(2025, 3, 4, 5, 0, 0, tzinfo=datetime.UTC), + "time_took_us": 367_000_000, + "score_str": "1337.000000", + "block_numbers": [100, 101], +} + +# Blobs as produced before field validation was added. Signatures already made across the network +# are over these exact bytes, so they must not change. +STARTED_BLOB = ( + '{"executor_class": "always_on.gpu-24gb", "is_organic": true, ' + '"job_uuid": "3342460e-4a99-438b-8757-795f4cb348dd", ' + '"miner_hotkey": "5FHneW46xGXgs5mUiveU4sbTyGBzmstUspZC92UhjJM694ty", ' + '"receipt_type": "JobStartedReceipt", "timestamp": "2025-03-04T05:06:07.123456Z", ' + '"ttl": 300, "validator_hotkey": "5HGjWAeFDfFCWPsjFQdVV2Msvz2XtMktvgocEZcCj68kUMaw"}' +) +ACCEPTED_BLOB = ( + '{"job_uuid": "3342460e-4a99-438b-8757-795f4cb348dd", ' + '"miner_hotkey": "5FHneW46xGXgs5mUiveU4sbTyGBzmstUspZC92UhjJM694ty", ' + '"receipt_type": "JobAcceptedReceipt", "time_accepted": "2025-03-04T05:06:08Z", ' + '"timestamp": "2025-03-04T05:06:07.123456Z", "ttl": 123, ' + '"validator_hotkey": "5HGjWAeFDfFCWPsjFQdVV2Msvz2XtMktvgocEZcCj68kUMaw"}' +) +FINISHED_BLOB = ( + '{"block_numbers": [100, 101], "job_uuid": "3342460e-4a99-438b-8757-795f4cb348dd", ' + '"miner_hotkey": "5FHneW46xGXgs5mUiveU4sbTyGBzmstUspZC92UhjJM694ty", ' + '"receipt_type": "JobFinishedReceipt", "score_str": "1337.000000", ' + '"time_started": "2025-03-04T05:00:00Z", "time_took_us": 367000000, ' + '"timestamp": "2025-03-04T05:06:07.123456Z", ' + '"validator_hotkey": "5HGjWAeFDfFCWPsjFQdVV2Msvz2XtMktvgocEZcCj68kUMaw"}' +) + + +@pytest.mark.parametrize( + "payload_cls,fields,expected_blob", + [ + (JobStartedReceiptPayload, STARTED_FIELDS, STARTED_BLOB), + (JobAcceptedReceiptPayload, ACCEPTED_FIELDS, ACCEPTED_BLOB), + (JobFinishedReceiptPayload, FINISHED_FIELDS, FINISHED_BLOB), + ], +) +def test_blob_for_signing_is_unchanged(payload_cls, fields, expected_blob): + assert payload_cls(**fields).blob_for_signing() == expected_blob + + +@pytest.mark.parametrize( + "payload_cls,fields,overrides", + [ + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"job_uuid": "not-a-uuid"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"job_uuid": ""}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"job_uuid": JOB_UUID.upper()}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"job_uuid": JOB_UUID.replace("-", "")}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"job_uuid": "{" + JOB_UUID + "}"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "abc"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "nan"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "inf"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "-1"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "+1"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "1e5"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "1_000"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "0x10"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": ""}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": " 1"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "1\n"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "1."}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": ".5"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "١"}), + ( + JobFinishedReceiptPayload, + FINISHED_FIELDS, + {"score_str": "9" * (MAX_SCORE_STR_LENGTH + 1)}, + ), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"time_took_us": -1}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"time_took_us": MAX_TIME_TOOK_US + 1}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"time_took_us": 2**70}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"block_numbers": [-1]}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"block_numbers": [MAX_BLOCK_NUMBER + 1]}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"block_numbers": [[1]]}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"block_numbers": ["a"]}), + ( + JobFinishedReceiptPayload, + FINISHED_FIELDS, + {"block_numbers": [1] * (MAX_BLOCK_NUMBERS + 1)}, + ), + ( + JobFinishedReceiptPayload, + FINISHED_FIELDS, + {"timestamp": datetime.datetime(1999, 12, 31, 23, 59, 59, tzinfo=datetime.UTC)}, + ), + ( + JobFinishedReceiptPayload, + FINISHED_FIELDS, + {"timestamp": datetime.datetime(2100, 1, 1, tzinfo=datetime.UTC)}, + ), + ( + JobFinishedReceiptPayload, + FINISHED_FIELDS, + {"time_started": datetime.datetime(9999, 12, 31, tzinfo=datetime.UTC)}, + ), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"timestamp": datetime.datetime(2025, 3, 4)}), + ( + JobFinishedReceiptPayload, + FINISHED_FIELDS, + {"time_started": datetime.datetime(2025, 3, 4)}, + ), + (JobStartedReceiptPayload, STARTED_FIELDS, {"ttl": -1}), + (JobStartedReceiptPayload, STARTED_FIELDS, {"ttl": MAX_TTL_SECONDS + 1}), + (JobStartedReceiptPayload, STARTED_FIELDS, {"ttl": 2**40}), + (JobStartedReceiptPayload, STARTED_FIELDS, {"job_uuid": "not-a-uuid"}), + (JobAcceptedReceiptPayload, ACCEPTED_FIELDS, {"ttl": -1}), + (JobAcceptedReceiptPayload, ACCEPTED_FIELDS, {"ttl": MAX_TTL_SECONDS + 1}), + ( + JobAcceptedReceiptPayload, + ACCEPTED_FIELDS, + {"time_accepted": datetime.datetime(1, 1, 1, tzinfo=datetime.UTC)}, + ), + ], + ids=lambda value: repr(value)[:60] if isinstance(value, dict) else None, +) +def test_rejects_out_of_bounds_field(payload_cls, fields, overrides): + with pytest.raises(ValidationError): + payload_cls(**{**fields, **overrides}) + + +@pytest.mark.parametrize( + "payload_cls,fields,overrides", + [ + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "0"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "0.5"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "007.250"}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"score_str": "9" * MAX_SCORE_STR_LENGTH}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"time_took_us": 0}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"time_took_us": MAX_TIME_TOOK_US}), + (JobFinishedReceiptPayload, FINISHED_FIELDS, {"block_numbers": []}), + ( + JobFinishedReceiptPayload, + FINISHED_FIELDS, + {"block_numbers": [MAX_BLOCK_NUMBER] * MAX_BLOCK_NUMBERS}, + ), + ( + JobFinishedReceiptPayload, + FINISHED_FIELDS, + {"timestamp": datetime.datetime(2000, 1, 1, tzinfo=datetime.UTC)}, + ), + ( + JobFinishedReceiptPayload, + FINISHED_FIELDS, + {"timestamp": datetime.datetime(2099, 12, 31, 23, 59, 59, tzinfo=datetime.UTC)}, + ), + (JobStartedReceiptPayload, STARTED_FIELDS, {"ttl": 0}), + (JobStartedReceiptPayload, STARTED_FIELDS, {"ttl": MAX_TTL_SECONDS}), + (JobAcceptedReceiptPayload, ACCEPTED_FIELDS, {"ttl": MAX_TTL_SECONDS}), + ], + ids=lambda value: repr(value)[:60] if isinstance(value, dict) else None, +) +def test_accepts_boundary_field(payload_cls, fields, overrides): + payload = payload_cls(**{**fields, **overrides}) + assert payload.model_dump() == {**fields, **overrides, "receipt_type": payload.receipt_type} + + +@pytest.mark.parametrize( + "score_str,expected", + [("0", 0.0), ("0.5", 0.5), ("1337.000000", 1337.0), ("007", 7.0)], +) +def test_parse_score_str(score_str, expected): + assert parse_score_str(score_str) == expected + + +@pytest.mark.parametrize("score_str", ["nan", "inf", "-1", "abc", "1e5", ""]) +def test_parse_score_str_rejects_what_the_wire_format_rejects(score_str): + with pytest.raises(ValueError): + parse_score_str(score_str) + + +def test_signed_receipt_survives_json_round_trip(): + miner_keypair = random_keypair() + validator_keypair = random_keypair() + payload = JobFinishedReceiptPayload( + **{ + **FINISHED_FIELDS, + "miner_hotkey": miner_keypair.ss58_address, + "validator_hotkey": validator_keypair.ss58_address, + } + ) + blob = payload.blob_for_signing() + receipt = Receipt( + payload=payload, + miner_signature=sign_blob(miner_keypair, blob), + validator_signature=sign_blob(validator_keypair, blob), + ) + + received = Receipt.model_validate_json(receipt.model_dump_json()) + + assert received == receipt + assert received.verify_miner_signature(throw=True) + assert received.verify_validator_signature(throw=True) + + +def test_undecodable_hotkeys_count_as_bad_signatures(): + payload = JobStartedReceiptPayload( + **{**STARTED_FIELDS, "miner_hotkey": "not-an-address", "validator_hotkey": "x" * 300} + ) + receipt = Receipt(payload=payload, miner_signature="0x00", validator_signature="not-hex") + + assert receipt.verify_miner_signature() is False + assert receipt.verify_validator_signature() is False + with pytest.raises(BadMinerReceiptSignature): + receipt.verify_miner_signature(throw=True) + with pytest.raises(BadValidatorReceiptSignature): + receipt.verify_validator_signature(throw=True) diff --git a/compute_horde/tests/test_receipts_transfer.py b/compute_horde/tests/test_receipts_transfer.py index 47d70de1e..62fafac5b 100644 --- a/compute_horde/tests/test_receipts_transfer.py +++ b/compute_horde/tests/test_receipts_transfer.py @@ -1,13 +1,22 @@ from asyncio import Semaphore +from datetime import timedelta from functools import partial from unittest.mock import AsyncMock, Mock, patch +from uuid import uuid4 import aiohttp import pytest +from django.db import DataError +from django.utils import timezone from pydantic import ValidationError -from compute_horde.receipts.models import JobStartedReceipt -from compute_horde.receipts.transfer import ReceiptsTransfer +from compute_horde.receipts.models import JobFinishedReceipt, JobStartedReceipt +from compute_horde.receipts.schemas import JobFinishedReceiptPayload, Receipt, ReceiptType +from compute_horde.receipts.transfer import ( + MAX_RECEIPT_TIMESTAMP_SKEW, + FutureDatedReceipt, + ReceiptsTransfer, +) from compute_horde.receipts.transfer_checkpoints import TransferCheckpointBackend from compute_horde.utils import sign_blob from tests.utils import random_keypair, random_receipt @@ -92,6 +101,127 @@ async def test_rejects_bad_signatures(miner_keypair, validator_keypair): assert await JobStartedReceipt.objects.acount() == 0 +def malformed_finished_receipt_line(miner_keypair, validator_keypair, **overrides) -> bytes: + """ + A JobFinishedReceipt line signed by both parties, built without validation so that it can + carry values the wire format rejects. + """ + payload = JobFinishedReceiptPayload.model_construct( + **{ + "receipt_type": ReceiptType.JobFinishedReceipt, + "job_uuid": str(uuid4()), + "miner_hotkey": miner_keypair.ss58_address, + "validator_hotkey": validator_keypair.ss58_address, + "timestamp": timezone.now(), + "time_started": timezone.now(), + "time_took_us": 1_000_000, + "score_str": "1.000000", + "block_numbers": [1, 2], + **overrides, + } + ) + blob = payload.blob_for_signing() + receipt = Receipt.model_construct( + payload=payload, + miner_signature=sign_blob(miner_keypair, blob), + validator_signature=sign_blob(validator_keypair, blob), + ) + return receipt.model_dump_json().encode() + + +async def stored_started_job_uuids() -> list[str]: + return sorted( + [str(u) async for u in JobStartedReceipt.objects.values_list("job_uuid", flat=True)] + ) + + +@pytest.mark.asyncio +@pytest.mark.django_db(transaction=True) +async def test_rejects_malformed_signed_receipts_and_keeps_the_rest( + miner_keypair, validator_keypair +): + good_receipts = [random_receipt(miner_keypair, validator_keypair) for _ in range(4)] + lines_to_serve = [ + good_receipts[0].model_dump_json().encode(), + malformed_finished_receipt_line(miner_keypair, validator_keypair, score_str="nan"), + good_receipts[1].model_dump_json().encode(), + malformed_finished_receipt_line(miner_keypair, validator_keypair, job_uuid="not-a-uuid"), + good_receipts[2].model_dump_json().encode(), + malformed_finished_receipt_line(miner_keypair, validator_keypair, time_took_us=2**70), + malformed_finished_receipt_line(miner_keypair, validator_keypair, block_numbers=[-1]), + good_receipts[3].model_dump_json().encode(), + ] + + session = Mock() + session.get = AsyncMock() + session.get.return_value.status = 200 + session.get.return_value.read = AsyncMock(return_value=b"\n".join(lines_to_serve)) + + result = await do_transfer(session=session) + + assert result.n_receipts == 4 + assert len(result.transfer_errors) == 0 + assert [type(e.cause) for e in result.line_errors] == [ValidationError] * 4 + assert await JobFinishedReceipt.objects.acount() == 0 + assert await stored_started_job_uuids() == sorted(r.payload.job_uuid for r in good_receipts) + + +@pytest.mark.asyncio +@pytest.mark.django_db(transaction=True) +async def test_rejects_receipts_dated_beyond_the_clock_skew(miner_keypair, validator_keypair): + now = timezone.now() + within_skew = random_receipt( + miner_keypair, validator_keypair, timestamp=now + MAX_RECEIPT_TIMESTAMP_SKEW / 2 + ) + beyond_skew = random_receipt( + miner_keypair, validator_keypair, timestamp=now + timedelta(days=3) + ) + + session = Mock() + session.get = AsyncMock() + session.get.return_value.status = 200 + session.get.return_value.read = AsyncMock( + return_value=b"\n".join(r.model_dump_json().encode() for r in [within_skew, beyond_skew]) + ) + + result = await do_transfer(session=session) + + assert result.n_receipts == 1 + assert len(result.transfer_errors) == 0 + assert [type(e.cause) for e in result.line_errors] == [FutureDatedReceipt] + assert await stored_started_job_uuids() == [within_skew.payload.job_uuid] + + +@pytest.mark.asyncio +@pytest.mark.django_db(transaction=True) +async def test_row_rejected_by_database_does_not_lose_the_batch(miner_keypair, validator_keypair): + receipts = [random_receipt(miner_keypair, validator_keypair) for _ in range(4)] + poisoned_job_uuid = receipts[1].payload.job_uuid + original_abulk_create = JobStartedReceipt.objects.abulk_create + + async def abulk_create(objs, **kwargs): + if any(str(obj.job_uuid) == poisoned_job_uuid for obj in objs): + raise DataError("integer out of range") + return await original_abulk_create(objs, **kwargs) + + session = Mock() + session.get = AsyncMock() + session.get.return_value.status = 200 + session.get.return_value.read = AsyncMock( + return_value=b"\n".join(r.model_dump_json().encode() for r in receipts) + ) + + with patch.object(JobStartedReceipt.objects, "abulk_create", abulk_create): + result = await do_transfer(session=session) + + assert result.n_receipts == 3 + assert len(result.transfer_errors) == 0 + assert [type(e.cause) for e in result.line_errors] == [DataError] + assert await stored_started_job_uuids() == sorted( + r.payload.job_uuid for r in receipts if r.payload.job_uuid != poisoned_job_uuid + ) + + @pytest.mark.asyncio @pytest.mark.django_db(transaction=True) async def test_handles_failures(miner_keypair, mocked_checkpoint_backend): diff --git a/compute_horde/tests/test_run_organic_job.py b/compute_horde/tests/test_run_organic_job.py index 187b72bec..740175005 100644 --- a/compute_horde/tests/test_run_organic_job.py +++ b/compute_horde/tests/test_run_organic_job.py @@ -1,6 +1,9 @@ +import datetime + import pytest from compute_horde_core.executor_class import ExecutorClass from compute_horde_core.streaming import StreamingDetails +from freezegun import freeze_time from pydantic import TypeAdapter from compute_horde.executor_class import EXECUTOR_CLASS, ExecutorClassSpec @@ -30,6 +33,7 @@ ValidatorAuthForMiner, ValidatorToMinerMessage, ) +from compute_horde.receipts.schemas import MAX_TTL_SECONDS, JobAcceptedReceiptPayload from compute_horde.transport import StubTransport from compute_horde.transport.base import TransportConnectionError @@ -55,24 +59,26 @@ async def start(self): raise self.exception_to_raise +@pytest.fixture +def successful_job_flow() -> list[str]: + return [ + V0AcceptJobRequest(job_uuid=JOB_UUID).model_dump_json(), + V0ExecutorReadyRequest(job_uuid=JOB_UUID).model_dump_json(), + V0VolumesReadyRequest(job_uuid=JOB_UUID).model_dump_json(), + V0ExecutionDoneRequest(job_uuid=JOB_UUID).model_dump_json(), + V0JobFinishedRequest( + job_uuid=JOB_UUID, + docker_process_stdout="stdout", + docker_process_stderr="stderr", + artifacts={}, + ).model_dump_json(), + ] + + @pytest.mark.asyncio @pytest.mark.django_db(transaction=True) -async def test_run_organic_job__success(keypair): - mock_transport = MinerStubTransport( - "mock", - [ - V0AcceptJobRequest(job_uuid=JOB_UUID).model_dump_json(), - V0ExecutorReadyRequest(job_uuid=JOB_UUID).model_dump_json(), - V0VolumesReadyRequest(job_uuid=JOB_UUID).model_dump_json(), - V0ExecutionDoneRequest(job_uuid=JOB_UUID).model_dump_json(), - V0JobFinishedRequest( - job_uuid=JOB_UUID, - docker_process_stdout="stdout", - docker_process_stderr="stderr", - artifacts={}, - ).model_dump_json(), - ], - ) +async def test_run_organic_job__success(keypair, successful_job_flow): + mock_transport = MinerStubTransport("mock", successful_job_flow) client = OrganicMinerClient( miner_hotkey="mock", miner_address="0.0.0.0", @@ -101,6 +107,66 @@ async def test_run_organic_job__success(keypair): V0JobRequest, V0JobFinishedReceiptRequest, ] + _, initial_job_request, _, job_request, _ = mock_transport.sent_models + assert initial_job_request.is_signed_by(keypair.ss58_address) + assert job_request == V0JobRequest( + job_uuid=JOB_UUID, + executor_class=ExecutorClass.always_on__llm__a6000, + docker_image="mock", + docker_run_options_preset="nvidia_all", + docker_run_cmd=[], + env={}, + volume=None, + output_upload=None, + artifacts_dir=None, + validator_signature=job_request.validator_signature, + ) + assert job_request.is_signed_by(keypair.ss58_address) + + +@pytest.mark.asyncio +@pytest.mark.django_db(transaction=True) +async def test_run_organic_job__clamps_job_accepted_receipt_ttl(keypair, successful_job_flow): + accepted_at = datetime.datetime(2025, 3, 4, 12, 0, tzinfo=datetime.UTC) + mock_transport = MinerStubTransport("mock", successful_job_flow) + client = OrganicMinerClient( + miner_hotkey="mock", + miner_address="0.0.0.0", + miner_port=1234, + job_uuid=JOB_UUID, + my_keypair=keypair, + transport=mock_transport, + ) + job_details = OrganicJobDetails( + job_uuid=JOB_UUID, + executor_class=ExecutorClass.always_on__llm__a6000, + docker_image="mock", + job_timing=OrganicJobDetails.TimingDetails( + allowed_leeway=5, + download_time_limit=5, + execution_time_limit=MAX_TTL_SECONDS, + upload_time_limit=5, + streaming_start_time_limit=0, + ), + ) + with freeze_time(accepted_at, real_asyncio=True): + await execute_organic_job_on_miner( + client, job_details, reservation_time_limit=2, executor_startup_time_limit=2 + ) + + [receipt] = [ + model + for model in mock_transport.sent_models + if isinstance(model, V0JobAcceptedReceiptRequest) + ] + assert receipt.payload == JobAcceptedReceiptPayload( + job_uuid=JOB_UUID, + miner_hotkey="mock", + validator_hotkey=keypair.ss58_address, + timestamp=accepted_at, + time_accepted=accepted_at, + ttl=MAX_TTL_SECONDS, + ) @pytest.mark.asyncio diff --git a/compute_horde/tests/test_transport.py b/compute_horde/tests/test_transport.py index c9c895487..66c120225 100644 --- a/compute_horde/tests/test_transport.py +++ b/compute_horde/tests/test_transport.py @@ -1,46 +1,28 @@ import asyncio +import contextlib +import itertools import pytest import pytest_asyncio import websockets from compute_horde.transport import WSTransport +from tests.ws_server import AbstractWSTestServer, ReconnectingWSTestServer, wait_until -class WSTestServer: - host = "localhost" - +class WSTestServer(AbstractWSTestServer): def __init__(self, port: int) -> None: - self.ws_server: websockets.Server | None = None + super().__init__(port) self.received: asyncio.Queue[str] = asyncio.Queue() - self.port = port - async def srv(self, ws: websockets.ServerConnection): + async def _serve(self, ws: websockets.ServerConnection) -> None: async for message in ws: await self.received.put(message) - async def start(self): - self.ws_server = await websockets.serve(self.srv, self.host, self.port) - - async def stop(self): - self.ws_server.close() - await self.ws_server.wait_closed() - - async def __aenter__(self): - await self.start() - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.stop() - @property def connection(self) -> websockets.ServerConnection | None: - assert len(self.ws_server.connections) <= 1 - - try: - return next(iter(self.ws_server.connections)) - except StopIteration: - return None + assert len(self.connections) <= 1 + return next(iter(self.connections), None) @property def is_connected(self) -> bool: @@ -55,12 +37,7 @@ async def server(unused_tcp_port): @pytest.fixture def ws_transport(server): - return WSTransport( - "test", - f"ws://{server.host}:{server.port}", - base_retry_delay=0.1, - retry_jitter=0.1, - ) + return WSTransport("test", server.url, base_retry_delay=0.1, retry_jitter=0.1) @pytest.mark.asyncio @@ -119,3 +96,82 @@ async def test_receive_reconnect(server: WSTestServer, ws_transport: WSTransport await server.connection.send("foo") assert await receive_task == "foo" + + +@pytest.fixture +def greeting(): + numbers = itertools.count(1) + return lambda: f"greeting {next(numbers)}" + + +def greeting_transport(server: ReconnectingWSTestServer, greeting) -> WSTransport: + return WSTransport( + "test", server.url, base_retry_delay=0.1, retry_jitter=0.1, reconnect_greeting=greeting + ) + + +@pytest.mark.asyncio +async def test_receive_reconnect_sends_greeting_first(unused_tcp_port, greeting): + async with ReconnectingWSTestServer(port=unused_tcp_port) as server: + ws_transport = greeting_transport(server, greeting) + await ws_transport.start() + + receive_task = asyncio.create_task(ws_transport.receive()) + assert await asyncio.wait_for(server.received.get(), 1) == (2, "greeting 1") + await server.connection.send("foo") + + assert await receive_task == "foo" + await ws_transport.stop() + + +@pytest.mark.asyncio +async def test_send_reconnect_sends_greeting_first(unused_tcp_port, greeting): + async with ReconnectingWSTestServer(port=unused_tcp_port) as server: + ws_transport = greeting_transport(server, greeting) + await ws_transport.start() + await wait_until(lambda: ws_transport.ws.state is not websockets.State.OPEN) + + await ws_transport.send("foo") + + assert await asyncio.wait_for(server.received.get(), 1) == (2, "greeting 1") + assert await asyncio.wait_for(server.received.get(), 1) == (2, "foo") + await ws_transport.stop() + + +@pytest.mark.asyncio +async def test_failing_greeting_closes_reopened_connection(unused_tcp_port): + def greeting() -> str: + raise ValueError("cannot sign the greeting") + + async with ReconnectingWSTestServer(port=unused_tcp_port) as server: + ws_transport = greeting_transport(server, greeting) + await ws_transport.start() + + with pytest.raises(ValueError, match="cannot sign the greeting"): + await ws_transport.receive() + await ws_transport.stop() + + assert server.connection_count == 2 + await wait_until(lambda: not server.connections) + + +@pytest.mark.asyncio +async def test_concurrent_reconnects_open_one_connection(unused_tcp_port, greeting): + async with ReconnectingWSTestServer(port=unused_tcp_port, handshake_delay=0.3) as server: + ws_transport = greeting_transport(server, greeting) + await ws_transport.start() + await wait_until(lambda: ws_transport.ws.state is not websockets.State.OPEN) + + receive_task = asyncio.create_task(ws_transport.receive()) + await wait_until(lambda: server.delayed_handshakes == 1) + await ws_transport.send("foo") + + receive_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await receive_task + await ws_transport.stop() + + assert server.connection_count == 2 + assert await asyncio.wait_for(server.received.get(), 1) == (2, "greeting 1") + assert await asyncio.wait_for(server.received.get(), 1) == (2, "foo") + await wait_until(lambda: not server.connections) diff --git a/compute_horde/tests/test_utils.py b/compute_horde/tests/test_utils.py index 90826229f..b29c10803 100644 --- a/compute_horde/tests/test_utils.py +++ b/compute_horde/tests/test_utils.py @@ -1,7 +1,7 @@ import pytest from freezegun import freeze_time -from compute_horde.utils import Timer +from compute_horde.utils import Timer, canonical_json, sign_blob, verify_blob_signature def test_timer_none(): @@ -17,3 +17,19 @@ def test_timer_timeout(): frozen_time.move_to("2024-01-01 00:00:30") assert timer.passed_time() == 30.0 assert timer.time_left() == 270.0 + + +def test_canonical_json_depends_only_on_content(): + assert canonical_json({"b": [1, {"d": None, "c": "x"}], "a": 1.5}) == ( + '{"a":1.5,"b":[1,{"c":"x","d":null}]}' + ) + assert canonical_json({"a": 1, "b": 2}) == canonical_json({"b": 2, "a": 1}) + + +def test_verify_blob_signature(keypair): + signature = sign_blob(keypair, "blob") + + assert verify_blob_signature(keypair.ss58_address, "blob", signature) + assert not verify_blob_signature(keypair.ss58_address, "other blob", signature) + assert not verify_blob_signature(keypair.ss58_address, "blob", "gibberish") + assert not verify_blob_signature("not an ss58 address", "blob", signature) diff --git a/compute_horde/tests/test_validator_signed_messages.py b/compute_horde/tests/test_validator_signed_messages.py new file mode 100644 index 000000000..472870b63 --- /dev/null +++ b/compute_horde/tests/test_validator_signed_messages.py @@ -0,0 +1,335 @@ +import datetime +import json +import uuid + +import pytest +from compute_horde_core.executor_class import ExecutorClass +from pydantic import TypeAdapter, create_model + +from compute_horde.executor_class import DEFAULT_EXECUTOR_CLASS +from compute_horde.protocol_messages import ( + GenericError, + V0InitialJobRequest, + V0JobAcceptedReceiptRequest, + V0JobFinishedReceiptRequest, + V0JobRequest, + V0MainHotkeyMessage, + ValidatorAuthForMiner, + ValidatorSignedMessage, + ValidatorToMinerMessage, + parse_validator_to_miner_message, +) +from compute_horde.receipts.schemas import JobStartedReceiptPayload +from compute_horde.utils import canonical_json, sign_blob + +JOB_UUID = "b4793a02-33a2-4a49-b4e2-4a7b903847e7" +OTHER_EXECUTOR_CLASS = next(c for c in ExecutorClass if c != DEFAULT_EXECUTOR_CLASS) + + +@pytest.fixture +def initial_job_request(validator_keypair, miner_keypair) -> V0InitialJobRequest: + receipt = JobStartedReceiptPayload( + job_uuid=JOB_UUID, + miner_hotkey=miner_keypair.ss58_address, + validator_hotkey=validator_keypair.ss58_address, + timestamp=datetime.datetime(2025, 1, 2, 3, 4, 5, 678901, tzinfo=datetime.UTC), + executor_class=DEFAULT_EXECUTOR_CLASS, + is_organic=True, + ttl=30, + ) + return V0InitialJobRequest( + job_uuid=JOB_UUID, + executor_class=DEFAULT_EXECUTOR_CLASS, + docker_image="image:v1", + timeout_seconds=60, + volume={"volume_type": "inline", "contents": "aW5wdXQ=", "relative_path": "input"}, + job_started_receipt_payload=receipt, + job_started_receipt_signature=sign_blob(validator_keypair, receipt.blob_for_signing()), + streaming_details={"public_key": "certificate"}, + executor_timing={ + "allowed_leeway": 1, + "download_time_limit": 2, + "execution_time_limit": 3, + "upload_time_limit": 4, + "streaming_start_time_limit": 5, + }, + ) + + +@pytest.fixture +def job_request() -> V0JobRequest: + return V0JobRequest( + job_uuid=JOB_UUID, + executor_class=DEFAULT_EXECUTOR_CLASS, + docker_image="image:v1", + raw_script="print(1)", + docker_run_options_preset="nvidia_all", + docker_run_cmd=["python", "-c", "print(1)"], + env={"B": "2", "A": "1"}, + volume={"volume_type": "inline", "contents": "aW5wdXQ="}, + output_upload={ + "output_upload_type": "zip_and_http_post", + "url": "http://localhost/upload", + "form_fields": {"key": "value"}, + }, + artifacts_dir="/artifacts", + ) + + +@pytest.fixture(params=["initial_job_request", "job_request"]) +def message(request) -> ValidatorSignedMessage: + return request.getfixturevalue(request.param) + + +INITIAL_JOB_REQUEST_CHANGES = { + "message_type": "V0JobRequest", + "job_uuid": str(uuid.uuid4()), + "executor_class": OTHER_EXECUTOR_CLASS, + "docker_image": "image:v2", + "timeout_seconds": 61, + "volume": None, + "job_started_receipt_payload": JobStartedReceiptPayload( + job_uuid=JOB_UUID, + miner_hotkey="miner", + validator_hotkey="validator", + timestamp=datetime.datetime(2025, 1, 2, 3, 4, 5, 678901, tzinfo=datetime.UTC), + executor_class=DEFAULT_EXECUTOR_CLASS, + is_organic=False, + ttl=30, + ), + "job_started_receipt_signature": "0x" + "00" * 64, + "streaming_details": None, + "executor_timing": None, +} +JOB_REQUEST_CHANGES = { + "message_type": "V0InitialJobRequest", + "job_uuid": str(uuid.uuid4()), + "executor_class": OTHER_EXECUTOR_CLASS, + "docker_image": "image:v2", + "raw_script": None, + "docker_run_options_preset": "none", + "docker_run_cmd": ["python", "-c", "print(2)"], + "env": {"A": "1", "B": "2", "C": "3"}, + "volume": None, + "output_upload": None, + "artifacts_dir": None, +} +FIELD_CHANGES = [ + pytest.param("initial_job_request", field, value, id=f"initial_job_request-{field}") + for field, value in INITIAL_JOB_REQUEST_CHANGES.items() +] + [ + pytest.param("job_request", field, value, id=f"job_request-{field}") + for field, value in JOB_REQUEST_CHANGES.items() +] + + +@pytest.mark.parametrize( + ("message_class", "changes"), + [ + (V0InitialJobRequest, INITIAL_JOB_REQUEST_CHANGES), + (V0JobRequest, JOB_REQUEST_CHANGES), + ], +) +def test_field_changes_cover_every_signed_field(message_class, changes): + assert set(changes) == set(message_class.model_fields) - {"validator_signature"} + + +def test_blob_for_signing_is_canonical_and_excludes_the_signature(message, validator_keypair): + blob = message.blob_for_signing() + content = json.loads(blob) + + assert "validator_signature" not in content + assert blob == json.dumps(content, sort_keys=True, separators=(",", ":")) + assert message.signed(validator_keypair).blob_for_signing() == blob + assert type(message).model_validate_json(message.model_dump_json()).blob_for_signing() == blob + + +def test_signature_covers_the_sent_json_object_without_the_signature(message, validator_keypair): + sent = json.loads(message.signed(validator_keypair).model_dump_json()) + del sent["validator_signature"] + + assert message.blob_for_signing() == canonical_json(sent) + + +def test_received_message_verifies_and_keeps_every_field(message, validator_keypair): + signed = message.signed(validator_keypair) + + received = parse_validator_to_miner_message(signed.model_dump_json()) + + assert received.is_signed_by(validator_keypair.ss58_address) + assert (type(received), received.model_dump()) == (type(signed), signed.model_dump()) + + +def test_received_message_with_a_field_unknown_to_the_receiver_verifies(message, validator_keypair): + newer_model = create_model( + f"Newer{type(message).__name__}", __base__=type(message), priority=(int, 0) + ) + sent = newer_model(**message.model_dump(), priority=7).signed(validator_keypair) + sent_content = json.loads(sent.model_dump_json()) + + received = parse_validator_to_miner_message(json.dumps(sent_content)) + stripped = parse_validator_to_miner_message( + json.dumps({key: value for key, value in sent_content.items() if key != "priority"}) + ) + + assert (type(received), received.is_signed_by(validator_keypair.ss58_address)) == ( + type(message), + True, + ) + assert not stripped.is_signed_by(validator_keypair.ss58_address) + + +@pytest.mark.parametrize( + ("message_fixture", "field"), + [("initial_job_request", "executor_timing"), ("job_request", "env")], +) +def test_received_message_without_a_field_unknown_to_the_sender_verifies( + request, validator_keypair, message_fixture, field +): + message = request.getfixturevalue(message_fixture) + sent_content = json.loads(message.model_dump_json(exclude={"validator_signature", field})) + signature = sign_blob(validator_keypair, canonical_json(sent_content)) + + received = parse_validator_to_miner_message( + json.dumps({**sent_content, "validator_signature": signature}) + ) + + assert received.is_signed_by(validator_keypair.ss58_address) + assert received.model_dump() == { + **message.model_dump(), + field: type(message).model_fields[field].get_default(call_default_factory=True), + "validator_signature": signature, + } + + +# Pairs of a signed key-value pair of the job_request fixture and a duplicate of its key. +DUPLICATE_KEYS = [ + ('"docker_image":"image:v1"', '"docker_image":"injected"'), + ('"url":"http://localhost/upload"', '"url":"http://injected"'), +] + + +def frames_with_duplicate_keys(wire: str, injected_first: bool) -> list[str]: + frames = [] + for signed_pair, injected_pair in DUPLICATE_KEYS: + assert wire.count(signed_pair) == 1 + pairs = (injected_pair, signed_pair) if injected_first else (signed_pair, injected_pair) + frames.append(wire.replace(signed_pair, ",".join(pairs))) + return frames + + +def received_job_request_values(frames: list[str], validator_hotkey: str) -> list[tuple]: + return [ + (msg.docker_image, msg.output_upload.url, msg.is_signed_by(validator_hotkey)) + for msg in map(parse_validator_to_miner_message, frames) + ] + + +def test_value_injected_before_a_signed_key_is_ignored(job_request, validator_keypair): + wire = job_request.signed(validator_keypair).model_dump_json() + + frames = frames_with_duplicate_keys(wire, injected_first=True) + + assert received_job_request_values(frames, validator_keypair.ss58_address) == [ + ("image:v1", "http://localhost/upload", True), + ("image:v1", "http://localhost/upload", True), + ] + + +def test_value_injected_after_a_signed_key_fails_the_signature(job_request, validator_keypair): + wire = job_request.signed(validator_keypair).model_dump_json() + + frames = frames_with_duplicate_keys(wire, injected_first=False) + + assert received_job_request_values(frames, validator_keypair.ss58_address) == [ + ("injected", "http://localhost/upload", False), + ("image:v1", "http://injected", False), + ] + + +def test_parsing_matches_pydantic_json_validation( + initial_job_request, job_request, validator_keypair, miner_keypair, receipts +): + _, accepted_receipt, finished_receipt = receipts + messages = [ + GenericError(details="error"), + ValidatorAuthForMiner( + validator_hotkey=validator_keypair.ss58_address, + miner_hotkey=miner_keypair.ss58_address, + timestamp=1_700_000_000, + signature="0x00", + ), + initial_job_request.signed(validator_keypair), + job_request.signed(validator_keypair), + V0JobAcceptedReceiptRequest( + payload=accepted_receipt.payload, signature=accepted_receipt.validator_signature + ), + V0JobFinishedReceiptRequest( + payload=finished_receipt.payload, signature=finished_receipt.validator_signature + ), + V0MainHotkeyMessage(main_hotkey=miner_keypair.ss58_address), + ] + wires = [message.model_dump_json() for message in messages] + + assert [ + (type(parsed), parsed.model_dump()) + for parsed in map(parse_validator_to_miner_message, wires) + ] == [ + (type(parsed), parsed.model_dump()) + for parsed in map(TypeAdapter(ValidatorToMinerMessage).validate_json, wires) + ] + + +@pytest.mark.parametrize( + "frame", + [ + "not json", + "[1, 2]", + '{"message_type": "Unknown"}', + pytest.param( + '{"message_type": "GenericError", "details": ' + "1" * 5000 + "}", id="huge_int" + ), + pytest.param("[" * 100_000 + "]" * 100_000, id="deep_nesting"), + ], +) +def test_malformed_frame_raises(frame): + with pytest.raises(ValueError): + parse_validator_to_miner_message(frame) + + +def test_signature_is_bound_to_the_signing_key(message, validator_keypair, miner_keypair): + assert not message.signed(validator_keypair).is_signed_by(miner_keypair.ss58_address) + + +@pytest.mark.parametrize("signature", [None, "", "gibberish", "0x1234"]) +def test_missing_or_malformed_signature_is_invalid(message, validator_keypair, signature): + tampered = message.model_copy(update={"validator_signature": signature}) + + assert not tampered.is_signed_by(validator_keypair.ss58_address) + + +@pytest.mark.parametrize(("message_fixture", "field", "value"), FIELD_CHANGES) +def test_signature_is_invalid_after_changing_any_field( + request, validator_keypair, message_fixture, field, value +): + signed = request.getfixturevalue(message_fixture).signed(validator_keypair) + assert signed.is_signed_by(validator_keypair.ss58_address) + + changed = signed.model_copy(update={field: value}) + + assert changed != signed + assert not changed.is_signed_by(validator_keypair.ss58_address) + + +@pytest.mark.parametrize( + ("message_fixture", "field", "value"), + # A frame whose message type changed is not parsed as the same message at all. + [change for change in FIELD_CHANGES if change.values[1] != "message_type"], +) +def test_received_signature_is_invalid_after_changing_any_field( + request, validator_keypair, message_fixture, field, value +): + signed = request.getfixturevalue(message_fixture).signed(validator_keypair) + frame = signed.model_copy(update={field: value}).model_dump_json() + + assert not parse_validator_to_miner_message(frame).is_signed_by(validator_keypair.ss58_address) diff --git a/compute_horde/tests/utils.py b/compute_horde/tests/utils.py index 83548c4b2..3eb6f4cb3 100644 --- a/compute_horde/tests/utils.py +++ b/compute_horde/tests/utils.py @@ -1,3 +1,4 @@ +import datetime import random from uuid import uuid4 @@ -14,6 +15,7 @@ def random_receipt( miner_keypair: Keypair = _USE_RANDOM, validator_keypair: Keypair = _USE_RANDOM, + timestamp: datetime.datetime | None = None, ): if miner_keypair is _USE_RANDOM: miner_keypair = random_keypair() @@ -25,7 +27,7 @@ def random_receipt( job_uuid=str(uuid4()), miner_hotkey=miner_keypair.ss58_address, validator_hotkey=validator_keypair.ss58_address, - timestamp=timezone.now(), + timestamp=timestamp or timezone.now(), receipt_type=ReceiptType.JobStartedReceipt, executor_class=ExecutorClass.always_on__gpu_24gb, ttl=100, diff --git a/compute_horde/tests/ws_server.py b/compute_horde/tests/ws_server.py new file mode 100644 index 000000000..70754a1a6 --- /dev/null +++ b/compute_horde/tests/ws_server.py @@ -0,0 +1,99 @@ +import asyncio +from abc import ABC, abstractmethod +from collections.abc import Callable + +import websockets + + +async def wait_until(condition: Callable[[], bool], timeout: float = 2) -> None: + async with asyncio.timeout(timeout): + while not condition(): + await asyncio.sleep(0.01) + + +class AbstractWSTestServer(ABC): + """A WebSocket server on localhost that `async with` starts and stops.""" + + host = "localhost" + + def __init__(self, port: int) -> None: + self.port = port + self.ws_server: websockets.Server | None = None + + async def start(self) -> None: + self.ws_server = await websockets.serve( + self._serve, self.host, self.port, process_request=self._process_request + ) + + async def stop(self) -> None: + self.ws_server.close() + await self.ws_server.wait_closed() + + async def __aenter__(self): + await self.start() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self.stop() + + @property + def url(self) -> str: + return f"ws://{self.host}:{self.port}" + + @property + def connections(self) -> set[websockets.ServerConnection]: + return self.ws_server.connections + + @abstractmethod + async def _serve(self, ws: websockets.ServerConnection) -> None: ... + + async def _process_request( + self, connection: websockets.ServerConnection, request: websockets.Request + ) -> None: + """Called during every opening handshake, before the connection is accepted.""" + + +class ReconnectingWSTestServer(AbstractWSTestServer): + """ + A WebSocket peer that closes the first connection after `drop_first_connection_after` + messages, so the client's next send() or receive() has to reconnect. Every later handshake + takes `handshake_delay` seconds, which keeps a reconnect in flight long enough for a + concurrent one to start. Received messages are queued as (connection number, message). + """ + + def __init__( + self, port: int, *, drop_first_connection_after: int = 0, handshake_delay: float = 0 + ) -> None: + super().__init__(port) + self.drop_first_connection_after = drop_first_connection_after + self.handshake_delay = handshake_delay + self.connection_count = 0 + self.delayed_handshakes = 0 + self.received: asyncio.Queue[tuple[int, str | bytes]] = asyncio.Queue() + + @property + def connection(self) -> websockets.ServerConnection: + (connection,) = self.connections + return connection + + async def _process_request( + self, connection: websockets.ServerConnection, request: websockets.Request + ) -> None: + if self.connection_count: + self.delayed_handshakes += 1 + await asyncio.sleep(self.handshake_delay) + self.delayed_handshakes -= 1 + + async def _serve(self, ws: websockets.ServerConnection) -> None: + self.connection_count += 1 + number = self.connection_count + limit = self.drop_first_connection_after if number == 1 else None + received = 0 + while limit is None or received < limit: + try: + message = await ws.recv() + except websockets.ConnectionClosed: + return + await self.received.put((number, message)) + received += 1 + # Returning closes the connection. diff --git a/compute_horde_sdk/changelog.d/+multi-upload-cancellation.fixed.md b/compute_horde_sdk/changelog.d/+multi-upload-cancellation.fixed.md new file mode 100644 index 000000000..ac911f50c --- /dev/null +++ b/compute_horde_sdk/changelog.d/+multi-upload-cancellation.fixed.md @@ -0,0 +1 @@ +When one file of a `MultiUpload` fails to upload, the uploads still in progress are cancelled before `OutputUploadFailed` is raised; uploads that already finished stay delivered. diff --git a/compute_horde_sdk/changelog.d/+output-upload-confinement.fixed.md b/compute_horde_sdk/changelog.d/+output-upload-confinement.fixed.md new file mode 100644 index 000000000..9046b9d90 --- /dev/null +++ b/compute_horde_sdk/changelog.d/+output-upload-confinement.fixed.md @@ -0,0 +1 @@ +Output uploads are confined to the output directory: `relative_path` of single-file uploads must be a canonical relative path with no component longer than 255 bytes, the uploader opens output files without following symlinks and rejects anything that is not a regular file or exceeds the size limit, zipped uploads leave out symlinks, other non-regular entries and empty directories, and the HTTP response of an upload receiver is cut off at the 10 MiB cap. diff --git a/compute_horde_sdk/changelog.d/+output-volume-path-validation.changed.md b/compute_horde_sdk/changelog.d/+output-volume-path-validation.changed.md new file mode 100644 index 000000000..be226f600 --- /dev/null +++ b/compute_horde_sdk/changelog.d/+output-volume-path-validation.changed.md @@ -0,0 +1 @@ +`HTTPOutputVolume` mount paths must be `/output/` followed by a canonical relative path; paths such as `/output//x`, `/output/../x` or `/output/x/` now raise `ValueError` when the job is created. diff --git a/compute_horde_sdk/changelog.d/+put-upload-chunks.fixed.md b/compute_horde_sdk/changelog.d/+put-upload-chunks.fixed.md new file mode 100644 index 000000000..2502c36e9 --- /dev/null +++ b/compute_horde_sdk/changelog.d/+put-upload-chunks.fixed.md @@ -0,0 +1 @@ +HTTP PUT output uploads send the file in chunks of at most 64 KiB instead of one chunk per line, so a file without newlines is no longer read into memory whole. diff --git a/compute_horde_sdk/changelog.d/+untrusted-files.added.md b/compute_horde_sdk/changelog.d/+untrusted-files.added.md new file mode 100644 index 000000000..e53d221a2 --- /dev/null +++ b/compute_horde_sdk/changelog.d/+untrusted-files.added.md @@ -0,0 +1 @@ +Added `compute_horde_core.untrusted_files` for reading files from, and writing files into, a directory populated by an untrusted process without following symlinks and with a size bound checked before reading. diff --git a/compute_horde_sdk/changelog.d/+zipped-output-upload-errors.fixed.md b/compute_horde_sdk/changelog.d/+zipped-output-upload-errors.fixed.md new file mode 100644 index 000000000..2363e06d8 --- /dev/null +++ b/compute_horde_sdk/changelog.d/+zipped-output-upload-errors.fixed.md @@ -0,0 +1 @@ +Zipped output uploads accept files modified after 2107 and fail with `OutputUploadFailed` instead of an unhandled exception when a file name is not valid UTF-8. An output directory with directories nested more than 100 levels deep also fails with `OutputUploadFailed`. `MultiUpload` builds its system output zip before it sends any file. diff --git a/compute_horde_sdk/src/compute_horde_core/output_upload/_models.py b/compute_horde_sdk/src/compute_horde_core/output_upload/_models.py index 61d34e3f5..7303cae89 100644 --- a/compute_horde_sdk/src/compute_horde_core/output_upload/_models.py +++ b/compute_horde_sdk/src/compute_horde_core/output_upload/_models.py @@ -5,11 +5,25 @@ from urllib.parse import urlparse import pydantic -from pydantic import Field +from pydantic import AfterValidator, Field + +from ..untrusted_files import split_relative_path SAFE_DOMAIN_REGEX = re.compile(r".*") +def _validate_relative_path(value: str) -> str: + split_relative_path(value) + return value + + +RelativePath = Annotated[str, AfterValidator(_validate_relative_path)] +""" +Path of a file inside the output directory, as sent over the wire: rejected unless it is +already in its canonical relative spelling (see `split_relative_path`), never rewritten. +""" + + class OutputUploadType(str, enum.Enum): zip_and_http_post = "zip_and_http_post" zip_and_http_put = "zip_and_http_put" @@ -39,7 +53,7 @@ class SingleFilePostUpload(pydantic.BaseModel): output_upload_type: Literal[OutputUploadType.single_file_post] = OutputUploadType.single_file_post url: str form_fields: Mapping[str, str] | None = None - relative_path: str + relative_path: RelativePath signed_headers: Mapping[str, str] | None = None def is_safe(self) -> bool: @@ -54,7 +68,7 @@ class SingleFilePutUpload(pydantic.BaseModel): url: str # TODO: PUT implementation does not support it - find out why and clean up # form_fields: Mapping[str, str] | None = None - relative_path: str + relative_path: RelativePath signed_headers: Mapping[str, str] | None = None def is_safe(self) -> bool: diff --git a/compute_horde_sdk/src/compute_horde_core/output_upload/_uploader.py b/compute_horde_sdk/src/compute_horde_core/output_upload/_uploader.py index 444a85e3d..2bae259e9 100644 --- a/compute_horde_sdk/src/compute_horde_core/output_upload/_uploader.py +++ b/compute_horde_sdk/src/compute_horde_core/output_upload/_uploader.py @@ -4,20 +4,29 @@ import asyncio import contextlib import logging +import os import pathlib +import shutil import tempfile +import time import zipfile -from collections.abc import AsyncIterable, Callable, Iterable, Iterator +from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator from functools import wraps -from typing import IO, Any +from typing import IO, Any, TypeVar import httpx +from ..untrusted_files import ( + OpenedFile, + UntrustedFileRejected, + open_directory, + open_regular_file, + walk_regular_files, +) from ._models import ( HttpOutputVolumeResponse, MultiUpload, OutputUpload, - OutputUploadType, SingleFilePostUpload, SingleFilePutUpload, ZipAndHttpPostUpload, @@ -26,9 +35,14 @@ logger = logging.getLogger(__name__) +_T = TypeVar("_T") + OUTPUT_UPLOAD_TIMEOUT_SECONDS = 300 MAX_NUMBER_OF_FILES = 1000 MAX_CONCURRENT_UPLOADS = 3 +_PUT_CHUNK_SIZE = 64 * 1024 +_ZIP_FIRST_DATE_TIME = (1980, 1, 1, 0, 0, 0) +_ZIP_LAST_DATE_TIME = (2107, 12, 31, 23, 59, 59) def retry( @@ -96,16 +110,17 @@ def handles_output_type(cls) -> type[OutputUpload]: return ZipAndHttpPostUpload async def upload(self, directory: pathlib.Path) -> dict[str, HttpOutputVolumeResponse]: - with zipped_directory(directory, max_size_bytes=self.max_size_bytes) as (file_size, fp): - async with self._semaphore: - file_name = "output.zip" - response = await upload_post( - fp, - file_name, - self.upload_output.url, - form_fields=self.upload_output.form_fields, - ) - return {file_name: response} + with open_directory(directory) as dir_fd: + with zipped_directory(dir_fd, max_size_bytes=self.max_size_bytes) as (file_size, fp): + async with self._semaphore: + file_name = "output.zip" + response = await upload_post( + fp, + file_name, + self.upload_output.url, + form_fields=self.upload_output.form_fields, + ) + return {file_name: response} class ZipAndHTTPPutOutputUploader(OutputUploader): @@ -120,10 +135,11 @@ def handles_output_type(cls) -> type[OutputUpload]: return ZipAndHttpPutUpload async def upload(self, directory: pathlib.Path) -> dict[str, HttpOutputVolumeResponse]: - with zipped_directory(directory, max_size_bytes=self.max_size_bytes) as (file_size, fp): - async with self._semaphore: - response = await upload_put(fp, file_size, self.upload_output.url) - return {"output.zip": response} + with open_directory(directory) as dir_fd: + with zipped_directory(dir_fd, max_size_bytes=self.max_size_bytes) as (file_size, fp): + async with self._semaphore: + response = await upload_put(fp, file_size, self.upload_output.url) + return {"output.zip": response} class MultiUploadOutputUploader(OutputUploader): @@ -138,95 +154,41 @@ def handles_output_type(cls) -> type[OutputUpload]: return MultiUpload async def upload(self, directory: pathlib.Path) -> dict[str, HttpOutputVolumeResponse]: - single_file_uploads = [] - tasks = [] - for upload in self.upload_output.uploads: - file_path = directory / upload.relative_path - if not file_path.exists(): - raise OutputUploadFailed(f"File not found: {file_path}") - - if upload.output_upload_type == OutputUploadType.single_file_post: - # we run those concurrently but for loop changes slots - we need to bind - async def _single_post_upload_task( - file_path: pathlib.Path, upload: SingleFilePostUpload - ) -> tuple[str, HttpOutputVolumeResponse]: - with file_path.open("rb") as fp: - response = await upload_post( - fp, - file_path.name, - upload.url, - form_fields=upload.form_fields, - headers=upload.signed_headers, - ) - return upload.relative_path, response + # Every file is opened and the system output is zipped before the first request, so a bad + # path fails the upload before anything is sent. When one request fails the others are + # cancelled, but uploads that already finished stay delivered. + with contextlib.ExitStack() as stack: + dir_fd = stack.enter_context(open_directory(directory)) + files = [ + stack.enter_context(open_output_file(dir_fd, upload.relative_path, self.max_size_bytes)) + for upload in self.upload_output.uploads + ] + system_output_upload = self.upload_output.system_output + if system_output_upload: + exclude = [upload.relative_path for upload in self.upload_output.uploads] + zip_size, zip_fp = stack.enter_context( + zipped_directory(dir_fd, exclude=exclude, max_size_bytes=self.max_size_bytes) + ) + tasks = [] + for file, upload in zip(files, self.upload_output.uploads): async with self._semaphore: - tasks.append(_single_post_upload_task(file_path, upload)) - single_file_uploads.append(upload.relative_path) - elif upload.output_upload_type == OutputUploadType.single_file_put: - # we run those concurrently but for loop changes slots - we need to bind - async def _single_put_upload_task( - file_path: pathlib.Path, upload: SingleFilePutUpload - ) -> tuple[str, HttpOutputVolumeResponse]: - with file_path.open("rb") as fp: - response = await upload_put( - fp, file_path.stat().st_size, upload.url, headers=upload.signed_headers - ) - return upload.relative_path, response - + tasks.append(upload_single_file(file, upload)) + if system_output_upload: async with self._semaphore: - tasks.append(_single_put_upload_task(file_path, upload)) - single_file_uploads.append(upload.relative_path) - else: - raise OutputUploadFailed(f"Unsupported upload type: {upload.output_upload_type}") - - system_output_upload = self.upload_output.system_output - if system_output_upload: - if isinstance(system_output_upload, ZipAndHttpPostUpload): - # we don't need to bind any vars because we don't run it in a loop - async def _output_post_upload_task( - upload: ZipAndHttpPostUpload, - ) -> tuple[str, HttpOutputVolumeResponse]: - with zipped_directory( - directory, exclude=single_file_uploads, max_size_bytes=self.max_size_bytes - ) as ( - file_size, - fp, - ): - response = await upload_post( - fp, - "output.zip", - upload.url, - form_fields=upload.form_fields, - ) - return "system_output", response + tasks.append(self._upload_system_output(system_output_upload, zip_fp, zip_size)) - async with self._semaphore: - tasks.append(_output_post_upload_task(system_output_upload)) - elif isinstance(system_output_upload, ZipAndHttpPutUpload): - # we don't need to bind any vars because we don't run it in a loop - async def _output_put_upload_task(upload: ZipAndHttpPutUpload) -> tuple[str, HttpOutputVolumeResponse]: - with zipped_directory( - directory, exclude=single_file_uploads, max_size_bytes=self.max_size_bytes - ) as ( - file_size, - fp, - ): - response = await upload_put( - fp, - file_size, - upload.url, - ) - return "system_output", response + responses = await _gather_or_cancel(tasks) + return dict(responses) - async with self._semaphore: - tasks.append(_output_put_upload_task(system_output_upload)) - else: - raise OutputUploadFailed( - f"Unsupported system output upload type: {system_output_upload.output_upload_type}" - ) - responses = await asyncio.gather(*tasks) - return {path: result for path, result in responses} + async def _upload_system_output( + self, upload: ZipAndHttpPostUpload | ZipAndHttpPutUpload, fp: IO[bytes], file_size: int + ) -> tuple[str, HttpOutputVolumeResponse]: + if isinstance(upload, ZipAndHttpPostUpload): + response = await upload_post(fp, "output.zip", upload.url, form_fields=upload.form_fields) + else: + response = await upload_put(fp, file_size, upload.url) + return "system_output", response class SingleFilePutOutputUploader(OutputUploader): @@ -239,14 +201,7 @@ def handles_output_type(cls) -> type[OutputUpload]: return SingleFilePutUpload async def upload(self, directory: pathlib.Path) -> dict[str, HttpOutputVolumeResponse]: - file_path = directory / self.upload_output.relative_path - if not file_path.exists(): - raise OutputUploadFailed(f"File not found: {file_path}") - with file_path.open("rb") as fp: - response = await upload_put( - fp, file_path.stat().st_size, self.upload_output.url, headers=self.upload_output.signed_headers - ) - return {self.upload_output.relative_path: response} + return await upload_single_file_from_directory(directory, self.upload_output, self.max_size_bytes) class SingleFilePostOutputUploader(OutputUploader): @@ -259,28 +214,81 @@ def handles_output_type(cls) -> type[OutputUpload]: return SingleFilePostUpload async def upload(self, directory: pathlib.Path) -> dict[str, HttpOutputVolumeResponse]: - file_path = directory / self.upload_output.relative_path - if not file_path.exists(): - raise OutputUploadFailed(f"File not found: {file_path}") - with file_path.open("rb") as fp: + return await upload_single_file_from_directory(directory, self.upload_output, self.max_size_bytes) + + +@contextlib.contextmanager +def open_output_file(dir_fd: int, relative_path: str, max_size_bytes: int) -> Iterator[OpenedFile]: + """ + Open the regular file `relative_path` beneath the output directory `dir_fd` for uploading. + + Raises OutputUploadFailed when the file is missing, is not a regular file reachable + without following symlinks, cannot be opened, or is larger than `max_size_bytes`. + """ + with contextlib.ExitStack() as stack: + try: + file = stack.enter_context(open_regular_file(dir_fd, relative_path)) + except FileNotFoundError: + raise OutputUploadFailed(f"File not found: {relative_path}") + except UntrustedFileRejected as e: + raise OutputUploadFailed(f"Output file {relative_path!r} cannot be uploaded: {e.reason}") + except OSError as e: + raise OutputUploadFailed(f"Output file {relative_path!r} cannot be uploaded: {e.strerror}") + if file.stat.st_size > max_size_bytes: + raise OutputUploadFailed( + f"Output file {relative_path!r} is {file.stat.st_size:,} bytes, " + f"more than the allowed {max_size_bytes:,} bytes" + ) + yield file + + +async def upload_single_file_from_directory( + directory: pathlib.Path, upload: SingleFilePostUpload | SingleFilePutUpload, max_size_bytes: int +) -> dict[str, HttpOutputVolumeResponse]: + with open_directory(directory) as dir_fd: + with open_output_file(dir_fd, upload.relative_path, max_size_bytes) as file: + relative_path, response = await upload_single_file(file, upload) + return {relative_path: response} + + +async def upload_single_file( + file: OpenedFile, upload: SingleFilePostUpload | SingleFilePutUpload +) -> tuple[str, HttpOutputVolumeResponse]: + with os.fdopen(file.fd, "rb", closefd=False) as fp: + if isinstance(upload, SingleFilePostUpload): response = await upload_post( fp, - file_path.name, - self.upload_output.url, - form_fields=self.upload_output.form_fields, - headers=self.upload_output.signed_headers, + pathlib.PurePosixPath(upload.relative_path).name, + upload.url, + form_fields=upload.form_fields, + headers=upload.signed_headers, ) - return {self.upload_output.relative_path: response} + else: + response = await upload_put(fp, file.stat.st_size, upload.url, headers=upload.signed_headers) + return upload.relative_path, response -async def make_iterator_async(it: Iterable[Any]) -> AsyncIterable[Any]: +async def _gather_or_cancel(coros: Iterable[Coroutine[Any, Any, _T]]) -> list[_T]: """ - Make an iterator async. + Await `coros` concurrently and return their results in order, like `asyncio.gather`. - This is stupid. + When one fails, or the caller is cancelled, the others are cancelled and then awaited. + If the caller is cancelled during that wait, this stops waiting and raises + `CancelledError`, and some of the others may still be unwinding. """ - for x in it: - yield x + tasks = [asyncio.create_task(coro) for coro in coros] + try: + return await asyncio.gather(*tasks) + finally: + for task in tasks: + task.cancel() + if tasks: + await asyncio.wait(tasks) + + +async def _read_in_chunks(fp: IO[bytes]) -> AsyncIterator[bytes]: + while chunk := fp.read(_PUT_CHUNK_SIZE): + yield chunk @retry(max_retries=3, exceptions=OutputUploadFailed) @@ -335,6 +343,9 @@ async def read_stream_with_cap(response: httpx.Response, max_allowed_size: int = """ Reads the response in a streaming fashion, enforcing a maximum allowed response size. + Reading stops at the first chunk that takes the total over the limit, so at most + `max_allowed_size` bytes are ever kept in memory. + Args: response: The httpx response object. max_allowed_size: Maximum allowed size in bytes for the response body, defaults to 10MB. @@ -355,7 +366,7 @@ async def read_stream_with_cap(response: httpx.Response, max_allowed_size: int = total_bytes = 0 chunks: list[bytes] = [] async for chunk in response.aiter_bytes(): - total_bytes += len(chunks) + total_bytes += len(chunk) if total_bytes > max_allowed_size: raise OutputUploadFailed("Response size exceeds allowed limit") chunks.append(chunk) @@ -391,7 +402,7 @@ async def upload_put( async with client.stream( method="PUT", url=url, - content=make_iterator_async(fp), + content=_read_in_chunks(fp), headers=headers, timeout=OUTPUT_UPLOAD_TIMEOUT_SECONDS, ) as response: @@ -407,30 +418,28 @@ async def upload_put( @contextlib.contextmanager def zipped_directory( - directory: pathlib.Path, exclude: list[str] | None = None, max_size_bytes: int = 2147483648 + dir_fd: int, exclude: list[str] | None = None, max_size_bytes: int = 2147483648 ) -> Iterator[tuple[int, IO[bytes]]]: """ - Context manager that creates a temporary zip file with the files from given directory. + Context manager that creates a temporary zip file with the regular files from the given directory. The temporary file is cleared after the context manager exits. - :param directory: The directory to zip. + Symlinks and other non-regular entries are left out, so the archive never contains + anything from outside the directory. Empty directories are left out as well. Any + failure to build the archive, such as a file name that is not valid UTF-8, raises + OutputUploadFailed. + + :param dir_fd: Descriptor of the directory to zip, from `open_directory`. :param exclude: A list of relative paths to exclude from the zip file. :param max_size_bytes: Maximum allowed size of the zip file in bytes. Defaults to 2147483648. :return: tuple of size and the file object of the zip file """ - files = list(directory.glob("**/*")) - exclude_set = set(exclude) if exclude else set() - - filtered_files = [file for file in files if str(file.relative_to(directory)) not in exclude_set] - - if len(filtered_files) > MAX_NUMBER_OF_FILES: - raise OutputUploadFailed("Attempting to upload too many files") - with tempfile.TemporaryFile() as fp: - with zipfile.ZipFile(fp, mode="w") as zipf: - for file in filtered_files: - zipf.write(filename=file, arcname=file.relative_to(directory)) + try: + _write_zip(fp, dir_fd, exclude=set(exclude or ())) + except (UntrustedFileRejected, OSError, UnicodeEncodeError) as e: + raise OutputUploadFailed(f"Output directory cannot be zipped: {e}") from e file_size = fp.tell() fp.seek(0) @@ -439,3 +448,33 @@ def zipped_directory( raise OutputUploadFailed("Attempting to upload too large file") yield file_size, fp + + +def _write_zip(fp: IO[bytes], dir_fd: int, exclude: set[str]) -> None: + relative_paths = [path for path in walk_regular_files(dir_fd) if path not in exclude] + + if len(relative_paths) > MAX_NUMBER_OF_FILES: + raise OutputUploadFailed("Attempting to upload too many files") + + with zipfile.ZipFile(fp, mode="w") as zipf: + for relative_path in relative_paths: + with open_regular_file(dir_fd, relative_path) as file: + _add_to_zip(zipf, file) + + +def _add_to_zip(zipf: zipfile.ZipFile, file: OpenedFile) -> None: + # Same metadata as ZipFile.write, taken from the open descriptor instead of the path. + info = zipfile.ZipInfo(file.path, date_time=_zip_date_time(file.stat.st_mtime)) + info.external_attr = (file.stat.st_mode & 0xFFFF) << 16 + info.file_size = file.stat.st_size + with os.fdopen(file.fd, "rb", closefd=False) as source, zipf.open(info, "w") as target: + shutil.copyfileobj(source, target) + + +def _zip_date_time(mtime: float) -> tuple[int, int, int, int, int, int]: + # Zip dates span 1980 to 2107; clamped as ZipInfo.from_file does with strict_timestamps=False. + try: + date_time = time.localtime(mtime)[:6] + except (OverflowError, OSError, ValueError): + return _ZIP_FIRST_DATE_TIME + return min(max(date_time, _ZIP_FIRST_DATE_TIME), _ZIP_LAST_DATE_TIME) diff --git a/compute_horde_sdk/src/compute_horde_core/untrusted_files.py b/compute_horde_sdk/src/compute_horde_core/untrusted_files.py new file mode 100644 index 000000000..07c481c86 --- /dev/null +++ b/compute_horde_sdk/src/compute_horde_core/untrusted_files.py @@ -0,0 +1,279 @@ +""" +Reading and writing files in a directory that an untrusted process has populated. + +A job container writes into host directories that are later read with the +privileges of the process collecting the results. Every entry is treated as +hostile: no path component is ever followed as a symlink, only regular files +are opened, and the size reported by the kernel is checked before any content +is read. A file is written there only as a new regular file that replaces +whatever entry had its name. Relative paths chosen by an untrusted party are +accepted only in a single canonical spelling, so a path can never name +anything above the directory it is resolved against. +""" + +import contextlib +import dataclasses +import errno +import os +import pathlib +import stat +from collections.abc import Iterator + +# O_NONBLOCK and O_NOCTTY: an entry swapped for a FIFO or a terminal between the +# type check and the open neither blocks the open nor becomes the controlling terminal. +_OPEN_FLAGS = { + stat.S_IFDIR: os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC, + stat.S_IFREG: os.O_RDONLY | os.O_NOFOLLOW | os.O_CLOEXEC | os.O_NONBLOCK | os.O_NOCTTY, +} +# O_EXCL: an entry planted under the name after it was unlinked fails the open instead of being opened. +_WRITE_FLAGS = os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW | os.O_CLOEXEC | os.O_NONBLOCK +_FILE_TYPE_NAMES = { + stat.S_IFREG: "regular file", + stat.S_IFDIR: "directory", + stat.S_IFLNK: "symlink", + stat.S_IFIFO: "FIFO", + stat.S_IFCHR: "character device", + stat.S_IFBLK: "block device", + stat.S_IFSOCK: "socket", +} +_READ_CHUNK_SIZE = 1 << 20 +_FORBIDDEN_COMPONENTS = frozenset({"", ".", ".."}) +_MAX_NAME_BYTES = 255 +# Every directory level being walked keeps a descriptor open and a generator frame alive. +_MAX_WALK_DEPTH = 100 + + +class UntrustedFileRejected(Exception): + """A path under an untrusted directory was not opened or read; `reason` says why.""" + + def __init__(self, path: str, reason: str) -> None: + super().__init__(f"{path}: {reason}") + self.path = path + self.reason = reason + + +class InvalidRelativePath(UntrustedFileRejected, ValueError): + """Also a ValueError, so pydantic validators can raise it as is.""" + + +class SymlinkRejected(UntrustedFileRejected): + pass + + +class UnexpectedFileType(UntrustedFileRejected): + pass + + +class FileTooLarge(UntrustedFileRejected): + pass + + +class FileSizeChanged(UntrustedFileRejected): + pass + + +class EntryNotReplaceable(UntrustedFileRejected): + pass + + +class DirectoryTooDeep(UntrustedFileRejected): + pass + + +@dataclasses.dataclass(frozen=True) +class OpenedFile: + path: str + fd: int + stat: os.stat_result + + +def split_relative_path(relative_path: str) -> list[str]: + """ + Components of a relative path in its single accepted spelling, or `InvalidRelativePath`. + + Meant for paths chosen by an untrusted party. Accepted: a path `open_regular_file` + accepts that contains no backslash and no component longer than 255 bytes in UTF-8. + Joining the components with '/' gives the input back, so a validated path needs no + normalization anywhere. + """ + if "\\" in relative_path: + raise InvalidRelativePath(relative_path, "must not contain backslashes") + components = _split_confined(relative_path) + if any(len(component.encode()) > _MAX_NAME_BYTES for component in components): + raise InvalidRelativePath(relative_path, f"must not have components longer than {_MAX_NAME_BYTES} bytes") + return components + + +@contextlib.contextmanager +def open_directory(path: str | os.PathLike[str]) -> Iterator[int]: + """Open `path` as a directory descriptor; `path` itself must not be a symlink.""" + with _open_component(None, path, os.fspath(path), stat.S_IFDIR) as (fd, _): + yield fd + + +def list_entries(dir_fd: int) -> list[str]: + """Names of the current entries in the directory, whatever their type.""" + # Rewind first: btrfs otherwise hides entries created after the descriptor was opened. + os.lseek(dir_fd, 0, os.SEEK_SET) + return os.listdir(dir_fd) + + +def walk_regular_files(dir_fd: int) -> Iterator[str]: + """ + Relative paths of the regular files below the directory `dir_fd`, in no particular order. + + Subdirectories are entered through descriptors opened without following + symlinks. Symlinks and every other non-regular entry are skipped. Names are + kept as `os.listdir` returns them, so they may contain backslashes or + surrogate escapes; `open_regular_file` accepts every path yielded here. + A directory nested more than 100 levels deep raises `DirectoryTooDeep`. + """ + yield from _walk(dir_fd, pathlib.PurePosixPath()) + + +@contextlib.contextmanager +def open_regular_file(dir_fd: int, relative_path: str) -> Iterator[OpenedFile]: + """ + Open `relative_path` beneath the directory `dir_fd` without following any symlink. + + `relative_path` must be non-empty and '/'-separated, without a leading '/', NUL + bytes or empty, '.' or '..' components, or `InvalidRelativePath` is raised; any + other character, a backslash included, is part of a name. Each component is + opened relative to the previous one, so a symlink or an unexpected file type + anywhere in the path is rejected instead of followed. + """ + *directories, name = _split_confined(relative_path) + with contextlib.ExitStack() as stack: + walked = pathlib.PurePosixPath() + for directory in directories: + walked /= directory + dir_fd, _ = stack.enter_context(_open_component(dir_fd, directory, str(walked), stat.S_IFDIR)) + fd, st = stack.enter_context(_open_component(dir_fd, name, relative_path, stat.S_IFREG)) + yield OpenedFile(path=relative_path, fd=fd, stat=st) + + +def read_bounded(file: OpenedFile, max_size: int) -> bytes: + """ + Read the whole content of an opened regular file of at most `max_size` bytes. + + The size reported by fstat is checked before reading and never more than + `max_size + 1` bytes are read, so an oversized file costs no memory. Content + that does not match the reported size, because the file is still being + written to, is rejected. + """ + expected_size = file.stat.st_size + if expected_size > max_size: + raise FileTooLarge(file.path, f"is {expected_size:,} bytes, more than the allowed {max_size:,} bytes") + content = _read_up_to(file.fd, expected_size + 1) + if len(content) != expected_size: + raise FileSizeChanged(file.path, f"read {len(content):,} bytes, expected {expected_size:,} bytes") + return content + + +def read_regular_file(dir_fd: int, relative_path: str, max_size: int) -> bytes: + with open_regular_file(dir_fd, relative_path) as file: + return read_bounded(file, max_size) + + +def _read_up_to(fd: int, limit: int) -> bytes: + chunks: list[bytes] = [] + remaining = limit + while remaining > 0: + chunk = os.read(fd, min(remaining, _READ_CHUNK_SIZE)) + if not chunk: + break + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) + + +def _split_confined(relative_path: str) -> list[str]: + if not relative_path: + raise InvalidRelativePath(relative_path, "must not be empty") + if "\0" in relative_path: + raise InvalidRelativePath(relative_path, "must not contain NUL bytes") + if relative_path.startswith("/"): + raise InvalidRelativePath(relative_path, "must be relative, not absolute") + components = relative_path.split("/") + if _FORBIDDEN_COMPONENTS.intersection(components): + raise InvalidRelativePath(relative_path, "must not contain empty, '.' or '..' components or end with '/'") + return components + + +def _walk(dir_fd: int, prefix: pathlib.PurePosixPath) -> Iterator[str]: + for name in list_entries(dir_fd): + path = prefix / name + file_type = stat.S_IFMT(os.stat(name, dir_fd=dir_fd, follow_symlinks=False).st_mode) + if file_type == stat.S_IFREG: + yield str(path) + elif file_type == stat.S_IFDIR: + if len(path.parts) > _MAX_WALK_DEPTH: + raise DirectoryTooDeep(str(path), f"is nested more than {_MAX_WALK_DEPTH} directories deep") + with _open_component(dir_fd, name, str(path), stat.S_IFDIR) as (sub_fd, _): + yield from _walk(sub_fd, path) + + +@contextlib.contextmanager +def _open_component( + dir_fd: int | None, name: str | os.PathLike[str], path: str, expected_type: int +) -> Iterator[tuple[int, os.stat_result]]: + # Check the type before opening: opening a device node can have side effects. + _check_type(os.stat(name, dir_fd=dir_fd, follow_symlinks=False), path, expected_type) + try: + fd = os.open(name, _OPEN_FLAGS[expected_type], dir_fd=dir_fd) + except OSError as e: + if e.errno == errno.ELOOP: + raise SymlinkRejected(path, "is a symlink") from e + raise + try: + # The open descriptor is what gets read, so its type is the one that counts. + st = os.fstat(fd) + _check_type(st, path, expected_type) + yield fd, st + finally: + os.close(fd) + + +def _check_type(st: os.stat_result, path: str, expected_type: int) -> None: + actual_type = stat.S_IFMT(st.st_mode) + if actual_type == expected_type: + return + if actual_type == stat.S_IFLNK: + raise SymlinkRejected(path, "is a symlink") + actual_name = _FILE_TYPE_NAMES.get(actual_type, "file of unknown type") + raise UnexpectedFileType(path, f"is a {actual_name}, expected a {_FILE_TYPE_NAMES[expected_type]}") + + +def write_regular_file(dir_fd: int, name: str, data: str | bytes) -> None: + """ + Create `name` directly beneath the directory `dir_fd` as a new regular file holding `data`. + + `name` must be a single file name; '', '.', '..' and names containing '/' or + a NUL byte raise `InvalidRelativePath`. Whatever the untrusted process left + under that name is unlinked first, so a symlink is never followed and a FIFO + is never opened. A directory cannot be unlinked and is rejected with + `EntryNotReplaceable`. `str` data is written as UTF-8. + """ + _check_file_name(name) + try: + os.unlink(name, dir_fd=dir_fd) + except FileNotFoundError: + pass + except OSError as e: + raise EntryNotReplaceable(name, f"cannot be replaced: {e.strerror}") from e + fd = os.open(name, _WRITE_FLAGS, 0o644, dir_fd=dir_fd) + try: + _write_all(fd, data.encode() if isinstance(data, str) else data) + finally: + os.close(fd) + + +def _check_file_name(name: str) -> None: + if name in ("", ".", "..") or "/" in name or "\0" in name: + raise InvalidRelativePath(name, "must be a single file name") + + +def _write_all(fd: int, data: bytes) -> None: + view = memoryview(data) + while view: + view = view[os.write(fd, view) :] diff --git a/compute_horde_sdk/src/compute_horde_sdk/_internal/models.py b/compute_horde_sdk/src/compute_horde_sdk/_internal/models.py index 23d8e776b..504cd0357 100644 --- a/compute_horde_sdk/src/compute_horde_sdk/_internal/models.py +++ b/compute_horde_sdk/src/compute_horde_sdk/_internal/models.py @@ -13,6 +13,7 @@ from compute_horde_core import output_upload as compute_horde_output_upload from compute_horde_core import volume as compute_horde_volume +from compute_horde_core.untrusted_files import InvalidRelativePath, split_relative_path if sys.version_info >= (3, 11): # noqa: UP036 from enum import StrEnum @@ -324,7 +325,14 @@ class HTTPOutputVolume(pydantic.BaseModel): def get_volume_relative_path(self, mount_path: str) -> str: if not mount_path.startswith(OUTPUT_MOUNT_PATH_PREFIX): raise ValueError(f"Output volume paths must start with {OUTPUT_MOUNT_PATH_PREFIX!r}") - return mount_path.removeprefix(OUTPUT_MOUNT_PATH_PREFIX) + relative_path = mount_path.removeprefix(OUTPUT_MOUNT_PATH_PREFIX) + try: + split_relative_path(relative_path) + except InvalidRelativePath as e: + raise ValueError( + f"Output volume path {mount_path!r}: the part after {OUTPUT_MOUNT_PATH_PREFIX!r} {e.reason}" + ) + return relative_path def to_compute_horde_output_upload(self, mount_path: str) -> compute_horde_output_upload.OutputUpload: relative_path = self.get_volume_relative_path(mount_path) diff --git a/compute_horde_sdk/tests/unit/core/test_output_upload.py b/compute_horde_sdk/tests/unit/core/test_output_upload.py index 8fc8d6a06..355881c8a 100644 --- a/compute_horde_sdk/tests/unit/core/test_output_upload.py +++ b/compute_horde_sdk/tests/unit/core/test_output_upload.py @@ -1,12 +1,20 @@ import asyncio +import calendar +import io +import os import pathlib +import zipfile from unittest import mock +import httpx +import pydantic import pytest +from pytest_httpx import IteratorStream from compute_horde_core.output_upload import ( MultiUpload, MultiUploadOutputUploader, + OutputUpload, OutputUploader, OutputUploadFailed, SingleFilePostUpload, @@ -16,6 +24,7 @@ ZipAndHTTPPutOutputUploader, ZipAndHttpPutUpload, ) +from compute_horde_core.output_upload._uploader import _zip_date_time, read_stream_with_cap class TestOutputUploader: @@ -360,6 +369,33 @@ async def test_upload_error_handling(self, upload_without_system_output, setup_t with pytest.raises(OutputUploadFailed): await uploader.upload(setup_test_dir) + @pytest.fixture + def upload_with_hanging_and_failing_url(self): + return MultiUpload( + uploads=[ + SingleFilePutUpload(url="https://example.com/hanging", relative_path="test_file.txt"), + SingleFilePutUpload(url="https://example.com/failing", relative_path="subdir/subdir_file.txt"), + ] + ) + + @pytest.fixture + def hanging_and_failing_receivers(self, httpx_mock): + async def never_respond(request): + await asyncio.Event().wait() + + httpx_mock.add_callback(never_respond, url="https://example.com/hanging") + httpx_mock.add_response(url="https://example.com/failing", status_code=500, is_reusable=True) + + @pytest.mark.asyncio + @pytest.mark.usefixtures("async_sleep_mock", "hanging_and_failing_receivers") + async def test_failed_upload_cancels_the_other_uploads(self, upload_with_hanging_and_failing_url, setup_test_dir): + uploader = MultiUploadOutputUploader(upload_with_hanging_and_failing_url) + + with pytest.raises(OutputUploadFailed): + await uploader.upload(setup_test_dir) + + assert [task for task in asyncio.all_tasks() if task is not asyncio.current_task()] == [] + @pytest.mark.asyncio async def test_concurrency_limit(self, upload_without_system_output, setup_test_dir, httpx_mock): """Test that concurrency limit is respected.""" @@ -418,6 +454,7 @@ async def __aexit__(self, *args): async def test_system_output_with_zip_http_put(self, uploads, setup_test_dir, httpx_mock): """Test uploading with ZipAndHttpPutUpload as system output.""" system_output = ZipAndHttpPutUpload(url="https://example.com/system_output_put") + (setup_test_dir / "extra.txt").write_bytes(b"not uploaded separately") upload = MultiUpload(uploads=uploads, system_output=system_output) @@ -436,6 +473,398 @@ async def test_system_output_with_zip_http_put(self, uploads, setup_test_dir, ht request = requests[0] assert request.method == "PUT" - # Content should be a ZIP file - content = request.read() - assert b"PK\x03\x04" in content # ZIP signature + # The zip holds only what the single-file uploads did not cover + archive = zipfile.ZipFile(io.BytesIO(request.read())) + assert {name: archive.read(name) for name in archive.namelist()} == {"extra.txt": b"not uploaded separately"} + + +VALID_RELATIVE_PATHS = [ + "result.json", + "dir/file.bin", + "a/b/c.txt", + "with space.txt", + ".hidden", + pytest.param("x" * 255, id="255-byte-name"), +] +INVALID_RELATIVE_PATHS = [ + "", + "/", + "/etc/passwd", + "/proc/self/environ", + "//x", + "../secret", + "a/../secret", + "./x", + "a/./b", + "a//b", + "dir/", + "a\\b", + "a\0b", + pytest.param("x" * 256, id="256-byte-name"), + pytest.param("dir/" + "é" * 128, id="256-byte-utf8-name"), +] + + +class TestSingleFileUploadModels: + @pytest.mark.parametrize("relative_path", VALID_RELATIVE_PATHS) + def test_valid_relative_path_is_serialized_unchanged(self, relative_path): + upload = SingleFilePutUpload(url="https://example.com/upload", relative_path=relative_path) + + assert upload.model_dump() == { + "output_upload_type": "single_file_put", + "url": "https://example.com/upload", + "relative_path": relative_path, + "signed_headers": None, + } + + @pytest.mark.parametrize("model", [SingleFilePostUpload, SingleFilePutUpload]) + @pytest.mark.parametrize("relative_path", INVALID_RELATIVE_PATHS) + def test_invalid_relative_path_is_rejected(self, model, relative_path): + with pytest.raises(pydantic.ValidationError, match="relative_path"): + model(url="https://example.com/upload", relative_path=relative_path) + + def test_invalid_relative_path_inside_multi_upload_is_rejected(self): + with pytest.raises(pydantic.ValidationError, match="relative_path"): + pydantic.TypeAdapter(OutputUpload).validate_python( + { + "output_upload_type": "multi_upload", + "uploads": [ + { + "output_upload_type": "single_file_put", + "url": "https://example.com/upload", + "relative_path": "/proc/self/environ", + } + ], + } + ) + + +class TestOutputDirectoryConfinement: + """The uploader must never read anything a job can reach only through the output directory.""" + + @pytest.fixture + def secret_file(self, tmp_path): + path = tmp_path / "secret" + path.write_bytes(b"EXECUTOR_TOKEN=do-not-leak") + return path + + @pytest.fixture + def output_dir(self, tmp_path): + path = tmp_path / "output" + path.mkdir() + (path / "result.json").write_bytes(b'{"ok": true}') + return path + + @pytest.fixture + def hostile_entries(self, output_dir, secret_file): + os.symlink(secret_file, output_dir / "leak") + os.symlink("/proc/self/environ", output_dir / "environ") + os.symlink(secret_file.parent, output_dir / "escape") + os.mkfifo(output_dir / "fifo") + (output_dir / "sub").mkdir() + (output_dir / "sub" / "nested.txt").write_bytes(b"nested") + (output_dir / "empty-dir").mkdir() + + @pytest.fixture(params=["single_file_put", "single_file_post"]) + def single_file_upload(self, request): + def make(relative_path): + return pydantic.TypeAdapter(OutputUpload).validate_python( + { + "output_upload_type": request.param, + "url": "https://example.com/upload", + "relative_path": relative_path, + } + ) + + return make + + @pytest.mark.asyncio + async def test_regular_file_is_uploaded(self, output_dir, single_file_upload, httpx_mock): + httpx_mock.add_response(url="https://example.com/upload", content=b"stored") + uploader = OutputUploader.for_upload_output(single_file_upload("result.json")) + uploader.max_size_bytes = len(b'{"ok": true}') + + result = await uploader.upload(output_dir) + + assert list(result) == ["result.json"] + assert result["result.json"].body == "stored" + assert b'{"ok": true}' in httpx_mock.get_requests()[0].read() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "relative_path,reason", + [ + ("leak", "is a symlink"), + ("environ", "is a symlink"), + ("escape/secret", "is a symlink"), + ("fifo", "is a FIFO"), + ("sub", "is a directory"), + ], + ) + async def test_non_regular_entry_is_rejected_before_any_request( + self, output_dir, hostile_entries, single_file_upload, httpx_mock, relative_path, reason + ): + uploader = OutputUploader.for_upload_output(single_file_upload(relative_path)) + + with pytest.raises(OutputUploadFailed, match=reason): + await uploader.upload(output_dir) + + assert httpx_mock.get_requests() == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize("relative_path", ["/proc/self/environ", "../secret", "/output/../secret"]) + async def test_unvalidated_escaping_path_is_rejected_before_any_request( + self, output_dir, secret_file, httpx_mock, relative_path + ): + # model_construct skips validation, as a hop running older code would. + upload = SingleFilePutUpload.model_construct(url="https://example.com/upload", relative_path=relative_path) + uploader = OutputUploader.for_upload_output(upload) + + with pytest.raises(OutputUploadFailed, match="must"): + await uploader.upload(output_dir) + + assert httpx_mock.get_requests() == [] + + @pytest.mark.asyncio + async def test_unvalidated_overlong_name_is_rejected_before_any_request(self, output_dir, httpx_mock): + upload = SingleFilePutUpload.model_construct(url="https://example.com/upload", relative_path="x" * 256) + uploader = OutputUploader.for_upload_output(upload) + + with pytest.raises(OutputUploadFailed, match="File name too long"): + await uploader.upload(output_dir) + + assert httpx_mock.get_requests() == [] + + @pytest.mark.asyncio + async def test_file_over_size_limit_is_rejected_before_any_request( + self, output_dir, single_file_upload, httpx_mock + ): + uploader = OutputUploader.for_upload_output(single_file_upload("result.json")) + uploader.max_size_bytes = len(b'{"ok": true}') - 1 + + with pytest.raises(OutputUploadFailed, match="more than the allowed"): + await uploader.upload(output_dir) + + assert httpx_mock.get_requests() == [] + + @pytest.mark.asyncio + async def test_missing_file_is_rejected_before_any_request(self, output_dir, single_file_upload, httpx_mock): + uploader = OutputUploader.for_upload_output(single_file_upload("missing.json")) + + with pytest.raises(OutputUploadFailed, match="File not found"): + await uploader.upload(output_dir) + + assert httpx_mock.get_requests() == [] + + @pytest.mark.asyncio + async def test_multi_upload_rejects_symlink_before_any_request(self, output_dir, hostile_entries, httpx_mock): + upload = MultiUpload( + uploads=[ + SingleFilePutUpload(url="https://example.com/upload1", relative_path="result.json"), + SingleFilePutUpload(url="https://example.com/upload2", relative_path="leak"), + ], + ) + uploader = MultiUploadOutputUploader(upload) + + with pytest.raises(OutputUploadFailed, match="is a symlink"): + await uploader.upload(output_dir) + + assert httpx_mock.get_requests() == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "upload,method", + [ + (ZipAndHttpPutUpload(url="https://example.com/upload"), "PUT"), + (ZipAndHttpPostUpload(url="https://example.com/upload"), "POST"), + ], + ) + async def test_zip_contains_only_regular_files(self, output_dir, hostile_entries, httpx_mock, upload, method): + httpx_mock.add_response(url="https://example.com/upload", method=method) + uploader = OutputUploader.for_upload_output(upload) + + await uploader.upload(output_dir) + + archive = zipfile.ZipFile(io.BytesIO(zip_payload(httpx_mock.get_requests()[0].read()))) + assert {name: archive.read(name) for name in sorted(archive.namelist())} == { + "result.json": b'{"ok": true}', + "sub/nested.txt": b"nested", + } + + @pytest.mark.asyncio + async def test_system_output_zip_leaves_out_single_file_uploads_and_symlinks( + self, output_dir, hostile_entries, httpx_mock + ): + httpx_mock.add_response(url="https://example.com/upload1", method="PUT") + httpx_mock.add_response(url="https://example.com/system_output", method="PUT") + upload = MultiUpload( + uploads=[SingleFilePutUpload(url="https://example.com/upload1", relative_path="result.json")], + system_output=ZipAndHttpPutUpload(url="https://example.com/system_output"), + ) + + await MultiUploadOutputUploader(upload).upload(output_dir) + + (request,) = [r for r in httpx_mock.get_requests() if r.url == "https://example.com/system_output"] + assert zipfile.ZipFile(io.BytesIO(request.read())).namelist() == ["sub/nested.txt"] + + @pytest.fixture + def backslash_named_file(self, output_dir): + (output_dir / "a\\b").write_bytes(b"backslash") + + @pytest.fixture + def non_utf8_named_file(self, output_dir): + (output_dir / os.fsdecode(b"\xff.bin")).write_bytes(b"undecodable") + + @pytest.mark.asyncio + async def test_zip_contains_file_with_backslash_in_its_name(self, output_dir, backslash_named_file, httpx_mock): + httpx_mock.add_response(url="https://example.com/upload", method="PUT") + + await ZipAndHTTPPutOutputUploader(ZipAndHttpPutUpload(url="https://example.com/upload")).upload(output_dir) + + archive = zipfile.ZipFile(io.BytesIO(httpx_mock.get_requests()[0].read())) + assert {name: archive.read(name) for name in sorted(archive.namelist())} == { + "a\\b": b"backslash", + "result.json": b'{"ok": true}', + } + + @pytest.mark.asyncio + async def test_system_output_zip_contains_file_with_backslash_in_its_name( + self, output_dir, backslash_named_file, httpx_mock + ): + httpx_mock.add_response(url="https://example.com/upload1", method="PUT") + httpx_mock.add_response(url="https://example.com/system_output", method="PUT") + upload = MultiUpload( + uploads=[SingleFilePutUpload(url="https://example.com/upload1", relative_path="result.json")], + system_output=ZipAndHttpPutUpload(url="https://example.com/system_output"), + ) + + result = await MultiUploadOutputUploader(upload).upload(output_dir) + + assert sorted(result) == ["result.json", "system_output"] + (request,) = [r for r in httpx_mock.get_requests() if r.url == "https://example.com/system_output"] + assert zipfile.ZipFile(io.BytesIO(request.read())).namelist() == ["a\\b"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "upload", + [ + ZipAndHttpPutUpload(url="https://example.com/upload"), + MultiUpload( + uploads=[SingleFilePutUpload(url="https://example.com/upload1", relative_path="result.json")], + system_output=ZipAndHttpPutUpload(url="https://example.com/system_output"), + ), + ], + ) + async def test_zip_with_name_that_is_not_utf8_fails_before_any_request( + self, output_dir, non_utf8_named_file, httpx_mock, upload + ): + uploader = OutputUploader.for_upload_output(upload) + + with pytest.raises(OutputUploadFailed, match="cannot be zipped"): + await uploader.upload(output_dir) + + assert httpx_mock.get_requests() == [] + + @pytest.fixture + def directory_nested_101_levels_deep(self, output_dir): + deepest = output_dir.joinpath(*["d"] * 101) + deepest.mkdir(parents=True) + (deepest / "file").write_bytes(b"deep") + + @pytest.mark.asyncio + async def test_zip_of_deeply_nested_directories_fails_before_any_request( + self, output_dir, directory_nested_101_levels_deep, httpx_mock + ): + uploader = ZipAndHTTPPutOutputUploader(ZipAndHttpPutUpload(url="https://example.com/upload")) + + with pytest.raises(OutputUploadFailed, match="is nested more than 100 directories deep"): + await uploader.upload(output_dir) + + assert httpx_mock.get_requests() == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "mtime,date_time", + [ + (0, (1980, 1, 1, 0, 0, 0)), + # Zip stores seconds in 2-second steps. + (calendar.timegm((2200, 1, 1, 0, 0, 0)), (2107, 12, 31, 23, 59, 58)), + ], + ) + async def test_file_with_mtime_outside_zip_date_range_is_zipped(self, output_dir, httpx_mock, mtime, date_time): + os.utime(output_dir / "result.json", (mtime, mtime)) + httpx_mock.add_response(url="https://example.com/upload", method="PUT") + + await ZipAndHTTPPutOutputUploader(ZipAndHttpPutUpload(url="https://example.com/upload")).upload(output_dir) + + archive = zipfile.ZipFile(io.BytesIO(httpx_mock.get_requests()[0].read())) + assert archive.getinfo("result.json").date_time == date_time + + @pytest.mark.parametrize("mtime", [1e300, float("nan")]) + def test_mtime_beyond_platform_time_range_gets_first_zip_date(self, mtime): + assert _zip_date_time(mtime) == (1980, 1, 1, 0, 0, 0) + + +def zip_payload(request_body: bytes) -> bytes: + """The zip file from a PUT body or from the `file` part of a multipart POST body.""" + start = request_body.find(b"PK\x03\x04") + end = request_body.rfind(b"PK\x05\x06") + return request_body[start : end + 22] + + +class TestPutUploadBody: + @pytest.fixture + def sent_chunk_sizes(self, monkeypatch): + """Sizes of the request body chunks in the order httpx passes them to the transport.""" + sizes = [] + + async def handle_async_request(transport, request): + sizes.extend([len(chunk) async for chunk in request.stream]) + return httpx.Response(200) + + monkeypatch.setattr(httpx.AsyncHTTPTransport, "handle_async_request", handle_async_request) + return sizes + + @pytest.fixture + def output_dir(self, tmp_path): + (tmp_path / "result.bin").write_bytes(b"x" * 150_000) + return tmp_path + + @pytest.mark.asyncio + async def test_file_without_newlines_is_sent_in_bounded_chunks(self, output_dir, sent_chunk_sizes): + upload = SingleFilePutUpload(url="https://example.com/upload", relative_path="result.bin") + + await OutputUploader.for_upload_output(upload).upload(output_dir) + + assert sent_chunk_sizes == [65_536, 65_536, 18_928] + + +class TestReadStreamWithCap: + @pytest.fixture + def chunks_consumed(self): + return [] + + @pytest.fixture + def endless_response(self, httpx_mock, chunks_consumed): + def chunks(): + for i in range(1000): + chunks_consumed.append(i) + yield b"x" * 60 + + httpx_mock.add_response(url="https://example.com/response", stream=IteratorStream(chunks())) + + @pytest.mark.asyncio + @pytest.mark.usefixtures("endless_response") + async def test_reading_stops_at_the_first_chunk_over_the_cap(self, chunks_consumed): + async with httpx.AsyncClient() as client, client.stream("GET", "https://example.com/response") as response: + with pytest.raises(OutputUploadFailed, match="exceeds allowed limit"): + await read_stream_with_cap(response, max_allowed_size=100) + + assert chunks_consumed == [0, 1] + + @pytest.mark.asyncio + async def test_body_within_the_cap_is_returned(self, httpx_mock): + httpx_mock.add_response(url="https://example.com/response", stream=IteratorStream([b"ab", b"cd"])) + + async with httpx.AsyncClient() as client, client.stream("GET", "https://example.com/response") as response: + assert await read_stream_with_cap(response, max_allowed_size=4) == "abcd" diff --git a/compute_horde_sdk/tests/unit/core/test_untrusted_files.py b/compute_horde_sdk/tests/unit/core/test_untrusted_files.py new file mode 100644 index 000000000..65d42ea89 --- /dev/null +++ b/compute_horde_sdk/tests/unit/core/test_untrusted_files.py @@ -0,0 +1,382 @@ +import os +import pathlib +import socket +import stat + +import pytest + +from compute_horde_core.untrusted_files import ( + DirectoryTooDeep, + EntryNotReplaceable, + FileSizeChanged, + FileTooLarge, + InvalidRelativePath, + SymlinkRejected, + UnexpectedFileType, + list_entries, + open_directory, + open_regular_file, + read_bounded, + read_regular_file, + split_relative_path, + walk_regular_files, + write_regular_file, +) + +MAX_SIZE = 1000 + + +@pytest.fixture +def secret_file(tmp_path): + path = tmp_path / "secret" + path.write_bytes(b"EXECUTOR_TOKEN=do-not-leak") + return path + + +@pytest.fixture +def directory(tmp_path): + path = tmp_path / "untrusted" + path.mkdir() + return path + + +@pytest.fixture +def dir_fd(directory): + with open_directory(directory) as fd: + yield fd + + +@pytest.fixture +def unix_socket(directory): + with socket.socket(socket.AF_UNIX) as sock: + sock.bind(str(directory / "socket")) + yield + + +def test_regular_file_is_read(directory, dir_fd): + (directory / "artifact.txt").write_bytes(b"artifact content") + + assert read_regular_file(dir_fd, "artifact.txt", max_size=MAX_SIZE) == b"artifact content" + + +def test_empty_file_is_read(directory, dir_fd): + (directory / "empty").write_bytes(b"") + + assert read_regular_file(dir_fd, "empty", max_size=MAX_SIZE) == b"" + + +def test_file_exactly_at_limit_is_read(directory, dir_fd): + (directory / "at-limit").write_bytes(b"x" * MAX_SIZE) + + assert read_regular_file(dir_fd, "at-limit", max_size=MAX_SIZE) == b"x" * MAX_SIZE + + +def test_file_over_limit_is_rejected(directory, dir_fd): + (directory / "over-limit").write_bytes(b"x" * (MAX_SIZE + 1)) + + with pytest.raises(FileTooLarge): + read_regular_file(dir_fd, "over-limit", max_size=MAX_SIZE) + + +def test_sparse_file_over_limit_is_rejected(directory, dir_fd): + with open(directory / "sparse", "wb") as f: + f.truncate(8 * 1024 * 1024) + + with pytest.raises(FileTooLarge): + read_regular_file(dir_fd, "sparse", max_size=MAX_SIZE) + + +def test_symlink_to_file_outside_directory_is_rejected(directory, dir_fd, secret_file): + os.symlink(secret_file, directory / "leak") + + with pytest.raises(SymlinkRejected): + read_regular_file(dir_fd, "leak", max_size=MAX_SIZE) + + +def test_symlink_to_proc_environ_is_rejected(directory, dir_fd): + os.symlink("/proc/self/environ", directory / "leak") + + with pytest.raises(SymlinkRejected): + read_regular_file(dir_fd, "leak", max_size=MAX_SIZE) + + +def test_symlink_to_file_inside_directory_is_rejected(directory, dir_fd): + (directory / "artifact.txt").write_bytes(b"artifact content") + os.symlink("artifact.txt", directory / "link") + + with pytest.raises(SymlinkRejected): + read_regular_file(dir_fd, "link", max_size=MAX_SIZE) + + +def test_dangling_symlink_is_rejected(directory, dir_fd): + os.symlink("missing", directory / "dangling") + + with pytest.raises(SymlinkRejected): + read_regular_file(dir_fd, "dangling", max_size=MAX_SIZE) + + +def test_directory_entry_is_rejected(directory, dir_fd): + (directory / "subdir").mkdir() + + with pytest.raises(UnexpectedFileType, match="is a directory"): + read_regular_file(dir_fd, "subdir", max_size=MAX_SIZE) + + +def test_fifo_is_rejected_without_blocking(directory, dir_fd): + os.mkfifo(directory / "fifo") + + with pytest.raises(UnexpectedFileType, match="is a FIFO"): + read_regular_file(dir_fd, "fifo", max_size=MAX_SIZE) + + +def test_socket_is_rejected(dir_fd, unix_socket): + with pytest.raises(UnexpectedFileType, match="is a socket"): + read_regular_file(dir_fd, "socket", max_size=MAX_SIZE) + + +@pytest.mark.skipif(os.geteuid() != 0, reason="creating device nodes requires root") +def test_device_node_is_rejected(directory, dir_fd): + os.mknod(directory / "null", mode=0o600 | os.makedev(1, 3)) + + with pytest.raises(UnexpectedFileType, match="is a character device"): + read_regular_file(dir_fd, "null", max_size=MAX_SIZE) + + +def test_missing_entry_raises_os_error(dir_fd): + with pytest.raises(FileNotFoundError): + read_regular_file(dir_fd, "missing", max_size=MAX_SIZE) + + +def test_multi_component_path_is_read(directory, dir_fd): + (directory / "a" / "b").mkdir(parents=True) + (directory / "a" / "b" / "file.txt").write_bytes(b"nested") + + assert read_regular_file(dir_fd, "a/b/file.txt", max_size=MAX_SIZE) == b"nested" + + +def test_symlinked_intermediate_directory_is_rejected(directory, dir_fd, secret_file): + os.symlink(secret_file.parent, directory / "escape") + + with pytest.raises(SymlinkRejected, match="^escape: "): + read_regular_file(dir_fd, f"escape/{secret_file.name}", max_size=MAX_SIZE) + + +def test_regular_file_as_intermediate_component_is_rejected(directory, dir_fd): + (directory / "file").write_bytes(b"not a directory") + + with pytest.raises(UnexpectedFileType, match="^file: is a regular file, expected a directory"): + read_regular_file(dir_fd, "file/child", max_size=MAX_SIZE) + + +MALFORMED_RELATIVE_PATHS = [ + "", + ".", + "/", + "/etc/passwd", + "//etc/passwd", + "../secret", + "a/../../secret", + "./file", + "a/./file", + "a//file", + "dir/", + "a\0file", +] + + +@pytest.mark.parametrize("relative_path", MALFORMED_RELATIVE_PATHS) +def test_malformed_relative_path_is_rejected(dir_fd, relative_path): + with pytest.raises(InvalidRelativePath): + read_regular_file(dir_fd, relative_path, max_size=MAX_SIZE) + + +@pytest.mark.parametrize( + "relative_path", + [ + *MALFORMED_RELATIVE_PATHS, + "a\\file", + pytest.param("x" * 256, id="256-byte-name"), + pytest.param("dir/" + "é" * 128, id="256-byte-utf8-name"), + ], +) +def test_invalid_relative_path_is_a_value_error(relative_path): + with pytest.raises(ValueError): + split_relative_path(relative_path) + + +def test_backslash_is_part_of_a_name(directory, dir_fd): + (directory / "a\\file").write_bytes(b"backslash") + + assert read_regular_file(dir_fd, "a\\file", max_size=MAX_SIZE) == b"backslash" + + +@pytest.mark.parametrize( + "relative_path,components", + [ + ("file", ["file"]), + ("a/b/file.txt", ["a", "b", "file.txt"]), + (".hidden", [".hidden"]), + ("a b/c", ["a b", "c"]), + pytest.param("x" * 255, ["x" * 255], id="255-byte-name"), + pytest.param("é" * 127, ["é" * 127], id="254-byte-utf8-name"), + ], +) +def test_split_relative_path(relative_path, components): + assert split_relative_path(relative_path) == components + + +def test_walk_regular_files_skips_everything_but_regular_files(directory, dir_fd, secret_file): + (directory / "a.txt").write_bytes(b"") + (directory / "sub" / "deeper").mkdir(parents=True) + (directory / "sub" / "b.txt").write_bytes(b"") + (directory / "sub" / "deeper" / "c.txt").write_bytes(b"") + (directory / "sub" / "back\\slash").write_bytes(b"") + (directory / "empty-dir").mkdir() + os.symlink(secret_file, directory / "leak") + os.symlink(secret_file.parent, directory / "escape") + os.symlink("a.txt", directory / "sub" / "link-inside") + os.mkfifo(directory / "fifo") + + assert sorted(walk_regular_files(dir_fd)) == ["a.txt", "sub/b.txt", "sub/back\\slash", "sub/deeper/c.txt"] + + +def test_walk_regular_files_enters_directories_100_levels_deep(directory, dir_fd): + deepest = directory.joinpath(*["d"] * 100) + deepest.mkdir(parents=True) + (deepest / "file").write_bytes(b"") + + assert list(walk_regular_files(dir_fd)) == ["d/" * 100 + "file"] + + +def test_walk_regular_files_rejects_directory_101_levels_deep(directory, dir_fd): + directory.joinpath(*["d"] * 101).mkdir(parents=True) + + with pytest.raises(DirectoryTooDeep, match=f"^{'d/' * 100}d: is nested more than 100 directories deep$"): + list(walk_regular_files(dir_fd)) + + +def test_file_grown_after_open_is_rejected(directory, dir_fd): + path = directory / "growing" + path.write_bytes(b"x" * 10) + + with open_regular_file(dir_fd, "growing") as file: + with path.open("ab") as f: + f.write(b"y" * 10) + with pytest.raises(FileSizeChanged): + read_bounded(file, max_size=MAX_SIZE) + + +def test_file_truncated_after_open_is_rejected(directory, dir_fd): + path = directory / "shrinking" + path.write_bytes(b"x" * 10) + + with open_regular_file(dir_fd, "shrinking") as file: + path.write_bytes(b"x" * 5) + with pytest.raises(FileSizeChanged): + read_bounded(file, max_size=MAX_SIZE) + + +def test_open_directory_rejects_symlink(tmp_path, directory): + os.symlink(directory, tmp_path / "link-to-dir") + + with pytest.raises(SymlinkRejected): + with open_directory(tmp_path / "link-to-dir"): + pass + + +def test_open_directory_rejects_regular_file(tmp_path): + (tmp_path / "file").write_bytes(b"") + + with pytest.raises(UnexpectedFileType, match="is a regular file, expected a directory"): + with open_directory(tmp_path / "file"): + pass + + +def test_list_entries_lists_everything_without_following(directory, dir_fd, secret_file): + (directory / "artifact.txt").write_bytes(b"") + os.symlink(secret_file, directory / "leak") + os.mkfifo(directory / "fifo") + + assert sorted(list_entries(dir_fd)) == ["artifact.txt", "fifo", "leak"] + + +@pytest.mark.parametrize("data", ["job output", b"job output"]) +def test_write_regular_file_creates_regular_file(directory, dir_fd, data): + write_regular_file(dir_fd, "stdout.txt", data) + + path = directory / "stdout.txt" + assert stat.S_ISREG(path.lstat().st_mode) + assert path.read_bytes() == b"job output" + + +def test_write_regular_file_replaces_existing_file(directory, dir_fd): + (directory / "stdout.txt").write_bytes(b"left by the job") + + write_regular_file(dir_fd, "stdout.txt", "job output") + + assert (directory / "stdout.txt").read_bytes() == b"job output" + + +def test_write_regular_file_replaces_symlink_without_following(directory, dir_fd, secret_file): + os.symlink(secret_file, directory / "stdout.txt") + + write_regular_file(dir_fd, "stdout.txt", "job output") + + path = directory / "stdout.txt" + assert stat.S_ISREG(path.lstat().st_mode) + assert path.read_bytes() == b"job output" + assert secret_file.read_bytes() == b"EXECUTOR_TOKEN=do-not-leak" + + +def test_write_regular_file_replaces_fifo_without_blocking(directory, dir_fd): + os.mkfifo(directory / "stdout.txt") + + write_regular_file(dir_fd, "stdout.txt", "job output") + + path = directory / "stdout.txt" + assert stat.S_ISREG(path.lstat().st_mode) + assert path.read_bytes() == b"job output" + + +def test_write_regular_file_rejects_directory(directory, dir_fd): + (directory / "stdout.txt").mkdir() + + with pytest.raises(EntryNotReplaceable, match="^stdout.txt: cannot be replaced"): + write_regular_file(dir_fd, "stdout.txt", "job output") + + assert (directory / "stdout.txt").is_dir() + + +@pytest.fixture +def outside_directory(tmp_path): + path = tmp_path / "outside" + path.mkdir() + (path / "x").write_bytes(b"outside content") + return path + + +@pytest.fixture +def symlink_to_outside_directory(directory, outside_directory): + os.symlink(outside_directory, directory / "sub") + + +def tree_contents(root: pathlib.Path) -> dict[str, bytes | str]: + """Regular files mapped to their content and symlinks to their target, keyed by path relative to `root`.""" + return { + str(path.relative_to(root)): os.readlink(path) if path.is_symlink() else path.read_bytes() + for path in root.rglob("*") + if path.is_symlink() or path.is_file() + } + + +@pytest.mark.parametrize("name", ["sub/x", "../x", "", ".", "..", "x/", "x\0"]) +def test_write_regular_file_rejects_anything_but_a_single_file_name( + tmp_path, dir_fd, outside_directory, symlink_to_outside_directory, name +): + with pytest.raises(InvalidRelativePath, match="must be a single file name$"): + write_regular_file(dir_fd, name, "job output") + + assert tree_contents(tmp_path) == { + "outside/x": b"outside content", + "untrusted/sub": str(outside_directory), + } diff --git a/compute_horde_sdk/tests/unit/internal/test_output_volume.py b/compute_horde_sdk/tests/unit/internal/test_output_volume.py new file mode 100644 index 000000000..d8c374e78 --- /dev/null +++ b/compute_horde_sdk/tests/unit/internal/test_output_volume.py @@ -0,0 +1,44 @@ +import pytest + +from compute_horde_core.output_upload import SingleFilePostUpload, SingleFilePutUpload +from compute_horde_sdk._internal.models import HTTPOutputVolume + +URL = "https://example.com/upload" + + +@pytest.mark.parametrize( + "mount_path,relative_path", + [ + ("/output/result.json", "result.json"), + ("/output/dir/file.bin", "dir/file.bin"), + ("/output/.hidden", ".hidden"), + ], +) +def test_valid_output_volume_path_is_converted(mount_path, relative_path): + assert HTTPOutputVolume(http_method="PUT", url=URL).to_compute_horde_output_upload( + mount_path + ) == SingleFilePutUpload(url=URL, relative_path=relative_path) + assert HTTPOutputVolume(http_method="POST", url=URL).to_compute_horde_output_upload( + mount_path + ) == SingleFilePostUpload(url=URL, relative_path=relative_path) + + +@pytest.mark.parametrize( + "mount_path", + [ + "/output", + "/output/", + "/output//proc/self/environ", + "/output/../secret", + "/output/dir/../../secret", + "/output/./result.json", + "/output/dir/", + "/output/a//b", + "/output/a\\b", + "/volume/result.json", + "output/result.json", + ], +) +def test_invalid_output_volume_path_is_rejected(mount_path): + with pytest.raises(ValueError, match="Output volume path"): + HTTPOutputVolume(http_method="PUT", url=URL).to_compute_horde_output_upload(mount_path) diff --git a/executor/app/src/compute_horde_executor/executor/job_runner.py b/executor/app/src/compute_horde_executor/executor/job_runner.py index f6f26f2b9..f646a4a7a 100644 --- a/executor/app/src/compute_horde_executor/executor/job_runner.py +++ b/executor/app/src/compute_horde_executor/executor/job_runner.py @@ -2,7 +2,6 @@ import base64 import json import logging -import os import pathlib import re import shlex @@ -26,6 +25,14 @@ start_nginx, ) from compute_horde_core.output_upload import OutputUploader, OutputUploadFailed +from compute_horde_core.untrusted_files import ( + FileTooLarge, + UntrustedFileRejected, + list_entries, + open_directory, + read_regular_file, + write_regular_file, +) from compute_horde_core.volume import ( HuggingfaceVolume, Volume, @@ -49,6 +56,8 @@ TRUNCATED_RESPONSE_PREFIX_LEN = MAX_RESULT_SIZE_IN_RESPONSE // 2 TRUNCATED_RESPONSE_SUFFIX_LEN = MAX_RESULT_SIZE_IN_RESPONSE // 2 MAX_ARTIFACT_SIZE = 1_000_000 +MAX_ARTIFACTS_TOTAL_SIZE = 10_000_000 +MAX_ARTIFACTS_COUNT = 1_000 DOCKER_STOP_TIMEOUT_SECONDS = 15 NGINX_CONF = """ @@ -324,28 +333,11 @@ async def upload_results(self) -> JobResult: ) # Save the streams in output volume and truncate them in response. - with open(self.output_volume_mount_dir / "stdout.txt", "w") as f: - f.write(self.execution_result.stdout) + self.save_output_streams() stdout = truncate(self.execution_result.stdout) - with open(self.output_volume_mount_dir / "stderr.txt", "w") as f: - f.write(self.execution_result.stderr) stderr = truncate(self.execution_result.stderr) - artifacts = {} - for artifact_filename in os.listdir(self.artifacts_mount_dir): - artifact_path = self.artifacts_mount_dir / artifact_filename - if os.path.isfile(artifact_path): - with open(artifact_path, "rb") as f: - content = f.read() - artifact_size = len(content) - if artifact_size < MAX_ARTIFACT_SIZE: - artifacts[f"{self.full_job_request.artifacts_dir}/{artifact_filename}"] = ( - base64.b64encode(content).decode() - ) - else: - logger.error(f"Artifact {artifact_filename} too large: {artifact_size:,} bytes") - else: - logger.error(f"Directory found in artifacts: {artifact_filename}") + artifacts = self.collect_artifacts() upload_results = {} @@ -372,6 +364,64 @@ async def upload_results(self) -> JobResult: ), ) + def save_output_streams(self) -> None: + """ + Write the job's stdout and stderr as stdout.txt and stderr.txt in the output directory, + replacing whatever the job left under those names without following it. + """ + assert self.execution_result is not None, "No execution result" + try: + with open_directory(self.output_volume_mount_dir) as dir_fd: + write_regular_file(dir_fd, "stdout.txt", self.execution_result.stdout) + write_regular_file(dir_fd, "stderr.txt", self.execution_result.stderr) + except UntrustedFileRejected as ex: + raise JobError( + f"Saving output streams failed: {ex}", JobFailureReason.UPLOAD_FAILED + ) from ex + + def collect_artifacts(self) -> dict[str, str]: + """ + Base64-encode the regular files the job left in the artifacts directory, keyed by + their path inside the job container. Anything else the job put there (symlinks, + directories, FIFOs, devices, files of MAX_ARTIFACT_SIZE bytes or more) is skipped. + Only the first MAX_ARTIFACTS_COUNT entries in name order are visited, and a file + that does not fit in what is left of the MAX_ARTIFACTS_TOTAL_SIZE bytes shared by + all artifacts is skipped without being read. + """ + assert self.full_job_request is not None, ( + "Full job request must be set. Call prepare_full() first." + ) + artifacts = {} + budget_left = MAX_ARTIFACTS_TOTAL_SIZE + with open_directory(self.artifacts_mount_dir) as dir_fd: + artifact_filenames = sorted(list_entries(dir_fd)) + if len(artifact_filenames) > MAX_ARTIFACTS_COUNT: + logger.error( + f"Artifacts skipped: only the first {MAX_ARTIFACTS_COUNT:,} of " + f"{len(artifact_filenames):,} entries in name order are collected" + ) + for artifact_filename in artifact_filenames[:MAX_ARTIFACTS_COUNT]: + max_size = min(MAX_ARTIFACT_SIZE - 1, budget_left) + try: + content = read_regular_file(dir_fd, artifact_filename, max_size=max_size) + except FileTooLarge as e: + if max_size == budget_left: + logger.error( + f"Artifact skipped: {artifact_filename}: would exceed the " + f"{MAX_ARTIFACTS_TOTAL_SIZE:,} bytes allowed for all artifacts together" + ) + else: + logger.error(f"Artifact skipped: {e}") + continue + except (UntrustedFileRejected, OSError) as e: + logger.error(f"Artifact skipped: {e}") + continue + budget_left -= len(content) + artifacts[f"{self.full_job_request.artifacts_dir}/{artifact_filename}"] = ( + base64.b64encode(content).decode() + ) + return artifacts + async def clean(self): # remove input/output directories with docker, to deal with funky file permissions root_for_remove = pathlib.Path("/temp_dir/") diff --git a/executor/app/src/compute_horde_executor/executor/tests/test_job_runner.py b/executor/app/src/compute_horde_executor/executor/tests/test_job_runner.py new file mode 100644 index 000000000..eaa7a8d33 --- /dev/null +++ b/executor/app/src/compute_horde_executor/executor/tests/test_job_runner.py @@ -0,0 +1,265 @@ +import base64 +import logging +import os +import shutil +import stat +import uuid + +import pytest +from asgiref.sync import async_to_sync +from compute_horde.job_errors import JobError +from compute_horde.protocol_consts import JobFailureReason +from compute_horde.protocol_messages import V0InitialJobRequest, V0JobRequest +from compute_horde_core import untrusted_files +from compute_horde_core.output_upload import SingleFilePutUpload + +from compute_horde_executor.executor import job_runner +from compute_horde_executor.executor.job_runner import MAX_ARTIFACT_SIZE, DefaultJobRunner +from compute_horde_executor.executor.miner_client import ExecutionResult, JobResult + +JOB_UUID = str(uuid.uuid4()) +EXECUTOR_CLASS = "spin_up-4min.gpu-24gb" +DOCKER_IMAGE = "backenddevelopersltd/compute-horde-job-echo:v0-latest" + + +@pytest.fixture +def runner(): + runner = DefaultJobRunner() + runner.initial_job_request = V0InitialJobRequest.model_validate( + { + "job_uuid": JOB_UUID, + "executor_class": EXECUTOR_CLASS, + "docker_image": DOCKER_IMAGE, + "timeout_seconds": 10, + "job_started_receipt_payload": { + "job_uuid": JOB_UUID, + "miner_hotkey": "miner_hotkey", + "validator_hotkey": "validator_hotkey", + "timestamp": "2025-01-01T00:00:00+00:00", + "executor_class": EXECUTOR_CLASS, + "is_organic": True, + "ttl": 5, + }, + "job_started_receipt_signature": "blah", + } + ) + runner.full_job_request = V0JobRequest( + job_uuid=JOB_UUID, + executor_class=EXECUTOR_CLASS, + docker_image=DOCKER_IMAGE, + docker_run_options_preset="none", + docker_run_cmd=[], + artifacts_dir="/artifacts", + ) + runner.execution_result = ExecutionResult( + timed_out=False, return_code=0, stdout="job stdout", stderr="job stderr" + ) + runner.output_volume_mount_dir.mkdir() + runner.artifacts_mount_dir.mkdir() + yield runner + shutil.rmtree(runner.temp_dir) + + +def job_runner_errors(caplog) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.name == job_runner.logger.name and record.levelno == logging.ERROR + ] + + +@pytest.fixture +def secret_file(tmp_path): + path = tmp_path / "secret" + path.write_bytes(b"EXECUTOR_TOKEN=do-not-leak") + return path + + +@pytest.fixture +def hostile_artifacts(runner, secret_file): + artifacts_dir = runner.artifacts_mount_dir + (artifacts_dir / "normal.txt").write_bytes(b"artifact content") + os.symlink(secret_file, artifacts_dir / "secret-link") + os.symlink("/proc/self/environ", artifacts_dir / "environ-link") + os.symlink("missing", artifacts_dir / "dangling-link") + os.mkfifo(artifacts_dir / "fifo") + (artifacts_dir / "subdir").mkdir() + (artifacts_dir / "too-large.bin").write_bytes(b"x" * MAX_ARTIFACT_SIZE) + with open(artifacts_dir / "huge-sparse.bin", "wb") as f: + f.truncate(8 * 1024 * 1024) + + +def test_upload_results_collects_only_regular_files(runner, hostile_artifacts, caplog): + result = async_to_sync(runner.upload_results)() + + assert result == JobResult( + exit_status=0, + timeout=False, + stdout="job stdout", + stderr="job stderr", + artifacts={"/artifacts/normal.txt": base64.b64encode(b"artifact content").decode()}, + upload_results={}, + ) + assert job_runner_errors(caplog) == [ + "Artifact skipped: dangling-link: is a symlink", + "Artifact skipped: environ-link: is a symlink", + "Artifact skipped: fifo: is a FIFO, expected a regular file", + "Artifact skipped: huge-sparse.bin: is 8,388,608 bytes, more than the allowed 999,999 bytes", + "Artifact skipped: secret-link: is a symlink", + "Artifact skipped: subdir: is a directory, expected a regular file", + "Artifact skipped: too-large.bin: is 1,000,000 bytes, more than the allowed 999,999 bytes", + ] + + +@pytest.fixture +def artifact_with_backslash_in_name(runner): + (runner.artifacts_mount_dir / "a\\b.txt").write_bytes(b"artifact content") + + +def test_upload_results_collects_artifact_with_backslash_in_name( + runner, artifact_with_backslash_in_name +): + result = async_to_sync(runner.upload_results)() + + assert result == JobResult( + exit_status=0, + timeout=False, + stdout="job stdout", + stderr="job stderr", + artifacts={"/artifacts/a\\b.txt": base64.b64encode(b"artifact content").decode()}, + upload_results={}, + ) + + +@pytest.fixture +def symlinked_stdout(runner, secret_file): + os.symlink(secret_file, runner.output_volume_mount_dir / "stdout.txt") + + +def test_upload_results_writes_stdout_without_following_symlink( + runner, secret_file, symlinked_stdout +): + async_to_sync(runner.upload_results)() + + stdout_path = runner.output_volume_mount_dir / "stdout.txt" + assert stat.S_ISREG(stdout_path.lstat().st_mode) + assert stdout_path.read_text() == "job stdout" + assert secret_file.read_bytes() == b"EXECUTOR_TOKEN=do-not-leak" + + +def test_upload_results_fails_when_stream_file_cannot_be_replaced(runner): + (runner.output_volume_mount_dir / "stdout.txt").mkdir() + + with pytest.raises(JobError, match="stdout.txt: cannot be replaced") as exc_info: + async_to_sync(runner.upload_results)() + + assert exc_info.value.reason == JobFailureReason.UPLOAD_FAILED + + +@pytest.fixture +def artifacts_over_budget(runner, monkeypatch): + monkeypatch.setattr(job_runner, "MAX_ARTIFACTS_TOTAL_SIZE", 10) + for name, size in [("a.bin", 4), ("b.bin", 4), ("c.bin", 4), ("d.bin", 2)]: + (runner.artifacts_mount_dir / name).write_bytes(b"x" * size) + + +def test_upload_results_skips_artifacts_that_do_not_fit_in_the_total_size_budget( + runner, artifacts_over_budget, caplog +): + result = async_to_sync(runner.upload_results)() + + assert result == JobResult( + exit_status=0, + timeout=False, + stdout="job stdout", + stderr="job stderr", + artifacts={ + "/artifacts/a.bin": base64.b64encode(b"xxxx").decode(), + "/artifacts/b.bin": base64.b64encode(b"xxxx").decode(), + "/artifacts/d.bin": base64.b64encode(b"xx").decode(), + }, + upload_results={}, + ) + assert job_runner_errors(caplog) == [ + "Artifact skipped: c.bin: would exceed the 10 bytes allowed for all artifacts together" + ] + + +@pytest.fixture +def artifacts_over_count(runner, monkeypatch): + monkeypatch.setattr(job_runner, "MAX_ARTIFACTS_COUNT", 2) + for name in ["c.txt", "a.txt", "b.txt"]: + (runner.artifacts_mount_dir / name).write_bytes(b"") + + +def test_upload_results_collects_only_the_first_artifacts_in_name_order( + runner, artifacts_over_count, caplog +): + result = async_to_sync(runner.upload_results)() + + assert result == JobResult( + exit_status=0, + timeout=False, + stdout="job stdout", + stderr="job stderr", + artifacts={"/artifacts/a.txt": "", "/artifacts/b.txt": ""}, + upload_results={}, + ) + assert job_runner_errors(caplog) == [ + "Artifacts skipped: only the first 2 of 3 entries in name order are collected" + ] + + +@pytest.fixture +def artifact_over_remaining_budget(runner, monkeypatch): + monkeypatch.setattr(job_runner, "MAX_ARTIFACTS_TOTAL_SIZE", 10) + (runner.artifacts_mount_dir / "a.bin").write_bytes(b"x" * 10) + (runner.artifacts_mount_dir / "b.bin").write_bytes(b"x" * 11) + + +@pytest.fixture +def read_chunk_sizes(monkeypatch): + sizes = [] + original_read = untrusted_files.os.read + + def recording_read(fd, length): + chunk = original_read(fd, length) + sizes.append(len(chunk)) + return chunk + + monkeypatch.setattr(untrusted_files.os, "read", recording_read) + return sizes + + +def test_upload_results_does_not_read_an_artifact_over_the_remaining_budget( + runner, artifact_over_remaining_budget, read_chunk_sizes +): + result = async_to_sync(runner.upload_results)() + + assert result == JobResult( + exit_status=0, + timeout=False, + stdout="job stdout", + stderr="job stderr", + artifacts={"/artifacts/a.bin": base64.b64encode(b"x" * 10).decode()}, + upload_results={}, + ) + assert read_chunk_sizes == [10, 0] + + +@pytest.fixture +def symlink_output_upload(runner, secret_file): + os.symlink(secret_file, runner.output_volume_mount_dir / "leak") + runner.full_job_request.output_upload = SingleFilePutUpload( + url="https://example.com/upload", relative_path="leak" + ) + + +def test_upload_results_does_not_follow_symlink_in_output_dir( + runner, symlink_output_upload, httpx_mock +): + with pytest.raises(JobError) as exc_info: + async_to_sync(runner.upload_results)() + + assert exc_info.value.reason == JobFailureReason.UPLOAD_FAILED + assert httpx_mock.get_requests() == [] diff --git a/facilitator/app/src/project/core/migrations/0055_clear_malformed_job_uploads.py b/facilitator/app/src/project/core/migrations/0055_clear_malformed_job_uploads.py new file mode 100644 index 000000000..ef879adcf --- /dev/null +++ b/facilitator/app/src/project/core/migrations/0055_clear_malformed_job_uploads.py @@ -0,0 +1,44 @@ +import pydantic +from compute_horde_core.output_upload import SingleFileUpload +from django.db import migrations +from structlog import get_logger + +log = get_logger(__name__) + +UPLOADS_ADAPTER = pydantic.TypeAdapter(list[SingleFileUpload]) + + +def is_valid_uploads(uploads) -> bool: + # The field hands back the raw JSON string when its own validation fails. + try: + if isinstance(uploads, str): + UPLOADS_ADAPTER.validate_json(uploads) + else: + UPLOADS_ADAPTER.validate_python(uploads) + except pydantic.ValidationError: + return False + return True + + +def clear_malformed_uploads(apps, schema_editor): + """ + Empty Job.uploads where it no longer validates, e.g. an absolute relative_path stored + before such paths were rejected, since loading such a job raises ValidationError. + """ + Job = apps.get_model("core", "Job") + malformed = [ + uuid for uuid, uploads in Job.objects.values_list("uuid", "uploads").iterator() if not is_valid_uploads(uploads) + ] + if malformed: + Job.objects.filter(uuid__in=malformed).update(uploads=[]) + log.warning("Cleared job uploads that fail validation", job_uuids=[str(uuid) for uuid in malformed]) + + +class Migration(migrations.Migration): + dependencies = [ + ("core", "0054_cheatedjobreport"), + ] + + operations = [ + migrations.RunPython(clear_malformed_uploads, migrations.RunPython.noop), + ] diff --git a/facilitator/app/src/project/core/tests/test_clear_malformed_job_uploads.py b/facilitator/app/src/project/core/tests/test_clear_malformed_job_uploads.py new file mode 100644 index 000000000..6ab097ea1 --- /dev/null +++ b/facilitator/app/src/project/core/tests/test_clear_malformed_job_uploads.py @@ -0,0 +1,91 @@ +import importlib +import json +from datetime import timedelta + +import pytest +from compute_horde_core.output_upload import SingleFilePutUpload +from django.apps import apps +from django.core.exceptions import ValidationError +from django.db import connection +from django.utils.timezone import now +from rest_framework.test import APIClient + +from project.core.eviction import JOBS_RETENTION_PERIOD, evict_jobs +from project.core.models import Job + +migration = importlib.import_module("project.core.migrations.0055_clear_malformed_job_uploads") + +VALID_UPLOAD = SingleFilePutUpload(url="https://example.com/upload", relative_path="result.json") +VALID_UPLOADS = [VALID_UPLOAD.model_dump(mode="json")] +MALFORMED_UPLOADS = [{**VALID_UPLOADS[0], "relative_path": "/proc/self/environ"}] + + +def store_raw_uploads(job: Job, uploads: list[dict]) -> None: + with connection.cursor() as cursor: + cursor.execute(f"UPDATE {Job._meta.db_table} SET uploads = %s WHERE uuid = %s", [json.dumps(uploads), job.uuid]) + + +@pytest.fixture +def make_job(user, validator, signature, dummy_job_params): + def make(created_at, uploads): + job = Job.objects.create( + user=user, + validator=validator, + target_validator_hotkey=validator.ss58_address, + signature=signature.model_dump(), + created_at=created_at, + **dummy_job_params, + ) + store_raw_uploads(job, uploads) + return job + + return make + + +@pytest.fixture +def recent_malformed_job(make_job): + return make_job(now(), MALFORMED_UPLOADS) + + +@pytest.fixture +def expired_malformed_job(make_job): + return make_job(now() - JOBS_RETENTION_PERIOD - timedelta(days=1), MALFORMED_UPLOADS) + + +@pytest.fixture +def recent_valid_job(make_job): + return make_job(now(), VALID_UPLOADS) + + +@pytest.fixture +def owner_api_client(user): + client = APIClient() + client.force_authenticate(user=user) + return client + + +@pytest.mark.django_db(transaction=True) +def test_malformed_uploads_are_cleared(recent_malformed_job, recent_valid_job): + with pytest.raises(ValidationError): + Job.objects.get(pk=recent_malformed_job.pk) + + migration.clear_malformed_uploads(apps, None) + + assert {job.uuid: job.uploads for job in Job.objects.all()} == { + recent_malformed_job.uuid: [], + recent_valid_job.uuid: [VALID_UPLOAD], + } + + +@pytest.mark.django_db(transaction=True) +def test_jobs_with_cleared_uploads_can_be_evicted_and_listed( + recent_malformed_job, expired_malformed_job, owner_api_client +): + migration.clear_malformed_uploads(apps, None) + + evict_jobs() + response = owner_api_client.get("/api/v1/jobs/") + + assert list(Job.objects.values_list("uuid", flat=True)) == [recent_malformed_job.uuid] + assert response.status_code == 200 + assert [job["uuid"] for job in response.data["results"]] == [str(recent_malformed_job.uuid)] diff --git a/miner/app/src/compute_horde_miner/miner/migrations/0016_consumedvalidatorauth.py b/miner/app/src/compute_horde_miner/miner/migrations/0016_consumedvalidatorauth.py new file mode 100644 index 000000000..cb0414fb2 --- /dev/null +++ b/miner/app/src/compute_horde_miner/miner/migrations/0016_consumedvalidatorauth.py @@ -0,0 +1,26 @@ +# Generated by Django 4.2.19 on 2026-10-19 16:02 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("miner", "0015_acceptedjob_upload_results"), + ] + + operations = [ + migrations.CreateModel( + name="ConsumedValidatorAuth", + fields=[ + ( + "id", + models.BigAutoField( + auto_created=True, primary_key=True, serialize=False, verbose_name="ID" + ), + ), + ("signature", models.CharField(max_length=128, unique=True)), + ("validator_hotkey", models.CharField(max_length=256)), + ("timestamp", models.DateTimeField(db_index=True)), + ], + ), + ] diff --git a/miner/app/src/compute_horde_miner/miner/miner_consumer/base_compute_horde_consumer.py b/miner/app/src/compute_horde_miner/miner/miner_consumer/base_compute_horde_consumer.py index 98c4e8973..a591100ff 100644 --- a/miner/app/src/compute_horde_miner/miner/miner_consumer/base_compute_horde_consumer.py +++ b/miner/app/src/compute_horde_miner/miner/miner_consumer/base_compute_horde_consumer.py @@ -6,7 +6,7 @@ from channels.generic.websocket import AsyncWebsocketConsumer from compute_horde.protocol_messages import GenericError -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel logger = logging.getLogger(__name__) @@ -33,7 +33,10 @@ async def wrapper(*args: Params.args, **kwargs: Params.kwargs) -> TResult: class BaseConsumer(AsyncWebsocketConsumer, Generic[TReceived], metaclass=abc.ABCMeta): @abc.abstractmethod def parse_message(self, raw_msg: str | bytes) -> TReceived: - """Parse raw message into a pydantic model""" + """ + Parses a raw message into a pydantic model; raises ValueError, such as + pydantic.ValidationError or json.JSONDecodeError, when it is malformed. + """ @abc.abstractmethod async def handle(self, msg: TReceived) -> None: ... @@ -43,11 +46,13 @@ async def connect(self): @log_errors_explicitly async def receive(self, text_data=None, bytes_data=None) -> None: + if text_data is None: + await self.reject_malformed_message("expected a text frame") + return try: msg = self.parse_message(text_data) - except ValidationError as ex: - logger.error(f"Malformed message: {str(ex)}") - await self.send(GenericError(details=f"Malformed message: {str(ex)}").model_dump_json()) + except ValueError as ex: + await self.reject_malformed_message(str(ex)) return if isinstance(msg, GenericError): @@ -58,3 +63,7 @@ async def receive(self, text_data=None, bytes_data=None) -> None: return await self.handle(msg) + + async def reject_malformed_message(self, reason: str) -> None: + logger.error(f"Malformed message: {reason}") + await self.send(GenericError(details=f"Malformed message: {reason}").model_dump_json()) diff --git a/miner/app/src/compute_horde_miner/miner/miner_consumer/layer_utils.py b/miner/app/src/compute_horde_miner/miner/miner_consumer/layer_utils.py index 5bd7d6e58..91b5fc599 100644 --- a/miner/app/src/compute_horde_miner/miner/miner_consumer/layer_utils.py +++ b/miner/app/src/compute_horde_miner/miner/miner_consumer/layer_utils.py @@ -1,5 +1,6 @@ import abc import logging +from collections.abc import Awaitable, Callable from typing import Any, TypeVar import pydantic @@ -48,60 +49,83 @@ def validate_event( class ValidatorInterfaceMixin(BaseMixin, abc.ABC): + validator_authenticated: bool + @classmethod def group_name(cls, executor_token: str): return f"validator_interface_{executor_token}" + async def relay_executor_event( + self, + event_type: str, + models_class: type[TModel], + event: dict[str, Any], + handler: Callable[[TModel], Awaitable[None]], + ) -> None: + """ + Forwards an executor-originated channel-layer event to the handler that reports it + to the validator, unless the validator has not authenticated yet; job results, specs + and streaming endpoints must never reach an unauthenticated socket. + """ + payload = self.validate_event(event_type, models_class, event) + if payload is None: + return + if not self.validator_authenticated: + logger.warning( + f"Dropping {event_type} message for job {event.get('job_uuid')}: " + f"validator socket is not authenticated" + ) + return + await handler(payload) + @log_errors_explicitly async def executor_ready(self, event: dict[str, Any]): - payload = self.validate_event("executor_ready", V0ExecutorReadyRequest, event) - if payload: - await self._executor_ready(payload) + await self.relay_executor_event( + "executor_ready", V0ExecutorReadyRequest, event, self._executor_ready + ) @abc.abstractmethod async def _executor_ready(self, msg: V0ExecutorReadyRequest): ... @log_errors_explicitly async def streaming_job_ready(self, event: dict[str, Any]): - payload = self.validate_event("streaming_job_ready", V0StreamingJobReadyRequest, event) - if payload: - await self._streaming_job_ready(payload) + await self.relay_executor_event( + "streaming_job_ready", V0StreamingJobReadyRequest, event, self._streaming_job_ready + ) @abc.abstractmethod async def _streaming_job_ready(self, msg: V0StreamingJobReadyRequest): ... @log_errors_explicitly async def executor_finished(self, event: dict[str, Any]): - payload = self.validate_event("executor_finished", V0JobFinishedRequest, event) - if payload: - await self._executor_finished(payload) + await self.relay_executor_event( + "executor_finished", V0JobFinishedRequest, event, self._executor_finished + ) @abc.abstractmethod - async def _executor_specs(self, event: V0MachineSpecsRequest): ... + async def _executor_finished(self, msg: V0JobFinishedRequest): ... @log_errors_explicitly async def executor_specs(self, event: dict[str, Any]): - payload = self.validate_event("executor_specs", V0MachineSpecsRequest, event) - if payload: - await self._executor_specs(payload) + await self.relay_executor_event( + "executor_specs", V0MachineSpecsRequest, event, self._executor_specs + ) @abc.abstractmethod - async def _executor_finished(self, msg: V0JobFinishedRequest): ... + async def _executor_specs(self, event: V0MachineSpecsRequest): ... @log_errors_explicitly async def job_failed(self, event: dict[str, Any]): - payload = self.validate_event("job_failed", V0JobFailedRequest, event) - if payload: - await self._job_failed(payload) + await self.relay_executor_event("job_failed", V0JobFailedRequest, event, self._job_failed) @abc.abstractmethod async def _job_failed(self, msg: V0JobFailedRequest): ... @log_errors_explicitly async def horde_failed(self, event: dict[str, Any]): - payload = self.validate_event("horde_failed", V0HordeFailedRequest, event) - if payload: - await self._horde_failed(payload) + await self.relay_executor_event( + "horde_failed", V0HordeFailedRequest, event, self._horde_failed + ) @abc.abstractmethod async def _horde_failed(self, msg: V0HordeFailedRequest): ... @@ -114,18 +138,18 @@ async def send_job_request(self, executor_token, job_request: V0JobRequest): @log_errors_explicitly async def volumes_ready(self, event: dict[str, Any]) -> None: - payload = self.validate_event("volumes_ready", V0VolumesReadyRequest, event) - if payload: - await self._volumes_ready(payload) + await self.relay_executor_event( + "volumes_ready", V0VolumesReadyRequest, event, self._volumes_ready + ) @abc.abstractmethod async def _volumes_ready(self, msg: V0VolumesReadyRequest): ... @log_errors_explicitly async def execution_done(self, event: dict[str, Any]) -> None: - payload = self.validate_event("execution_done", V0ExecutionDoneRequest, event) - if payload: - await self._execution_done(payload) + await self.relay_executor_event( + "execution_done", V0ExecutionDoneRequest, event, self._execution_done + ) @abc.abstractmethod async def _execution_done(self, msg: V0ExecutionDoneRequest): ... diff --git a/miner/app/src/compute_horde_miner/miner/miner_consumer/validator_interface.py b/miner/app/src/compute_horde_miner/miner/miner_consumer/validator_interface.py index ec5795d4c..efc42afcb 100644 --- a/miner/app/src/compute_horde_miner/miner/miner_consumer/validator_interface.py +++ b/miner/app/src/compute_horde_miner/miner/miner_consumer/validator_interface.py @@ -1,15 +1,18 @@ import asyncio import datetime as dt import logging -import time import uuid from datetime import timedelta from functools import cached_property from typing import Protocol, assert_never -import bittensor from compute_horde.executor_class import EXECUTOR_CLASS -from compute_horde.protocol_consts import HordeFailureReason, JobParticipantType, JobRejectionReason +from compute_horde.protocol_consts import ( + TRUSTED_MINER_FAKE_KEY, + HordeFailureReason, + JobParticipantType, + JobRejectionReason, +) from compute_horde.protocol_messages import ( GenericError, V0AcceptJobRequest, @@ -30,13 +33,15 @@ V0VolumesReadyRequest, ValidatorAuthForMiner, ValidatorToMinerMessage, + parse_validator_to_miner_message, ) from compute_horde.receipts.models import JobAcceptedReceipt, JobFinishedReceipt, JobStartedReceipt -from compute_horde.receipts.schemas import JobStartedReceiptPayload, ReceiptPayload +from compute_horde.receipts.schemas import ReceiptPayload +from compute_horde.utils import verify_blob_signature from django.conf import settings +from django.db import IntegrityError from django.db.models import DateTimeField, ExpressionWrapper, F from django.utils import timezone -from pydantic import TypeAdapter from compute_horde_miner.miner.executor_manager import current from compute_horde_miner.miner.executor_manager.base import ( @@ -54,10 +59,20 @@ ValidatorBlacklist, ) from compute_horde_miner.miner.receipts import current_store +from compute_horde_miner.miner.validator_auth import ( + AUTH_MESSAGE_MAX_CLOCK_SKEW, + auth_signature_bytes, + auth_timestamp_error, + consume_auth_signature, +) logger = logging.getLogger(__name__) -AUTH_MESSAGE_MAX_AGE = 10 +# Seconds a connected validator has to authenticate before the socket is closed. +AUTH_TIMEOUT = 5 +# A validator dates the receipt right before sending the job request, but reconnecting to the +# miner can delay the request by up to about 5 minutes. +JOB_STARTED_RECEIPT_MAX_AGE = dt.timedelta(minutes=10) DONT_CHECK = "DONT_CHECK" @@ -80,15 +95,17 @@ class MissingExecutorToken(Exception): def __init__(self, *a, **kw): super().__init__(*a, **kw) - if settings.DEBUG_TURN_AUTHENTICATION_OFF or settings.IS_LOCAL_MINER: + if settings.DEBUG_TURN_AUTHENTICATION_OFF: self.my_hotkey = DONT_CHECK + elif settings.IS_LOCAL_MINER: + # Validators address their trusted miner by this key, not by its wallet's hotkey. + self.my_hotkey = TRUSTED_MINER_FAKE_KEY else: self.my_hotkey = settings.BITTENSOR_WALLET().get_hotkey().ss58_address self.validator_authenticated = False self.msg_queue = [] - self.defer_saving_jobs = [] - self.defer_executor_ready = [] self.pending_jobs: dict[str, AcceptedJob] = {} + self._auth_timeout_task: asyncio.Task[None] | None = None # Validator is populated only after the connection initialization succeeds self._maybe_validator: Validator | None = None @@ -118,7 +135,28 @@ async def connect(self): await self.close_with_error_msg(f"Inactive validator: {self.validator_key}") return + # The validator hotkey in the URL is public, so until the socket authenticates it must + # not be subscribed to any job nor receive anything; see handle_authentication. + self._auth_timeout_task = asyncio.create_task(self._close_on_auth_timeout()) + + async def _close_on_auth_timeout(self): + await asyncio.sleep(AUTH_TIMEOUT) + await self.close_with_error_msg( + f"Validator {self.validator_key} did not authenticate within {AUTH_TIMEOUT}s" + ) + + def _cancel_auth_timeout(self): + if self._auth_timeout_task is not None: + self._auth_timeout_task.cancel() + self._auth_timeout_task = None + + async def _resume_pending_jobs(self) -> list[AcceptedJob]: + """ + Re-subscribes to the validator's in-progress jobs after a reconnect, failing the ones + that went stale; returns the jobs whose executor is still waiting for the payload. + """ self.pending_jobs = await AcceptedJob.get_for_validator(self.validator) + jobs_awaiting_payload = [] for job in self.pending_jobs.values(): if job.status != AcceptedJob.Status.WAITING_FOR_PAYLOAD: # TODO: this actually works only for temporary connection issue between validator and miner; @@ -130,53 +168,46 @@ async def connect(self): # this only handles this one particular case of broken connection between miner and validator if timezone.now() - job.updated_at > dt.timedelta(days=1): job.status = AcceptedJob.Status.FAILED - # we don't want to block accepting connection - we defer it until after authorized - self.defer_saving_jobs.append(job) + await job.asave() logger.debug(f"Give up on job {job.job_uuid} after no status change for a day.") else: await self.group_add(job.executor_token) continue if timezone.now() - job.updated_at > dt.timedelta(minutes=10): job.status = AcceptedJob.Status.FAILED - # we don't want to block accepting connection - we defer it until after authorized - self.defer_saving_jobs.append(job) + await job.asave() logger.debug( f"Give up on job {job.job_uuid} after not receiving payload after 10 minutes" ) else: - logger.debug( - f"Deferring executor ready msg for job {job.job_uuid}: {job.executor_token}" - ) await self.group_add(job.executor_token) - # we don't send anything until we get authorization confirmation - self.defer_executor_ready.append(job) + jobs_awaiting_payload.append(job) + return jobs_awaiting_payload async def close_with_error_msg(self, msg: str): await self.send(GenericError(details=msg).model_dump_json()) logger.info(msg) await self.close(1000) - def verify_auth_msg(self, msg: ValidatorAuthForMiner) -> tuple[bool, str]: - if msg.timestamp < time.time() - AUTH_MESSAGE_MAX_AGE: - return False, "msg too old" - if not settings.IS_LOCAL_MINER and msg.miner_hotkey != self.my_hotkey: + async def verify_auth_msg(self, msg: ValidatorAuthForMiner) -> tuple[bool, str]: + if timestamp_error := auth_timestamp_error(msg): + return False, timestamp_error + if msg.miner_hotkey != self.my_hotkey: return False, f"wrong miner hotkey ({self.my_hotkey}!={msg.miner_hotkey})" if msg.validator_hotkey != self.validator_key: return ( False, f"wrong validator hotkey ({self.validator_key}!={msg.validator_hotkey})", ) - - keypair = bittensor.Keypair(ss58_address=self.validator_key) - if keypair.verify(msg.blob_for_signing(), msg.signature): - return True, "" - - return False, "Signature mismatches" + if (signature := auth_signature_bytes(msg)) is None: + return False, "malformed signature" + if not verify_blob_signature(self.validator_key, msg.blob_for_signing(), msg.signature): + return False, "Signature mismatches" + if not await consume_auth_signature(msg, signature): + return False, "auth message already used" + return True, "" def verify_receipt_payload(self, payload: ReceiptPayload, signature: str) -> bool: - if settings.IS_LOCAL_MINER: - return True - if self.my_hotkey != DONT_CHECK and payload.miner_hotkey != self.my_hotkey: logger.warning( f"Miner hotkey mismatch in receipt for job_uuid {payload.job_uuid} ({payload.miner_hotkey!r} != {self.my_hotkey!r})" @@ -188,21 +219,49 @@ def verify_receipt_payload(self, payload: ReceiptPayload, signature: str) -> boo ) return False - keypair = bittensor.Keypair(ss58_address=self.validator_key) - if keypair.verify(payload.blob_for_signing(), signature): + if verify_blob_signature(self.validator_key, payload.blob_for_signing(), signature): return True logger.warning(f"Validator signature mismatch in receipt for job_uuid {payload.job_uuid}") return False + def job_request_signature_error(self, msg: V0InitialJobRequest | V0JobRequest) -> str | None: + """ + Why the request cannot be attributed to the authenticated validator, or None. + A missing signature is tolerated only when REQUIRE_SIGNED_JOB_REQUESTS is off. + """ + if msg.validator_signature is None: + if settings.REQUIRE_SIGNED_JOB_REQUESTS: + return "missing validator signature" + logger.warning( + f"Accepting unsigned {msg.message_type} for job {msg.job_uuid} from validator " + f"{self.validator_key} because REQUIRE_SIGNED_JOB_REQUESTS is off" + ) + return None + if not msg.is_signed_by(self.validator_key): + return "invalid validator signature" + return None + + async def decline_job(self, job_uuid: str, reason: JobRejectionReason, message: str) -> None: + logger.info(f"Declining job {job_uuid} from validator {self.validator_key}: {message}") + await self.send( + V0DeclineJobRequest(job_uuid=job_uuid, reason=reason, message=message).model_dump_json() + ) + async def handle_authentication(self, msg: ValidatorAuthForMiner): + if self.validator_authenticated: + logger.warning( + f"Ignoring repeated authentication message from validator {self.validator_key}" + ) + return + if settings.DEBUG_TURN_AUTHENTICATION_OFF: logger.critical( f"Validator {self.validator_key} passed authentication without checking, because " f'"DEBUG_TURN_AUTHENTICATION_OFF" is on' ) else: - authenticated, auth_error_msg = self.verify_auth_msg(msg) + authenticated, auth_error_msg = await self.verify_auth_msg(msg) if not authenticated: close_error_msg = ( f"Validator {self.validator_key} not authenticated due to: {auth_error_msg}" @@ -211,9 +270,12 @@ async def handle_authentication(self, msg: ValidatorAuthForMiner): return self.validator_authenticated = True + self._cancel_auth_timeout() # (!) We should not send any messages until the validator authorizes itself. + jobs_awaiting_payload = await self._resume_pending_jobs() + manifest = await current.executor_manager.get_manifest() await self.send(V0ExecutorManifestRequest(manifest=manifest).model_dump_json()) @@ -256,20 +318,14 @@ async def handle_authentication(self, msg: ValidatorAuthForMiner): # job.result_reported_to_validator = timezone.now() # await job.asave() - while self.defer_executor_ready: - job = self.defer_executor_ready.pop() + for job in jobs_awaiting_payload: await self.send(V0ExecutorReadyRequest(job_uuid=str(job.job_uuid)).model_dump_json()) logger.debug( f"Readiness for job {job.job_uuid} reported to validator {self.validator_key}" ) - # we could do this anywhere, but this sounds like a good enough place - while self.defer_saving_jobs: - job = self.defer_saving_jobs.pop() - await job.asave() - def parse_message(self, raw_msg: str | bytes) -> ValidatorToMinerMessage: - return TypeAdapter(ValidatorToMinerMessage).validate_json(raw_msg) + return parse_validator_to_miner_message(raw_msg) async def handle(self, msg: ValidatorToMinerMessage) -> None: if self._maybe_validator is None: @@ -290,7 +346,12 @@ async def handle(self, msg: ValidatorToMinerMessage) -> None: self.msg_queue.append(msg) return - if isinstance(msg, V0InitialJobRequest) or isinstance(msg, V0JobRequest): + if isinstance(msg, V0InitialJobRequest | V0JobRequest): + if signature_error := self.job_request_signature_error(msg): + await self.decline_job( + msg.job_uuid, JobRejectionReason.INVALID_SIGNATURE, signature_error + ) + return # Proactively check volume safety in both requests that may contain a volume if msg.volume and not msg.volume.is_safe(): error_msg = f"Received JobRequest with unsafe volume: {msg.volume}" @@ -330,20 +391,14 @@ async def handle_initial_job_request(self, msg: V0InitialJobRequest): validator=self.validator ).aexists() if validator_blacklisted: - logger.info( - f"Declining job {msg.job_uuid} from blacklisted validator: {self.validator_key}" - ) - await self.send( - V0DeclineJobRequest( - job_uuid=msg.job_uuid, - reason=JobRejectionReason.VALIDATOR_BLACKLISTED, - ).model_dump_json() + await self.decline_job( + msg.job_uuid, JobRejectionReason.VALIDATOR_BLACKLISTED, "validator is blacklisted" ) return - await self.handle_job_started_receipt( - msg.job_started_receipt_payload, msg.job_started_receipt_signature - ) + if receipt_error := await self.accept_job_started_receipt(msg): + await self.decline_job(msg.job_uuid, JobRejectionReason.INVALID_RECEIPT, receipt_error) + return # TODO add rate limiting per validator key here executor_token = f"{msg.job_uuid}-{uuid.uuid4()}" @@ -476,30 +531,54 @@ async def handle_job_request(self, msg: V0JobRequest): job.full_job_details = msg.model_dump() await job.asave() - async def handle_job_started_receipt(self, payload: JobStartedReceiptPayload, signature: str): + async def accept_job_started_receipt(self, msg: V0InitialJobRequest) -> str | None: + """ + Verifies the job-started receipt carried by the request and records it, unless this is a + local miner, which keeps no receipts; returns why the receipt was refused, or None. + Each receipt is accepted once and only while it is recent, so that a captured request + cannot start its job again. + """ + payload = msg.job_started_receipt_payload + signature = msg.job_started_receipt_signature logger.info( f"Received job started receipt for" f" job_uuid={payload.job_uuid} validator_hotkey={payload.validator_hotkey}" ) - if settings.IS_LOCAL_MINER: - return - + if payload.job_uuid != msg.job_uuid: + return f"job started receipt is for job {payload.job_uuid}" + if payload.executor_class != msg.executor_class: + return f"job started receipt is for executor class {payload.executor_class}" if not self.verify_receipt_payload(payload, signature): - return + return "job started receipt could not be verified" + now = timezone.now() + if payload.timestamp < now - JOB_STARTED_RECEIPT_MAX_AGE: + return "job started receipt is too old" + if payload.timestamp > now + dt.timedelta(seconds=AUTH_MESSAGE_MAX_CLOCK_SKEW): + return "job started receipt is dated in the future" + if settings.IS_LOCAL_MINER: + if await AcceptedJob.objects.filter( + validator=self.validator, job_uuid=msg.job_uuid + ).aexists(): + return "job was already accepted" + return None - receipt = await JobStartedReceipt.objects.acreate( - job_uuid=payload.job_uuid, - validator_hotkey=payload.validator_hotkey, - miner_hotkey=payload.miner_hotkey, - validator_signature=signature, - miner_signature=get_miner_signature(payload), - timestamp=payload.timestamp, - executor_class=payload.executor_class, - is_organic=payload.is_organic, - ttl=payload.ttl, - ) + try: + receipt = await JobStartedReceipt.objects.acreate( + job_uuid=payload.job_uuid, + validator_hotkey=payload.validator_hotkey, + miner_hotkey=payload.miner_hotkey, + validator_signature=signature, + miner_signature=get_miner_signature(payload), + timestamp=payload.timestamp, + executor_class=payload.executor_class, + is_organic=payload.is_organic, + ttl=payload.ttl, + ) + except IntegrityError: + return "job started receipt was already used" (await current_store()).store([receipt.to_receipt()]) + return None async def handle_job_accepted_receipt(self, msg: V0JobAcceptedReceiptRequest): logger.info( @@ -658,4 +737,5 @@ async def _execution_done(self, msg: V0ExecutionDoneRequest): await self.send(msg.model_dump_json()) async def disconnect(self, close_code): + self._cancel_auth_timeout() logger.info(f"Validator {self.validator_key} disconnected") diff --git a/miner/app/src/compute_horde_miner/miner/models.py b/miner/app/src/compute_horde_miner/miner/models.py index 9087fe0de..d92bfb8ea 100644 --- a/miner/app/src/compute_horde_miner/miner/models.py +++ b/miner/app/src/compute_horde_miner/miner/models.py @@ -124,6 +124,21 @@ async def get_not_reported(cls, validator: Validator) -> Iterable[Self]: ] +class ConsumedValidatorAuth(models.Model): + """ + Signature of a validator authentication frame the miner has accepted, as 128 lowercase hex + digits. The same frame presented again is a replay; keying on the signature keeps separate, + legitimately signed frames apart even when they share a timestamp. + """ + + signature = models.CharField(max_length=128, unique=True) + validator_hotkey = models.CharField(max_length=256) + timestamp = models.DateTimeField(db_index=True) + + def __str__(self): + return f"validator hotkey: {self.validator_hotkey} - {self.timestamp}" + + class ClusterMiner(models.Model): hotkey = models.CharField(max_length=48, primary_key=True) block = models.BigIntegerField(null=True) diff --git a/miner/app/src/compute_horde_miner/miner/tests/executor_manager.py b/miner/app/src/compute_horde_miner/miner/tests/executor_manager.py index 584f668d8..e7d2aaab5 100644 --- a/miner/app/src/compute_horde_miner/miner/tests/executor_manager.py +++ b/miner/app/src/compute_horde_miner/miner/tests/executor_manager.py @@ -5,14 +5,18 @@ from channels.testing import WebsocketCommunicator from compute_horde.executor_class import DEFAULT_EXECUTOR_CLASS from compute_horde.protocol_messages import V0ExecutionDoneRequest, V0VolumesReadyRequest +from compute_horde.test_wallet import get_test_miner_wallet, get_test_validator_wallet -from compute_horde_miner import asgi from compute_horde_miner.miner.executor_manager import v1 WEBSOCKET_TIMEOUT = 100000 async def fake_executor(token, streaming=False): + # The test settings make the executor manager load this module while the consumers that + # make up the ASGI application are still being imported, so asgi cannot be imported above. + from compute_horde_miner import asgi + communicator = WebsocketCommunicator(asgi.application, f"v0.1/executor_interface/{token}") connected, _ = await communicator.connect() assert connected @@ -27,14 +31,16 @@ async def fake_executor(token, streaming=False): "job_started_receipt_payload": { "receipt_type": "JobStartedReceipt", "job_uuid": fake_executor.job_uuid, - "miner_hotkey": "miner_hotkey", - "validator_hotkey": "some_public_key", - "timestamp": "2020-01-01T00:00:00Z", + "miner_hotkey": fake_executor.miner_hotkey + or get_test_miner_wallet().hotkey.ss58_address, + "validator_hotkey": get_test_validator_wallet().hotkey.ss58_address, + "timestamp": mock.ANY, "executor_class": DEFAULT_EXECUTOR_CLASS, "is_organic": True, - "ttl": 5, + "ttl": 60, }, - "job_started_receipt_signature": "gibberish", + "job_started_receipt_signature": mock.ANY, + "validator_signature": mock.ANY, "executor_timing": None, "streaming_details": None if not streaming @@ -63,6 +69,7 @@ async def fake_executor(token, streaming=False): "volume": {"volume_type": "inline", "contents": "nonsense", "relative_path": None}, "output_upload": mock.ANY, "artifacts_dir": None, + "validator_signature": mock.ANY, }, response if streaming: await communicator.send_json_to( @@ -96,6 +103,8 @@ async def fake_executor(token, streaming=False): fake_executor.job_uuid = None +# Miner hotkey the receipt names, when the validator does not address the test miner. +fake_executor.miner_hotkey = None class StubExecutorManager(v1.BaseExecutorManager): diff --git a/miner/app/src/compute_horde_miner/miner/tests/integration/conftest.py b/miner/app/src/compute_horde_miner/miner/tests/integration/conftest.py new file mode 100644 index 000000000..7e7308777 --- /dev/null +++ b/miner/app/src/compute_horde_miner/miner/tests/integration/conftest.py @@ -0,0 +1,27 @@ +import pytest +import pytest_asyncio +from pytest_mock import MockerFixture + +from compute_horde_miner.miner.models import Validator +from compute_horde_miner.miner.tests.executor_manager import StubExecutorManager +from compute_horde_miner.miner.tests.validator import validator_hotkey + + +@pytest.fixture +def authentication_on(settings): + settings.DEBUG_TURN_AUTHENTICATION_OFF = False + + +@pytest_asyncio.fixture +async def validator() -> Validator: + return await Validator.objects.acreate(public_key=validator_hotkey(), active=True) + + +# Somehow the regular dependency mechanism doesn't work with multiple test cases +# Explicit patching with a new instance each time is a temporary workaround +@pytest.fixture +def mock_executor_manager_class(mocker: MockerFixture): + mocker.patch( + "compute_horde_miner.miner.executor_manager.current.executor_manager", + StubExecutorManager(), + ) diff --git a/miner/app/src/compute_horde_miner/miner/tests/integration/test_job_request_gates.py b/miner/app/src/compute_horde_miner/miner/tests/integration/test_job_request_gates.py new file mode 100644 index 000000000..1ee48b192 --- /dev/null +++ b/miner/app/src/compute_horde_miner/miner/tests/integration/test_job_request_gates.py @@ -0,0 +1,654 @@ +import contextlib +import datetime as dt +import json +import uuid +from collections.abc import AsyncIterator, Callable +from typing import Any +from unittest import mock +from unittest.mock import AsyncMock + +import pytest +from channels.testing import WebsocketCommunicator +from compute_horde.executor_class import DEFAULT_EXECUTOR_CLASS +from compute_horde.protocol_consts import TRUSTED_MINER_FAKE_KEY, JobRejectionReason +from compute_horde.protocol_messages import ( + V0DeclineJobRequest, + V0InitialJobRequest, + V0JobRequest, + ValidatorSignedMessage, + parse_validator_to_miner_message, +) +from compute_horde.receipts.models import JobStartedReceipt +from compute_horde.test_wallet import get_test_misc_wallet +from compute_horde.utils import canonical_json, sign_blob +from compute_horde_core.executor_class import ExecutorClass +from django.utils import timezone +from pytest_mock import MockerFixture + +from compute_horde_miner.miner.miner_consumer.validator_interface import MinerValidatorConsumer +from compute_horde_miner.miner.models import AcceptedJob, Validator +from compute_horde_miner.miner.tests.validator import ( + MANIFEST_MESSAGE, + WEBSOCKET_TIMEOUT, + authenticate, + fake_validator, + initial_job_request, + job_request, + job_started_receipt, + miner_hotkey, + organic_client_reconnect_greeting, + received_output, + validator_keypair, + wait_until, +) + +pytestmark = [ + pytest.mark.asyncio, + pytest.mark.django_db(transaction=True), + pytest.mark.usefixtures("authentication_on"), +] + +OTHER_EXECUTOR_CLASS = next(c for c in ExecutorClass if c != DEFAULT_EXECUTOR_CLASS) +STALE_RECEIPT_AGE = dt.timedelta(minutes=11) +NO_STATE: dict[str, Any] = { + "accepted_jobs": [], + "job_started_receipts": [], + "reserved_executor_classes": [], + "joined_executor_groups": 0, +} + + +@pytest.fixture +def executor_manager(mocker: MockerFixture) -> AsyncMock: + manager = AsyncMock() + manager.get_manifest.return_value = {DEFAULT_EXECUTOR_CLASS: 1} + manager.reserve_executor_class.return_value = None + manager.wait_for_executor_reservation.return_value = None + manager.get_executor_public_address.return_value = None + mocker.patch("compute_horde_miner.miner.executor_manager.current.executor_manager", manager) + return manager + + +@pytest.fixture +def group_add(mocker: MockerFixture) -> AsyncMock: + return mocker.patch.object(MinerValidatorConsumer, "group_add", AsyncMock()) + + +@pytest.fixture +def send_job_request(mocker: MockerFixture) -> AsyncMock: + return mocker.patch.object(MinerValidatorConsumer, "send_job_request", AsyncMock()) + + +@pytest.fixture +def job_uuid() -> str: + return str(uuid.uuid4()) + + +@pytest.fixture +def local_miner(settings) -> None: + settings.IS_LOCAL_MINER = True + + +async def miner_state(executor_manager: AsyncMock, group_add: AsyncMock) -> dict[str, Any]: + return { + "accepted_jobs": [ + job + async for job in AcceptedJob.objects.values("job_uuid", "status", "full_job_details") + ], + "job_started_receipts": [ + job_uuid + async for job_uuid in JobStartedReceipt.objects.values_list("job_uuid", flat=True) + ], + "reserved_executor_classes": [ + call.args[1] for call in executor_manager.reserve_executor_class.await_args_list + ], + "joined_executor_groups": len(group_add.await_args_list), + } + + +def accepted_job_state(job_uuid: str, full_job_details: dict[str, Any] | None) -> dict[str, Any]: + return { + "accepted_jobs": [ + { + "job_uuid": uuid.UUID(job_uuid), + "status": AcceptedJob.Status.WAITING_FOR_EXECUTOR + if full_job_details is None + else AcceptedJob.Status.RUNNING, + "full_job_details": full_job_details, + } + ], + "job_started_receipts": [uuid.UUID(job_uuid)], + "reserved_executor_classes": [DEFAULT_EXECUTOR_CLASS], + "joined_executor_groups": 1, + } + + +def accept(job_uuid: str) -> dict[str, Any]: + return {"message_type": "V0AcceptJobRequest", "job_uuid": job_uuid} + + +def decline(job_uuid: str, reason: JobRejectionReason, message: str) -> dict[str, Any]: + return V0DeclineJobRequest(job_uuid=job_uuid, reason=reason, message=message).model_dump( + mode="json" + ) + + +async def wait_for_forwarding(send_job_request: AsyncMock) -> None: + async def forwarded() -> bool: + return bool(send_job_request.await_args_list) + + await wait_until(forwarded) + + +@contextlib.asynccontextmanager +async def authenticated_validator( + to_miner: str | None = None, +) -> AsyncIterator[WebsocketCommunicator]: + async with fake_validator() as communicator: + await authenticate(communicator, to_miner) + assert await communicator.receive_json_from(timeout=WEBSOCKET_TIMEOUT) == MANIFEST_MESSAGE + yield communicator + + +@contextlib.asynccontextmanager +async def accepted_job(job_uuid: str) -> AsyncIterator[WebsocketCommunicator]: + async with authenticated_validator() as communicator: + await communicator.send_to(initial_job_request(job_uuid).model_dump_json()) + assert await received_output(communicator, 1) == [accept(job_uuid)] + yield communicator + + +def signed_content(msg: ValidatorSignedMessage) -> dict[str, Any]: + return json.loads(msg.model_dump_json(exclude={"validator_signature"})) + + +def signed_frame(content: dict[str, Any]) -> str: + """Frame of the JSON object signed the way a validator signs a job request.""" + signature = sign_blob(validator_keypair(), canonical_json(content)) + return json.dumps({**content, "validator_signature": signature}) + + +def reserved_executor_token(executor_manager: AsyncMock) -> str: + [executor_token] = [ + call.args[0] for call in executor_manager.reserve_executor_class.await_args_list + ] + return executor_token + + +def received(frame: str) -> ValidatorSignedMessage: + """The message the miner makes of the frame, signed content included.""" + msg = parse_validator_to_miner_message(frame) + assert isinstance(msg, ValidatorSignedMessage) + return msg + + +def receipt_with_bad_signature(job_uuid: str) -> tuple[dict[str, Any], str]: + return ( + {"job_started_receipt_signature": "0x" + "00" * 64}, + "job started receipt could not be verified", + ) + + +def receipt_with_malformed_signature(job_uuid: str) -> tuple[dict[str, Any], str]: + return { + "job_started_receipt_signature": "gibberish" + }, "job started receipt could not be verified" + + +def receipt_signed_by_another_key(job_uuid: str) -> tuple[dict[str, Any], str]: + payload, signature = job_started_receipt(job_uuid, signer=get_test_misc_wallet().hotkey) + return ( + {"job_started_receipt_payload": payload, "job_started_receipt_signature": signature}, + "job started receipt could not be verified", + ) + + +def receipt_for_another_miner(job_uuid: str) -> tuple[dict[str, Any], str]: + payload, signature = job_started_receipt( + job_uuid, miner_hotkey=get_test_misc_wallet().hotkey.ss58_address + ) + return ( + {"job_started_receipt_payload": payload, "job_started_receipt_signature": signature}, + "job started receipt could not be verified", + ) + + +def receipt_for_another_job(job_uuid: str) -> tuple[dict[str, Any], str]: + other_job_uuid = str(uuid.uuid4()) + payload, signature = job_started_receipt(other_job_uuid) + return ( + {"job_started_receipt_payload": payload, "job_started_receipt_signature": signature}, + f"job started receipt is for job {other_job_uuid}", + ) + + +def receipt_for_another_executor_class(job_uuid: str) -> tuple[dict[str, Any], str]: + payload, signature = job_started_receipt(job_uuid, executor_class=OTHER_EXECUTOR_CLASS) + return ( + {"job_started_receipt_payload": payload, "job_started_receipt_signature": signature}, + f"job started receipt is for executor class {OTHER_EXECUTOR_CLASS}", + ) + + +def receipt_dated_ago( + job_uuid: str, age: dt.timedelta, to_miner: str | None = None +) -> dict[str, Any]: + """Initial job request fields carrying a receipt dated the given time before now.""" + payload, signature = job_started_receipt( + job_uuid, miner_hotkey=to_miner or miner_hotkey(), timestamp=timezone.now() - age + ) + return {"job_started_receipt_payload": payload, "job_started_receipt_signature": signature} + + +def receipt_too_old(job_uuid: str) -> tuple[dict[str, Any], str]: + return receipt_dated_ago(job_uuid, STALE_RECEIPT_AGE), "job started receipt is too old" + + +def receipt_from_the_future(job_uuid: str) -> tuple[dict[str, Any], str]: + return ( + receipt_dated_ago(job_uuid, -dt.timedelta(minutes=1)), + "job started receipt is dated in the future", + ) + + +@pytest.mark.parametrize( + "bad_receipt", + [ + receipt_with_bad_signature, + receipt_with_malformed_signature, + receipt_signed_by_another_key, + receipt_for_another_miner, + receipt_for_another_job, + receipt_for_another_executor_class, + receipt_too_old, + receipt_from_the_future, + ], + ids=lambda f: f.__name__, +) +async def test_initial_job_request_with_bad_receipt_is_declined_without_state( + validator: Validator, + job_uuid: str, + executor_manager: AsyncMock, + group_add: AsyncMock, + bad_receipt: Callable[[str], tuple[dict[str, Any], str]], +): + overrides, expected_message = bad_receipt(job_uuid) + + async with authenticated_validator() as communicator: + await communicator.send_to(initial_job_request(job_uuid, **overrides).model_dump_json()) + assert await received_output(communicator, 1) == [ + decline(job_uuid, JobRejectionReason.INVALID_RECEIPT, expected_message) + ] + + assert await miner_state(executor_manager, group_add) == NO_STATE + + +async def test_initial_job_request_delayed_by_reconnects_is_accepted( + validator: Validator, job_uuid: str, executor_manager: AsyncMock, group_add: AsyncMock +): + overrides = receipt_dated_ago(job_uuid, dt.timedelta(minutes=5)) + + async with authenticated_validator() as communicator: + await communicator.send_to(initial_job_request(job_uuid, **overrides).model_dump_json()) + assert await received_output(communicator, 1) == [accept(job_uuid)] + + assert await miner_state(executor_manager, group_add) == accepted_job_state(job_uuid, None) + + +async def test_repeated_initial_job_request_is_declined_without_new_state( + validator: Validator, job_uuid: str, executor_manager: AsyncMock, group_add: AsyncMock +): + frame = initial_job_request(job_uuid).model_dump_json() + + async with authenticated_validator() as communicator: + await communicator.send_to(frame) + assert await received_output(communicator, 1) == [accept(job_uuid)] + await communicator.send_to(frame) + assert await received_output(communicator, 1) == [ + decline( + job_uuid, JobRejectionReason.INVALID_RECEIPT, "job started receipt was already used" + ) + ] + + assert await miner_state(executor_manager, group_add) == accepted_job_state(job_uuid, None) + + +@pytest.mark.usefixtures("local_miner") +async def test_local_miner_accepts_receipt_for_the_trusted_miner_without_storing_it( + validator: Validator, job_uuid: str, executor_manager: AsyncMock, group_add: AsyncMock +): + async with authenticated_validator(TRUSTED_MINER_FAKE_KEY) as communicator: + await communicator.send_to( + initial_job_request(job_uuid, TRUSTED_MINER_FAKE_KEY).model_dump_json() + ) + assert await received_output(communicator, 1) == [accept(job_uuid)] + + assert await miner_state(executor_manager, group_add) == { + **accepted_job_state(job_uuid, None), + "job_started_receipts": [], + } + + +@pytest.mark.parametrize( + ("bad_receipt", "expected_message"), + [ + pytest.param( + lambda job_uuid: {}, + "job started receipt could not be verified", + id="receipt_for_another_miner", + ), + pytest.param( + lambda job_uuid: { + "to_miner": TRUSTED_MINER_FAKE_KEY, + "job_started_receipt_signature": "0x" + "00" * 64, + }, + "job started receipt could not be verified", + id="receipt_with_bad_signature", + ), + pytest.param( + lambda job_uuid: receipt_dated_ago(job_uuid, STALE_RECEIPT_AGE, TRUSTED_MINER_FAKE_KEY), + "job started receipt is too old", + id="receipt_too_old", + ), + ], +) +@pytest.mark.usefixtures("local_miner") +async def test_local_miner_declines_bad_receipt_without_state( + validator: Validator, + job_uuid: str, + executor_manager: AsyncMock, + group_add: AsyncMock, + bad_receipt: Callable[[str], dict[str, Any]], + expected_message: str, +): + request = initial_job_request(job_uuid, **bad_receipt(job_uuid)) + + async with authenticated_validator(TRUSTED_MINER_FAKE_KEY) as communicator: + await communicator.send_to(request.model_dump_json()) + assert await received_output(communicator, 1) == [ + decline(job_uuid, JobRejectionReason.INVALID_RECEIPT, expected_message) + ] + + assert await miner_state(executor_manager, group_add) == NO_STATE + + +@pytest.mark.usefixtures("local_miner") +async def test_local_miner_declines_repeated_initial_job_request( + validator: Validator, job_uuid: str, executor_manager: AsyncMock, group_add: AsyncMock +): + frame = initial_job_request(job_uuid, TRUSTED_MINER_FAKE_KEY).model_dump_json() + + async with authenticated_validator(TRUSTED_MINER_FAKE_KEY) as communicator: + await communicator.send_to(frame) + assert await received_output(communicator, 1) == [accept(job_uuid)] + await communicator.send_to(frame) + assert await received_output(communicator, 1) == [ + decline(job_uuid, JobRejectionReason.INVALID_RECEIPT, "job was already accepted") + ] + + assert await miner_state(executor_manager, group_add) == { + **accepted_job_state(job_uuid, None), + "job_started_receipts": [], + } + + +def unsigned(msg: ValidatorSignedMessage) -> ValidatorSignedMessage: + return msg.model_copy(update={"validator_signature": None}) + + +def signed_by_another_key(msg: ValidatorSignedMessage) -> ValidatorSignedMessage: + return msg.signed(get_test_misc_wallet().hotkey) + + +def changed_after_signing(msg: ValidatorSignedMessage) -> ValidatorSignedMessage: + return msg.model_copy(update={"docker_image": "something else"}) + + +BAD_SIGNATURES = [ + pytest.param(unsigned, "missing validator signature", id="unsigned"), + pytest.param(signed_by_another_key, "invalid validator signature", id="signed_by_another_key"), + pytest.param(changed_after_signing, "invalid validator signature", id="changed_after_signing"), +] + + +@pytest.mark.parametrize(("tamper", "expected_message"), BAD_SIGNATURES) +async def test_initial_job_request_without_valid_signature_is_declined_without_state( + validator: Validator, + job_uuid: str, + executor_manager: AsyncMock, + group_add: AsyncMock, + tamper: Callable[[V0InitialJobRequest], V0InitialJobRequest], + expected_message: str, +): + async with authenticated_validator() as communicator: + await communicator.send_to(tamper(initial_job_request(job_uuid)).model_dump_json()) + assert await received_output(communicator, 1) == [ + decline(job_uuid, JobRejectionReason.INVALID_SIGNATURE, expected_message) + ] + + assert await miner_state(executor_manager, group_add) == NO_STATE + + +async def test_signed_initial_job_request_is_accepted( + validator: Validator, job_uuid: str, executor_manager: AsyncMock, group_add: AsyncMock +): + async with accepted_job(job_uuid): + pass + + assert await miner_state(executor_manager, group_add) == accepted_job_state(job_uuid, None) + + +async def test_unsigned_initial_job_request_is_accepted_when_signatures_are_not_required( + validator: Validator, + job_uuid: str, + executor_manager: AsyncMock, + group_add: AsyncMock, + settings, +): + settings.REQUIRE_SIGNED_JOB_REQUESTS = False + + async with authenticated_validator() as communicator: + await communicator.send_to(unsigned(initial_job_request(job_uuid)).model_dump_json()) + assert await received_output(communicator, 1) == [accept(job_uuid)] + + assert await miner_state(executor_manager, group_add) == accepted_job_state(job_uuid, None) + + +async def test_wrongly_signed_initial_job_request_is_declined_even_when_signatures_are_not_required( + validator: Validator, + job_uuid: str, + executor_manager: AsyncMock, + group_add: AsyncMock, + settings, +): + settings.REQUIRE_SIGNED_JOB_REQUESTS = False + + async with authenticated_validator() as communicator: + await communicator.send_to( + signed_by_another_key(initial_job_request(job_uuid)).model_dump_json() + ) + assert await received_output(communicator, 1) == [ + decline(job_uuid, JobRejectionReason.INVALID_SIGNATURE, "invalid validator signature") + ] + + assert await miner_state(executor_manager, group_add) == NO_STATE + + +@pytest.mark.parametrize(("tamper", "expected_message"), BAD_SIGNATURES) +async def test_job_request_without_valid_signature_is_not_forwarded( + validator: Validator, + job_uuid: str, + executor_manager: AsyncMock, + group_add: AsyncMock, + send_job_request: AsyncMock, + tamper: Callable[[V0JobRequest], V0JobRequest], + expected_message: str, +): + async with accepted_job(job_uuid) as communicator: + await communicator.send_to(tamper(job_request(job_uuid)).model_dump_json()) + assert await received_output(communicator, 1) == [ + decline(job_uuid, JobRejectionReason.INVALID_SIGNATURE, expected_message) + ] + + assert send_job_request.await_args_list == [] + assert await miner_state(executor_manager, group_add) == accepted_job_state(job_uuid, None) + + +async def test_job_request_for_unknown_job_is_not_forwarded( + validator: Validator, + job_uuid: str, + executor_manager: AsyncMock, + group_add: AsyncMock, + send_job_request: AsyncMock, +): + other_job_uuid = str(uuid.uuid4()) + + async with accepted_job(job_uuid) as communicator: + await communicator.send_to(job_request(other_job_uuid).model_dump_json()) + assert await received_output(communicator, 1) == [ + { + "message_type": "GenericError", + "details": f"Received JobRequest for unknown job_uuid: {other_job_uuid}", + } + ] + + assert send_job_request.await_args_list == [] + assert await miner_state(executor_manager, group_add) == accepted_job_state(job_uuid, None) + + +async def test_signed_job_request_is_forwarded_unchanged( + validator: Validator, + job_uuid: str, + executor_manager: AsyncMock, + group_add: AsyncMock, + send_job_request: AsyncMock, +): + frame = job_request(job_uuid).model_dump_json() + + async with accepted_job(job_uuid) as communicator: + await communicator.send_to(frame) + await wait_for_forwarding(send_job_request) + assert await received_output(communicator) == [] + + assert send_job_request.await_args_list == [ + mock.call(reserved_executor_token(executor_manager), received(frame)) + ] + assert await miner_state(executor_manager, group_add) == accepted_job_state( + job_uuid, json.loads(frame) + ) + + +async def test_job_request_with_a_field_unknown_to_the_sender_is_forwarded( + validator: Validator, + job_uuid: str, + executor_manager: AsyncMock, + group_add: AsyncMock, + send_job_request: AsyncMock, +): + content = signed_content(job_request(job_uuid, env={})) + del content["env"] + frame = signed_frame(content) + + async with accepted_job(job_uuid) as communicator: + await communicator.send_to(frame) + await wait_for_forwarding(send_job_request) + assert await received_output(communicator) == [] + + assert send_job_request.await_args_list == [ + mock.call(reserved_executor_token(executor_manager), received(frame)) + ] + + +async def test_initial_job_request_with_a_field_unknown_to_the_miner_is_accepted( + validator: Validator, job_uuid: str, executor_manager: AsyncMock, group_add: AsyncMock +): + frame = signed_frame({**signed_content(initial_job_request(job_uuid)), "priority": 7}) + + async with authenticated_validator() as communicator: + await communicator.send_to(frame) + assert await received_output(communicator, 1) == [accept(job_uuid)] + + assert await miner_state(executor_manager, group_add) == accepted_job_state(job_uuid, None) + + +async def test_job_request_with_a_duplicate_key_runs_the_signed_value( + validator: Validator, + job_uuid: str, + executor_manager: AsyncMock, + group_add: AsyncMock, + send_job_request: AsyncMock, +): + signed = job_request(job_uuid, docker_image="signed").model_dump_json() + # pydantic's JSON validation picks the first of these two keys, json.loads the last one. + frame = signed.replace( + '"docker_image":"signed"', '"docker_image":"injected","docker_image":"signed"' + ) + assert frame != signed + + async with accepted_job(job_uuid) as communicator: + await communicator.send_to(frame) + await wait_for_forwarding(send_job_request) + assert await received_output(communicator) == [] + + assert send_job_request.await_args_list == [ + mock.call(reserved_executor_token(executor_manager), received(signed)) + ] + + +@pytest.mark.parametrize( + ("frame", "error"), + [ + pytest.param( + {"text_data": "not json"}, "Expecting value: line 1 column 1 (char 0)", id="text" + ), + pytest.param({"bytes_data": b"{}"}, "expected a text frame", id="binary"), + pytest.param( + {"text_data": '{"message_type": "GenericError", "details": ' + "1" * 5000 + "}"}, + "Exceeds the limit (4300 digits) for integer string conversion: value has 5000" + " digits; use sys.set_int_max_str_digits() to increase the limit", + id="huge_int", + ), + pytest.param( + {"text_data": "[" * 100_000 + "]" * 100_000}, + "JSON is nested too deeply", + id="deep_nesting", + ), + ], +) +async def test_malformed_frame_is_reported(validator: Validator, frame: dict[str, Any], error: str): + async with authenticated_validator() as communicator: + await communicator.send_to(**frame) + assert await received_output(communicator, 1) == [ + {"message_type": "GenericError", "details": f"Malformed message: {error}"} + ] + + +async def test_signed_job_request_is_forwarded_on_a_reopened_socket( + validator: Validator, + job_uuid: str, + executor_manager: AsyncMock, + group_add: AsyncMock, + send_job_request: AsyncMock, +): + greeting = organic_client_reconnect_greeting() + request = job_request(job_uuid) + + async with fake_validator() as dropped: + await dropped.send_to(greeting()) + await dropped.send_to(initial_job_request(job_uuid).model_dump_json()) + assert await received_output(dropped, 2) == [MANIFEST_MESSAGE, accept(job_uuid)] + + async with fake_validator() as reopened: + await reopened.send_to(greeting()) + await reopened.send_to(request.model_dump_json()) + await wait_for_forwarding(send_job_request) + assert await received_output(reopened, 1) == [MANIFEST_MESSAGE] + + assert send_job_request.await_args_list == [ + mock.call(reserved_executor_token(executor_manager), received(request.model_dump_json())) + ] + assert await miner_state(executor_manager, group_add) == { + **accepted_job_state(job_uuid, request.model_dump(mode="json")), + # The reopened socket subscribes to the job again. + "joined_executor_groups": 2, + } diff --git a/miner/app/src/compute_horde_miner/miner/tests/integration/test_mocked_executor_manager.py b/miner/app/src/compute_horde_miner/miner/tests/integration/test_mocked_executor_manager.py index f4cdfa389..ac66d5d35 100644 --- a/miner/app/src/compute_horde_miner/miner/tests/integration/test_mocked_executor_manager.py +++ b/miner/app/src/compute_horde_miner/miner/tests/integration/test_mocked_executor_manager.py @@ -1,32 +1,31 @@ -import contextlib import time import uuid from unittest.mock import MagicMock import pytest -import pytest_asyncio -from channels.testing import WebsocketCommunicator -from compute_horde.executor_class import DEFAULT_EXECUTOR_CLASS +from compute_horde.protocol_consts import TRUSTED_MINER_FAKE_KEY +from compute_horde_core.streaming import StreamingDetails from pytest_mock import MockerFixture -from compute_horde_miner import asgi from compute_horde_miner.miner.models import Validator from compute_horde_miner.miner.tests.executor_manager import ( - StubExecutorManager, StubStreamingExecutorManager, fake_executor, ) +from compute_horde_miner.miner.tests.validator import ( + MANIFEST_MESSAGE, + WEBSOCKET_TIMEOUT, + authenticate, + fake_validator, + initial_job_request, + job_request, +) -pytestmark = [pytest.mark.asyncio, pytest.mark.django_db(transaction=True)] - -WEBSOCKET_TIMEOUT = 3 - - -@pytest.fixture -def mock_keypair(mocker: MockerFixture): - return mocker.patch( - "compute_horde_miner.miner.miner_consumer.validator_interface.bittensor.Keypair" - ) +pytestmark = [ + pytest.mark.asyncio, + pytest.mark.django_db(transaction=True), + pytest.mark.usefixtures("authentication_on"), +] @pytest.fixture @@ -37,16 +36,6 @@ def mock_executor_ip(mocker: MockerFixture): ) -# Somehow the regular dependency mechanism doesn't work with multiple test cases -# Explicit patching with a new instance each time is a temporary workaround -@pytest.fixture -def mock_executor_manager_class(mocker: MockerFixture): - mocker.patch( - "compute_horde_miner.miner.executor_manager.current.executor_manager", - StubExecutorManager(), - ) - - @pytest.fixture def mock_streaming_executor_manager_class(mocker: MockerFixture): mocker.patch( @@ -63,68 +52,19 @@ def job_uuid(): fake_executor.job_uuid = None -@pytest.fixture -def validator_hotkey(): - return "some_public_key" - - -@pytest_asyncio.fixture -async def validator(validator_hotkey: str): - return await Validator.objects.acreate(public_key=validator_hotkey, active=True) - - -@contextlib.asynccontextmanager -async def make_communicator(validator_key: str): - communicator = WebsocketCommunicator( - asgi.application, f"v0.1/validator_interface/{validator_key}" - ) - connected, _ = await communicator.connect() - assert connected - yield communicator - await communicator.disconnect() - - -async def run_regular_flow_test(validator_key: str, job_uuid: str, streaming: bool = False): - async with make_communicator(validator_key) as communicator: - await communicator.send_json_to( - { - "message_type": "ValidatorAuthForMiner", - "validator_hotkey": validator_key, - "miner_hotkey": "some key", - "timestamp": int(time.time()), - "signature": "gibberish", - } - ) +async def run_regular_flow_test( + job_uuid: str, streaming: bool = False, to_miner: str | None = None +): + streaming_details = StreamingDetails(public_key="some_public_key") if streaming else None + async with fake_validator() as communicator: + await authenticate(communicator, to_miner) response = await communicator.receive_json_from(timeout=WEBSOCKET_TIMEOUT) - assert response == { - "message_type": "V0ExecutorManifestRequest", - "manifest": {DEFAULT_EXECUTOR_CLASS: 1}, - } - msg = { - "message_type": "V0InitialJobRequest", - "job_uuid": job_uuid, - "executor_class": DEFAULT_EXECUTOR_CLASS, - "docker_image": "it's teeeeests", - "timeout_seconds": 60, - "volume": None, - "job_started_receipt_payload": { - "receipt_type": "JobStartedReceipt", - "job_uuid": job_uuid, - "miner_hotkey": "miner_hotkey", - "validator_hotkey": validator_key, - "timestamp": "2020-01-01T00:00:00Z", - "executor_class": DEFAULT_EXECUTOR_CLASS, - "is_organic": True, - "ttl": 5, - }, - "job_started_receipt_signature": "gibberish", - } - if streaming: - msg["streaming_details"] = { - "public_key": "some_public_key", - "executor_ip": None, - } - await communicator.send_json_to(msg) + assert response == MANIFEST_MESSAGE + await communicator.send_to( + initial_job_request( + job_uuid, to_miner, streaming_details=streaming_details + ).model_dump_json() + ) response = await communicator.receive_json_from(timeout=WEBSOCKET_TIMEOUT) assert response == { @@ -137,23 +77,7 @@ async def run_regular_flow_test(validator_key: str, job_uuid: str, streaming: bo "job_uuid": job_uuid, "executor_token": None, } - await communicator.send_json_to( - { - "message_type": "V0JobRequest", - "job_uuid": job_uuid, - "executor_class": DEFAULT_EXECUTOR_CLASS, - "docker_image": "it's teeeeests again", - "docker_run_cmd": [], - "env": {"SOME_ENV_VAR": "some value"}, - "docker_run_options_preset": "none", - "volume": {"volume_type": "inline", "contents": "nonsense"}, - "output_upload": { - "output_upload_type": "zip_and_http_post", - "url": "http://localhost/output/upload", - "form_fields": {}, - }, - } - ) + await communicator.send_to(job_request(job_uuid).model_dump_json()) if streaming: response = await communicator.receive_json_from(timeout=WEBSOCKET_TIMEOUT) response.pop("miner_signature") # ignore miner_signature field @@ -179,47 +103,44 @@ async def run_regular_flow_test(validator_key: str, job_uuid: str, streaming: bo } -async def test_main_loop( - validator: Validator, job_uuid: str, mock_keypair: MagicMock, mock_executor_manager_class -): - await run_regular_flow_test(validator.public_key, job_uuid) +async def test_main_loop(validator: Validator, job_uuid: str, mock_executor_manager_class): + await run_regular_flow_test(job_uuid) async def test_streaming_main_loop( validator: Validator, job_uuid: str, - mock_keypair: MagicMock, mock_streaming_executor_manager_class, mock_executor_ip: MagicMock, ): - await run_regular_flow_test(validator.public_key, job_uuid, streaming=True) + await run_regular_flow_test(job_uuid, streaming=True) + + +@pytest.fixture +def trusted_miner_receipts(): + fake_executor.miner_hotkey = TRUSTED_MINER_FAKE_KEY + yield + fake_executor.miner_hotkey = None async def test_local_miner( validator: Validator, job_uuid: str, - mock_keypair: MagicMock, mock_executor_manager_class, + trusted_miner_receipts, settings, ): settings.IS_LOCAL_MINER = True - settings.DEBUG_TURN_AUTHENTICATION_OFF = False - - await run_regular_flow_test(validator.public_key, job_uuid) - mock_keypair.assert_called_once_with(ss58_address=validator.public_key) - mock_keypair.return_value.verify.assert_called_once() + await run_regular_flow_test(job_uuid, to_miner=TRUSTED_MINER_FAKE_KEY) -async def test_local_miner_unknown_validator( - mock_keypair: MagicMock, mock_executor_manager_class, settings -): +async def test_local_miner_unknown_validator(mock_executor_manager_class, settings): settings.IS_LOCAL_MINER = True - settings.DEBUG_TURN_AUTHENTICATION_OFF = False validator_key = "unknown_validator" - async with make_communicator(validator_key) as communicator: + async with fake_validator(validator_key) as communicator: await communicator.send_json_to( { "message_type": "ValidatorAuthForMiner", diff --git a/miner/app/src/compute_horde_miner/miner/tests/integration/test_validator_auth_replay.py b/miner/app/src/compute_horde_miner/miner/tests/integration/test_validator_auth_replay.py new file mode 100644 index 000000000..7f0f3ddac --- /dev/null +++ b/miner/app/src/compute_horde_miner/miner/tests/integration/test_validator_auth_replay.py @@ -0,0 +1,139 @@ +import time +from collections.abc import Callable +from typing import Any + +import pytest +from compute_horde.protocol_consts import TRUSTED_MINER_FAKE_KEY + +from compute_horde_miner.miner.models import Validator +from compute_horde_miner.miner.tests.validator import ( + MANIFEST_MESSAGE, + auth_message, + fake_validator, + miner_hotkey, + organic_client_reconnect_greeting, + received_output, + validator_hotkey, +) + +pytestmark = [ + pytest.mark.asyncio, + pytest.mark.django_db(transaction=True), + pytest.mark.usefixtures("authentication_on", "mock_executor_manager_class"), +] + + +def rejected(reason: str) -> list[dict[str, Any]]: + return [ + { + "message_type": "GenericError", + "details": f"Validator {validator_hotkey()} not authenticated due to: {reason}", + }, + {"type": "websocket.close", "code": 1000}, + ] + + +async def authentication_response(frame: str) -> list[dict[str, Any]]: + async with fake_validator() as communicator: + await communicator.send_to(frame) + return await received_output(communicator, 1) + + +async def test_replayed_auth_frame_is_rejected(validator: Validator): + frame = auth_message().model_dump_json() + async with fake_validator() as first, fake_validator() as second: + await first.send_to(frame) + assert await received_output(first, 1) == [MANIFEST_MESSAGE] + + await second.send_to(frame) + assert await received_output(second, 1) == rejected("auth message already used") + + +@pytest.mark.parametrize( + ("respell", "reason"), + [ + pytest.param( + lambda signature: "0x" + signature[2:].upper(), "auth message already used", id="upper" + ), + pytest.param(lambda signature: "0x" + signature, "malformed signature", id="0x0x"), + pytest.param( + lambda signature: "0x" * 63 + signature, "malformed signature", id="256_characters" + ), + pytest.param( + lambda signature: "0x" * 70 + signature, "malformed signature", id="270_characters" + ), + ], +) +async def test_respelled_signature_of_used_auth_frame_is_rejected( + validator: Validator, respell: Callable[[str], str], reason: str +): + frame = auth_message() + respelled = frame.model_copy(update={"signature": respell(frame.signature)}) + + async with fake_validator() as first, fake_validator() as second: + await first.send_to(frame.model_dump_json()) + assert await received_output(first, 1) == [MANIFEST_MESSAGE] + + await second.send_to(respelled.model_dump_json()) + assert await received_output(second, 1) == rejected(reason) + + +async def test_auth_frames_with_same_timestamp_are_both_accepted(validator: Validator): + timestamp = int(time.time()) + first_frame, second_frame = auth_message(timestamp), auth_message(timestamp) + assert first_frame.signature != second_frame.signature + + async with fake_validator() as first, fake_validator() as second: + await first.send_to(first_frame.model_dump_json()) + await second.send_to(second_frame.model_dump_json()) + assert (await received_output(first, 1), await received_output(second, 1)) == ( + [MANIFEST_MESSAGE], + [MANIFEST_MESSAGE], + ) + + +@pytest.mark.parametrize( + ("timestamp_offset", "reason"), + [ + pytest.param(60, "msg timestamp is in the future", id="future"), + pytest.param(-60, "msg too old", id="stale"), + ], +) +async def test_auth_frame_outside_validity_window_is_rejected( + validator: Validator, timestamp_offset: int, reason: str +): + frame = auth_message(int(time.time()) + timestamp_offset).model_dump_json() + assert await authentication_response(frame) == rejected(reason) + + +async def test_reconnect_greeting_authenticates_each_reopened_socket_once(validator: Validator): + greeting = organic_client_reconnect_greeting() + + first_frame = greeting() + assert await authentication_response(first_frame) == [MANIFEST_MESSAGE] + reconnect_frame = greeting() + assert await authentication_response(reconnect_frame) == [MANIFEST_MESSAGE] + + assert [ + await authentication_response(replayed) for replayed in (first_frame, reconnect_frame) + ] == [rejected("auth message already used")] * 2 + + +async def test_local_miner_accepts_auth_frame_addressed_to_the_trusted_miner( + validator: Validator, settings +): + settings.IS_LOCAL_MINER = True + frame = auth_message(to_miner=TRUSTED_MINER_FAKE_KEY).model_dump_json() + + assert await authentication_response(frame) == [MANIFEST_MESSAGE] + + +async def test_local_miner_rejects_auth_frame_addressed_to_another_miner( + validator: Validator, settings +): + settings.IS_LOCAL_MINER = True + frame = auth_message().model_dump_json() + + assert await authentication_response(frame) == rejected( + f"wrong miner hotkey ({TRUSTED_MINER_FAKE_KEY}!={miner_hotkey()})" + ) diff --git a/miner/app/src/compute_horde_miner/miner/tests/integration/test_validator_authentication.py b/miner/app/src/compute_horde_miner/miner/tests/integration/test_validator_authentication.py new file mode 100644 index 000000000..a143f9165 --- /dev/null +++ b/miner/app/src/compute_horde_miner/miner/tests/integration/test_validator_authentication.py @@ -0,0 +1,224 @@ +import asyncio +import contextlib +import time +import uuid +from collections.abc import AsyncIterator +from typing import Any +from unittest import mock + +import pytest +import pytest_asyncio +from channels.layers import get_channel_layer +from channels.testing import WebsocketCommunicator +from compute_horde.protocol_messages import V0JobFinishedRequest, V0StreamingJobReadyRequest +from pytest_mock import MockerFixture + +from compute_horde_miner import asgi +from compute_horde_miner.miner.miner_consumer.layer_utils import ValidatorInterfaceMixin +from compute_horde_miner.miner.miner_consumer.validator_interface import MinerValidatorConsumer +from compute_horde_miner.miner.models import AcceptedJob, Validator +from compute_horde_miner.miner.tests.validator import ( + MANIFEST_MESSAGE, + WEBSOCKET_TIMEOUT, + received_output, +) + +pytestmark = [pytest.mark.asyncio, pytest.mark.django_db(transaction=True)] + +SHORT_AUTH_TIMEOUT = 1 + + +def _patch_auth_timeout(mocker: MockerFixture, seconds: float) -> float: + mocker.patch( + "compute_horde_miner.miner.miner_consumer.validator_interface.AUTH_TIMEOUT", seconds + ) + return seconds + + +@pytest.fixture +def generous_auth_timeout(mocker: MockerFixture) -> float: + return _patch_auth_timeout(mocker, 30) + + +@pytest.fixture +def short_auth_timeout(mocker: MockerFixture) -> float: + return _patch_auth_timeout(mocker, SHORT_AUTH_TIMEOUT) + + +@pytest.fixture +def connect_finished(mocker: MockerFixture) -> asyncio.Event: + """ + Set once MinerValidatorConsumer.connect() has run to completion; the websocket handshake + alone finishes before connect() does. + """ + finished = asyncio.Event() + original_connect = MinerValidatorConsumer.connect + + async def connect(self: MinerValidatorConsumer) -> None: + await original_connect(self) + finished.set() + + mocker.patch.object(MinerValidatorConsumer, "connect", connect) + return finished + + +@pytest_asyncio.fixture +async def running_job(validator: Validator) -> AcceptedJob: + return await AcceptedJob.objects.acreate( + validator=validator, + job_uuid=uuid.uuid4(), + executor_token=f"running-{uuid.uuid4()}", + status=AcceptedJob.Status.RUNNING, + initial_job_details={}, + ) + + +@pytest_asyncio.fixture +async def job_awaiting_payload(validator: Validator) -> AcceptedJob: + return await AcceptedJob.objects.acreate( + validator=validator, + job_uuid=uuid.uuid4(), + executor_token=f"awaiting-{uuid.uuid4()}", + status=AcceptedJob.Status.WAITING_FOR_PAYLOAD, + initial_job_details={}, + ) + + +@contextlib.asynccontextmanager +async def validator_socket(validator_key: str) -> AsyncIterator[WebsocketCommunicator]: + communicator = WebsocketCommunicator( + asgi.application, f"v0.1/validator_interface/{validator_key}" + ) + connected, _ = await communicator.connect(timeout=WEBSOCKET_TIMEOUT) + assert connected + try: + yield communicator + finally: + await communicator.disconnect() + + +async def authenticate(communicator: WebsocketCommunicator, validator_key: str) -> None: + await communicator.send_json_to( + { + "message_type": "ValidatorAuthForMiner", + "validator_hotkey": validator_key, + "miner_hotkey": "some key", + "timestamp": int(time.time()), + "signature": "gibberish", + } + ) + + +def executor_messages( + job: AcceptedJob, +) -> list[tuple[str, V0StreamingJobReadyRequest | V0JobFinishedRequest]]: + job_uuid = str(job.job_uuid) + return [ + ( + "streaming_job.ready", + V0StreamingJobReadyRequest( + job_uuid=job_uuid, + executor_token=job.executor_token, + public_key="streaming server certificate", + ip="10.0.0.1", + port=4321, + ), + ), + ( + "executor.finished", + V0JobFinishedRequest( + job_uuid=job_uuid, + docker_process_stdout="customer stdout", + docker_process_stderr="customer stderr", + artifacts={"/artifacts/result": "cmVzdWx0"}, + upload_results={}, + ), + ), + ] + + +def relayed_executor_messages(job: AcceptedJob) -> list[dict[str, Any]]: + streaming_ready, finished = (msg for _, msg in executor_messages(job)) + return [ + { + **streaming_ready.model_dump(), + "executor_token": None, + "miner_signature": mock.ANY, + }, + finished.model_dump(), + ] + + +async def send_executor_messages(job: AcceptedJob) -> None: + group = ValidatorInterfaceMixin.group_name(job.executor_token) + for event_type, msg in executor_messages(job): + await get_channel_layer().group_send(group, {"type": event_type, **msg.model_dump()}) + + +def executor_ready_message(job: AcceptedJob) -> dict[str, Any]: + return { + "message_type": "V0ExecutorReadyRequest", + "job_uuid": str(job.job_uuid), + "executor_token": None, + } + + +async def test_unauthenticated_socket_receives_nothing( + validator: Validator, + running_job: AcceptedJob, + job_awaiting_payload: AcceptedJob, + generous_auth_timeout: float, + connect_finished: asyncio.Event, + mock_executor_manager_class, +): + async with validator_socket(validator.public_key) as communicator: + await asyncio.wait_for(connect_finished.wait(), WEBSOCKET_TIMEOUT) + await send_executor_messages(running_job) + assert await received_output(communicator) == [] + + +async def test_authenticated_socket_receives_executor_messages( + validator: Validator, + running_job: AcceptedJob, + job_awaiting_payload: AcceptedJob, + generous_auth_timeout: float, + mock_executor_manager_class, +): + async with validator_socket(validator.public_key) as communicator: + await authenticate(communicator, validator.public_key) + assert await received_output(communicator, 2) == [ + MANIFEST_MESSAGE, + executor_ready_message(job_awaiting_payload), + ] + await send_executor_messages(running_job) + assert await received_output(communicator, 2) == relayed_executor_messages(running_job) + + +async def test_unauthenticated_socket_is_closed_after_auth_timeout( + validator: Validator, + short_auth_timeout: float, + mock_executor_manager_class, +): + async with validator_socket(validator.public_key) as communicator: + assert await received_output(communicator, quiet_period=short_auth_timeout * 2) == [ + { + "message_type": "GenericError", + "details": ( + f"Validator {validator.public_key} did not authenticate" + f" within {short_auth_timeout}s" + ), + }, + {"type": "websocket.close", "code": 1000}, + ] + + +async def test_authenticated_socket_stays_open_past_auth_timeout( + validator: Validator, + short_auth_timeout: float, + mock_executor_manager_class, +): + async with validator_socket(validator.public_key) as communicator: + await authenticate(communicator, validator.public_key) + assert await received_output(communicator, quiet_period=short_auth_timeout * 2) == [ + MANIFEST_MESSAGE + ] diff --git a/miner/app/src/compute_horde_miner/miner/tests/test_busy.py b/miner/app/src/compute_horde_miner/miner/tests/test_busy.py index 1dd78eade..a97b24071 100644 --- a/miner/app/src/compute_horde_miner/miner/tests/test_busy.py +++ b/miner/app/src/compute_horde_miner/miner/tests/test_busy.py @@ -1,14 +1,10 @@ -from datetime import datetime, timedelta +from datetime import timedelta from unittest.mock import patch from uuid import uuid4 import pytest from compute_horde.executor_class import DEFAULT_EXECUTOR_CLASS from compute_horde.protocol_consts import JobRejectionReason -from compute_horde.protocol_messages import ( - V0InitialJobRequest, - ValidatorAuthForMiner, -) from compute_horde.receipts import Receipt from compute_horde.receipts.models import JobStartedReceipt from compute_horde.receipts.schemas import JobStartedReceiptPayload @@ -24,7 +20,13 @@ from compute_horde_miner.miner.executor_manager.v0 import AllExecutorsBusy from compute_horde_miner.miner.models import Validator from compute_horde_miner.miner.tests.executor_manager import fake_executor -from compute_horde_miner.miner.tests.validator import fake_validator +from compute_horde_miner.miner.tests.validator import ( + WEBSOCKET_TIMEOUT, + authenticate, + fake_validator, + initial_job_request, + wait_until, +) pytestmark = [pytest.mark.asyncio, pytest.mark.django_db(transaction=True)] @@ -37,48 +39,12 @@ def job_uuid(): fake_executor.job_uuid = None -async def _authenticate(channel, validator, miner_hotkey): - auth_payload = ValidatorAuthForMiner( - validator_hotkey=validator.hotkey.ss58_address, - miner_hotkey=miner_hotkey, - timestamp=int(datetime.now().timestamp()), - signature="", - ) - auth_payload.signature = sign_blob(validator.hotkey, auth_payload.blob_for_signing()) - await channel.send_to(auth_payload.model_dump_json()) - - async def _send_initial_job_request(validator_channel, job_uuid): - validator_wallet = get_test_validator_wallet() - miner_wallet = get_test_miner_wallet() - job_started_receipt_payload = JobStartedReceiptPayload( - job_uuid=job_uuid, - is_organic=True, - miner_hotkey=miner_wallet.hotkey.ss58_address, - validator_hotkey=validator_wallet.hotkey.ss58_address, - timestamp=timezone.now(), - executor_class=DEFAULT_EXECUTOR_CLASS, - ttl=60, - ) - job_started_receipt_signature = sign_blob( - validator_wallet.hotkey, - job_started_receipt_payload.blob_for_signing(), - ) - await validator_channel.send_to( - V0InitialJobRequest( - job_uuid=job_uuid, - executor_class=DEFAULT_EXECUTOR_CLASS, - docker_image="it's teeeeests", - timeout_seconds=60, - job_started_receipt_payload=job_started_receipt_payload, - job_started_receipt_signature=job_started_receipt_signature, - ).model_dump_json() - ) + await validator_channel.send_to(initial_job_request(job_uuid).model_dump_json()) async def test_receipt_stored_before_executor_reserved(job_uuid): validator_wallet = get_test_validator_wallet() - miner_wallet = get_test_miner_wallet() receipt_created_before_executor_reserved = False @@ -86,16 +52,18 @@ async def _on_start_new_executor(*args, **kwargs): nonlocal receipt_created_before_executor_reserved receipt_created_before_executor_reserved = await JobStartedReceipt.objects.acount() == 1 + async def executor_reserved() -> bool: + return reserve_executor_class.called + with patch.object(executor_manager, "reserve_executor_class") as reserve_executor_class: reserve_executor_class.side_effect = _on_start_new_executor await Validator.objects.acreate( public_key=validator_wallet.hotkey.ss58_address, active=True ) async with fake_validator() as validator_channel: - await _authenticate( - validator_channel, validator_wallet, miner_wallet.hotkey.ss58_address - ) + await authenticate(validator_channel) await _send_initial_job_request(validator_channel, job_uuid) + await wait_until(executor_reserved) reserve_executor_class.assert_called() assert receipt_created_before_executor_reserved @@ -150,13 +118,11 @@ async def test_reject_as_busy_when_busy(job_uuid): wait_for_executor_reservation.side_effect = AllExecutorsBusy() async with fake_validator() as validator_channel: - await _authenticate( - validator_channel, validator_wallet, miner_wallet.hotkey.ss58_address - ) - await validator_channel.receive_from() # this gets the manifest + await authenticate(validator_channel) + await validator_channel.receive_from(timeout=WEBSOCKET_TIMEOUT) # the manifest await _send_initial_job_request(validator_channel, job_uuid) - response = await validator_channel.receive_json_from() + response = await validator_channel.receive_json_from(timeout=WEBSOCKET_TIMEOUT) assert response["message_type"] == "V0DeclineJobRequest" assert response["reason"] == JobRejectionReason.BUSY.value diff --git a/miner/app/src/compute_horde_miner/miner/tests/test_receipts.py b/miner/app/src/compute_horde_miner/miner/tests/test_receipts.py index f9f78b964..16934e6c8 100644 --- a/miner/app/src/compute_horde_miner/miner/tests/test_receipts.py +++ b/miner/app/src/compute_horde_miner/miner/tests/test_receipts.py @@ -1,4 +1,5 @@ from datetime import datetime +from functools import partial from unittest.mock import AsyncMock from uuid import uuid4 @@ -24,7 +25,7 @@ from compute_horde_miner.miner.models import Validator from compute_horde_miner.miner.tests.executor_manager import fake_executor -from compute_horde_miner.miner.tests.validator import fake_validator +from compute_horde_miner.miner.tests.validator import fake_validator, wait_until @pytest.fixture @@ -111,7 +112,9 @@ async def test_receipt_is_saved( timeout_seconds=60, job_started_receipt_payload=job_started_receipt_payload, job_started_receipt_signature=job_started_receipt_signature, - ).model_dump_json() + ) + .signed(validator_wallet.hotkey) + .model_dump_json() ) # skip doing the job, and only send receipts @@ -150,6 +153,7 @@ async def test_receipt_is_saved( ), ).model_dump_json() ) + await wait_until(partial(JobFinishedReceipt.objects.filter(job_uuid=job_uuid).aexists)) assert await JobStartedReceipt.objects.filter( job_uuid=job_uuid, is_organic=organic_job diff --git a/miner/app/src/compute_horde_miner/miner/tests/test_validator_auth.py b/miner/app/src/compute_horde_miner/miner/tests/test_validator_auth.py new file mode 100644 index 000000000..f199a906e --- /dev/null +++ b/miner/app/src/compute_horde_miner/miner/tests/test_validator_auth.py @@ -0,0 +1,80 @@ +import datetime as dt +import time + +import pytest + +from compute_horde_miner.miner.models import ConsumedValidatorAuth +from compute_horde_miner.miner.tests.validator import auth_message +from compute_horde_miner.miner.validator_auth import ( + AUTH_MESSAGE_MAX_AGE, + AUTH_MESSAGE_MAX_CLOCK_SKEW, + CONSUMED_AUTH_RETENTION, + auth_signature_bytes, + auth_timestamp_error, + consume_auth_signature, +) + +SIGNATURE_HEX = "0123456789abcdef" * 8 + + +@pytest.mark.parametrize( + ("timestamp_offset", "error"), + [ + (0, None), + (-AUTH_MESSAGE_MAX_AGE + 1, None), + (-AUTH_MESSAGE_MAX_AGE - 1, "msg too old"), + (AUTH_MESSAGE_MAX_CLOCK_SKEW - 1, None), + (AUTH_MESSAGE_MAX_CLOCK_SKEW + 1, "msg timestamp is in the future"), + ], +) +def test_auth_timestamp_error(timestamp_offset: int, error: str | None): + assert auth_timestamp_error(auth_message(int(time.time()) + timestamp_offset)) == error + + +@pytest.mark.parametrize( + ("signature", "expected"), + [ + pytest.param("0x" + SIGNATURE_HEX, bytes.fromhex(SIGNATURE_HEX), id="canonical"), + pytest.param("0x" + SIGNATURE_HEX.upper(), bytes.fromhex(SIGNATURE_HEX), id="upper_case"), + pytest.param("0x0x" + SIGNATURE_HEX, None, id="double_prefix"), + pytest.param("0x" * 64 + SIGNATURE_HEX, None, id="repeated_prefix"), + pytest.param("0X" + SIGNATURE_HEX, None, id="upper_case_prefix"), + pytest.param(SIGNATURE_HEX, None, id="no_prefix"), + pytest.param("0x" + SIGNATURE_HEX + "\n", None, id="trailing_newline"), + pytest.param("0x" + SIGNATURE_HEX[:-2], None, id="too_short"), + pytest.param("0x" + SIGNATURE_HEX + "00", None, id="too_long"), + ], +) +def test_auth_signature_bytes(signature: str, expected: bytes | None): + msg = auth_message().model_copy(update={"signature": signature}) + + assert auth_signature_bytes(msg) == expected + + +@pytest.mark.django_db(transaction=True) +@pytest.mark.asyncio +async def test_consume_auth_signature_is_single_use_and_prunes_expired_rows(): + expired = dt.datetime.now(dt.UTC) - CONSUMED_AUTH_RETENTION - dt.timedelta(seconds=1) + await ConsumedValidatorAuth.objects.acreate( + signature="expired", validator_hotkey="hotkey", timestamp=expired + ) + msg = auth_message() + signature = auth_signature_bytes(msg) + assert signature is not None + + assert await consume_auth_signature(msg, signature) is True + assert await consume_auth_signature(msg, signature) is False + + consumed = [ + row + async for row in ConsumedValidatorAuth.objects.values( + "signature", "validator_hotkey", "timestamp" + ) + ] + assert consumed == [ + { + "signature": signature.hex(), + "validator_hotkey": msg.validator_hotkey, + "timestamp": dt.datetime.fromtimestamp(msg.timestamp, dt.UTC), + } + ] diff --git a/miner/app/src/compute_horde_miner/miner/tests/test_validator_relay_guard.py b/miner/app/src/compute_horde_miner/miner/tests/test_validator_relay_guard.py new file mode 100644 index 000000000..576189024 --- /dev/null +++ b/miner/app/src/compute_horde_miner/miner/tests/test_validator_relay_guard.py @@ -0,0 +1,110 @@ +import logging +import uuid +from unittest.mock import AsyncMock, MagicMock, call + +import pytest +from compute_horde.protocol_consts import HordeFailureReason, JobParticipantType +from compute_horde.protocol_messages import ( + V0ExecutionDoneRequest, + V0ExecutorReadyRequest, + V0HordeFailedRequest, + V0JobFailedRequest, + V0JobFinishedRequest, + V0MachineSpecsRequest, + V0StreamingJobReadyRequest, + V0VolumesReadyRequest, +) +from compute_horde.utils import MachineSpecs + +from compute_horde_miner.miner.miner_consumer.validator_interface import MinerValidatorConsumer + +JOB_UUID = str(uuid.uuid4()) + +RELAYED_EVENTS = [ + ( + "executor_ready", + "_executor_ready", + V0ExecutorReadyRequest(job_uuid=JOB_UUID, executor_token="token"), + ), + ( + "streaming_job_ready", + "_streaming_job_ready", + V0StreamingJobReadyRequest( + job_uuid=JOB_UUID, executor_token="token", public_key="key", ip="10.0.0.1", port=1234 + ), + ), + ( + "executor_finished", + "_executor_finished", + V0JobFinishedRequest( + job_uuid=JOB_UUID, docker_process_stdout="out", docker_process_stderr="err" + ), + ), + ( + "executor_specs", + "_executor_specs", + V0MachineSpecsRequest(job_uuid=JOB_UUID, specs=MachineSpecs(specs={"gpu": "fake"})), + ), + ("job_failed", "_job_failed", V0JobFailedRequest(job_uuid=JOB_UUID)), + ( + "horde_failed", + "_horde_failed", + V0HordeFailedRequest( + job_uuid=JOB_UUID, + reported_by=JobParticipantType.EXECUTOR, + message="boom", + reason=HordeFailureReason.GENERIC_ERROR, + ), + ), + ("volumes_ready", "_volumes_ready", V0VolumesReadyRequest(job_uuid=JOB_UUID)), + ("execution_done", "_execution_done", V0ExecutionDoneRequest(job_uuid=JOB_UUID)), +] + + +@pytest.fixture +def consumer() -> MinerValidatorConsumer: + consumer = MinerValidatorConsumer(MagicMock()) + consumer.send = AsyncMock() + return consumer + + +@pytest.fixture(params=RELAYED_EVENTS, ids=[event_type for event_type, _, _ in RELAYED_EVENTS]) +def relayed_event(request, consumer: MinerValidatorConsumer): + event_type, handler_name, message = request.param + handler = AsyncMock() + setattr(consumer, handler_name, handler) + event = {"type": event_type.replace("_", "."), **message.model_dump()} + return event_type, handler, message, event + + +@pytest.mark.asyncio +async def test_executor_event_dropped_before_authentication( + consumer: MinerValidatorConsumer, relayed_event, caplog: pytest.LogCaptureFixture +): + event_type, handler, _, event = relayed_event + caplog.set_level(logging.WARNING) + + await getattr(consumer, event_type)(event) + + assert handler.await_args_list == [] + assert consumer.send.await_args_list == [] + assert caplog.record_tuples == [ + ( + "compute_horde_miner.miner.miner_consumer.layer_utils", + logging.WARNING, + f"Dropping {event_type} message for job {JOB_UUID}: " + f"validator socket is not authenticated", + ) + ] + + +@pytest.mark.asyncio +async def test_executor_event_relayed_after_authentication( + consumer: MinerValidatorConsumer, relayed_event +): + event_type, handler, message, event = relayed_event + consumer.validator_authenticated = True + + await getattr(consumer, event_type)(event) + + assert handler.await_args_list == [call(message)] diff --git a/miner/app/src/compute_horde_miner/miner/tests/validator.py b/miner/app/src/compute_horde_miner/miner/tests/validator.py index 45784125e..c4318c7c6 100644 --- a/miner/app/src/compute_horde_miner/miner/tests/validator.py +++ b/miner/app/src/compute_horde_miner/miner/tests/validator.py @@ -1,19 +1,191 @@ +import asyncio +import json +import time +import uuid +from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager +from typing import Any +import bittensor from channels.testing import WebsocketCommunicator -from compute_horde.test_wallet import get_test_validator_wallet +from compute_horde.executor_class import DEFAULT_EXECUTOR_CLASS +from compute_horde.miner_client.organic import OrganicMinerClient +from compute_horde.protocol_messages import ( + V0InitialJobRequest, + V0JobRequest, + ValidatorAuthForMiner, +) +from compute_horde.receipts.schemas import JobStartedReceiptPayload +from compute_horde.test_wallet import get_test_miner_wallet, get_test_validator_wallet +from compute_horde.transport import WSTransport +from compute_horde.utils import sign_blob +from django.utils import timezone from compute_horde_miner import asgi +WEBSOCKET_TIMEOUT = 3 +# What the miner sends an authenticated validator when it has one executor of the default class. +MANIFEST_MESSAGE = { + "message_type": "V0ExecutorManifestRequest", + "manifest": {DEFAULT_EXECUTOR_CLASS: 1}, +} + + +def validator_hotkey() -> str: + return get_test_validator_wallet().hotkey.ss58_address + + +def miner_hotkey() -> str: + return get_test_miner_wallet().hotkey.ss58_address + + +def validator_keypair() -> bittensor.Keypair: + return get_test_validator_wallet().hotkey + @asynccontextmanager -async def fake_validator() -> WebsocketCommunicator: - validator_wallet = get_test_validator_wallet() +async def fake_validator( + validator_key: str | None = None, +) -> AsyncIterator[WebsocketCommunicator]: + """ + Validator socket to the miner, opened under the test validator wallet's hotkey unless + another key is given. + """ communicator = WebsocketCommunicator( application=asgi.application, - path=f"v0.1/validator_interface/{validator_wallet.hotkey.ss58_address}", + path=f"v0.1/validator_interface/{validator_key or validator_hotkey()}", ) - connected, _ = await communicator.connect() + connected, _ = await communicator.connect(timeout=WEBSOCKET_TIMEOUT) assert connected - yield communicator - await communicator.disconnect() + try: + yield communicator + finally: + await communicator.disconnect() + + +def auth_message( + timestamp: int | None = None, to_miner: str | None = None +) -> ValidatorAuthForMiner: + """ + Authentication frame of the test validator, signed for real, for the test miner unless + another miner hotkey is given. + """ + msg = ValidatorAuthForMiner( + validator_hotkey=validator_hotkey(), + miner_hotkey=to_miner or miner_hotkey(), + timestamp=int(time.time()) if timestamp is None else timestamp, + signature="", + ) + msg.signature = sign_blob(validator_keypair(), msg.blob_for_signing()) + return msg + + +async def authenticate(channel: WebsocketCommunicator, to_miner: str | None = None) -> None: + await channel.send_to(auth_message(to_miner=to_miner).model_dump_json()) + + +def organic_client_reconnect_greeting() -> Callable[[], str]: + """ + What the test validator's OrganicMinerClient sends first on every connection to the test + miner that its transport reopens. Needs a running event loop. + """ + client = OrganicMinerClient( + miner_hotkey=miner_hotkey(), + miner_address="localhost", + miner_port=8000, + job_uuid=str(uuid.uuid4()), + my_keypair=validator_keypair(), + ) + assert isinstance(client.transport, WSTransport) + assert client.transport.reconnect_greeting is not None + return client.transport.reconnect_greeting + + +def job_started_receipt( + job_uuid: str, + signer: bittensor.Keypair | None = None, + **overrides: Any, +) -> tuple[JobStartedReceiptPayload, str]: + """Receipt of the test validator for the test miner; overrides replace payload fields.""" + payload = JobStartedReceiptPayload( + **{ + "job_uuid": job_uuid, + "miner_hotkey": miner_hotkey(), + "validator_hotkey": validator_hotkey(), + "timestamp": timezone.now(), + "executor_class": DEFAULT_EXECUTOR_CLASS, + "is_organic": True, + "ttl": 60, + **overrides, + } + ) + return payload, sign_blob(signer or validator_keypair(), payload.blob_for_signing()) + + +def initial_job_request( + job_uuid: str, to_miner: str | None = None, **overrides: Any +) -> V0InitialJobRequest: + """ + Signed by the test validator, with a receipt for the test miner unless another miner hotkey + is given; overrides replace fields before signing. + """ + payload, signature = job_started_receipt(job_uuid, miner_hotkey=to_miner or miner_hotkey()) + return V0InitialJobRequest( + **{ + "job_uuid": job_uuid, + "executor_class": DEFAULT_EXECUTOR_CLASS, + "docker_image": "it's teeeeests", + "timeout_seconds": 60, + "job_started_receipt_payload": payload, + "job_started_receipt_signature": signature, + **overrides, + } + ).signed(validator_keypair()) + + +def job_request(job_uuid: str, **overrides: Any) -> V0JobRequest: + """Signed by the test validator; overrides replace fields before signing.""" + return V0JobRequest( + **{ + "job_uuid": job_uuid, + "executor_class": DEFAULT_EXECUTOR_CLASS, + "docker_image": "it's teeeeests again", + "docker_run_cmd": [], + "env": {"SOME_ENV_VAR": "some value"}, + "docker_run_options_preset": "none", + "volume": {"volume_type": "inline", "contents": "nonsense"}, + "output_upload": { + "output_upload_type": "zip_and_http_post", + "url": "http://localhost/output/upload", + "form_fields": {}, + }, + **overrides, + } + ).signed(validator_keypair()) + + +async def received_output( + communicator: WebsocketCommunicator, expected_count: int = 0, quiet_period: float = 0.5 +) -> list[dict[str, Any]]: + """ + The next expected_count frames the socket receives, waiting up to WEBSOCKET_TIMEOUT for + each, plus everything else that arrives until nothing has for quiet_period seconds: sent + frames as parsed JSON, other events (such as websocket.close) as raw ASGI messages. + """ + outputs: list[dict[str, Any]] = [] + while len(outputs) < expected_count or not await communicator.receive_nothing( + timeout=quiet_period + ): + output = await communicator.receive_output(timeout=WEBSOCKET_TIMEOUT) + outputs.append(json.loads(output["text"]) if output["type"] == "websocket.send" else output) + return outputs + + +async def wait_until( + condition: Callable[[], Awaitable[bool]], timeout: float = WEBSOCKET_TIMEOUT +) -> None: + """Polls the condition until it holds; fails when it still does not after the timeout.""" + deadline = time.monotonic() + timeout + while not await condition(): + assert time.monotonic() < deadline, "condition not met in time" + await asyncio.sleep(0.05) diff --git a/miner/app/src/compute_horde_miner/miner/validator_auth.py b/miner/app/src/compute_horde_miner/miner/validator_auth.py new file mode 100644 index 000000000..fc04907ca --- /dev/null +++ b/miner/app/src/compute_horde_miner/miner/validator_auth.py @@ -0,0 +1,58 @@ +""" +Replay protection for validator authentication frames: a frame is accepted only when its +timestamp is current and its signature has not been accepted before. +""" + +import datetime as dt +import re +import time + +from compute_horde.protocol_messages import ValidatorAuthForMiner +from django.db import IntegrityError + +from compute_horde_miner.miner.models import ConsumedValidatorAuth + +AUTH_MESSAGE_MAX_AGE = 10 +AUTH_MESSAGE_MAX_CLOCK_SKEW = 5 +# Consumed signatures are kept well past the age limit, after which the timestamp check alone +# rejects the frame. +CONSUMED_AUTH_RETENTION = dt.timedelta(minutes=1) +# How sign_blob writes a signature, in any letter case. The verifier also strips repeated "0x" +# prefixes, which would give one signature many more spellings. +SIGNATURE_FORMAT = re.compile(r"0x[0-9a-fA-F]{128}") + + +def auth_timestamp_error(msg: ValidatorAuthForMiner) -> str | None: + now = time.time() + if msg.timestamp < now - AUTH_MESSAGE_MAX_AGE: + return "msg too old" + if msg.timestamp > now + AUTH_MESSAGE_MAX_CLOCK_SKEW: + return "msg timestamp is in the future" + return None + + +def auth_signature_bytes(msg: ValidatorAuthForMiner) -> bytes | None: + """The frame's 64-byte signature, or None when it is not spelled as "0x" and 128 hex digits.""" + if not SIGNATURE_FORMAT.fullmatch(msg.signature): + return None + return bytes.fromhex(msg.signature[2:]) + + +async def consume_auth_signature(msg: ValidatorAuthForMiner, signature: bytes) -> bool: + """ + Records the frame's signature, as returned by auth_signature_bytes, as used; False means it + was used before, so the frame is a replay. Call it only for a frame whose signature + verified, so that unauthenticated clients cannot fill the table. + """ + await ConsumedValidatorAuth.objects.filter( + timestamp__lt=dt.datetime.now(dt.UTC) - CONSUMED_AUTH_RETENTION + ).adelete() + try: + await ConsumedValidatorAuth.objects.acreate( + signature=signature.hex(), + validator_hotkey=msg.validator_hotkey, + timestamp=dt.datetime.fromtimestamp(msg.timestamp, dt.UTC), + ) + except IntegrityError: + return False + return True diff --git a/miner/app/src/compute_horde_miner/settings.py b/miner/app/src/compute_horde_miner/settings.py index 2ae8567da..4b42bf617 100644 --- a/miner/app/src/compute_horde_miner/settings.py +++ b/miner/app/src/compute_horde_miner/settings.py @@ -483,6 +483,10 @@ def wrapped(*args, **kwargs): BITTENSOR_WALLET_NAME = env.str("BITTENSOR_WALLET_NAME") BITTENSOR_WALLET_HOTKEY_NAME = env.str("BITTENSOR_WALLET_HOTKEY_NAME") DEBUG_TURN_AUTHENTICATION_OFF = env.bool("DEBUG_TURN_AUTHENTICATION_OFF", default=False) +# Decline V0InitialJobRequest and V0JobRequest messages that carry no signature of the +# authenticated validator. Turn off only while validators that do not sign yet are in use; +# a signature that is present is verified either way. +REQUIRE_SIGNED_JOB_REQUESTS = env.bool("REQUIRE_SIGNED_JOB_REQUESTS", default=True) DYNAMIC_CONFIG_ENV = env.str("DYNAMIC_CONFIG_ENV", default="prod") CONFIG_CONTRACT_ADDRESS = env.str( diff --git a/validator/app/src/compute_horde_validator/validator/eviction.py b/validator/app/src/compute_horde_validator/validator/eviction.py index 7271856b6..932bdf9b4 100644 --- a/validator/app/src/compute_horde_validator/validator/eviction.py +++ b/validator/app/src/compute_horde_validator/validator/eviction.py @@ -3,6 +3,7 @@ from compute_horde.receipts.models import JobAcceptedReceipt, JobFinishedReceipt, JobStartedReceipt from django.conf import settings +from django.db.models import Q from django.utils.timezone import now from .models import ( @@ -25,6 +26,7 @@ def evict_all() -> None: evict_organic_jobs() evict_receipts() + evict_malformed_receipts() evict_system_events() evict_miner_blacklist() evict_miner_preliminary_reservations() @@ -43,11 +45,26 @@ def evict_organic_jobs() -> None: def evict_receipts() -> None: - logger.info("Evicting old receipts") - cutoff = now() - RECEIPTS_RETENTION_PERIOD - JobStartedReceipt.objects.filter(timestamp__lt=cutoff).delete() - JobAcceptedReceipt.objects.filter(timestamp__lt=cutoff).delete() - JobFinishedReceipt.objects.filter(timestamp__lt=cutoff).delete() + logger.info("Evicting receipts outside the retention window") + current_time = now() + # Receipts are dated by the issuing validator; a transferred receipt dated further ahead than + # the retention period would otherwise outlive it. + outside_retention_window = Q(timestamp__lt=current_time - RECEIPTS_RETENTION_PERIOD) | Q( + timestamp__gt=current_time + RECEIPTS_RETENTION_PERIOD + ) + for receipt_model in (JobStartedReceipt, JobAcceptedReceipt, JobFinishedReceipt): + receipt_model.objects.filter(outside_retention_window).delete() + + +def evict_malformed_receipts() -> None: + """ + Deletes job finished receipts whose score cannot be read. + + Such rows were transferred from miners before receipt fields were validated; scoring skips + them, so they only take up space and produce warnings. + """ + logger.info("Evicting job finished receipts with a malformed score") + JobFinishedReceipt.with_malformed_score().delete() def evict_miner_blacklist() -> None: diff --git a/validator/app/src/compute_horde_validator/validator/management/commands/debug_run_mock_streaming_job_to_miner.py b/validator/app/src/compute_horde_validator/validator/management/commands/debug_run_mock_streaming_job_to_miner.py index 136de749a..e09931dae 100644 --- a/validator/app/src/compute_horde_validator/validator/management/commands/debug_run_mock_streaming_job_to_miner.py +++ b/validator/app/src/compute_horde_validator/validator/management/commands/debug_run_mock_streaming_job_to_miner.py @@ -12,7 +12,7 @@ OrganicJobDetails, OrganicMinerClient, ) -from compute_horde.protocol_consts import HordeFailureReason +from compute_horde.protocol_consts import TRUSTED_MINER_FAKE_KEY, HordeFailureReason from compute_horde.protocol_messages import ( V0InitialJobRequest, V0JobRequest, @@ -86,7 +86,7 @@ async def run_streaming_job(options, wait_timeout: int = 300): job_started_receipt_payload=receipt_payload, job_started_receipt_signature=receipt_signature, streaming_details=StreamingDetails(public_key=public_key), - ), + ).signed(client.my_keypair), ) try: @@ -125,7 +125,7 @@ async def run_streaming_job(options, wait_timeout: int = 300): volume=job_details.volume, output_upload=job_details.output, artifacts_dir=job_details.artifacts_dir, - ) + ).signed(client.my_keypair) ) try: @@ -192,7 +192,12 @@ async def run_streaming_job(options, wait_timeout: int = 300): class Command(BaseCommand): def add_arguments(self, parser): - parser.add_argument("--miner_hotkey", type=str, help="Miner hotkey", required=True) + parser.add_argument( + "--miner_hotkey", + type=str, + help=f"Miner hotkey; {TRUSTED_MINER_FAKE_KEY} for the validator's trusted miner", + required=True, + ) parser.add_argument( "--miner_address", type=str, help="Miner IPv4 address", default="127.0.0.1" ) diff --git a/validator/app/src/compute_horde_validator/validator/migrations/0095_systemevent_receipt_malformed.py b/validator/app/src/compute_horde_validator/validator/migrations/0095_systemevent_receipt_malformed.py new file mode 100644 index 000000000..eba1cb34b --- /dev/null +++ b/validator/app/src/compute_horde_validator/validator/migrations/0095_systemevent_receipt_malformed.py @@ -0,0 +1,90 @@ +# Generated by Django 4.2.19 on 2026-10-19 13:58 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("validator", "0094_auto_20251126_1345"), + ] + + operations = [ + migrations.AlterField( + model_name="systemevent", + name="subtype", + field=models.CharField( + choices=[ + ("SUCCESS", "Success"), + ("FAILURE", "Failure"), + ("SUBTENSOR_CONNECTIVITY_ERROR", "Subtensor Connectivity Error"), + ("COMMIT_WEIGHTS_SUCCESS", "Commit Weights Success"), + ("COMMIT_WEIGHTS_ERROR", "Commit Weights Error"), + ("COMMIT_WEIGHTS_UNREVEALED_ERROR", "Commit Weights Unrevealed Error"), + ("REVEAL_WEIGHTS_ERROR", "Reveal Weights Error"), + ("REVEAL_WEIGHTS_SUCCESS", "Reveal Weights Success"), + ("SET_WEIGHTS_SUCCESS", "Set Weights Success"), + ("SET_WEIGHTS_ERROR", "Set Weights Error"), + ("GENERIC_ERROR", "Generic Error"), + ("WRITING_TO_CHAIN_TIMEOUT", "Writing To Chain Timeout"), + ("WRITING_TO_CHAIN_GENERIC_ERROR", "Writing To Chain Generic Error"), + ("GIVING_UP", "Giving Up"), + ("MANIFEST_ERROR", "Manifest Error"), + ("MANIFEST_TIMEOUT", "Manifest Timeout"), + ("MINER_CONNECTION_ERROR", "Miner Connection Error"), + ("MINER_SEND_ERROR", "Miner Send Error"), + ("MINER_SCORING_ERROR", "Miner Scoring Error"), + ("GENERIC_JOB_FAILURE", "Generic Job Failure"), + ("JOB_NOT_STARTED", "Job Not Started"), + ("JOB_REJECTED", "Job Rejected"), + ("JOB_EXCUSED", "Job Excused"), + ("JOB_CHEATED", "Job Cheated"), + ("MINER_BLACKLISTED", "Miner Blacklisted"), + ("JOB_TIMEOUT", "Job Timeout"), + ("JOB_EXECUTION_TIMEOUT", "Job Execution Timeout"), + ("RECEIPT_FETCH_ERROR", "Receipt Fetch Error"), + ("RECEIPT_SEND_ERROR", "Receipt Send Error"), + ("RECEIPT_MALFORMED", "Receipt Malformed"), + ("SPECS_SENDING_ERROR", "Specs Send Error"), + ("HEARTBEAT_ERROR", "Heartbeat Error"), + ("UNEXPECTED_MESSAGE", "Unexpected Message"), + ("UNAUTHORIZED", "Unauthorized"), + ("SYNTHETIC_BATCH", "Synthetic Batch"), + ("SYNTHETIC_JOB", "Synthetic Job"), + ("CHECKPOINT", "Checkpoint"), + ("OVERSLEPT", "Overslept"), + ("WARNING", "Warning"), + ("FAILED_TO_WAIT", "Failed To Wait"), + ("TRUSTED_MINER_NOT_CONFIGURED", "Trusted Miner Not Configured"), + ("LLM_PROMPT_ANSWERS_MISSING", "Llm Prompt Answers Missing"), + ("INSUFFICIENT_PROMPTS", "Insufficient Prompts"), + ("UNPROCESSED_WORKLOADS", "Unprocessed Workloads"), + ("PROMPT_GENERATION_STARTED", "Prompt Generation Started"), + ("PROMPT_SAMPLING_SKIPPED", "Prompt Sampling Skipped"), + ("NEW_WORKLOADS_CREATED", "New Workloads Created"), + ("ERROR_UPLOADING_TO_S3", "Error Uploading To S3"), + ("ERROR_DOWNLOADING_FROM_S3", "Error Downloading From S3"), + ("ERROR_DOWNLOADING_FROM_HUGGINGFACE", "Error Downloading From Huggingface"), + ("ERROR_DURING_SECURITY_CHECK", "Error Failed Security Check"), + ("ERROR_EXECUTOR_REPORTED_TIMEOUT", "Error Executor Reported Timeout"), + ("ERROR_VALIDATOR_REPORTED_TIMEOUT", "Error Validator Reported Timeout"), + ("JOB_PROCESS_NONZERO_EXIT_CODE", "Job Process Nonzero Exit Code"), + ("JOB_VOLUME_DOWNLOAD_FAILED", "Job Volume Download Failed"), + ("JOB_RESULT_UPLOAD_FAILED", "Job Result Upload Failed"), + ( + "LLM_PROMPT_ANSWERS_DOWNLOAD_WORKER_FAILED", + "Llm Prompt Answers Download Worker Failed", + ), + ("APPLIED_BURNING", "Applied Burning"), + ("NO_BURNING", "No Burning"), + ("GETTING_MINER_COLLATERAL_FAILED", "Getting Miner Collateral Failed"), + ( + "BLOCK_ALREADY_INSERTED_PROBABLY_BLOCK_CACHE_JITTER_IGNORE_ME_FOR_NOW", + "Block Already Inserted Probably Block Cache Jitter Ignore Me For Now", + ), + ("JOB_ROUTING_SUCCESS", "Job Routing Success"), + ("JOB_ROUTING_FAILURE", "Job Routing Failure"), + ], + max_length=255, + ), + ), + ] diff --git a/validator/app/src/compute_horde_validator/validator/models/spaghetti.py b/validator/app/src/compute_horde_validator/validator/models/spaghetti.py index 1e925307f..28713b8c7 100644 --- a/validator/app/src/compute_horde_validator/validator/models/spaghetti.py +++ b/validator/app/src/compute_horde_validator/validator/models/spaghetti.py @@ -79,6 +79,7 @@ class EventSubType(models.TextChoices): JOB_EXECUTION_TIMEOUT = "JOB_EXECUTION_TIMEOUT" RECEIPT_FETCH_ERROR = "RECEIPT_FETCH_ERROR" RECEIPT_SEND_ERROR = "RECEIPT_SEND_ERROR" + RECEIPT_MALFORMED = "RECEIPT_MALFORMED" SPECS_SEND_ERROR = "SPECS_SENDING_ERROR" HEARTBEAT_ERROR = "HEARTBEAT_ERROR" UNEXPECTED_MESSAGE = "UNEXPECTED_MESSAGE" diff --git a/validator/app/src/compute_horde_validator/validator/organic_jobs/miner_driver.py b/validator/app/src/compute_horde_validator/validator/organic_jobs/miner_driver.py index 14a7ae007..f0f5109fa 100644 --- a/validator/app/src/compute_horde_validator/validator/organic_jobs/miner_driver.py +++ b/validator/app/src/compute_horde_validator/validator/organic_jobs/miner_driver.py @@ -28,6 +28,7 @@ execute_organic_job_on_miner, ) from compute_horde.protocol_consts import ( + TRUSTED_MINER_FAKE_KEY, HordeFailureReason, JobFailureReason, JobParticipantType, @@ -56,7 +57,6 @@ from compute_horde_validator.validator.organic_jobs.miner_client import MinerClient from compute_horde_validator.validator.routing.default import routing from compute_horde_validator.validator.routing.types import JobRoute, MinerIncidentType -from compute_horde_validator.validator.utils import TRUSTED_MINER_FAKE_KEY logger = logging.getLogger(__name__) diff --git a/validator/app/src/compute_horde_validator/validator/organic_jobs/miner_driver_sync.py b/validator/app/src/compute_horde_validator/validator/organic_jobs/miner_driver_sync.py index a7c36e65f..843e66966 100644 --- a/validator/app/src/compute_horde_validator/validator/organic_jobs/miner_driver_sync.py +++ b/validator/app/src/compute_horde_validator/validator/organic_jobs/miner_driver_sync.py @@ -27,6 +27,7 @@ ) from compute_horde.job_errors import HordeError from compute_horde.protocol_consts import ( + TRUSTED_MINER_FAKE_KEY, HordeFailureReason, JobFailureReason, JobParticipantType, @@ -60,6 +61,7 @@ JobAcceptedReceiptPayload, JobFinishedReceiptPayload, JobStartedReceiptPayload, + clamp_ttl, ) from compute_horde.utils import Timer, sign_blob from django.conf import settings @@ -77,7 +79,6 @@ ) from compute_horde_validator.validator.routing.default import routing from compute_horde_validator.validator.routing.types import JobRoute, MinerIncidentType -from compute_horde_validator.validator.utils import TRUSTED_MINER_FAKE_KEY logger = logging.getLogger(__name__) @@ -637,6 +638,21 @@ def run(self) -> None: self._undo_allowance_if_failed() self.miner_client.close() + def _job_finished_receipt_payload( + self, score_str: str, block_numbers: list[int] + ) -> JobFinishedReceiptPayload: + return JobFinishedReceiptPayload( + job_uuid=str(self.job.job_uuid), + miner_hotkey=self.miner_hotkey, + validator_hotkey=self.my_keypair.ss58_address, + timestamp=datetime.datetime.now(datetime.UTC), + time_started=datetime.datetime.fromtimestamp(self._started_at, datetime.UTC), + # Wall-clock durations go negative after a backwards clock step; the payload rejects that. + time_took_us=max(0, int((time.time() - self._started_at) * 1_000_000)), + block_numbers=block_numbers, + score_str=score_str, + ) + def _send_job_finished_receipt_if_failed(self) -> None: if self._state == DriverState.COMPLETE: return @@ -645,16 +661,8 @@ def _send_job_finished_receipt_if_failed(self) -> None: return try: - started_timestamp = datetime.datetime.fromtimestamp(self._started_at, datetime.UTC) - receipt_payload = JobFinishedReceiptPayload( - job_uuid=str(self.job.job_uuid), - miner_hotkey=self.miner_hotkey, - validator_hotkey=self.my_keypair.ss58_address, - timestamp=datetime.datetime.now(datetime.UTC), - time_started=started_timestamp, - time_took_us=int((time.time() - self._started_at) * 1_000_000), - block_numbers=[], - score_str="0.000000", + receipt_payload = self._job_finished_receipt_payload( + score_str="0.000000", block_numbers=[] ) receipt_signature = sign_blob(self.my_keypair, receipt_payload.blob_for_signing()) self.miner_client.send( @@ -776,7 +784,7 @@ def _reserve_executor(self) -> DriverState: job_started_receipt_signature=receipt_signature, streaming_details=self.request.streaming_details, executor_timing=executor_timing, - ) + ).signed(self.my_keypair) try: self.miner_client.send(msg) @@ -840,7 +848,7 @@ def _send_job_accepted_receipt(self) -> DriverState: validator_hotkey=self.my_keypair.ss58_address, timestamp=datetime.datetime.now(datetime.UTC), time_accepted=datetime.datetime.now(datetime.UTC), - ttl=job_accepted_receipt_ttl, + ttl=clamp_ttl(job_accepted_receipt_ttl), ) receipt_signature = sign_blob(self.my_keypair, receipt_payload.blob_for_signing()) receipt_msg = V0JobAcceptedReceiptRequest( @@ -910,7 +918,7 @@ def _prepare_volumes(self) -> DriverState: volume=None, # already sent in V0InitialJobRequest output_upload=self.request.output_upload, artifacts_dir=self.request.artifacts_dir, - ) + ).signed(self.my_keypair) try: self.miner_client.send(msg) @@ -1045,16 +1053,9 @@ def _collect_results(self) -> DriverState: ) try: - started_timestamp = datetime.datetime.fromtimestamp(self._started_at, datetime.UTC) - receipt_payload = JobFinishedReceiptPayload( - job_uuid=str(self.job.job_uuid), - miner_hotkey=self.miner_hotkey, - validator_hotkey=self.my_keypair.ss58_address, - timestamp=datetime.datetime.now(datetime.UTC), - time_started=started_timestamp, - time_took_us=int((time.time() - self._started_at) * 1_000_000), - block_numbers=self.job.allowance_blocks or [], + receipt_payload = self._job_finished_receipt_payload( score_str=f"{float(self.job.allowance_job_value or 0):.6f}", + block_numbers=self.job.allowance_blocks or [], ) receipt_signature = sign_blob(self.my_keypair, receipt_payload.blob_for_signing()) self.miner_client.send( diff --git a/validator/app/src/compute_horde_validator/validator/receipts/README.md b/validator/app/src/compute_horde_validator/validator/receipts/README.md index fba30e37d..58afd9df7 100644 --- a/validator/app/src/compute_horde_validator/validator/receipts/README.md +++ b/validator/app/src/compute_horde_validator/validator/receipts/README.md @@ -66,7 +66,7 @@ rows: list[JobSpendingInfo] = receipts().get_finished_jobs_for_block_range( - validator_hotkey: str - miner_hotkey: str - executor_class: ExecutorClass -- executor_seconds_cost: int +- executor_seconds_cost: float - paid_with_blocks: list[int] - started_at: datetime """ diff --git a/validator/app/src/compute_horde_validator/validator/receipts/default.py b/validator/app/src/compute_horde_validator/validator/receipts/default.py index f83a4e6aa..3c689be7b 100644 --- a/validator/app/src/compute_horde_validator/validator/receipts/default.py +++ b/validator/app/src/compute_horde_validator/validator/receipts/default.py @@ -31,6 +31,7 @@ from compute_horde_validator.validator.allowance.default import allowance from compute_horde_validator.validator.allowance.utils.supertensor import supertensor from compute_horde_validator.validator.dynamic_config import aget_config +from compute_horde_validator.validator.models import SystemEvent from compute_horde_validator.validator.models.allowance.internal import Block from .base import ReceiptsBase @@ -251,18 +252,45 @@ def get_finished_jobs_for_block_range( # TODO(new scoring): this could be inefficient finished_qs = finished_qs.annotate(started_at=Subquery(starts.values("timestamp")[:1])) - return [ - JobSpendingInfo( - job_uuid=str(r.job_uuid), - validator_hotkey=r.validator_hotkey, - miner_hotkey=r.miner_hotkey, - executor_class=executor_class, - executor_seconds_cost=r.score(), - paid_with_blocks=r.block_numbers, - started_at=r.started_at, - ) - for r in finished_qs - ] + spendings: list[JobSpendingInfo] = [] + for r in finished_qs: + try: + spendings.append( + JobSpendingInfo( + job_uuid=str(r.job_uuid), + validator_hotkey=r.validator_hotkey, + miner_hotkey=r.miner_hotkey, + executor_class=executor_class, + executor_seconds_cost=r.score(), + paid_with_blocks=r.block_numbers, + started_at=r.started_at, + ) + ) + except ValueError as e: + # Rows stored before receipt fields were validated on transfer may hold anything; + # one such row must not stop the scoring of all the others. + self._report_malformed_receipt(r, e) + return spendings + + def _report_malformed_receipt(self, receipt: JobFinishedReceipt, error: Exception) -> None: + logger.warning( + "Skipping malformed job finished receipt job_uuid=%s miner_hotkey=%s " + "validator_hotkey=%s: %s", + receipt.job_uuid, + receipt.miner_hotkey, + receipt.validator_hotkey, + error, + ) + SystemEvent.objects.using(settings.DEFAULT_DB_ALIAS).create( + type=SystemEvent.EventType.RECEIPT_FAILURE, + subtype=SystemEvent.EventSubType.RECEIPT_MALFORMED, + long_description=str(error), + data={ + "job_uuid": str(receipt.job_uuid), + "miner_hotkey": receipt.miner_hotkey, + "validator_hotkey": receipt.validator_hotkey, + }, + ) def get_busy_executor_count( self, executor_class: ExecutorClass, at_time: datetime.datetime diff --git a/validator/app/src/compute_horde_validator/validator/receipts/tests/test_receipts.py b/validator/app/src/compute_horde_validator/validator/receipts/tests/test_receipts.py index 00651920b..032cf82b6 100644 --- a/validator/app/src/compute_horde_validator/validator/receipts/tests/test_receipts.py +++ b/validator/app/src/compute_horde_validator/validator/receipts/tests/test_receipts.py @@ -1,5 +1,6 @@ import datetime import uuid +from typing import Any import bittensor_wallet import pytest @@ -15,10 +16,11 @@ from compute_horde_core.executor_class import ExecutorClass from django.utils import timezone -from compute_horde_validator.validator.models import Miner +from compute_horde_validator.validator.models import Miner, SystemEvent from compute_horde_validator.validator.models.allowance.internal import Block from compute_horde_validator.validator.receipts import receipts -from compute_horde_validator.validator.tests.helpers import patch_constance +from compute_horde_validator.validator.receipts.types import JobSpendingInfo +from compute_horde_validator.validator.tests.helpers import patch_constance, persist_finished_job @pytest.mark.asyncio @@ -639,3 +641,68 @@ def test_get_busy_executor_count__two_miners_saturated(): busy_map = receipts().get_busy_executor_count(executor_class=executor, at_time=now) for miner, count in miners_counts.items(): assert busy_map.get(miner) == count + + +@pytest.mark.django_db(databases=["default", "default_alias"], transaction=True) +def test_get_finished_jobs_for_block_range_skips_malformed_receipts(settings): + start_block, end_block = 100, 105 + start_ts = datetime.datetime(2025, 1, 1, tzinfo=datetime.UTC) + end_ts = start_ts + datetime.timedelta(minutes=10) + Block.objects.create(block_number=start_block, creation_timestamp=start_ts) + Block.objects.create(block_number=end_block, creation_timestamp=end_ts) + validator_hotkey = settings.BITTENSOR_WALLET().get_hotkey().ss58_address + started_at = start_ts - datetime.timedelta(seconds=30) + finished_at = start_ts + datetime.timedelta(minutes=5) + + good_job_uuid = persist_finished_job( + "honest_miner", + validator_hotkey, + started_at, + finished_at, + score_str="3.000000", + block_numbers=[100, 101], + ) + malformed_fields: list[dict[str, Any]] = [ + {"score_str": "abc"}, + {"score_str": "nan"}, + {"score_str": "inf"}, + {"score_str": "-1"}, + {"score_str": "1e3"}, + {"score_str": "", "block_numbers": [100]}, + {"score_str": "1.000000", "block_numbers": [[100]]}, + {"score_str": "1.000000", "block_numbers": {"100": 1}}, + {"score_str": "1.000000", "block_numbers": "100"}, + ] + malformed_job_uuids = [ + persist_finished_job("rogue_miner", validator_hotkey, started_at, finished_at, **fields) + for fields in malformed_fields + ] + + rows = receipts().get_finished_jobs_for_block_range( + start_block, end_block, executor_class=ExecutorClass.always_on__gpu_24gb + ) + + assert rows == [ + JobSpendingInfo( + job_uuid=good_job_uuid, + validator_hotkey=validator_hotkey, + miner_hotkey="honest_miner", + executor_class=ExecutorClass.always_on__gpu_24gb, + executor_seconds_cost=3, + paid_with_blocks=[100, 101], + started_at=started_at, + ) + ] + reported = sorted( + (event.type, event.subtype, event.data["job_uuid"], event.data["miner_hotkey"]) + for event in SystemEvent.objects.using(settings.DEFAULT_DB_ALIAS).all() + ) + assert reported == sorted( + ( + SystemEvent.EventType.RECEIPT_FAILURE, + SystemEvent.EventSubType.RECEIPT_MALFORMED, + job_uuid, + "rogue_miner", + ) + for job_uuid in malformed_job_uuids + ) diff --git a/validator/app/src/compute_horde_validator/validator/receipts/types.py b/validator/app/src/compute_horde_validator/validator/receipts/types.py index cf1faecbe..2d9bac70a 100644 --- a/validator/app/src/compute_horde_validator/validator/receipts/types.py +++ b/validator/app/src/compute_horde_validator/validator/receipts/types.py @@ -13,6 +13,6 @@ class JobSpendingInfo(pydantic.BaseModel): validator_hotkey: str miner_hotkey: str executor_class: ExecutorClass - executor_seconds_cost: int + executor_seconds_cost: float paid_with_blocks: list[int] started_at: datetime diff --git a/validator/app/src/compute_horde_validator/validator/routing/default.py b/validator/app/src/compute_horde_validator/validator/routing/default.py index 6c9863a26..c98dc81f6 100644 --- a/validator/app/src/compute_horde_validator/validator/routing/default.py +++ b/validator/app/src/compute_horde_validator/validator/routing/default.py @@ -7,6 +7,7 @@ OrganicJobRequest, V2JobRequest, ) +from compute_horde.protocol_consts import TRUSTED_MINER_FAKE_KEY from compute_horde_core.executor_class import ExecutorClass from django.conf import settings from django.db.models import Count @@ -33,7 +34,6 @@ NotEnoughCollateralException, ) from compute_horde_validator.validator.routing.utils import weighted_shuffle -from compute_horde_validator.validator.utils import TRUSTED_MINER_FAKE_KEY logger = logging.getLogger(__name__) diff --git a/validator/app/src/compute_horde_validator/validator/routing/tests/test_job_routing.py b/validator/app/src/compute_horde_validator/validator/routing/tests/test_job_routing.py index ceb388cc3..34e2e4400 100644 --- a/validator/app/src/compute_horde_validator/validator/routing/tests/test_job_routing.py +++ b/validator/app/src/compute_horde_validator/validator/routing/tests/test_job_routing.py @@ -6,6 +6,7 @@ import pytest from compute_horde.executor_class import DEFAULT_EXECUTOR_CLASS from compute_horde.fv_protocol.facilitator_requests import V2JobRequest +from compute_horde.protocol_consts import TRUSTED_MINER_FAKE_KEY from compute_horde.receipts.models import JobFinishedReceipt, JobStartedReceipt from django.utils import timezone @@ -20,7 +21,6 @@ AllMinersBusy, NotEnoughCollateralException, ) -from compute_horde_validator.validator.utils import TRUSTED_MINER_FAKE_KEY JOB_REQUEST = V2JobRequest( uuid=str(uuid.uuid4()), diff --git a/validator/app/src/compute_horde_validator/validator/scoring/main_hotkey_querying.py b/validator/app/src/compute_horde_validator/validator/scoring/main_hotkey_querying.py index 7c3808498..9e1083703 100644 --- a/validator/app/src/compute_horde_validator/validator/scoring/main_hotkey_querying.py +++ b/validator/app/src/compute_horde_validator/validator/scoring/main_hotkey_querying.py @@ -72,17 +72,20 @@ async def _query_single_miner(miner: Miner) -> str | None: Returns: Main hotkey (or None) """ + + def greeting() -> str: + return _create_auth_message(miner.hotkey).model_dump_json() + transport = None try: transport = WSTransport( f"main_hotkey_query_{miner.hotkey}", f"ws://{miner.address}:{miner.port}/v0.1/validator_interface/{settings.BITTENSOR_WALLET().get_hotkey().ss58_address}", max_retries=2, + reconnect_greeting=greeting, ) await transport.start() - - auth_msg = _create_auth_message(miner.hotkey) - await transport.send(auth_msg.model_dump_json()) + await transport.send(greeting()) request = V0MainHotkeyMessage() await transport.send(request.model_dump_json()) diff --git a/validator/app/src/compute_horde_validator/validator/scoring/tests/test_paid_job_scores.py b/validator/app/src/compute_horde_validator/validator/scoring/tests/test_paid_job_scores.py index 2ea6ce73b..3cc464b1f 100644 --- a/validator/app/src/compute_horde_validator/validator/scoring/tests/test_paid_job_scores.py +++ b/validator/app/src/compute_horde_validator/validator/scoring/tests/test_paid_job_scores.py @@ -11,7 +11,7 @@ from compute_horde_validator.validator.scoring.calculations import ( calculate_allowance_paid_job_scores, ) -from compute_horde_validator.validator.tests.helpers import patch_constance +from compute_horde_validator.validator.tests.helpers import patch_constance, persist_finished_job pytestmark = [ pytest.mark.django_db, @@ -245,3 +245,38 @@ def finished_jobs(_, __, executor_class, organic_only): assert scores == {} else: assert scores == {EXECUTOR_CLASS: {MINER: expected_scores}} + + +@pytest.mark.django_db(databases=["default", "default_alias"], transaction=True) +def test_paid_job_scores__malformed_persisted_receipts_do_not_break_scoring(): + """ + Receipts stored before wire-format validation may carry anything; scoring must skip them + and still score the well-formed ones. + """ + started_at = _block_time(50) + persist_finished_job( + MINER, + VALIDATOR, + started_at, + _block_time(55), + score_str="3.000000", + block_numbers=[40, 41, 42], + ) + persist_finished_job( + "rogue_miner", VALIDATOR, started_at, _block_time(56), score_str="nan", block_numbers=[43] + ) + persist_finished_job( + "rogue_miner", VALIDATOR, started_at, _block_time(57), score_str="abc", block_numbers=[44] + ) + persist_finished_job( + "rogue_miner", + VALIDATOR, + started_at, + _block_time(58), + score_str="1.000000", + block_numbers=[[45]], + ) + + assert calculate_allowance_paid_job_scores(CYCLE_START, CYCLE_END) == { + EXECUTOR_CLASS: {MINER: 3} + } diff --git a/validator/app/src/compute_horde_validator/validator/scoring/tests/test_scoring_engine.py b/validator/app/src/compute_horde_validator/validator/scoring/tests/test_scoring_engine.py index e09ebd740..8e191fad7 100644 --- a/validator/app/src/compute_horde_validator/validator/scoring/tests/test_scoring_engine.py +++ b/validator/app/src/compute_horde_validator/validator/scoring/tests/test_scoring_engine.py @@ -118,7 +118,7 @@ def _reset_transports(self): def _get_mock_wstransport_constructor(self): """Get a mock WSTransport constructor that returns simulation transports.""" - def mock_wstransport_constructor(name, url, max_retries=2): + def mock_wstransport_constructor(name, url, **kwargs): if ":8081" in url: return self.transport1 elif ":8082" in url: diff --git a/validator/app/src/compute_horde_validator/validator/tests/helpers.py b/validator/app/src/compute_horde_validator/validator/tests/helpers.py index babc883ae..de89f9c5b 100644 --- a/validator/app/src/compute_horde_validator/validator/tests/helpers.py +++ b/validator/app/src/compute_horde_validator/validator/tests/helpers.py @@ -4,7 +4,8 @@ import os import shlex import subprocess -from datetime import timedelta +import uuid +from datetime import datetime, timedelta from pathlib import Path from time import monotonic from typing import Any @@ -25,7 +26,9 @@ V0VolumesReadyRequest, ValidatorToMinerMessage, ) +from compute_horde.receipts.models import JobFinishedReceipt, JobStartedReceipt from compute_horde.utils import ValidatorInfo +from compute_horde_core.executor_class import ExecutorClass from compute_horde_core.signature import Signature from django.conf import settings from pydantic import TypeAdapter @@ -362,6 +365,45 @@ def __init__( self.block = MockBlock(block_num) +def persist_finished_job( + miner_hotkey: str, + validator_hotkey: str, + started_at: datetime, + finished_at: datetime, + executor_class: ExecutorClass = ExecutorClass.always_on__gpu_24gb, + **finished_fields: Any, +) -> str: + """ + Stores a job started receipt and the matching job finished receipt, returning the job uuid. + + finished_fields are written to the finished receipt unchecked, so a test can store values + the wire format would reject, such as a malformed score_str. + """ + job_uuid = str(uuid.uuid4()) + JobStartedReceipt.objects.create( + job_uuid=job_uuid, + miner_hotkey=miner_hotkey, + validator_hotkey=validator_hotkey, + validator_signature="sig", + timestamp=started_at, + executor_class=str(executor_class), + is_organic=True, + ttl=60, + ) + JobFinishedReceipt.objects.create( + job_uuid=job_uuid, + miner_hotkey=miner_hotkey, + validator_hotkey=validator_hotkey, + validator_signature="sig", + miner_signature="sig", + timestamp=finished_at, + time_started=started_at, + time_took_us=1, + **finished_fields, + ) + return job_uuid + + def check_system_events( type: SystemEvent.EventType, subtype: SystemEvent.EventSubType, count: int = 1 ): diff --git a/validator/app/src/compute_horde_validator/validator/tests/test_organic_jobs/test_facilitator_client.py b/validator/app/src/compute_horde_validator/validator/tests/test_organic_jobs/test_facilitator_client.py index 9a295db12..31fef7b03 100644 --- a/validator/app/src/compute_horde_validator/validator/tests/test_organic_jobs/test_facilitator_client.py +++ b/validator/app/src/compute_horde_validator/validator/tests/test_organic_jobs/test_facilitator_client.py @@ -18,6 +18,7 @@ V0AuthenticationRequest, V0MachineSpecsUpdate, ) +from compute_horde.protocol_consts import TRUSTED_MINER_FAKE_KEY from django.utils import timezone from compute_horde_validator.validator.allowance.default import allowance @@ -34,7 +35,7 @@ from compute_horde_validator.validator.organic_jobs.facilitator_client import ( FacilitatorClient, ) -from compute_horde_validator.validator.utils import MACHINE_SPEC_CHANNEL, TRUSTED_MINER_FAKE_KEY +from compute_horde_validator.validator.utils import MACHINE_SPEC_CHANNEL from ..helpers import ( MockFaillingMinerClient, diff --git a/validator/app/src/compute_horde_validator/validator/tests/test_organic_jobs/test_miner_driver_sync.py b/validator/app/src/compute_horde_validator/validator/tests/test_organic_jobs/test_miner_driver_sync.py index a8f3705f4..1503a5fe6 100644 --- a/validator/app/src/compute_horde_validator/validator/tests/test_organic_jobs/test_miner_driver_sync.py +++ b/validator/app/src/compute_horde_validator/validator/tests/test_organic_jobs/test_miner_driver_sync.py @@ -1,3 +1,4 @@ +import datetime import uuid from collections.abc import Callable from functools import partial @@ -17,13 +18,22 @@ V0ExecutorManifestRequest, V0ExecutorReadyRequest, V0HordeFailedRequest, + V0InitialJobRequest, V0JobAcceptedReceiptRequest, V0JobFailedRequest, V0JobFinishedReceiptRequest, V0JobFinishedRequest, + V0JobRequest, V0VolumesReadyRequest, ValidatorToMinerMessage, ) +from compute_horde.receipts.models import JobFinishedReceipt +from compute_horde.receipts.schemas import ( + MAX_TTL_SECONDS, + JobAcceptedReceiptPayload, + JobFinishedReceiptPayload, +) +from freezegun import freeze_time from compute_horde_validator.validator.models import Miner from compute_horde_validator.validator.organic_jobs.facilitator_client import OrganicJob @@ -82,10 +92,10 @@ def query_sent_models( return result -def create_job_request(job_uuid: str) -> OrganicJobRequest: +def create_job_request(job_uuid: str, **overrides) -> OrganicJobRequest: """Convert V2JobRequest to OrganicJobRequest for the sync driver.""" v2_request = get_dummy_job_request_v2(job_uuid) - return OrganicJobRequest( + fields = dict( type="job.new", uuid=job_uuid, docker_image=v2_request.docker_image, @@ -99,6 +109,60 @@ def create_job_request(job_uuid: str) -> OrganicJobRequest: streaming_start_time_limit=v2_request.streaming_start_time_limit, upload_time_limit=v2_request.upload_time_limit, ) + return OrganicJobRequest(**{**fields, **overrides}) + + +SUCCESSFUL_FLOW = [ + V0AcceptJobRequest, + V0ExecutorReadyRequest, + V0VolumesReadyRequest, + V0ExecutionDoneRequest, + partial(V0JobFinishedRequest, docker_process_stdout="", docker_process_stderr="", artifacts={}), +] + + +def _driver_for( + settings, job_request: OrganicJobRequest, messages_from_miner: list, **job_fields +) -> tuple[SyncOrganicJobDriver, MockSyncMinerClient]: + miner, _ = Miner.objects.get_or_create(hotkey="miner_client") + job = OrganicJob.objects.create( + job_uuid=job_request.uuid, + miner=miner, + miner_address="irrelevant", + miner_address_ip_version=4, + miner_port=9999, + executor_class=DEFAULT_EXECUTOR_CLASS, + job_description="User job from facilitator", + block=42, + **job_fields, + ) + my_keypair = settings.BITTENSOR_WALLET().hotkey + miner_client = MockSyncMinerClient( + url="ws://miner:9999/v0.1/validator_interface/test", + miner_hotkey="miner_hotkey", + validator_keypair=my_keypair, + messages_to_return=[message(job_uuid=job_request.uuid) for message in messages_from_miner], + ) + driver = SyncOrganicJobDriver( + miner_client, + job, + job_request, + miner_hotkey="miner_hotkey", + my_keypair=my_keypair, + allowed_leeway=5, + reservation_time_limit=10, + executor_startup_time_limit=10, + max_overall_time_limit=60, + ) + return driver, miner_client + + +def _run(driver: SyncOrganicJobDriver) -> None: + with mock.patch( + "compute_horde_validator.validator.organic_jobs.miner_driver_sync.supertensor" + ) as mock_supertensor: + mock_supertensor.return_value.list_validators.return_value = [] + driver.run() @pytest.mark.django_db(databases=["default", "default_alias"], transaction=True) @@ -287,6 +351,104 @@ def track_job_status_updates(x: JobStatusUpdate): assert miner_client.query_sent_models(model_class=V0JobFinishedReceiptRequest) +@pytest.mark.django_db(databases=["default", "default_alias"], transaction=True) +def test_sync_miner_driver_signs_job_requests(settings): + miner, _ = Miner.objects.get_or_create(hotkey="miner_client") + job_uuid = str(uuid.uuid4()) + job_request = create_job_request(job_uuid) + job = OrganicJob.objects.create( + job_uuid=job_uuid, + miner=miner, + miner_address="irrelevant", + miner_address_ip_version=4, + miner_port=9999, + executor_class=DEFAULT_EXECUTOR_CLASS, + job_description="User job from facilitator", + block=42, + ) + my_keypair = settings.BITTENSOR_WALLET().hotkey + miner_client = MockSyncMinerClient( + url="ws://miner:9999/v0.1/validator_interface/test", + miner_hotkey="miner_hotkey", + validator_keypair=my_keypair, + messages_to_return=[ + V0AcceptJobRequest(job_uuid=job_uuid), + V0ExecutorReadyRequest(job_uuid=job_uuid), + V0VolumesReadyRequest(job_uuid=job_uuid), + V0ExecutionDoneRequest(job_uuid=job_uuid), + V0JobFinishedRequest( + job_uuid=job_uuid, + docker_process_stdout="mocked stdout", + docker_process_stderr="mocked stderr", + artifacts={}, + ), + ], + ) + + with mock.patch( + "compute_horde_validator.validator.organic_jobs.miner_driver_sync.supertensor" + ) as mock_supertensor: + mock_supertensor.return_value.list_validators.return_value = [] + SyncOrganicJobDriver( + miner_client, + job, + job_request, + miner_hotkey="miner_hotkey", + my_keypair=my_keypair, + allowed_leeway=5, + reservation_time_limit=10, + executor_startup_time_limit=10, + max_overall_time_limit=60, + ).run() + + [initial_job_request] = miner_client.query_sent_models(model_class=V0InitialJobRequest) + [full_job_request] = miner_client.query_sent_models(model_class=V0JobRequest) + assert initial_job_request.model_dump(mode="json") == { + "message_type": "V0InitialJobRequest", + "job_uuid": job_uuid, + "executor_class": DEFAULT_EXECUTOR_CLASS, + "docker_image": "nvidia", + "timeout_seconds": None, + "volume": job_request.volume.model_dump(mode="json"), + "job_started_receipt_payload": { + "receipt_type": "JobStartedReceipt", + "job_uuid": job_uuid, + "miner_hotkey": "miner_hotkey", + "validator_hotkey": my_keypair.ss58_address, + "timestamp": mock.ANY, + "executor_class": DEFAULT_EXECUTOR_CLASS, + "is_organic": True, + "ttl": 10, + }, + "job_started_receipt_signature": mock.ANY, + "streaming_details": None, + "executor_timing": { + "allowed_leeway": 5, + "download_time_limit": 1, + "execution_time_limit": 1, + "upload_time_limit": 1, + "streaming_start_time_limit": 1, + }, + "validator_signature": mock.ANY, + } + assert initial_job_request.is_signed_by(my_keypair.ss58_address) + assert full_job_request.model_dump(mode="json") == { + "message_type": "V0JobRequest", + "job_uuid": job_uuid, + "executor_class": DEFAULT_EXECUTOR_CLASS, + "docker_image": "nvidia", + "raw_script": None, + "docker_run_options_preset": "nvidia_all", + "docker_run_cmd": [], + "env": {}, + "volume": None, + "output_upload": job_request.output_upload.model_dump(mode="json"), + "artifacts_dir": None, + "validator_signature": mock.ANY, + } + assert full_job_request.is_signed_by(my_keypair.ss58_address) + + @pytest.mark.django_db(databases=["default", "default_alias"], transaction=True) def test_sync_miner_driver_connection_failure(settings): """Test that connection failures are handled properly.""" @@ -723,3 +885,62 @@ def test_sync_miner_driver_fails_on_fatal_generic_error(settings): job.refresh_from_db() assert job.status == OrganicJob.Status.FAILED assert [u.status for u in job_status_updates] == [JobStatus.ACCEPTED, JobStatus.FAILED] + + +@pytest.mark.django_db(databases=["default", "default_alias"], transaction=True) +@pytest.mark.parametrize( + "messages_from_miner,expected_score_str,expected_block_numbers", + [ + pytest.param([V0DeclineJobRequest], "0.000000", [], id="failed_job"), + pytest.param(SUCCESSFUL_FLOW, "3.000000", [7, 8], id="completed_job"), + ], +) +def test_sync_miner_driver_job_finished_receipt_survives_a_backwards_clock_step( + messages_from_miner, expected_score_str, expected_block_numbers, settings +): + started_at = datetime.datetime(2025, 3, 4, 12, 0, tzinfo=datetime.UTC) + finished_at = started_at - datetime.timedelta(minutes=5) + job_request = create_job_request(str(uuid.uuid4())) + with freeze_time(started_at): + driver, miner_client = _driver_for( + settings, + job_request, + messages_from_miner, + allowance_job_value=3, + allowance_blocks=[7, 8], + ) + with freeze_time(finished_at): + _run(driver) + + [receipt] = miner_client.query_sent_models(model_class=V0JobFinishedReceiptRequest) + assert receipt.payload == JobFinishedReceiptPayload( + job_uuid=job_request.uuid, + miner_hotkey="miner_hotkey", + validator_hotkey=settings.BITTENSOR_WALLET().hotkey.ss58_address, + timestamp=finished_at, + time_started=started_at, + time_took_us=0, + block_numbers=expected_block_numbers, + score_str=expected_score_str, + ) + assert JobFinishedReceipt.objects.get(job_uuid=job_request.uuid).time_took_us == 0 + + +@pytest.mark.django_db(databases=["default", "default_alias"], transaction=True) +def test_sync_miner_driver_clamps_job_accepted_receipt_ttl(settings): + accepted_at = datetime.datetime(2025, 3, 4, 12, 0, tzinfo=datetime.UTC) + job_request = create_job_request(str(uuid.uuid4()), execution_time_limit=MAX_TTL_SECONDS) + with freeze_time(accepted_at): + driver, miner_client = _driver_for(settings, job_request, SUCCESSFUL_FLOW) + _run(driver) + + [receipt] = miner_client.query_sent_models(model_class=V0JobAcceptedReceiptRequest) + assert receipt.payload == JobAcceptedReceiptPayload( + job_uuid=job_request.uuid, + miner_hotkey="miner_hotkey", + validator_hotkey=settings.BITTENSOR_WALLET().hotkey.ss58_address, + timestamp=accepted_at, + time_accepted=accepted_at, + ttl=MAX_TTL_SECONDS, + ) + assert OrganicJob.objects.get(job_uuid=job_request.uuid).status == OrganicJob.Status.COMPLETED diff --git a/validator/app/src/compute_horde_validator/validator/tests/test_other/test_eviction.py b/validator/app/src/compute_horde_validator/validator/tests/test_other/test_eviction.py new file mode 100644 index 000000000..bff4b7fea --- /dev/null +++ b/validator/app/src/compute_horde_validator/validator/tests/test_other/test_eviction.py @@ -0,0 +1,77 @@ +import datetime +import uuid + +import pytest +from compute_horde.receipts.models import JobAcceptedReceipt, JobFinishedReceipt, JobStartedReceipt +from compute_horde_core.executor_class import ExecutorClass + +from compute_horde_validator.validator.eviction import evict_malformed_receipts, evict_receipts + + +def _finished_receipt(score_str: str) -> JobFinishedReceipt: + now = datetime.datetime.now(datetime.UTC) + return JobFinishedReceipt( + job_uuid=str(uuid.uuid4()), + miner_hotkey="miner", + validator_hotkey="validator", + validator_signature="sig", + timestamp=now, + time_started=now, + time_took_us=1, + score_str=score_str, + ) + + +@pytest.mark.django_db +def test_evict_malformed_receipts_keeps_parsable_scores(): + JobFinishedReceipt.objects.bulk_create( + [ + _finished_receipt(score_str) + for score_str in ["0", "1.000000", "12.5", "nan", "inf", "abc", "-1", "1e5", "", " 1"] + ] + ) + + evict_malformed_receipts() + + assert sorted(JobFinishedReceipt.objects.values_list("score_str", flat=True)) == [ + "0", + "1.000000", + "12.5", + ] + + +def _receipts_dated( + timestamp: datetime.datetime, +) -> list[JobStartedReceipt | JobAcceptedReceipt | JobFinishedReceipt]: + common_fields = dict( + job_uuid=str(uuid.uuid4()), + miner_hotkey="miner", + validator_hotkey="validator", + validator_signature="sig", + timestamp=timestamp, + ) + return [ + JobStartedReceipt( + **common_fields, + executor_class=str(ExecutorClass.always_on__gpu_24gb), + is_organic=True, + ttl=60, + ), + JobAcceptedReceipt(**common_fields, time_accepted=timestamp, ttl=60), + JobFinishedReceipt(**common_fields, time_started=timestamp, time_took_us=1, score_str="1"), + ] + + +@pytest.mark.django_db +def test_evict_receipts_drops_rows_dated_outside_the_retention_window(): + now = datetime.datetime.now(datetime.UTC) + kept = [now - datetime.timedelta(days=1), now + datetime.timedelta(minutes=10)] + dropped = [now - datetime.timedelta(days=3), now + datetime.timedelta(days=3)] + for timestamp in kept + dropped: + for receipt in _receipts_dated(timestamp): + receipt.save() + + evict_receipts() + + for receipt_model in (JobStartedReceipt, JobAcceptedReceipt, JobFinishedReceipt): + assert sorted(receipt_model.objects.values_list("timestamp", flat=True)) == sorted(kept) diff --git a/validator/app/src/compute_horde_validator/validator/tests/test_other/test_job_excuses.py b/validator/app/src/compute_horde_validator/validator/tests/test_other/test_job_excuses.py index 006641101..3af7e45f3 100644 --- a/validator/app/src/compute_horde_validator/validator/tests/test_other/test_job_excuses.py +++ b/validator/app/src/compute_horde_validator/validator/tests/test_other/test_job_excuses.py @@ -43,7 +43,7 @@ def test_filter_valid_excuse_receipts_valid(): receipt = schemas.Receipt( payload=schemas.JobStartedReceiptPayload( - job_uuid="job-uuid-1", + job_uuid="00000000-0000-0000-0000-000000000001", miner_hotkey=miner, validator_hotkey=validator, is_organic=True, @@ -58,7 +58,7 @@ def test_filter_valid_excuse_receipts_valid(): valid_receipts = job_excuses.filter_valid_excuse_receipts( receipts_to_check=[receipt], check_time=now, - declined_job_uuid="declined", + declined_job_uuid="00000000-0000-0000-0000-00000000000b", declined_job_executor_class=executor_class, declined_job_is_synthetic=True, miner_hotkey=miner, @@ -78,7 +78,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): receipts = [ schemas.Receipt( payload=schemas.JobStartedReceiptPayload( - job_uuid="job-with-duplicated-uuid-1", + job_uuid="00000000-0000-0000-0000-000000000002", miner_hotkey=miner, validator_hotkey=validator, is_organic=True, @@ -91,7 +91,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): ), schemas.Receipt( payload=schemas.JobStartedReceiptPayload( - job_uuid="job-with-duplicated-uuid-1", + job_uuid="00000000-0000-0000-0000-000000000002", miner_hotkey=miner, validator_hotkey=validator, is_organic=True, @@ -107,7 +107,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): valid_receipts = job_excuses.filter_valid_excuse_receipts( receipts_to_check=receipts, check_time=now, - declined_job_uuid="declined", + declined_job_uuid="00000000-0000-0000-0000-00000000000b", declined_job_executor_class=executor_class, declined_job_is_synthetic=True, miner_hotkey=miner, @@ -116,7 +116,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): ) assert len(valid_receipts) == 1 - assert valid_receipts[0].payload.job_uuid == "job-with-duplicated-uuid-1" + assert valid_receipts[0].payload.job_uuid == "00000000-0000-0000-0000-000000000002" @pytest.mark.parametrize( @@ -127,7 +127,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): lambda now, miner, validator, executor_class: [ schemas.Receipt( payload=schemas.JobAcceptedReceiptPayload( - job_uuid="job-accepted-uuid-1", + job_uuid="00000000-0000-0000-0000-000000000003", miner_hotkey=miner, validator_hotkey=validator, timestamp=now - timedelta(seconds=1), @@ -144,7 +144,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): lambda now, miner, validator, executor_class: [ schemas.Receipt( payload=schemas.JobStartedReceiptPayload( - job_uuid="job-synthetic-uuid-2", + job_uuid="00000000-0000-0000-0000-000000000004", miner_hotkey=miner, validator_hotkey=validator, is_organic=False, @@ -162,7 +162,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): lambda now, miner, validator, executor_class: [ schemas.Receipt( payload=schemas.JobStartedReceiptPayload( - job_uuid="job-from-other-miner-uuid-3", + job_uuid="00000000-0000-0000-0000-000000000005", miner_hotkey="other_miner", validator_hotkey=validator, is_organic=True, @@ -180,7 +180,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): lambda now, miner, validator, executor_class: [ schemas.Receipt( payload=schemas.JobStartedReceiptPayload( - job_uuid="declined", + job_uuid="00000000-0000-0000-0000-00000000000b", miner_hotkey=miner, validator_hotkey=validator, is_organic=True, @@ -198,7 +198,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): lambda now, miner, validator, executor_class: [ schemas.Receipt( payload=schemas.JobStartedReceiptPayload( - job_uuid="job-from-other-validator-uuid-4", + job_uuid="00000000-0000-0000-0000-000000000006", miner_hotkey=miner, validator_hotkey="validator_not_allowed", is_organic=True, @@ -216,7 +216,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): lambda now, miner, validator, executor_class: [ schemas.Receipt( payload=schemas.JobStartedReceiptPayload( - job_uuid="job-from-other-executor-class-uuid-5", + job_uuid="00000000-0000-0000-0000-000000000007", miner_hotkey=miner, validator_hotkey=validator, is_organic=True, @@ -234,7 +234,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): lambda now, miner, validator, executor_class: [ schemas.Receipt( payload=schemas.JobStartedReceiptPayload( - job_uuid="job-from-future-uuid-6", + job_uuid="00000000-0000-0000-0000-000000000008", miner_hotkey=miner, validator_hotkey=validator, is_organic=True, @@ -252,7 +252,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): lambda now, miner, validator, executor_class: [ schemas.Receipt( payload=schemas.JobStartedReceiptPayload( - job_uuid="job-expired-uuid-7", + job_uuid="00000000-0000-0000-0000-000000000009", miner_hotkey=miner, validator_hotkey=validator, is_organic=True, @@ -270,7 +270,7 @@ def test_filter_valid_excuse_receipts_with_duplicate_receipts(): lambda now, miner, validator, executor_class: [ schemas.Receipt( payload=schemas.JobStartedReceiptPayload( - job_uuid="job-from-validator-with-invalid-signature-uuid-8", + job_uuid="00000000-0000-0000-0000-00000000000a", miner_hotkey=miner, validator_hotkey=validator, is_organic=True, @@ -300,7 +300,7 @@ def test_filter_valid_excuse_receipts_failures(case_name, make_receipts): res = job_excuses.filter_valid_excuse_receipts( receipts_to_check=receipts, check_time=now, - declined_job_uuid="declined", + declined_job_uuid="00000000-0000-0000-0000-00000000000b", declined_job_executor_class=executor_class, declined_job_is_synthetic=True, miner_hotkey=miner, diff --git a/validator/app/src/compute_horde_validator/validator/utils.py b/validator/app/src/compute_horde_validator/validator/utils.py index f64bfa31f..58cd788b6 100644 --- a/validator/app/src/compute_horde_validator/validator/utils.py +++ b/validator/app/src/compute_horde_validator/validator/utils.py @@ -4,7 +4,6 @@ from functools import cache MACHINE_SPEC_CHANNEL = "machine_spec_sending" -TRUSTED_MINER_FAKE_KEY = "0" * 48 def single_file_zip(filename: str, contents: str | bytes) -> str: