Fix: use UUID types for request/execution IDs

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

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(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,

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,

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

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

src/aleph_vrf/types.py

@@ -0,0 +1,8 @@
+from typing import NewType
+from uuid import UUID
+
+from pydantic import UUID4
+
+Nonce = NewType("Nonce", int)
+RequestId = NewType("RequestId", UUID4)
+ExecutionId = NewType("ExecutionId", UUID4)

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

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
@@ -23,6 +24,7 @@
     PublishedVRFResponseHash,
     PublishedVRFRandomBytes,
 )
+from aleph_vrf.types import Nonce, RequestId
 from aleph_vrf.utils import binary_to_bytes, verify
 
 
@@ -69,13 +71,13 @@ def make_post_message(
 
 @pytest.fixture
 def mock_vrf_request() -> VRFRequest:
-    request_id = "513eb52c-cb74-463a-b40e-0e2adedafb8b"
+    request_id = RequestId(UUID("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",
     )