move parse_verification_key function to dispatch.signature module

Signed-off-by: Achille Roussel <achille.roussel@gmail.com>

Author: achille-roussel

src/dispatch/__init__.py

@@ -10,19 +10,27 @@
 from dispatch.status import Status
 
 __all__ = [
+    "Call",
     "Client",
-    "DispatchID",
     "DEFAULT_API_URL",
+    "DispatchID",
+    "Error",
     "Input",
     "Output",
-    "Call",
-    "Error",
+    "Registry",
     "Reset",
     "Status",
-    "call",
-    "gather",
     "all",
     "any",
+    "call",
+    "function",
+    "gather",
     "race",
-    "Registry",
+    "run",
 ]
+
+function = None
+primitive_function = None
+
+
+def run(): ...

src/dispatch/fastapi.py

@@ -1,4 +1,4 @@
-"""Integration of Dispatch programmable endpoints for FastAPI.
+"""Integration of Dispatch functions with FastAPI.
 
 Example:
 
@@ -18,7 +18,6 @@ def read_root():
     """
 
 import asyncio
-import base64
 import logging
 import os
 from datetime import timedelta
@@ -36,8 +35,7 @@ def read_root():
     CaseInsensitiveDict,
     Ed25519PublicKey,
     Request,
-    public_key_from_bytes,
-    public_key_from_pem,
+    parse_verification_key,
     verify_request,
 )
 from dispatch.status import Status
@@ -46,9 +44,7 @@ def read_root():
 
 
 class Dispatch(Registry):
-    """A Dispatch programmable endpoint, powered by FastAPI."""
-
-    __slots__ = ("client",)
+    """A Dispatch instance, powered by FastAPI."""
 
     def __init__(
         self,
@@ -65,9 +61,9 @@ def __init__(
         Args:
             app: The FastAPI app to configure.
 
-            endpoint: Full URL of the application the Dispatch programmable
-                endpoint will be running on. Uses the value of the
-                DISPATCH_ENDPOINT_URL environment variable by default.
+            endpoint: Full URL of the application the Dispatch instance will
+                be running on. Uses the value of the DISPATCH_ENDPOINT_URL
+                environment variable by default.
 
             verification_key: Key to use when verifying signed requests. Uses
                 the value of the DISPATCH_VERIFICATION_KEY environment variable
@@ -108,55 +104,13 @@ def __init__(
                 f"{endpoint_from} must be a full URL with protocol and domain (e.g., https://example.com)"
             )
 
-        verification_key = parse_verification_key(verification_key)
-        if verification_key:
-            base64_key = base64.b64encode(verification_key.public_bytes_raw()).decode()
-            logger.info("verifying request signatures using key %s", base64_key)
-        elif parsed_url.scheme != "bridge":
-            logger.warning(
-                "request verification is disabled because DISPATCH_VERIFICATION_KEY is not set"
-            )
-
         super().__init__(endpoint, api_key=api_key, api_url=api_url)
 
+        verification_key = parse_verification_key(verification_key, url_scheme=parsed_url.scheme)
         function_service = _new_app(self, verification_key)
         app.mount("/dispatch.sdk.v1.FunctionService", function_service)
 
 
-def parse_verification_key(
-    verification_key: Optional[Union[Ed25519PublicKey, str, bytes]],
-) -> Optional[Ed25519PublicKey]:
-    if isinstance(verification_key, Ed25519PublicKey):
-        return verification_key
-
-    from_env = False
-    if not verification_key:
-        try:
-            verification_key = os.environ["DISPATCH_VERIFICATION_KEY"]
-        except KeyError:
-            return None
-        from_env = True
-
-    if isinstance(verification_key, bytes):
-        verification_key = verification_key.decode()
-
-    # Be forgiving when accepting keys in PEM format, which may span
-    # multiple lines. Users attempting to pass a PEM key via an environment
-    # variable may accidentally include literal "\n" bytes rather than a
-    # newline char (0xA).
-    try:
-        return public_key_from_pem(verification_key.replace("\\n", "\n"))
-    except ValueError:
-        pass
-
-    try:
-        return public_key_from_bytes(base64.b64decode(verification_key.encode()))
-    except ValueError:
-        if from_env:
-            raise ValueError(f"invalid DISPATCH_VERIFICATION_KEY '{verification_key}'")
-        raise ValueError(f"invalid verification key '{verification_key}'")
-
-
 class _ConnectResponse(fastapi.Response):
     media_type = "application/grpc+proto"
 

src/dispatch/grpc.py

@@ -0,0 +1,7 @@
+"""Integration of Dispatch functions with gRPC."""
+
+from dispatch.function import Batch, Registry
+
+
+class Dispatch(Registry):
+    """A Dispatch instance to be serviced by a gRPC server."""

src/dispatch/signature/__init__.py

@@ -1,6 +1,8 @@
+import base64
 import logging
+import os
 from datetime import datetime, timedelta
-from typing import Sequence, Set, cast
+from typing import Optional, Sequence, Set, Union, cast
 
 import http_sfv
 from cryptography.hazmat.primitives.asymmetric.ed25519 import (
@@ -123,3 +125,61 @@ def extract_covered_components(result: VerifyResult) -> Set[str]:
         covered_components.add(item.value)
 
     return covered_components
+
+
+def parse_verification_key(
+    verification_key: Optional[Union[Ed25519PublicKey, str, bytes]],
+    url_scheme: str | None = None,
+) -> Optional[Ed25519PublicKey]:
+    # This function depends a lot on global context like enviornment variables
+    # and logging configuration. It's not ideal for testing, but it's useful to
+    # unify the behavior of the Dispatch class everywhere the signature module
+    # is used.
+    if isinstance(verification_key, Ed25519PublicKey):
+        return verification_key
+
+    # Keep track of whether the key was obtained from the environment, so that
+    # we can tweak the error messages accordingly.
+    from_env = False
+    if not verification_key:
+        try:
+            verification_key = os.environ["DISPATCH_VERIFICATION_KEY"]
+        except KeyError:
+            return None
+        from_env = verification_key is not None
+
+    if isinstance(verification_key, bytes):
+        verification_key = verification_key.decode()
+
+    # Be forgiving when accepting keys in PEM format, which may span
+    # multiple lines. Users attempting to pass a PEM key via an environment
+    # variable may accidentally include literal "\n" bytes rather than a
+    # newline char (0xA).
+    public_key: Optional[Ed25519PublicKey] = None
+    try:
+        public_key = public_key_from_pem(verification_key.replace("\\n", "\n"))
+    except ValueError:
+        pass
+
+    # If the key is not in PEM format, try to decode it as base64 string.
+    if not public_key:
+        try:
+            public_key = public_key_from_bytes(
+                base64.b64decode(verification_key.encode())
+            )
+        except ValueError:
+            if from_env:
+                raise ValueError(
+                    f"invalid DISPATCH_VERIFICATION_KEY '{verification_key}'"
+                )
+            raise ValueError(f"invalid verification key '{verification_key}'")
+
+    # Print diagostic information about the key, this is useful for debugging.
+    if public_key:
+        base64_key = base64.b64encode(public_key.public_bytes_raw()).decode()
+        logger.info("verifying request signatures using key %s", base64_key)
+    elif url_scheme != "bridge":
+        logger.warning(
+            "request verification is disabled because DISPATCH_VERIFICATION_KEY is not set"
+        )
+    return public_key