coda-node 0.1.0__py3-none-any.whl

coda_node/__init__.py

@@ -0,0 +1,18 @@
+"""Coda node: a standalone server runtime for quantum processing units.
+
+This package provides everything needed to connect an execution backend
+(QPU or simulator) to the Coda cloud platform.  It handles JWT-based
+authentication, VPN tunnel management, Redis job consumption, and signed
+webhook delivery so that backend authors only need to implement the
+:class:`~coda_node.server.executor.JobExecutor` protocol.
+
+Typical usage::
+
+    uv run coda-node start --token <node-token>
+"""
+
+from coda_node.errors import CodaError
+from coda_node.server import app, create_app
+
+__version__ = "0.1.0"
+__all__ = ["CodaError", "__version__", "app", "create_app"]

coda_node/errors.py

@@ -0,0 +1,44 @@
+"""Shared exception hierarchy for coda-node.
+
+All domain-specific exceptions inherit from :class:`CodaError` so that
+callers can distinguish expected operational errors from unexpected bugs
+with a single ``except CodaError`` clause.
+"""
+
+__all__ = [
+    "AuthError",
+    "CodaError",
+    "ConfigError",
+    "ExecutorError",
+    "NodeError",
+    "VPNError",
+    "WebhookError",
+]
+
+
+class CodaError(Exception):
+    """Base exception for all coda-node errors."""
+
+
+class ConfigError(CodaError):
+    """Invalid or missing configuration."""
+
+
+class AuthError(CodaError):
+    """JWT authentication failure."""
+
+
+class VPNError(CodaError):
+    """VPN tunnel or health check failure."""
+
+
+class NodeError(CodaError):
+    """Node provisioning or reconnect failure."""
+
+
+class ExecutorError(CodaError):
+    """Executor loading or job execution failure."""
+
+
+class WebhookError(CodaError):
+    """Webhook delivery failure."""

coda_node/server/__init__.py

@@ -0,0 +1,9 @@
+"""Server runtime: FastAPI app, configuration, job consumption, and webhooks.
+
+Re-exports the application factory :func:`create_app` and the default
+``app`` instance used by the ``coda start`` CLI command.
+"""
+
+from coda_node.server.app import app, create_app
+
+__all__ = ["app", "create_app"]

coda_node/server/app.py

@@ -0,0 +1,191 @@
+"""FastAPI application for the standalone Coda-connected node server.
+
+The application is constructed via :func:`create_app`, which wires up a
+lifespan that boots the VPN guard, Redis consumer, heartbeat reporter,
+and webhook client.  A module-level ``app`` instance is provided for
+``uvicorn`` import-string references (``coda_node.server.app:app``).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager, suppress
+
+import redis.asyncio as aioredis
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
+
+from coda_node.server.config import Settings
+from coda_node.server.consumer import RedisConsumer
+from coda_node.server.executor import JobExecutor, load_executor
+from coda_node.server.heartbeat import HeartbeatClient
+from coda_node.server.webhook import WebhookClient
+from coda_node.vpn import (
+    ServiceState,
+    VPNGuard,
+    connect_settings,
+    kill_openvpn_daemon,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_connectivity(device_spec: object | None) -> list[list[int]] | None:
+    """Extract qubit connectivity from a device spec, preferring directed edges."""
+    if device_spec is None:
+        return None
+    directed = getattr(device_spec, "directed_edges", None)
+    if directed:
+        return [list(e) for e in directed]
+    return [list(e) for e in device_spec.logical_edges]  # type: ignore[attr-defined]
+
+
+async def _on_vpn_state_change(state: ServiceState) -> None:
+    """Log VPN state transitions at WARNING level."""
+    logger.warning("VPN state changed: %s", state.value)
+
+
+def create_app(executor: JobExecutor | None = None) -> FastAPI:
+    """Build a fully-wired FastAPI application.
+
+    The returned app uses an async lifespan that:
+
+    1. Loads ``Settings`` from environment / persisted config.
+    2. Connects to the Coda cloud (node or JWT reconnect).
+    3. Runs a VPN preflight check and starts background monitoring.
+    4. Opens a Redis consumer loop that dispatches jobs to *executor*.
+    5. On shutdown, drains in-flight work and tears down resources.
+
+    Args:
+        executor: Custom execution backend.  When ``None``, the executor
+            is resolved from ``CODA_EXECUTOR_FACTORY`` or falls back to
+            :class:`~coda_node.server.executor.NoopExecutor`.
+
+    Returns:
+        A configured :class:`~fastapi.FastAPI` instance with ``/health``
+        and ``/ready`` endpoints.
+    """
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
+        """Boot services on startup and tear them down on shutdown."""
+        settings = Settings()
+        await connect_settings(settings)
+
+        guard = VPNGuard(
+            probe_targets=settings.vpn_probe_urls,
+            interface_hint=settings.vpn_interface_hint,
+            check_interval_sec=settings.vpn_check_interval_sec,
+            vpn_required=settings.vpn_required,
+            extra_headers=settings.node_connect_headers,
+        )
+        vpn_status = await guard.preflight()
+        if (
+            not vpn_status.ok
+            and settings.vpn_required
+            and not settings.allow_degraded_startup
+        ):
+            raise RuntimeError(f"VPN preflight failed: {vpn_status.reason}")
+
+        redis_client = aioredis.from_url(settings.redis_url, decode_responses=True)
+        runner = executor or load_executor(settings)
+
+        device_spec = getattr(runner, "device", None)
+        connectivity = _resolve_connectivity(device_spec)
+
+        webhook = WebhookClient(
+            qpu_id=settings.qpu_id,
+            jwt_private_key=settings.jwt_private_key,
+            jwt_key_id=settings.jwt_key_id,
+            extra_headers=settings.node_connect_headers,
+        )
+        consumer = RedisConsumer(
+            redis=redis_client,
+            runner=runner,
+            webhook=webhook,
+            qpu_id=settings.qpu_id,
+        )
+
+        heartbeat = HeartbeatClient(
+            heartbeat_url=settings.heartbeat_url,
+            qpu_id=settings.qpu_id,
+            jwt_private_key=settings.jwt_private_key,
+            jwt_key_id=settings.jwt_key_id,
+            consumer=consumer,
+            interval=settings.heartbeat_interval_sec,
+            connectivity=connectivity,
+            extra_headers=settings.node_connect_headers,
+        )
+
+        watch_task = asyncio.create_task(guard.watch(_on_vpn_state_change))
+        consumer_task = asyncio.create_task(consumer.consume_loop())
+        heartbeat_task = asyncio.create_task(heartbeat.run())
+
+        app.state.settings = settings
+        app.state.guard = guard
+        app.state.consumer = consumer
+        app.state.webhook = webhook
+        app.state.heartbeat = heartbeat
+
+        yield
+
+        guard.stop()
+        watch_task.cancel()
+        with suppress(asyncio.CancelledError):
+            await watch_task
+
+        await heartbeat.close()
+        heartbeat_task.cancel()
+        with suppress(asyncio.CancelledError):
+            await heartbeat_task
+
+        drained = await consumer.drain(timeout=settings.shutdown_drain_timeout_sec)
+        if not drained:
+            logger.warning("Drain timeout expired, cancelling in-flight job")
+        consumer_task.cancel()
+        with suppress(asyncio.CancelledError):
+            await consumer_task
+
+        await webhook.close()
+        await redis_client.aclose()
+        kill_openvpn_daemon()
+
+    app = FastAPI(title="Coda Node", lifespan=lifespan)
+
+    @app.get("/health")
+    async def health() -> dict[str, str]:
+        """Liveness probe — returns 200 if the process is running."""
+        return {"status": "ok"}
+
+    async def _check_readiness() -> dict[str, object]:
+        """Gather composite readiness state from VPN guard and Redis consumer."""
+        guard: VPNGuard = app.state.guard
+        consumer: RedisConsumer = app.state.consumer
+        vpn_ok = guard.is_ready
+        redis_ok = consumer.redis_healthy
+        return {
+            "ready": vpn_ok and redis_ok,
+            "vpn_state": guard.state.value,
+            "redis_healthy": redis_ok,
+            "current_job": consumer.current_job_id,
+        }
+
+    @app.get("/ready")
+    async def ready() -> JSONResponse:
+        """Readiness probe — returns 200 when healthy, 503 otherwise."""
+        try:
+            result = await asyncio.wait_for(_check_readiness(), timeout=5.0)
+            status_code = 200 if result["ready"] else 503
+            return JSONResponse(content=result, status_code=status_code)
+        except TimeoutError:
+            return JSONResponse(
+                content={"ready": False, "reason": "health check timeout"},
+                status_code=503,
+            )
+
+    return app
+
+
+app = create_app()

coda_node/server/auth.py

@@ -0,0 +1,192 @@
+"""RS256 JWT authentication between the node runtime and the Coda cloud.
+
+Every request from the node to the Coda API (webhooks, heartbeats,
+reconnect) is authenticated with a short-lived JWT signed with an RSA
+private key.  The cloud verifies the signature using the corresponding
+public key registered during node provisioning.
+
+Key concepts:
+
+* **KeyPair** -- an RSA keypair generated locally or received from the
+  node endpoint.
+* **sign_token** -- creates a signed JWT for outbound requests.
+* **verify_token** / **verify_token_with_key** -- validate inbound JWTs
+  (used primarily in tests and potential future inbound auth).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import UTC, datetime, timedelta
+
+import jwt
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+__all__ = [
+    "DEFAULT_ISSUER",
+    "DEFAULT_TTL",
+    "KeyPair",
+    "generate_keypair",
+    "sign_token",
+    "verify_token",
+    "verify_token_with_key",
+]
+
+DEFAULT_ISSUER = "coda"
+DEFAULT_TTL = timedelta(hours=1)
+KEY_SIZE = 2048
+
+
+@dataclass(frozen=True, slots=True)
+class KeyPair:
+    """RSA keypair for JWT authentication."""
+
+    private_key_pem: str
+    public_key_pem: str
+    key_id: str
+
+
+def generate_keypair(key_id: str) -> KeyPair:
+    """Generate a fresh PEM-encoded RSA-2048 keypair.
+
+    Args:
+        key_id: Logical identifier embedded in JWTs as the ``kid`` header.
+            Typically the QPU node identifier.
+
+    Returns:
+        A :class:`KeyPair` containing PEM-encoded private and public keys.
+    """
+    private_key = rsa.generate_private_key(public_exponent=65537, key_size=KEY_SIZE)
+
+    private_pem = private_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    ).decode()
+
+    public_pem = (
+        private_key.public_key()
+        .public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo,
+        )
+        .decode()
+    )
+
+    return KeyPair(
+        private_key_pem=private_pem,
+        public_key_pem=public_pem,
+        key_id=key_id,
+    )
+
+
+def sign_token(
+    subject: str,
+    private_key_pem: str,
+    *,
+    issuer: str = DEFAULT_ISSUER,
+    ttl: timedelta = DEFAULT_TTL,
+    key_id: str | None = None,
+) -> str:
+    """Create a short-lived RS256-signed JWT.
+
+    The token carries ``sub``, ``iss``, ``iat``, and ``exp`` claims and a
+    ``kid`` header so the verifier can look up the correct public key.
+
+    Args:
+        subject: The ``sub`` claim -- usually the QPU node identifier.
+        private_key_pem: PEM-encoded RSA private key used to sign.
+        issuer: The ``iss`` claim.
+        ttl: Token lifetime from now.
+        key_id: Explicit ``kid`` header value.  Defaults to *subject*.
+
+    Returns:
+        A compact-serialised JWT string.
+    """
+    now = datetime.now(UTC)
+    payload = {
+        "sub": subject,
+        "iss": issuer,
+        "iat": now,
+        "exp": now + ttl,
+    }
+    headers = {"kid": key_id or subject}
+    return jwt.encode(payload, private_key_pem, algorithm="RS256", headers=headers)
+
+
+def verify_token(
+    token: str,
+    get_public_key: Callable[[str], str | None],
+    *,
+    issuer: str = DEFAULT_ISSUER,
+) -> dict[str, object]:
+    """Verify and decode a JWT by resolving its ``kid`` to a public key.
+
+    Args:
+        token: The compact-serialised JWT to verify.
+        get_public_key: Callable that maps a ``kid`` string to a
+            PEM-encoded public key, or ``None`` if the key is unknown.
+        issuer: Expected ``iss`` claim value.
+
+    Returns:
+        The decoded payload as a dictionary.
+
+    Raises:
+        jwt.InvalidTokenError: If the token is malformed, the ``kid`` is
+            missing or unknown, or signature / claim validation fails.
+    """
+    try:
+        header = jwt.get_unverified_header(token)
+    except jwt.DecodeError as exc:
+        raise jwt.InvalidTokenError(f"Malformed token header: {exc}") from exc
+
+    kid = header.get("kid")
+    if not kid:
+        raise jwt.InvalidTokenError("Token missing 'kid' header")
+
+    public_key_pem = get_public_key(kid)
+    if not public_key_pem:
+        raise jwt.InvalidTokenError(f"Unknown key_id: {kid}")
+
+    result: dict[str, object] = jwt.decode(
+        token,
+        public_key_pem,
+        algorithms=["RS256"],
+        issuer=issuer,
+        options={"require": ["sub", "iss", "exp", "iat"]},
+    )
+    return result
+
+
+def verify_token_with_key(
+    token: str,
+    public_key_pem: str,
+    *,
+    issuer: str = DEFAULT_ISSUER,
+) -> dict[str, object]:
+    """Verify and decode a JWT when the public key is already known.
+
+    Unlike :func:`verify_token`, this skips the ``kid`` lookup and uses
+    the provided key directly.
+
+    Args:
+        token: The compact-serialised JWT to verify.
+        public_key_pem: PEM-encoded RSA public key.
+        issuer: Expected ``iss`` claim value.
+
+    Returns:
+        The decoded payload as a dictionary.
+
+    Raises:
+        jwt.InvalidTokenError: If signature or claim validation fails.
+    """
+    result: dict[str, object] = jwt.decode(
+        token,
+        public_key_pem,
+        algorithms=["RS256"],
+        issuer=issuer,
+        options={"require": ["sub", "iss", "exp", "iat"]},
+    )
+    return result