Implement Confidential message type (#99)

* Problem: A user cannot create a Confidential VM because we don't support that kind of messages, that have a different schema from Instances and Programs.

Solution: Create a Confidential message schema with concrete fields.

* Fix: Removed old implementation and Added new one using a field instead of a new message type.

* Fix: Solve code quality issues.

* Fix: Solved test issues after last changes.

* Fix: Solved code quality issues on tests caused by different black versions.

* Fix: Solve mypy issue on CI.

* Fix: Added missing firmware field for confidentials.

* Add class TrustedExecutionEnvironment

* Add comment on retro-compatibility

* Improve tests

---------

Co-authored-by: Andres D. Molins <[email protected]>
Co-authored-by: Hugo Herter <[email protected]>

.github/workflows/tests.yml

@@ -31,7 +31,7 @@ jobs:
 
       - name: Test with ruff in the Docker image
         run: |
-          docker run aleph-message:${GITHUB_REF##*/} ruff aleph_message
+          docker run aleph-message:${GITHUB_REF##*/} ruff check aleph_message
 
       - name: Pytest in the Docker image
         run: |

aleph_message/models/execution/__init__.py

@@ -2,4 +2,8 @@
 from .instance import InstanceContent
 from .program import ProgramContent
 
-__all__ = ["BaseExecutableContent", "InstanceContent", "ProgramContent"]
+__all__ = [
+    "BaseExecutableContent",
+    "InstanceContent",
+    "ProgramContent",
+]

aleph_message/models/execution/abstract.py

@@ -1,13 +1,18 @@
 from __future__ import annotations
 
 from abc import ABC
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import Field
 
 from ..abstract import BaseContent, HashableModel
 from .base import Payment
-from .environment import FunctionEnvironment, HostRequirements, MachineResources
+from .environment import (
+    FunctionEnvironment,
+    HostRequirements,
+    InstanceEnvironment,
+    MachineResources,
+)
 from .volume import MachineVolume
 
 
@@ -22,7 +27,7 @@ class BaseExecutableContent(HashableModel, BaseContent, ABC):
     variables: Optional[Dict[str, str]] = Field(
         default=None, description="Environment variables available in the VM"
     )
-    environment: FunctionEnvironment = Field(
+    environment: Union[FunctionEnvironment, InstanceEnvironment] = Field(
         description="Properties of the execution environment"
     )
     resources: MachineResources = Field(description="System resources required")

aleph_message/models/execution/environment.py

@@ -3,10 +3,11 @@
 from enum import Enum
 from typing import List, Literal, Optional, Union
 
-from pydantic import Extra, Field
+from pydantic import Extra, Field, validator
 
 from ...utils import Mebibytes
 from ..abstract import HashableModel
+from ..item_hash import ItemHash
 
 
 class Subscription(HashableModel):
@@ -98,14 +99,72 @@ class FunctionEnvironment(HashableModel):
     internet: bool = False
     aleph_api: bool = False
     shared_cache: bool = False
-    hypervisor: Optional[HypervisorType]
+
+
+class AMDSEVPolicy(int, Enum):
+    """AMD Guest Policy for SEV-ES and SEV.
+
+    The firmware maintains a guest policy provided by the guest owner. This policy is enforced by the
+    firmware and restricts what configuration and operational commands can be performed on this
+    guest by the hypervisor. The policy also requires a minimum firmware level.
+
+    The policy comprises a set of flags that can be combined with bitwise OR.
+
+    See https://github.com/virtee/sev/blob/fbfed998930a0d1e6126462b371890b9f8d77148/src/launch/sev.rs#L245 for reference.
+    """
+
+    NO_DBG = 0b1  # Debugging of the guest is disallowed
+    NO_KS = 0b10  # Sharing keys with other guests is disallowed
+    SEV_ES = 0b100  # SEV-ES is required
+    NO_SEND = 0b1000  # Sending the guest to another platform is disallowed
+    DOMAIN = 0b10000  # The guest must not be transmitted to another platform that is not in the domain
+    SEV = 0b100000  # The guest must not be transmitted to another platform that is not SEV capable
+
+
+class TrustedExecutionEnvironment(HashableModel):
+    """Trusted Execution Environment properties."""
+
+    firmware: Optional[ItemHash] = Field(
+        default=None, description="Confidential OVMF firmware to use"
+    )
+    policy: int = Field(
+        default=AMDSEVPolicy.NO_DBG,
+        description="Policy of the TEE. Default value is 0x01 for SEV without debugging.",
+    )
+
+    class Config:
+        extra = Extra.allow
+
+
+class InstanceEnvironment(HashableModel):
+    internet: bool = False
+    aleph_api: bool = False
+    hypervisor: Optional[HypervisorType] = Field(
+        default=None, description="Hypervisor application to use. Default value is QEmu"
+    )
+    trusted_execution: Optional[TrustedExecutionEnvironment] = Field(
+        default=None,
+        description="Trusted Execution Environment properties. Defaults to no TEE.",
+    )
+    # The following fields are kept for retro-compatibility.
+    reproducible: bool = False
+    shared_cache: bool = False
+
+    @validator("trusted_execution", pre=True)
+    def check_hypervisor(cls, v, values):
+        if v and values.get("hypervisor") != HypervisorType.qemu:
+            raise ValueError("Trusted Execution Environment is only supported for QEmu")
+        return v
 
 
 class NodeRequirements(HashableModel):
     owner: Optional[str] = Field(default=None, description="Address of the node owner")
     address_regex: Optional[str] = Field(
         default=None, description="Node address must match this regular expression"
     )
+    node_hash: Optional[ItemHash] = Field(
+        default=None, description="Hash of the compute resource node that must be used"
+    )
 
     class Config:
         extra = Extra.forbid

aleph_message/models/execution/instance.py

@@ -5,6 +5,7 @@
 from aleph_message.models.abstract import HashableModel
 
 from .abstract import BaseExecutableContent
+from .environment import InstanceEnvironment
 from .volume import ParentVolume, PersistentVolumeSizeMib, VolumePersistence
 
 
@@ -25,6 +26,9 @@ class RootfsVolume(HashableModel):
 class InstanceContent(BaseExecutableContent):
     """Message content for scheduling a VM instance on the network."""
 
+    environment: InstanceEnvironment = Field(
+        description="Properties of the instance execution environment"
+    )
     rootfs: RootfsVolume = Field(
         description="Root filesystem of the system, will be booted by the kernel"
     )

aleph_message/tests/messages/instance_confidential_machine.json

@@ -0,0 +1,102 @@
+{
+    "_id": {
+        "$oid": "6080402d7f44efefd611dc1e"
+    },
+    "chain": "ETH",
+    "sender": "0x9319Ad3B7A8E0eE24f2E639c40D8eD124C5520Ba",
+    "type": "INSTANCE",
+    "channel": "Fun-dApps",
+    "confirmed": true,
+    "content": {
+        "address": "0x9319Ad3B7A8E0eE24f2E639c40D8eD124C5520Ba",
+        "allow_amend": false,
+        "variables": {
+            "VM_CUSTOM_VARIABLE": "SOMETHING",
+            "VM_CUSTOM_VARIABLE_2": "32"
+        },
+        "environment": {
+            "reproducible": true,
+            "internet": false,
+            "aleph_api": false,
+            "shared_cache": false,
+            "hypervisor": "qemu",
+            "trusted_execution": {
+                "policy": 1,
+                "firmware": "e258d248fda94c63753607f7c4494ee0fcbe92f1a76bfdac795c9d84101eb317"
+            }
+        },
+        "resources": {
+            "vcpus": 1,
+            "memory": 128,
+            "seconds": 30
+        },
+        "requirements": {
+            "cpu": {
+                "architecture": "x86_64"
+            },
+            "node": {
+                "node_hash": "4d4db19afca380fdf06ba7f916153d0f740db9de9eee23ad26ba96a90d8a2920"
+            }
+        },
+        "rootfs": {
+            "parent": {
+              "ref": "549ec451d9b099cad112d4aaa2c00ac40fb6729a92ff252ff22eef0b5c3cb613",
+              "use_latest": true
+            },
+            "persistence": "host",
+            "size_mib": 20000
+        },
+        "authorized_keys": [
+            "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIGULT6A41Msmw2KEu0R9MvUjhuWNAsbdeZ0DOwYbt4Qt user@example",
+            "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIH0jqdc5dmt75QhTrWqeHDV9xN8vxbgFyOYs2fuQl7CI",
+            "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQDRsrQV1HVrcnskNhyH0may8TG9fHCPawpAi3ZgAWU6V/R7ezvZOHnZdaFeIsOpFbPbt/l67Fur3qniSXllI2kvuh2D4BBJ9PwwlB2sgWzFDF34ADsfLQf+C/vpwrWemEEE91Tpj0dWbnf219i3mZLxy/+5Sv6kUy9YJlzWnDEUbaMAZK2CXrlK90b9Ns7mT82h6h3x9dLF/oCjBAKOSxbH2X+KgsDEZT0soxyluDqKNgKflkav+pvKFyD4J9IWM4j36r80yW+OPGsHqWoWleEhprfNb60RJPwKAYCDiBiSg6wCq5P+kS15O79Ko45wPaYDUwhRoNTcrWeadvTaCZgz9X3KDHgrX6wzdKqzQwtQeabhCaIGLFRMNl1Oy/BR8VozPbIe/mY28IN84An50UYkbve7nOGJucKc4hKxZKEVPpnVpRtIoWGwBJY2fi6C6wy2pBa8UX4C4t9NLJjNQSwFBzYOrphLu3ZW9A+267nogQHGnsJ5xnQ/MXximP3BlwM= user@example"
+        ],
+        "volumes": [
+            {
+                "comment": "Python libraries. Read-only since a 'ref' is specified.",
+                "mount": "/opt/venv",
+                "ref": "5f31b0706f59404fad3d0bff97ef89ddf24da4761608ea0646329362c662ba51",
+                "use_latest": false
+            },
+            {
+                "comment": "Ephemeral storage, read-write but will not persist after the VM stops",
+                "mount": "/var/cache",
+                "ephemeral": true,
+                "size_mib": 5
+            },
+            {
+                "comment": "Working data persisted on the VM supervisor, not available on other nodes",
+                "mount": "/var/lib/sqlite",
+                "name": "sqlite-data",
+                "persistence": "host",
+                "size_mib": 10
+            },
+            {
+                "comment": "Working data persisted on the Aleph network. New VMs will try to use the latest version of this volume, with no guarantee against conflicts",
+                "mount": "/var/lib/statistics",
+                "name": "statistics",
+                "persistence": "store",
+                "size_mib": 10
+            },
+            {
+                "comment": "Raw drive to use by a process, do not mount it",
+                "name": "raw-data",
+                "persistence": "host",
+                "size_mib": 10
+            }
+        ],
+        "replaces": "0x9319Ad3B7A8E0eE24f2E639c40D8eD124C5520Ba",
+        "time": 1619017773.8950517
+    },
+    "item_type": "inline",
+    "signature": "0x372da8230552b8c3e65c05b31a0ff3a24666d66c575f8e11019f62579bf48c2b7fe2f0bbe907a2a5bf8050989cdaf8a59ff8a1cbcafcdef0656c54279b4aa0c71b",
+    "size": 749,
+    "time": 1619017773.8950577,
+    "confirmations": [
+        {
+            "chain": "ETH",
+            "height": 12284734,
+            "hash": "0x67f2f3cde5e94e70615c92629c70d22dc959a118f46e9411b29659c2fce87cdc"
+        }
+    ]
+}

aleph_message/tests/test_models.py

@@ -26,6 +26,7 @@
     create_new_message,
     parse_message,
 )
+from aleph_message.models.execution.environment import AMDSEVPolicy
 from aleph_message.tests.download_messages import MESSAGES_STORAGE_PATH
 
 console = Console(color_system="windows")
@@ -130,7 +131,7 @@ def test_post_content():
 
 
 def test_message_machine():
-    path = Path(os.path.abspath(os.path.join(__file__, "../messages/machine.json")))
+    path = Path(__file__).parent / "messages/machine.json"
     message = create_message_from_file(path, factory=ProgramMessage)
 
     assert isinstance(message, ProgramMessage)
@@ -140,13 +141,47 @@ def test_message_machine():
 
 
 def test_instance_message_machine():
-    path = Path(
-        os.path.abspath(os.path.join(__file__, "../messages/instance_machine.json"))
-    )
+    path = Path(__file__).parent / "messages/instance_machine.json"
+    message = create_message_from_file(path, factory=InstanceMessage)
+
+    assert isinstance(message, InstanceMessage)
+    assert hash(message.content)
+
+
+def test_instance_message_machine_with_confidential_options():
+    path = Path(__file__).parent / "messages/instance_confidential_machine.json"
     message = create_message_from_file(path, factory=InstanceMessage)
 
     assert isinstance(message, InstanceMessage)
     assert hash(message.content)
+    assert message.content.environment.trusted_execution
+    assert message.content.environment.trusted_execution.policy == AMDSEVPolicy.NO_DBG
+    assert (
+        message.content.environment.trusted_execution.firmware
+        == "e258d248fda94c63753607f7c4494ee0fcbe92f1a76bfdac795c9d84101eb317"
+    )
+    assert message.content.requirements and message.content.requirements.node
+    assert (
+        message.content.requirements.node.node_hash
+        == "4d4db19afca380fdf06ba7f916153d0f740db9de9eee23ad26ba96a90d8a2920"
+    )
+
+
+def test_validation_on_confidential_options():
+    """Ensure that a trusted environment is only allowed for QEmu."""
+    path = Path(__file__).parent / "messages/instance_confidential_machine.json"
+    message_dict = json.loads(path.read_text())
+    # Patch the hypervisor to be something other than QEmu
+    message_dict["content"]["environment"]["hypervisor"] = "firecracker"
+    try:
+        _ = create_new_message(message_dict, factory=InstanceMessage)
+        raise AssertionError("An exception should have been raised before this point.")
+    except ValidationError as e:
+        assert e.errors()[0]["loc"] == ("content", "environment", "trusted_execution")
+        assert (
+            e.errors()[0]["msg"]
+            == "Trusted Execution Environment is only supported for QEmu"
+        )
 
 
 def test_message_machine_port_mapping():