feat(cloud): add edge host enrollment

2026-07-13 13:54:16 +08:00
parent cd56facbbf
commit e61dcca801
40 changed files with 2302 additions and 48 deletions
@@ -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] = []
@@ -6,6 +6,27 @@ from typing import Any, Literal
 from pydantic import BaseModel, Field


+class HostEnrollmentRequest(BaseModel):
+    agent_instance_id: str = Field(min_length=1, max_length=256)
+    host_token: str = Field(min_length=32, max_length=512)
+    display_name: str | None = Field(default=None, max_length=256)
+
+
+class HostEnrollmentResponse(BaseModel):
+    host_id: str
+
+
+class DeviceEnrollmentRequest(BaseModel):
+    local_device_id: str = Field(min_length=1, max_length=256)
+    driver_type: str = Field(min_length=1, max_length=128)
+    name: str | None = Field(default=None, max_length=256)
+    capability_tags: list[str] = Field(default_factory=list)
+
+
+class DeviceEnrollmentResponse(BaseModel):
+    device_id: str
+
+
 class DeviceSnapshotModel(BaseModel):
    device_id: str = Field(min_length=1)
    driver_type: str = Field(min_length=1)