feat(cloud): add edge host enrollment

packages/cloud-platform/cloud/internal_api/api.py

@@ -5,26 +5,38 @@ from collections.abc import Awaitable, Callable
 from datetime import timedelta
 from time import monotonic
 from typing import TYPE_CHECKING
+from uuid import uuid4
 
 from fastapi import APIRouter, HTTPException, Request, status
 from fastapi.responses import JSONResponse
 
 from cloud.auth import (
     AuthProvider,
+    ConfiguredEnrollmentTokenProvider,
     HostAuthorizationError,
+    digest_token,
 )
 from cloud.internal_api.models import (
     AssignmentModel,
     ClaimRequest,
     ClaimResponse,
+    DeviceEnrollmentRequest,
+    DeviceEnrollmentResponse,
     HeartbeatRequest,
     HeartbeatResponse,
+    HostEnrollmentRequest,
+    HostEnrollmentResponse,
     LeaseRenewalRequest,
     LeaseRenewalResponse,
     StaleLeaseConflict,
     TerminalResultRequest,
     TerminalResultResponse,
 )
+from cloud.repository import (
+    DeviceEnrollmentConflictError,
+    EnrollmentTokenConflictError,
+    HostEnrollmentConflictError,
+)
 from core.models import Device, utc_now
 
 if TYPE_CHECKING:
@@ -35,6 +47,7 @@ def create_internal_router(
     *,
     pool: DevicePool,
     auth_provider: AuthProvider,
+    enrollment_auth_provider: ConfiguredEnrollmentTokenProvider | None = None,
     version_prefix: str = "/internal/v1",
     claim_poll_interval_seconds: float = 0.1,
     lease_duration_seconds: float = 60.0,
@@ -45,6 +58,7 @@ def create_internal_router(
     if lease_duration_seconds <= 0:
         raise ValueError("lease_duration_seconds must be greater than zero")
     router = APIRouter(prefix=version_prefix, tags=["host-agent"])
+    enrollment_auth = enrollment_auth_provider or ConfiguredEnrollmentTokenProvider(())
 
     def authorize_host(request: Request, host_id: str) -> None:
         principal = auth_provider.authenticate(request)
@@ -62,6 +76,66 @@ def create_internal_router(
                 detail=str(exc),
             ) from exc
 
+    @router.post(
+        "/enrollments",
+        response_model=HostEnrollmentResponse,
+        status_code=status.HTTP_201_CREATED,
+    )
+    def enroll_host(
+        payload: HostEnrollmentRequest,
+        request: Request,
+    ) -> HostEnrollmentResponse:
+        enrollment_principal = enrollment_auth.authenticate(request)
+        if enrollment_principal is None:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="unauthorized",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+        try:
+            enrollment = pool.store.enroll_host(
+                host_id=f"host-{uuid4().hex}",
+                agent_instance_id=payload.agent_instance_id,
+                credential_digest=digest_token(payload.host_token),
+                enrollment_token_digest=enrollment_principal.token_digest,
+                display_name=payload.display_name,
+                enrolled_at=utc_now(),
+            )
+        except (EnrollmentTokenConflictError, HostEnrollmentConflictError) as exc:
+            raise HTTPException(
+                status_code=status.HTTP_409_CONFLICT,
+                detail=str(exc),
+            ) from exc
+        return HostEnrollmentResponse(host_id=enrollment.host_id)
+
+    @router.post(
+        "/hosts/{host_id}/devices/enroll",
+        response_model=DeviceEnrollmentResponse,
+        status_code=status.HTTP_201_CREATED,
+    )
+    def enroll_device(
+        host_id: str,
+        payload: DeviceEnrollmentRequest,
+        request: Request,
+    ) -> DeviceEnrollmentResponse:
+        authorize_host(request, host_id)
+        try:
+            enrollment = pool.store.enroll_device(
+                device_id=f"device-{uuid4().hex}",
+                host_id=host_id,
+                local_device_id=payload.local_device_id,
+                driver_type=payload.driver_type,
+                name=payload.name,
+                capability_tags=list(payload.capability_tags),
+                enrolled_at=utc_now(),
+            )
+        except DeviceEnrollmentConflictError as exc:
+            raise HTTPException(
+                status_code=status.HTTP_409_CONFLICT,
+                detail=str(exc),
+            ) from exc
+        return DeviceEnrollmentResponse(device_id=enrollment.device_id)
+
     @router.put(
         "/hosts/{host_id}/heartbeat",
         response_model=HeartbeatResponse,
@@ -250,6 +324,25 @@ def _validate_snapshot(
             detail="heartbeat snapshot contains duplicate device ids",
         )
 
+    if pool.store.is_enrollment_managed_host(host_id):
+        enrollments = {
+            enrollment.device_id: enrollment
+            for enrollment in pool.store.list_device_enrollments(host_id)
+            if enrollment.revoked_at is None
+        }
+        invalid = [
+            device.device_id
+            for device in payload.devices
+            if device.device_id not in enrollments
+            or enrollments[device.device_id].driver_type != device.driver_type
+        ]
+        if invalid:
+            raise HTTPException(
+                status_code=status.HTTP_409_CONFLICT,
+                detail=f"device enrollment conflict: {sorted(set(invalid))}",
+            )
+        return False
+
     now = utc_now()
     hosts = {host.host_id: host for host in pool.store.list_hosts()}
     conflicts: list[str] = []