Add caps to free-form executable content fields

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",
     )
 

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"
     )

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

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):

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),
+        )