feat(cloud): add edge host enrollment
This commit is contained in:
@@ -8,10 +8,12 @@ from cloud.internal_api.models import AssignmentModel
|
||||
from device.manager import DeviceManager
|
||||
from driver.registry import build_driver_factory
|
||||
from host_agent.assignment import AssignmentExecutor
|
||||
from host_agent.client import HostAgentClient
|
||||
from host_agent.client import HostAgentClient, HostAgentEnrollmentClient
|
||||
from host_agent.config import HostAgentConfig, load_host_agent_config
|
||||
from host_agent.enrollment import resolve_host_identity
|
||||
from host_agent.execution import create_execution_factories
|
||||
from host_agent.heartbeat import HeartbeatSynchronizer
|
||||
from host_agent.identity import HostIdentityStore
|
||||
from host_agent.lease import ActiveAssignmentRunner
|
||||
from host_agent.processor import AssignmentProcessingResult, AssignmentProcessor
|
||||
from storage.device_config import DeviceConfigStore
|
||||
@@ -97,9 +99,29 @@ def create_application(
|
||||
config: HostAgentConfig | None = None,
|
||||
manager: DeviceManager | None = None,
|
||||
device_config_store: DeviceConfigStore | None = None,
|
||||
identity_store: HostIdentityStore | None = None,
|
||||
enrollment_client: HostAgentEnrollmentClient | None = None,
|
||||
) -> HostAgentApplication:
|
||||
resolved_config = config or load_host_agent_config()
|
||||
resolved_manager = manager or _configured_device_manager(device_config_store)
|
||||
startup_config = config or load_host_agent_config()
|
||||
config_store = device_config_store or DeviceConfigStore()
|
||||
owned_enrollment_client = enrollment_client is None
|
||||
bootstrap_client = enrollment_client or HostAgentEnrollmentClient(startup_config)
|
||||
try:
|
||||
resolved_config = resolve_host_identity(
|
||||
startup_config,
|
||||
identity_store=identity_store
|
||||
or HostIdentityStore(startup_config.identity_path),
|
||||
client=bootstrap_client,
|
||||
)
|
||||
bootstrap_client.config = resolved_config
|
||||
resolved_manager = manager or _configured_device_manager(
|
||||
config_store,
|
||||
config=resolved_config,
|
||||
enrollment_client=bootstrap_client,
|
||||
)
|
||||
finally:
|
||||
if owned_enrollment_client:
|
||||
bootstrap_client.close()
|
||||
client = HostAgentClient(resolved_config)
|
||||
heartbeat = HeartbeatSynchronizer(resolved_manager, client, resolved_config)
|
||||
executor = AssignmentExecutor(create_execution_factories(resolved_manager))
|
||||
@@ -112,13 +134,25 @@ def create_application(
|
||||
|
||||
|
||||
def _configured_device_manager(
|
||||
config_store: DeviceConfigStore | None,
|
||||
config_store: DeviceConfigStore,
|
||||
*,
|
||||
config: HostAgentConfig,
|
||||
enrollment_client: HostAgentEnrollmentClient,
|
||||
) -> DeviceManager:
|
||||
manager = DeviceManager()
|
||||
store = config_store or DeviceConfigStore()
|
||||
for device in store.list():
|
||||
for device in config_store.list():
|
||||
runtime_device_id = device["device_id"]
|
||||
if config.enrollment_managed:
|
||||
enrollment = enrollment_client.enroll_device(
|
||||
local_device_id=device["device_id"],
|
||||
driver_type=device["driver_type"],
|
||||
name=device["name"],
|
||||
capability_tags=[],
|
||||
)
|
||||
runtime_device_id = enrollment.device_id
|
||||
config_store.set_cloud_device_id(device["device_id"], runtime_device_id)
|
||||
manager.register_device(
|
||||
device["device_id"],
|
||||
runtime_device_id,
|
||||
build_driver_factory(
|
||||
device["driver_type"],
|
||||
device["connection_info"],
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
@@ -9,8 +10,10 @@ import httpx
|
||||
from cloud.internal_api.models import (
|
||||
AssignmentModel,
|
||||
ClaimResponse,
|
||||
DeviceEnrollmentResponse,
|
||||
DeviceSnapshotModel,
|
||||
HeartbeatResponse,
|
||||
HostEnrollmentResponse,
|
||||
LeaseRenewalResponse,
|
||||
TerminalResultResponse,
|
||||
)
|
||||
@@ -30,6 +33,99 @@ class StaleLeaseError(HostAgentAPIError):
|
||||
pass
|
||||
|
||||
|
||||
class HostAgentEnrollmentClient:
|
||||
def __init__(
|
||||
self,
|
||||
config: HostAgentConfig,
|
||||
*,
|
||||
http_client: httpx.Client | None = None,
|
||||
sleep: Callable[[float], None] = time.sleep,
|
||||
) -> None:
|
||||
self.config = config
|
||||
self._sleep = sleep
|
||||
self._owns_client = http_client is None
|
||||
self._client = http_client or httpx.Client(base_url=config.control_plane_url)
|
||||
|
||||
def enroll_host(
|
||||
self,
|
||||
*,
|
||||
agent_instance_id: str,
|
||||
host_token: str,
|
||||
display_name: str | None,
|
||||
) -> HostEnrollmentResponse:
|
||||
if not self.config.enrollment_token:
|
||||
raise HostAgentAPIError(0, "Host enrollment token is unavailable")
|
||||
response = self._request(
|
||||
"POST",
|
||||
"/internal/v1/enrollments",
|
||||
token=self.config.enrollment_token,
|
||||
json={
|
||||
"agent_instance_id": agent_instance_id,
|
||||
"host_token": host_token,
|
||||
"display_name": display_name,
|
||||
},
|
||||
)
|
||||
return HostEnrollmentResponse.model_validate(response.json())
|
||||
|
||||
def enroll_device(
|
||||
self,
|
||||
*,
|
||||
local_device_id: str,
|
||||
driver_type: str,
|
||||
name: str | None,
|
||||
capability_tags: list[str],
|
||||
) -> DeviceEnrollmentResponse:
|
||||
if not self.config.host_id or not self.config.token:
|
||||
raise HostAgentAPIError(0, "Host identity is unresolved")
|
||||
response = self._request(
|
||||
"POST",
|
||||
f"/internal/v1/hosts/{self.config.host_id}/devices/enroll",
|
||||
token=self.config.token,
|
||||
json={
|
||||
"local_device_id": local_device_id,
|
||||
"driver_type": driver_type,
|
||||
"name": name,
|
||||
"capability_tags": list(capability_tags),
|
||||
},
|
||||
)
|
||||
return DeviceEnrollmentResponse.model_validate(response.json())
|
||||
|
||||
def close(self) -> None:
|
||||
if self._owns_client:
|
||||
self._client.close()
|
||||
|
||||
def _request(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
*,
|
||||
token: str,
|
||||
json: dict[str, Any],
|
||||
) -> httpx.Response:
|
||||
backoff = self.config.retry_backoff_seconds
|
||||
for attempt in range(1, self.config.max_retry_attempts + 1):
|
||||
try:
|
||||
response = self._client.request(
|
||||
method,
|
||||
path,
|
||||
json=json,
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
)
|
||||
except httpx.TransportError:
|
||||
if attempt == self.config.max_retry_attempts:
|
||||
raise
|
||||
else:
|
||||
if response.status_code < 500:
|
||||
if response.is_success:
|
||||
return response
|
||||
_raise_api_error(response)
|
||||
if attempt == self.config.max_retry_attempts:
|
||||
_raise_api_error(response)
|
||||
self._sleep(backoff)
|
||||
backoff = min(backoff * 2, self.config.max_retry_backoff_seconds)
|
||||
raise AssertionError("retry loop exited unexpectedly")
|
||||
|
||||
|
||||
class HostAgentClient:
|
||||
def __init__(
|
||||
self,
|
||||
@@ -149,22 +245,24 @@ class HostAgentClient:
|
||||
if response.status_code < 500:
|
||||
if response.is_success:
|
||||
return response
|
||||
self._raise_api_error(response)
|
||||
_raise_api_error(response, stale_lease=True)
|
||||
if attempt == self.config.max_retry_attempts:
|
||||
self._raise_api_error(response)
|
||||
_raise_api_error(response, stale_lease=True)
|
||||
|
||||
await self._sleep(backoff)
|
||||
backoff = min(backoff * 2, self.config.max_retry_backoff_seconds)
|
||||
raise AssertionError("retry loop exited unexpectedly")
|
||||
|
||||
@staticmethod
|
||||
def _raise_api_error(response: httpx.Response) -> None:
|
||||
try:
|
||||
payload = response.json()
|
||||
except ValueError:
|
||||
payload = {}
|
||||
detail = payload.get("detail") or payload.get("code") or "request rejected"
|
||||
error_type = (
|
||||
StaleLeaseError if response.status_code == 409 else HostAgentAPIError
|
||||
)
|
||||
raise error_type(response.status_code, str(detail))
|
||||
|
||||
def _raise_api_error(response: httpx.Response, *, stale_lease: bool = False) -> None:
|
||||
try:
|
||||
payload = response.json()
|
||||
except ValueError:
|
||||
payload = {}
|
||||
detail = payload.get("detail") or payload.get("code") or "request rejected"
|
||||
error_type = (
|
||||
StaleLeaseError
|
||||
if stale_lease and response.status_code == 409
|
||||
else HostAgentAPIError
|
||||
)
|
||||
raise error_type(response.status_code, str(detail))
|
||||
|
||||
@@ -2,7 +2,8 @@ from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
||||
@@ -13,8 +14,12 @@ class HostAgentConfigurationError(ValueError):
|
||||
@dataclass(frozen=True)
|
||||
class HostAgentConfig:
|
||||
control_plane_url: str
|
||||
host_id: str
|
||||
token: str
|
||||
host_id: str = ""
|
||||
token: str = field(default="", repr=False)
|
||||
enrollment_token: str = field(default="", repr=False)
|
||||
identity_path: Path = Path("tasks/host_identity.json")
|
||||
enrollment_managed: bool = False
|
||||
display_name: str | None = None
|
||||
heartbeat_interval_seconds: float = 30.0
|
||||
poll_timeout_seconds: float = 20.0
|
||||
retry_backoff_seconds: float = 1.0
|
||||
@@ -41,16 +46,29 @@ def load_host_agent_config(
|
||||
)
|
||||
|
||||
host_id = values.get("HOST_AGENT_HOST_ID", "").strip()
|
||||
if not host_id:
|
||||
raise HostAgentConfigurationError("HOST_AGENT_HOST_ID is required")
|
||||
token = values.get("HOST_AGENT_TOKEN", "").strip()
|
||||
if not token:
|
||||
raise HostAgentConfigurationError("HOST_AGENT_TOKEN is required")
|
||||
if bool(host_id) != bool(token):
|
||||
raise HostAgentConfigurationError(
|
||||
"HOST_AGENT_HOST_ID and HOST_AGENT_TOKEN must be configured together"
|
||||
)
|
||||
enrollment_token = values.get("HOST_AGENT_ENROLLMENT_TOKEN", "").strip()
|
||||
identity_path = Path(
|
||||
values.get("HOST_AGENT_IDENTITY_PATH", "tasks/host_identity.json").strip()
|
||||
)
|
||||
if not host_id and not enrollment_token and not identity_path.is_file():
|
||||
raise HostAgentConfigurationError(
|
||||
"explicit Host credentials, an enrollment token, or existing identity state "
|
||||
"is required"
|
||||
)
|
||||
|
||||
config = HostAgentConfig(
|
||||
control_plane_url=control_plane_url,
|
||||
host_id=host_id,
|
||||
token=token,
|
||||
enrollment_token=enrollment_token,
|
||||
identity_path=identity_path,
|
||||
enrollment_managed=not bool(host_id),
|
||||
display_name=values.get("HOST_AGENT_DISPLAY_NAME") or None,
|
||||
heartbeat_interval_seconds=_positive_float(
|
||||
values,
|
||||
"HOST_AGENT_HEARTBEAT_INTERVAL_SECONDS",
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import replace
|
||||
|
||||
from host_agent.client import HostAgentEnrollmentClient
|
||||
from host_agent.config import HostAgentConfig, HostAgentConfigurationError
|
||||
from host_agent.identity import HostIdentityStore
|
||||
|
||||
|
||||
def resolve_host_identity(
|
||||
config: HostAgentConfig,
|
||||
*,
|
||||
identity_store: HostIdentityStore,
|
||||
client: HostAgentEnrollmentClient,
|
||||
) -> HostAgentConfig:
|
||||
if config.host_id and config.token:
|
||||
return config
|
||||
state = identity_store.load_or_create()
|
||||
if state.host_id is None:
|
||||
if not config.enrollment_token:
|
||||
raise HostAgentConfigurationError(
|
||||
"HOST_AGENT_ENROLLMENT_TOKEN is required to complete enrollment"
|
||||
)
|
||||
response = client.enroll_host(
|
||||
agent_instance_id=state.agent_instance_id,
|
||||
host_token=state.token,
|
||||
display_name=config.display_name,
|
||||
)
|
||||
state = identity_store.complete(state, response.host_id)
|
||||
return replace(
|
||||
config,
|
||||
host_id=state.host_id or "",
|
||||
token=state.token,
|
||||
enrollment_managed=True,
|
||||
)
|
||||
@@ -0,0 +1,101 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass, field, replace
|
||||
from pathlib import Path
|
||||
from secrets import token_urlsafe
|
||||
from uuid import uuid4
|
||||
|
||||
|
||||
class HostIdentityStateError(RuntimeError):
|
||||
"""Raised when persisted Host enrollment identity is missing or invalid."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HostIdentityState:
|
||||
agent_instance_id: str
|
||||
token: str = field(repr=False)
|
||||
host_id: str | None = None
|
||||
|
||||
|
||||
class HostIdentityStore:
|
||||
def __init__(self, path: str | Path) -> None:
|
||||
self.path = Path(path)
|
||||
|
||||
def load(self) -> HostIdentityState | None:
|
||||
if not self.path.exists():
|
||||
return None
|
||||
try:
|
||||
payload = json.loads(self.path.read_text(encoding="utf-8"))
|
||||
except (OSError, ValueError, json.JSONDecodeError) as exc:
|
||||
raise HostIdentityStateError("Host identity state is unreadable") from exc
|
||||
if not isinstance(payload, dict):
|
||||
raise HostIdentityStateError("Host identity state must be an object")
|
||||
agent_instance_id = payload.get("agent_instance_id")
|
||||
token = payload.get("token")
|
||||
host_id = payload.get("host_id")
|
||||
if (
|
||||
not isinstance(agent_instance_id, str)
|
||||
or not agent_instance_id
|
||||
or not isinstance(token, str)
|
||||
or len(token) < 32
|
||||
or (host_id is not None and (not isinstance(host_id, str) or not host_id))
|
||||
):
|
||||
raise HostIdentityStateError("Host identity state is invalid")
|
||||
return HostIdentityState(
|
||||
agent_instance_id=agent_instance_id,
|
||||
token=token,
|
||||
host_id=host_id,
|
||||
)
|
||||
|
||||
def load_or_create(self) -> HostIdentityState:
|
||||
existing = self.load()
|
||||
if existing is not None:
|
||||
return existing
|
||||
state = HostIdentityState(
|
||||
agent_instance_id=f"agent-{uuid4().hex}",
|
||||
token=token_urlsafe(48),
|
||||
)
|
||||
self._write(state)
|
||||
return state
|
||||
|
||||
def complete(self, state: HostIdentityState, host_id: str) -> HostIdentityState:
|
||||
if not host_id:
|
||||
raise ValueError("host_id must not be empty")
|
||||
current = self.load()
|
||||
if current is not None and (
|
||||
current.agent_instance_id != state.agent_instance_id
|
||||
or current.token != state.token
|
||||
):
|
||||
raise HostIdentityStateError("Host identity changed during enrollment")
|
||||
completed = replace(state, host_id=host_id)
|
||||
self._write(completed)
|
||||
return completed
|
||||
|
||||
def _write(self, state: HostIdentityState) -> None:
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = self.path.with_name(f".{self.path.name}.{uuid4().hex}.tmp")
|
||||
payload = {
|
||||
"agent_instance_id": state.agent_instance_id,
|
||||
"token": state.token,
|
||||
"host_id": state.host_id,
|
||||
}
|
||||
try:
|
||||
temporary.write_text(
|
||||
json.dumps(payload, ensure_ascii=True, indent=2) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
_restrict_permissions(temporary)
|
||||
os.replace(temporary, self.path)
|
||||
_restrict_permissions(self.path)
|
||||
finally:
|
||||
if temporary.exists():
|
||||
temporary.unlink()
|
||||
|
||||
|
||||
def _restrict_permissions(path: Path) -> None:
|
||||
try:
|
||||
path.chmod(0o600)
|
||||
except OSError:
|
||||
return
|
||||
Reference in New Issue
Block a user