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 70080c8..e8d58bc 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(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..f112a75 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=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..f2bfeb4 --- /dev/null +++ b/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) 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..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 @@ -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", )