From 7d21d7367fae2e0c85caa22028d4c021e19285c6 Mon Sep 17 00:00:00 2001 From: Olivier Desenfans Date: Fri, 22 Sep 2023 15:33:09 +0200 Subject: [PATCH 1/2] Internal: use NewType for ID objects Problem: ID objects (request ID, execution ID, nonce) all use raw string/int objects. This could lead to using the wrong object in the wrong place without noticing. Solution: defined dedicated types for each object type. --- src/aleph_vrf/coordinator/vrf.py | 9 ++++--- src/aleph_vrf/executor/main.py | 9 ++++--- src/aleph_vrf/models.py | 40 ++++++++++++++++-------------- src/aleph_vrf/settings.py | 8 +++--- src/aleph_vrf/types.py | 9 +++++++ src/aleph_vrf/utils.py | 8 +++--- tests/executor/test_integration.py | 7 +++--- 7 files changed, 54 insertions(+), 36 deletions(-) create mode 100644 src/aleph_vrf/types.py diff --git a/src/aleph_vrf/coordinator/vrf.py b/src/aleph_vrf/coordinator/vrf.py index 70080c8..016a295 100644 --- a/src/aleph_vrf/coordinator/vrf.py +++ b/src/aleph_vrf/coordinator/vrf.py @@ -24,6 +24,7 @@ PublishedVRFRandomBytes, ) from aleph_vrf.settings import settings +from aleph_vrf.types import RequestId, Nonce from aleph_vrf.utils import ( binary_to_bytes, bytes_to_int, @@ -132,7 +133,7 @@ async def generate_vrf(account: ETHAccount) -> VRFResponse: nb_executors=nb_executors, nonce=nonce, vrf_function=ItemHash(settings.FUNCTION), - request_id=str(uuid4()), + request_id=RequestId(str(uuid4())), node_list_hash=sha3_256(selected_node_list).hexdigest(), ) @@ -179,8 +180,8 @@ async def generate_vrf(account: ETHAccount) -> VRFResponse: async def send_generate_requests( selected_nodes: List[Node], - request_item_hash: str, - request_id: str, + request_item_hash: ItemHash, + request_id: RequestId, ) -> Dict[str, PublishedVRFResponseHash]: generate_tasks = [] nodes: List[str] = [] @@ -233,7 +234,7 @@ async def send_publish_requests( def generate_final_vrf( nb_executors: int, - nonce: int, + nonce: Nonce, vrf_generated_result: Dict[str, PublishedVRFResponseHash], vrf_publish_result: Dict[str, PublishedVRFRandomBytes], vrf_request: VRFRequest, diff --git a/src/aleph_vrf/executor/main.py b/src/aleph_vrf/executor/main.py index aca5559..5995f63 100644 --- a/src/aleph_vrf/executor/main.py +++ b/src/aleph_vrf/executor/main.py @@ -5,6 +5,7 @@ from aleph.sdk.exceptions import MessageNotFoundError, MultipleMessagesError from aleph_vrf.settings import settings +from aleph_vrf.types import ExecutionId, RequestId logger = logging.getLogger(__name__) @@ -36,8 +37,8 @@ GENERATE_MESSAGE_REF_PATH = "hash" # TODO: Use another method to save the data -ANSWERED_REQUESTS: Set[str] = set() -SAVED_GENERATED_BYTES: Dict[str, bytes] = {} +ANSWERED_REQUESTS: Set[RequestId] = set() +SAVED_GENERATED_BYTES: Dict[ExecutionId, bytes] = {} http_app = FastAPI() app = AlephApp(http_app=http_app) @@ -91,7 +92,7 @@ async def receive_generate( generated_bytes, hashed_bytes = generate( generation_request.nb_bytes, generation_request.nonce ) - SAVED_GENERATED_BYTES[str(generation_request.execution_id)] = generated_bytes + SAVED_GENERATED_BYTES[generation_request.execution_id] = generated_bytes ANSWERED_REQUESTS.add(generation_request.request_id) response_hash = VRFResponseHash( @@ -137,7 +138,7 @@ async def receive_publish( status_code=404, detail="The random number has already been published" ) - random_bytes: bytes = SAVED_GENERATED_BYTES.pop(str(response_hash.execution_id)) + random_bytes: bytes = SAVED_GENERATED_BYTES.pop(response_hash.execution_id) response_bytes = VRFRandomBytes( request_id=response_hash.request_id, diff --git a/src/aleph_vrf/models.py b/src/aleph_vrf/models.py index 4334223..ce877bc 100644 --- a/src/aleph_vrf/models.py +++ b/src/aleph_vrf/models.py @@ -7,6 +7,8 @@ from pydantic import BaseModel, ValidationError, Field from pydantic.generics import GenericModel +from aleph_vrf.types import Nonce, RequestId, ExecutionId + class Node(BaseModel): hash: str @@ -17,17 +19,17 @@ class Node(BaseModel): class VRFRequest(BaseModel): nb_bytes: int nb_executors: int - nonce: int + nonce: Nonce vrf_function: ItemHash - request_id: str + request_id: RequestId node_list_hash: str class VRFGenerationRequest(BaseModel): nb_bytes: int - nonce: int - request_id: str - execution_id: str = Field(default_factory=lambda: str(uuid4())) + nonce: Nonce + request_id: RequestId + execution_id: ExecutionId = Field(default_factory=lambda: str(uuid4())) vrf_function: ItemHash @@ -44,9 +46,9 @@ def generate_request_from_message(message: PostMessage) -> VRFGenerationRequest: class VRFResponseHash(BaseModel): nb_bytes: int - nonce: int - request_id: str - execution_id: str + nonce: Nonce + request_id: RequestId + execution_id: ExecutionId vrf_request: ItemHash random_bytes_hash: str @@ -61,7 +63,7 @@ class PublishedVRFResponseHash(VRFResponseHash): @classmethod def from_vrf_response_hash( - cls, vrf_response_hash: VRFResponseHash, message_hash: ItemHash + cls, vrf_response_hash: VRFResponseHash, message_hash: ItemHash ) -> "PublishedVRFResponseHash": return cls( nb_bytes=vrf_response_hash.nb_bytes, @@ -75,7 +77,7 @@ def from_vrf_response_hash( def generate_response_hash_from_message( - message: PostMessage, + message: PostMessage, ) -> PublishedVRFResponseHash: content = message.content.content try: @@ -92,8 +94,8 @@ def generate_response_hash_from_message( class VRFRandomBytes(BaseModel): - request_id: str - execution_id: str + request_id: RequestId + execution_id: ExecutionId vrf_request: ItemHash random_bytes: str random_bytes_hash: str @@ -105,7 +107,7 @@ class PublishedVRFRandomBytes(VRFRandomBytes): @classmethod def from_vrf_random_bytes( - cls, vrf_random_bytes: VRFRandomBytes, message_hash: ItemHash + cls, vrf_random_bytes: VRFRandomBytes, message_hash: ItemHash ) -> "PublishedVRFRandomBytes": return cls( request_id=vrf_random_bytes.request_id, @@ -120,23 +122,23 @@ def from_vrf_random_bytes( class CRNVRFResponse(BaseModel): url: str - execution_id: str + execution_id: ExecutionId random_number: str random_bytes: str random_bytes_hash: str - generation_message_hash: str - publish_message_hash: str + generation_message_hash: ItemHash + publish_message_hash: ItemHash class VRFResponse(BaseModel): nb_bytes: int nb_executors: int - nonce: int + nonce: Nonce vrf_function: ItemHash - request_id: str + request_id: RequestId nodes: List[CRNVRFResponse] random_number: str - message_hash: Optional[str] = None + message_hash: Optional[ItemHash] = None M = TypeVar("M", bound=BaseModel) diff --git a/src/aleph_vrf/settings.py b/src/aleph_vrf/settings.py index b563d11..20a20c4 100644 --- a/src/aleph_vrf/settings.py +++ b/src/aleph_vrf/settings.py @@ -1,8 +1,8 @@ -from pydantic import BaseSettings, Field +from pydantic import BaseSettings, Field, HttpUrl class Settings(BaseSettings): - API_HOST: str = Field( + API_HOST: HttpUrl = Field( default="https://api2.aleph.im", description="URL of the reference aleph.im Core Channel Node.", ) @@ -17,7 +17,9 @@ class Settings(BaseSettings): default="4992b4127d296b240bbb73058daea9bca09f717fa94767d6f4dc3ef53b4ef5ce", description="VRF function to use.", ) - NB_EXECUTORS: int = Field(default=32, description="Number of executors to use.") + NB_EXECUTORS: int = Field( + default=32, description="Number of executors to use." + ) NB_BYTES: int = Field( default=32, description="Number of bytes of the generated random number." ) diff --git a/src/aleph_vrf/types.py b/src/aleph_vrf/types.py new file mode 100644 index 0000000..6ad0dc8 --- /dev/null +++ b/src/aleph_vrf/types.py @@ -0,0 +1,9 @@ +from typing import NewType +from uuid import UUID + +from pydantic import UUID4 + +Nonce = NewType("Nonce", int) +# TODO: use UUID once https://github.com/aleph-im/aleph-sdk-python/pull/61 is merged +RequestId = NewType("RequestId", str) +ExecutionId = NewType("ExecutionId", str) diff --git a/src/aleph_vrf/utils.py b/src/aleph_vrf/utils.py index 49438dd..7771385 100644 --- a/src/aleph_vrf/utils.py +++ b/src/aleph_vrf/utils.py @@ -4,6 +4,8 @@ from utilitybelt import dev_urandom_entropy +from aleph_vrf.types import Nonce + def xor_all(x: List[bytes]) -> bytes: """XORs all the bytes in the list together.""" @@ -33,12 +35,12 @@ def binary_to_bytes(s: str): return int(s, 2).to_bytes((len(s) + 7) // 8, byteorder="big") -def generate_nonce() -> int: +def generate_nonce() -> Nonce: """Generates pseudo-random nonce number.""" - return randint(0, 100000000) + return Nonce(randint(0, 100000000)) -def generate(n: int, nonce: int) -> Tuple[bytes, str]: +def generate(n: int, nonce: Nonce) -> Tuple[bytes, str]: """Generates a number of random bytes and hashes them with the nonce.""" random_bytes: bytes = dev_urandom_entropy(n) random_hash = sha3_256(random_bytes + int_to_bytes(nonce)).hexdigest() diff --git a/tests/executor/test_integration.py b/tests/executor/test_integration.py index 0d0b60d..53ad5bb 100644 --- a/tests/executor/test_integration.py +++ b/tests/executor/test_integration.py @@ -23,6 +23,7 @@ PublishedVRFResponseHash, PublishedVRFRandomBytes, ) +from aleph_vrf.types import Nonce, RequestId from aleph_vrf.utils import binary_to_bytes, verify @@ -69,13 +70,13 @@ def make_post_message( @pytest.fixture def mock_vrf_request() -> VRFRequest: - request_id = "513eb52c-cb74-463a-b40e-0e2adedafb8b" + request_id = RequestId("513eb52c-cb74-463a-b40e-0e2adedafb8b") vrf_request = VRFRequest( nb_bytes=32, nb_executors=4, - nonce=42, - vrf_function="deca" * 16, + nonce=Nonce(42), + vrf_function=ItemHash("deca" * 16), request_id=request_id, node_list_hash="1234", ) From 57b6ffe0edbe08d428759ea9a6f8bcee14f7921a Mon Sep 17 00:00:00 2001 From: Olivier Desenfans Date: Fri, 22 Sep 2023 18:02:02 +0200 Subject: [PATCH 2/2] Fix: use UUID types for request/execution IDs Problem: older versions of aleph-sdk-python do not support serializing UUID objects. Solution: use the latest version of aleph-sdk-python. --- setup.cfg | 2 +- src/aleph_vrf/coordinator/vrf.py | 2 +- src/aleph_vrf/models.py | 2 +- src/aleph_vrf/types.py | 5 ++--- tests/executor/test_integration.py | 3 ++- 5 files changed, 7 insertions(+), 7 deletions(-) diff --git a/setup.cfg b/setup.cfg index 5498224..d993155 100644 --- a/setup.cfg +++ b/setup.cfg @@ -49,7 +49,7 @@ package_dir = # For more information, check out https://semver.org/. install_requires = aiohttp - aleph-sdk-python==0.7.0 + aleph-sdk-python@git+https://github.com/aleph-im/aleph-sdk-python@28a10519a5dce4d1f25fb5d9d08c077eb60690e5 fastapi importlib-metadata; python_version<"3.8" utilitybelt diff --git a/src/aleph_vrf/coordinator/vrf.py b/src/aleph_vrf/coordinator/vrf.py index 016a295..e8d58bc 100644 --- a/src/aleph_vrf/coordinator/vrf.py +++ b/src/aleph_vrf/coordinator/vrf.py @@ -133,7 +133,7 @@ async def generate_vrf(account: ETHAccount) -> VRFResponse: nb_executors=nb_executors, nonce=nonce, vrf_function=ItemHash(settings.FUNCTION), - request_id=RequestId(str(uuid4())), + request_id=RequestId(uuid4()), node_list_hash=sha3_256(selected_node_list).hexdigest(), ) diff --git a/src/aleph_vrf/models.py b/src/aleph_vrf/models.py index ce877bc..f112a75 100644 --- a/src/aleph_vrf/models.py +++ b/src/aleph_vrf/models.py @@ -29,7 +29,7 @@ class VRFGenerationRequest(BaseModel): nb_bytes: int nonce: Nonce request_id: RequestId - execution_id: ExecutionId = Field(default_factory=lambda: str(uuid4())) + execution_id: ExecutionId = Field(default_factory=uuid4) vrf_function: ItemHash diff --git a/src/aleph_vrf/types.py b/src/aleph_vrf/types.py index 6ad0dc8..f2bfeb4 100644 --- a/src/aleph_vrf/types.py +++ b/src/aleph_vrf/types.py @@ -4,6 +4,5 @@ from pydantic import UUID4 Nonce = NewType("Nonce", int) -# TODO: use UUID once https://github.com/aleph-im/aleph-sdk-python/pull/61 is merged -RequestId = NewType("RequestId", str) -ExecutionId = NewType("ExecutionId", str) +RequestId = NewType("RequestId", UUID4) +ExecutionId = NewType("ExecutionId", UUID4) diff --git a/tests/executor/test_integration.py b/tests/executor/test_integration.py index 53ad5bb..b8a783b 100644 --- a/tests/executor/test_integration.py +++ b/tests/executor/test_integration.py @@ -1,6 +1,7 @@ import datetime as dt from hashlib import sha256 from typing import Dict, Any, Union, Tuple +from uuid import UUID import aiohttp import pytest @@ -70,7 +71,7 @@ def make_post_message( @pytest.fixture def mock_vrf_request() -> VRFRequest: - request_id = RequestId("513eb52c-cb74-463a-b40e-0e2adedafb8b") + request_id = RequestId(UUID("513eb52c-cb74-463a-b40e-0e2adedafb8b")) vrf_request = VRFRequest( nb_bytes=32,