Fix: Operator views were incomplete

This uses pydantic classes in control endpoints.

The reboot functionnality of VMs is now implemented.

Authentication schemes are now fully used and functional.

src/aleph/vm/orchestrator/supervisor.py

@@ -32,7 +32,13 @@
     status_public_config,
     update_allocations,
 )
-from .views.operator import operate_erase, operate_expire, operate_stop, stream_logs
+from .views.operator import (
+    operate_erase,
+    operate_expire,
+    operate_reboot,
+    operate_stop,
+    stream_logs,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -52,6 +58,20 @@ async def server_version_middleware(
 
 app = web.Application(middlewares=[server_version_middleware])
 
+
+async def allow_cors_on_endpoint(request: web.Request):
+    """Allow CORS on endpoints that VM owners use to control their machine."""
+    return web.Response(
+        status=200,
+        headers={
+            "Access-Control-Allow-Headers": "*",
+            "Access-Control-Allow-Methods": "*",
+            "Access-Control-Allow-Origin": "*",
+            "Allow": "POST",
+        },
+    )
+
+
 app.add_routes(
     [
         web.get("/about/login", about_login),
@@ -64,6 +84,11 @@ async def server_version_middleware(
         web.post("/control/machine/{ref}/expire", operate_expire),
         web.post("/control/machine/{ref}/stop", operate_stop),
         web.post("/control/machine/{ref}/erase", operate_erase),
+        web.post("/control/machine/{ref}/reboot", operate_reboot),
+        web.options(
+            "/control/machine/{ref}/{view:.*}",
+            allow_cors_on_endpoint,
+        ),
         web.get("/status/check/fastapi", status_check_fastapi),
         web.get("/status/check/version", status_check_version),
         web.get("/status/config", status_public_config),

src/aleph/vm/orchestrator/views/authentication.py

@@ -0,0 +1,233 @@
+import functools
+import json
+import logging
+from datetime import datetime, timedelta, timezone
+from typing import Any, Awaitable, Callable, Coroutine, Literal, Union
+
+import pydantic
+from aiohttp import web
+from eth_account import Account
+from eth_account.messages import encode_defunct
+from jwskate import Jwk
+from pydantic import BaseModel, ValidationError, root_validator, validator
+
+from aleph.vm.conf import settings
+
+logger = logging.getLogger(__name__)
+
+
+def is_token_still_valid(timestamp):
+    """
+    Checks if a token has exprired based on its timestamp
+    """
+    current_datetime = datetime.now(tz=timezone.utc)
+    target_datetime = datetime.fromisoformat(timestamp)
+
+    return target_datetime > current_datetime
+
+
+def verify_wallet_signature(signature, message, address):
+    """
+    Verifies a signature issued by a wallet
+    """
+    enc_msg = encode_defunct(hexstr=message)
+    computed_address = Account.recover_message(enc_msg, signature=signature)
+    return computed_address.lower() == address.lower()
+
+
+class SignedPubKeyPayload(BaseModel):
+    """This payload is signed by the wallet of the user to authorize an ephemeral key to act on his behalf."""
+
+    pubkey: dict[str, Any]
+    # {'pubkey': {'alg': 'ES256', 'crv': 'P-256', 'ext': True, 'key_ops': ['verify'], 'kty': 'EC',
+    #  'x': '4blJBYpltvQLFgRvLE-2H7dsMr5O0ImHkgOnjUbG2AU', 'y': '5VHnq_hUSogZBbVgsXMs0CjrVfMy4Pa3Uv2BEBqfrN4'}
+    # alg: Literal["ECDSA"]
+    domain: str
+    address: str
+    expires: str
+
+    @property
+    def json_web_key(self) -> Jwk:
+        """Return the ephemeral public key as Json Web Key"""
+        return Jwk(self.pubkey)
+
+
+class SignedPubKeyHeader(BaseModel):
+    signature: bytes
+    payload: bytes
+
+    @validator("signature")
+    def signature_must_be_hex(cls, v: bytes) -> bytes:
+        """Convert the signature from hexadecimal to bytes"""
+        return bytes.fromhex(v.removeprefix(b"0x").decode())
+
+    @validator("payload")
+    def payload_must_be_hex(cls, v: bytes) -> bytes:
+        """Convert the payload from hexadecimal to bytes"""
+        return bytes.fromhex(v.decode())
+
+    @root_validator(pre=False, skip_on_failure=True)
+    def check_expiry(cls, values):
+        """Check that the token has not expired"""
+        payload: bytes = values["payload"]
+        content = SignedPubKeyPayload.parse_raw(payload)
+        if not is_token_still_valid(content.expires):
+            msg = "Token expired"
+            raise ValueError(msg)
+        return values
+
+    @root_validator(pre=False, skip_on_failure=True)
+    def check_signature(cls, values):
+        """Check that the signature is valid"""
+        signature: bytes = values["signature"]
+        payload: bytes = values["payload"]
+        content = SignedPubKeyPayload.parse_raw(payload)
+        if not verify_wallet_signature(signature, payload.hex(), content.address):
+            msg = "Invalid signature"
+            raise ValueError(msg)
+        return values
+
+    @property
+    def content(self) -> SignedPubKeyPayload:
+        """Return the content of the header"""
+        return SignedPubKeyPayload.parse_raw(self.payload)
+
+
+class SignedOperationPayload(BaseModel):
+    time: datetime
+    method: Union[Literal["POST"], Literal["GET"]]
+    path: str
+    # body_sha256: str  # disabled since there is no body
+
+    @validator("time")
+    def time_is_current(cls, v: datetime) -> datetime:
+        """Check that the time is current and the payload is not a replay attack."""
+        max_past = datetime.now(tz=timezone.utc) - timedelta(minutes=2)
+        max_future = datetime.now(tz=timezone.utc) + timedelta(minutes=2)
+        if v < max_past:
+            raise ValueError("Time is too far in the past")
+        if v > max_future:
+            raise ValueError("Time is too far in the future")
+        return v
+
+
+class SignedOperation(BaseModel):
+    """This payload is signed by the ephemeral key authorized above."""
+
+    signature: bytes
+    payload: bytes
+
+    @validator("signature")
+    def signature_must_be_hex(cls, v) -> bytes:
+        """Convert the signature from hexadecimal to bytes"""
+        try:
+            return bytes.fromhex(v.removeprefix(b"0x").decode())
+        except pydantic.ValidationError as error:
+            print(v)
+            logger.warning(v)
+            raise error
+
+    @validator("payload")
+    def payload_must_be_hex(cls, v) -> bytes:
+        """Convert the payload from hexadecimal to bytes"""
+        v = bytes.fromhex(v.decode())
+        _ = SignedOperationPayload.parse_raw(v)
+        return v
+
+    @property
+    def content(self) -> SignedOperationPayload:
+        """Return the content of the header"""
+        return SignedOperationPayload.parse_raw(self.payload)
+
+
+def get_signed_pubkey(request: web.Request) -> SignedPubKeyHeader:
+    """Get the ephemeral public key that is signed by the wallet from the request headers."""
+    signed_pubkey_header = request.headers.get("X-SignedPubKey")
+    if not signed_pubkey_header:
+        raise web.HTTPBadRequest(reason="Missing X-SignedPubKey header")
+
+    try:
+        return SignedPubKeyHeader.parse_raw(signed_pubkey_header)
+    except KeyError as error:
+        logger.debug(f"Missing X-SignedPubKey header: {error}")
+        raise web.HTTPBadRequest(reason="Invalid X-SignedPubKey fields") from error
+    except json.JSONDecodeError as error:
+        raise web.HTTPBadRequest(reason="Invalid X-SignedPubKey format") from error
+    except ValueError as error:
+        if error.args == ("Token expired",):
+            raise web.HTTPUnauthorized(reason="Token expired") from error
+        elif error.args == ("Invalid signature",):
+            raise web.HTTPUnauthorized(reason="Invalid signature") from error
+        else:
+            raise error
+
+
+def get_signed_operation(request: web.Request) -> SignedOperation:
+    """Get the signed operation public key that is signed by the ephemeral key from the request headers."""
+    try:
+        signed_operation = request.headers["X-SignedOperation"]
+        return SignedOperation.parse_raw(signed_operation)
+    except KeyError as error:
+        raise web.HTTPBadRequest(reason="Missing X-SignedOperation header") from error
+    except json.JSONDecodeError as error:
+        raise web.HTTPBadRequest(reason="Invalid X-SignedOperation format") from error
+    except ValidationError as error:
+        logger.debug(f"Invalid X-SignedOperation fields: {error}")
+        raise web.HTTPBadRequest(reason="Invalid X-SignedOperation fields") from error
+
+
+def verify_signed_operation(signed_operation: SignedOperation, signed_pubkey: SignedPubKeyHeader) -> str:
+    """Verify that the operation is signed by the ephemeral key authorized by the wallet."""
+    if signed_pubkey.content.json_web_key.verify(
+        data=signed_operation.payload,
+        signature=signed_operation.signature,
+        alg="ES256",
+    ):
+        logger.debug("Signature verified")
+        return signed_pubkey.content.address
+    else:
+        raise web.HTTPUnauthorized(reason="Signature could not verified")
+
+
+async def authenticate_jwk(request: web.Request) -> str:
+    """Authenticate a request using the X-SignedPubKey and X-SignedOperation headers."""
+    signed_pubkey = get_signed_pubkey(request)
+    signed_operation = get_signed_operation(request)
+    if signed_pubkey.content.domain != settings.DOMAIN_NAME:
+        logger.debug(f"Invalid domain '{signed_pubkey.content.domain}' != '{settings.DOMAIN_NAME}'")
+        raise web.HTTPUnauthorized(reason="Invalid domain")
+    if signed_operation.content.path != request.path:
+        logger.debug(f"Invalid path '{signed_operation.content.path}' != '{request.path}'")
+        raise web.HTTPUnauthorized(reason="Invalid path")
+    if signed_operation.content.method != request.method:
+        logger.debug(f"Invalid method '{signed_operation.content.method}' != '{request.method}'")
+        raise web.HTTPUnauthorized(reason="Invalid method")
+    return verify_signed_operation(signed_operation, signed_pubkey)
+
+
+async def authenicate_websocket_message(message) -> str:
+    """Authenticate a websocket message since JS cannot configure headers on WebSockets."""
+    signed_pubkey = SignedPubKeyHeader.parse_obj(message["X-SignedPubKey"])
+    signed_operation = SignedOperation.parse_obj(message["X-SignedOperation"])
+    if signed_pubkey.content.domain != settings.DOMAIN_NAME:
+        logger.debug(f"Invalid domain '{signed_pubkey.content.domain}' != '{settings.DOMAIN_NAME}'")
+        raise web.HTTPUnauthorized(reason="Invalid domain")
+    return verify_signed_operation(signed_operation, signed_pubkey)
+
+
+def require_jwk_authentication(
+    handler: Callable[[web.Request, str], Coroutine[Any, Any, web.StreamResponse]]
+) -> Callable[[web.Response], Awaitable[web.StreamResponse]]:
+    @functools.wraps(handler)
+    async def wrapper(request):
+        try:
+            authenticated_sender: str = await authenticate_jwk(request)
+        except web.HTTPException as e:
+            return web.json_response(data={"error": e.reason}, status=e.status)
+
+        response = await handler(request, authenticated_sender)
+        # Allow browser clients to access the body of the response
+        response.headers.update({"Access-Control-Allow-Origin": request.headers.get("Origin", "")})
+        return response
+
+    return wrapper