diff --git a/src/aleph/sdk/vm/app.py b/src/aleph/sdk/vm/app.py index 88bf3eb8..b6e3fbd6 100644 --- a/src/aleph/sdk/vm/app.py +++ b/src/aleph/sdk/vm/app.py @@ -1,3 +1,5 @@ +import base64 +import socket from dataclasses import dataclass from typing import ( Any, @@ -85,3 +87,22 @@ async def send_handler_result(): def __getattr__(self, name): # Default all calls to the HTTP handler return getattr(self.http_app, name) + + @property + def vm_hash(self): + """Returns the hash of the VM that is running this app.""" + # Get hostname from environment + hostname = socket.gethostname() + + # Add padding if necessary + padding_length = len(hostname) % 8 + if padding_length != 0: + hostname += "=" * (8 - padding_length) + + # Convert the hostname back to its original binary form + item_hash_binary = base64.b32decode(hostname.upper()) + + # Convert the binary form to the original vm_hash + vm_hash = base64.b16encode(item_hash_binary).decode().lower() + + return vm_hash diff --git a/tests/unit/test_vm_app.py b/tests/unit/test_vm_app.py index c9ed5aa9..4578f20a 100644 --- a/tests/unit/test_vm_app.py +++ b/tests/unit/test_vm_app.py @@ -1,4 +1,6 @@ import asyncio +import base64 +from unittest.mock import patch import pytest from fastapi.testclient import TestClient @@ -31,3 +33,16 @@ def test_app_http(): response = client.get("/") assert response.status_code == 200 assert response.json() == {"index": "/"} + + +@patch("socket.gethostname") +def test_get_vm_hash(mock_gethostname): + vm_hash = "deadbeef" * 8 + # Uses the same logic as + # https://github.com/aleph-im/aleph-vm/blob/main/runtimes/aleph-debian-11-python/init1.py#L488 + item_hash_binary: bytes = base64.b16decode(vm_hash.encode().upper()) + hostname = base64.b32encode(item_hash_binary).decode().strip("=").lower() + + mock_gethostname.return_value = hostname + + assert app.vm_hash == vm_hash