diff --git a/aleph_message/models/execution/abstract.py b/aleph_message/models/execution/abstract.py index 7707e69..d335ad2 100644 --- a/aleph_message/models/execution/abstract.py +++ b/aleph_message/models/execution/abstract.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import ABC -from typing import Any, Dict, List, Optional, Sequence, Union +from typing import Annotated, Any, Dict, List, Optional, Sequence, Union from pydantic import Field @@ -16,17 +16,38 @@ ) from .volume import MachineVolume +MAX_METADATA_ENTRIES = 256 +MAX_AUTHORIZED_KEYS = 256 +MAX_AUTHORIZED_KEY_LENGTH = 8192 +MAX_VARIABLE_ENTRIES = 256 +MAX_VARIABLE_KEY_LENGTH = 128 +MAX_VARIABLE_VALUE_LENGTH = 4096 +MAX_VOLUMES = 256 +MAX_REPLACES_LENGTH = 128 + +VariableKey = Annotated[str, Field(max_length=MAX_VARIABLE_KEY_LENGTH)] +VariableValue = Annotated[str, Field(max_length=MAX_VARIABLE_VALUE_LENGTH)] +AuthorizedKey = Annotated[str, Field(max_length=MAX_AUTHORIZED_KEY_LENGTH)] + class BaseExecutableContent(HashableModel, BaseContent, ABC): """Abstract content for execution messages (Instances, Programs).""" allow_amend: bool = Field(description="Allow amends to update this function") - metadata: Optional[Dict[str, Any]] = Field(description="Metadata of the VM") - authorized_keys: Optional[List[str]] = Field( + metadata: Optional[Dict[str, Any]] = Field( + default=None, + max_length=MAX_METADATA_ENTRIES, + description="Metadata of the VM", + ) + authorized_keys: Optional[List[AuthorizedKey]] = Field( + default=None, + max_length=MAX_AUTHORIZED_KEYS, description="SSH public keys authorized to connect to the VM", ) - variables: Optional[Dict[str, str]] = Field( - default=None, description="Environment variables available in the VM" + variables: Optional[Dict[VariableKey, VariableValue]] = Field( + default=None, + max_length=MAX_VARIABLE_ENTRIES, + description="Environment variables available in the VM", ) environment: Union[FunctionEnvironment, InstanceEnvironment] = Field( description="Properties of the execution environment" @@ -37,10 +58,13 @@ class BaseExecutableContent(HashableModel, BaseContent, ABC): default=None, description="System properties required" ) volumes: List[MachineVolume] = Field( - default=[], description="Volumes to mount on the filesystem" + default=[], + max_length=MAX_VOLUMES, + description="Volumes to mount on the filesystem", ) replaces: Optional[str] = Field( default=None, + max_length=MAX_REPLACES_LENGTH, description="Previous version to replace. Must be signed by the same address", ) diff --git a/aleph_message/models/execution/instance.py b/aleph_message/models/execution/instance.py index a2affc0..4620cb2 100644 --- a/aleph_message/models/execution/instance.py +++ b/aleph_message/models/execution/instance.py @@ -31,11 +31,7 @@ class RootfsVolume(HashableModel): class InstanceContent(BaseExecutableContent): """Message content for scheduling a VM instance on the network.""" - metadata: Optional[dict] = None payment: Optional[Payment] = None - authorized_keys: Optional[List[str]] = Field( - default=None, description="List of authorized SSH keys" - ) environment: InstanceEnvironment = Field( description="Properties of the instance execution environment" ) diff --git a/aleph_message/models/execution/program.py b/aleph_message/models/execution/program.py index 07ec2e7..3bb11a0 100644 --- a/aleph_message/models/execution/program.py +++ b/aleph_message/models/execution/program.py @@ -73,6 +73,4 @@ class ProgramContent(BaseExecutableContent): ) on: FunctionTriggers = Field(description="Signals that trigger an execution") - metadata: Optional[dict] = None - authorized_keys: Optional[List[str]] = None payment: Optional[Payment] = None diff --git a/aleph_message/models/execution/volume.py b/aleph_message/models/execution/volume.py index 9f0702d..e70043a 100644 --- a/aleph_message/models/execution/volume.py +++ b/aleph_message/models/execution/volume.py @@ -10,10 +10,12 @@ from ..abstract import HashableModel from ..item_hash import ItemHash +MAX_VOLUME_LABEL_LENGTH = 256 + class AbstractVolume(HashableModel, ABC): - comment: Optional[str] = None - mount: Optional[str] = None + comment: Optional[str] = Field(default=None, max_length=MAX_VOLUME_LABEL_LENGTH) + mount: Optional[str] = Field(default=None, max_length=MAX_VOLUME_LABEL_LENGTH) @abstractmethod def is_read_only(self): ... @@ -75,7 +77,7 @@ class VolumePersistence(str, Enum): class PersistentVolume(AbstractVolume): parent: Optional[ParentVolume] = None persistence: Optional[VolumePersistence] = None - name: Optional[str] = None + name: Optional[str] = Field(default=None, max_length=MAX_VOLUME_LABEL_LENGTH) size_mib: PersistentVolumeSizeMib def is_read_only(self): diff --git a/aleph_message/tests/test_models.py b/aleph_message/tests/test_models.py index 1a4ad47..347334f 100644 --- a/aleph_message/tests/test_models.py +++ b/aleph_message/tests/test_models.py @@ -32,12 +32,26 @@ create_new_message, parse_message, ) +from aleph_message.models.execution.abstract import ( + MAX_AUTHORIZED_KEY_LENGTH, + MAX_AUTHORIZED_KEYS, + MAX_METADATA_ENTRIES, + MAX_REPLACES_LENGTH, + MAX_VARIABLE_ENTRIES, + MAX_VARIABLE_KEY_LENGTH, + MAX_VARIABLE_VALUE_LENGTH, + MAX_VOLUMES, +) from aleph_message.models.execution.environment import ( MAX_ADDRESS_REGEX_LENGTH, AMDSEVPolicy, HypervisorType, NodeRequirements, ) +from aleph_message.models.execution.volume import ( + MAX_VOLUME_LABEL_LENGTH, + EphemeralVolume, +) from aleph_message.tests.download_messages import MESSAGES_STORAGE_PATH console = Console(color_system="windows") @@ -691,3 +705,114 @@ def test_address_regex_length_boundary(): def test_address_regex_none_passes(): assert NodeRequirements(address_regex=None).address_regex is None + + +def _load_instance_fixture() -> dict: + path = Path(__file__).parent / "messages/instance_machine.json" + return json.loads(path.read_text()) + + +def test_metadata_entries_boundary(): + message_dict = _load_instance_fixture() + message_dict["content"]["metadata"] = { + str(i): i for i in range(MAX_METADATA_ENTRIES) + } + create_new_message(message_dict, factory=InstanceMessage) + + message_dict["content"]["metadata"] = { + str(i): i for i in range(MAX_METADATA_ENTRIES + 1) + } + with pytest.raises(ValidationError): + create_new_message(message_dict, factory=InstanceMessage) + + +def test_authorized_keys_count_boundary(): + message_dict = _load_instance_fixture() + message_dict["content"]["authorized_keys"] = ["key"] * MAX_AUTHORIZED_KEYS + create_new_message(message_dict, factory=InstanceMessage) + + message_dict["content"]["authorized_keys"] = ["key"] * (MAX_AUTHORIZED_KEYS + 1) + with pytest.raises(ValidationError): + create_new_message(message_dict, factory=InstanceMessage) + + +def test_authorized_key_length_boundary(): + message_dict = _load_instance_fixture() + message_dict["content"]["authorized_keys"] = ["a" * MAX_AUTHORIZED_KEY_LENGTH] + create_new_message(message_dict, factory=InstanceMessage) + + message_dict["content"]["authorized_keys"] = ["a" * (MAX_AUTHORIZED_KEY_LENGTH + 1)] + with pytest.raises(ValidationError): + create_new_message(message_dict, factory=InstanceMessage) + + +def test_variables_entries_boundary(): + message_dict = _load_instance_fixture() + message_dict["content"]["variables"] = { + str(i): str(i) for i in range(MAX_VARIABLE_ENTRIES) + } + create_new_message(message_dict, factory=InstanceMessage) + + message_dict["content"]["variables"] = { + str(i): str(i) for i in range(MAX_VARIABLE_ENTRIES + 1) + } + with pytest.raises(ValidationError): + create_new_message(message_dict, factory=InstanceMessage) + + +def test_variable_key_length_boundary(): + message_dict = _load_instance_fixture() + message_dict["content"]["variables"] = {"a" * MAX_VARIABLE_KEY_LENGTH: "v"} + create_new_message(message_dict, factory=InstanceMessage) + + message_dict["content"]["variables"] = {"a" * (MAX_VARIABLE_KEY_LENGTH + 1): "v"} + with pytest.raises(ValidationError): + create_new_message(message_dict, factory=InstanceMessage) + + +def test_variable_value_length_boundary(): + message_dict = _load_instance_fixture() + message_dict["content"]["variables"] = {"k": "v" * MAX_VARIABLE_VALUE_LENGTH} + create_new_message(message_dict, factory=InstanceMessage) + + message_dict["content"]["variables"] = {"k": "v" * (MAX_VARIABLE_VALUE_LENGTH + 1)} + with pytest.raises(ValidationError): + create_new_message(message_dict, factory=InstanceMessage) + + +def test_volumes_count_boundary(): + message_dict = _load_instance_fixture() + volume = {"ephemeral": True, "mount": "/tmp/a", "size_mib": 1} + message_dict["content"]["volumes"] = [volume] * MAX_VOLUMES + create_new_message(message_dict, factory=InstanceMessage) + + message_dict["content"]["volumes"] = [volume] * (MAX_VOLUMES + 1) + with pytest.raises(ValidationError): + create_new_message(message_dict, factory=InstanceMessage) + + +def test_replaces_length_boundary(): + message_dict = _load_instance_fixture() + message_dict["content"]["replaces"] = "x" * MAX_REPLACES_LENGTH + create_new_message(message_dict, factory=InstanceMessage) + + message_dict["content"]["replaces"] = "x" * (MAX_REPLACES_LENGTH + 1) + with pytest.raises(ValidationError): + create_new_message(message_dict, factory=InstanceMessage) + + +def test_volume_label_length_boundary(): + # comment/mount/name share MAX_VOLUME_LABEL_LENGTH; exercising comment is enough. + EphemeralVolume( + ephemeral=True, + mount="/tmp/a", + size_mib=1, + comment="c" * MAX_VOLUME_LABEL_LENGTH, + ) + with pytest.raises(ValidationError): + EphemeralVolume( + ephemeral=True, + mount="/tmp/a", + size_mib=1, + comment="c" * (MAX_VOLUME_LABEL_LENGTH + 1), + )