diff --git a/alembic/versions/0007_host_service_account.py b/alembic/versions/0007_host_service_account.py new file mode 100644 index 0000000..124b4e0 --- /dev/null +++ b/alembic/versions/0007_host_service_account.py @@ -0,0 +1,36 @@ +"""Record the service account that holds each host. + +Revision ID: 0007_host_service_account +Revises: 0006_service_accounts +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +revision: str = "0007_host_service_account" +down_revision: str | None = "0006_service_accounts" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + with op.batch_alter_table("service_accounts") as service_accounts: + service_accounts.alter_column("fingerprint", nullable=True) + op.execute(sa.text("INSERT INTO service_accounts (name) VALUES ('admin')")) + op.add_column("hosts", sa.Column("service_account", sa.String(64))) + # Before this revision only admin keys could create or claim a host. + op.execute( + sa.text( + "UPDATE hosts SET service_account = 'admin' " + "WHERE NOT pool_member OR claimed_at IS NOT NULL" + ) + ) + + +def downgrade() -> None: + op.drop_column("hosts", "service_account") + op.execute(sa.text("DELETE FROM service_accounts WHERE name = 'admin'")) + with op.batch_alter_table("service_accounts") as service_accounts: + service_accounts.alter_column("fingerprint", nullable=False) diff --git a/api-tests/tests/full-api.spec.js b/api-tests/tests/full-api.spec.js index b45c27d..cc8123d 100644 --- a/api-tests/tests/full-api.spec.js +++ b/api-tests/tests/full-api.spec.js @@ -28,6 +28,7 @@ const EXPECTED_OPENAPI_OPERATIONS = [ const HOST_KEYS = [ "id", "name", + "service_account", "status", "provider", "image", diff --git a/docs/api.md b/docs/api.md index 2c2d2dc..11a2dae 100644 --- a/docs/api.md +++ b/docs/api.md @@ -18,7 +18,8 @@ service accounts. A missing token returns `401`, a rejected one `403`. and the SHA-256 fingerprint of the token. Names are 1–64 lowercase letters, digits, or hyphens. A duplicate name returns `409`. `DELETE /service-accounts/ci` revokes the token on the next request, or -returns `404` for an unknown name. +returns `404` for an unknown name. The `admin` account exists from the +migration, has no token, and cannot be removed. Admin keys act as it. ## Endpoints @@ -33,6 +34,11 @@ returns `404` for an unknown name. - `GET /doctor` — read-only dependency diagnostics - `GET /healthz` — unauthenticated liveness probe +Host responses carry `service_account`: the service account that created +or claimed the host, `admin` for an admin key, or `null` for an unclaimed +warm host. Callers cannot set it. An `Idempotency-Key` belongs to the +service account that first used it. Another one reusing it gets `409`. + ## The secrets exchange The exchange is a second process, `python -m secrets_exchange`, on a private diff --git a/src/conftest.py b/src/conftest.py index 0f5a1e7..78b97b6 100644 --- a/src/conftest.py +++ b/src/conftest.py @@ -39,14 +39,17 @@ def load_test_env() -> None: @pytest.fixture(autouse=True) async def reset_database() -> AsyncGenerator[None]: + from sqlalchemy import insert + from core.database import Base, engine from hosts import models # noqa: F401 - from service_accounts import models as service_account_models # noqa: F401 + from service_accounts.models import ServiceAccount from templates import models as template_models # noqa: F401 async with engine.begin() as connection: await connection.run_sync(Base.metadata.drop_all) await connection.run_sync(Base.metadata.create_all) + await connection.execute(insert(ServiceAccount).values(name=ServiceAccount.ADMIN)) yield diff --git a/src/hosts/api.py b/src/hosts/api.py index 7550f0d..86de2f5 100644 --- a/src/hosts/api.py +++ b/src/hosts/api.py @@ -2,14 +2,14 @@ import uuid from typing import Annotated -from fastapi import APIRouter, Depends, Header, HTTPException, Response, status +from fastapi import APIRouter, Body, Depends, Header, HTTPException, Response, status from sqlalchemy.exc import SQLAlchemyError from hosts.auth import require_auth from hosts.deps import get_host_service from hosts.exceptions import HostTeardownError from hosts.models import Host -from hosts.schemas import HostCreate, HostOut, HostRenew +from hosts.schemas import ExpiresAt, HostCreate, HostOut from hosts.service import HostService from networking.tailscale import NetworkError from providers.exceptions import ProviderError, UnknownProviderError, UnsupportedSizingError @@ -24,6 +24,7 @@ @router.post("", response_model=HostOut, status_code=status.HTTP_201_CREATED) async def create_host( service: HostServiceDep, + service_account: Annotated[str | None, Depends(require_auth)], payload: HostCreate | None = None, idempotency_key: Annotated[ str | None, @@ -47,6 +48,7 @@ async def create_host( try: return await service.get_or_create_host( + service_account=service_account, env=host_create.env, secrets={name: entry.to_storage() for name, entry in host_create.secrets.items()}, image=host_create.image, @@ -83,10 +85,18 @@ async def get_host(host_id: uuid.UUID, service: HostServiceDep) -> Host: async def renew_host( host_id: uuid.UUID, service: HostServiceDep, - payload: HostRenew | None = None, + expires_at: Annotated[ + ExpiresAt, + Body( + embed=True, + description=( + "Omit or null: extend by LEASE_DEFAULT_TTL from now. " + "Renewal never makes a host permanent." + ), + ), + ] = None, ) -> Host: - host_renew = payload or HostRenew() - return await service.renew_host(host_id, expires_at=host_renew.expires_at) + return await service.renew_host(host_id, expires_at=expires_at) @router.delete("/{host_id}", status_code=status.HTTP_204_NO_CONTENT) diff --git a/src/hosts/auth.py b/src/hosts/auth.py index e795d8f..7347430 100644 --- a/src/hosts/auth.py +++ b/src/hosts/auth.py @@ -28,19 +28,17 @@ def is_admin_key(token: str) -> bool: async def require_auth( credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer_scheme)], session: Annotated[AsyncSession, Depends(get_session)], -) -> None: +) -> str | None: if not credentials: raise HTTPException(status_code=401, detail="admin key or service account token required") if is_admin_key(credentials.credentials): - return + return ServiceAccount.ADMIN try: - if await ServiceAccount.authenticate(session, credentials.credentials): - return + return await ServiceAccount.authenticate(session, credentials.credentials) except SQLAlchemyError: logger.exception("service account token lookup failed") raise HTTPException(status_code=503, detail="service account tokens unavailable") from None - raise HTTPException(status_code=403, detail="admin key or service account token rejected") def require_admin_auth( diff --git a/src/hosts/exceptions.py b/src/hosts/exceptions.py index a04c14c..c2c558c 100644 --- a/src/hosts/exceptions.py +++ b/src/hosts/exceptions.py @@ -6,6 +6,11 @@ class HostStateError(AppException): error_code = "HOST_STATE" +class IdempotencyKeyConflictError(AppException): + status_code = 409 + error_code = "IDEMPOTENCY_KEY_CONFLICT" + + class HostTeardownError(AppException): status_code = 503 error_code = "HOST_TEARDOWN" diff --git a/src/hosts/models.py b/src/hosts/models.py index c9f152d..bf7b4d8 100644 --- a/src/hosts/models.py +++ b/src/hosts/models.py @@ -74,6 +74,7 @@ class Host(Base): env: Mapped[dict[str, str]] = mapped_column(_JSONType, default=dict) secrets: Mapped[SecretsMapping] = EncryptedJsonField() name: Mapped[str] = mapped_column(String(100), unique=True, index=True) + service_account: Mapped[str | None] = mapped_column(String(64)) status: Mapped[str] = mapped_column(String(32), default=HostStatus.PROVISIONING.value) provider: Mapped[str] = mapped_column(String(20), default="exe") image: Mapped[str] = mapped_column(Text) diff --git a/src/hosts/schemas.py b/src/hosts/schemas.py index f242359..f9bef97 100644 --- a/src/hosts/schemas.py +++ b/src/hosts/schemas.py @@ -1,8 +1,9 @@ import re import uuid from datetime import UTC, datetime +from typing import Annotated -from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator +from pydantic import AfterValidator, BaseModel, ConfigDict, Field, ValidationInfo, field_validator from host_secrets import catalog from host_secrets.schemas import SECRET_NAME_PATTERN, SecretEntry @@ -13,15 +14,18 @@ def _expires_at_must_be_future_and_tz_aware(expires_at: datetime | None) -> datetime | None: - if expires_at is None: - return None - if expires_at.tzinfo is None: + if not expires_at: + return + if not expires_at.tzinfo: raise ValueError("expires_at must include a timezone offset") if expires_at <= datetime.now(UTC): raise ValueError("expires_at must be in the future") return expires_at +ExpiresAt = Annotated[datetime | None, AfterValidator(_expires_at_must_be_future_and_tz_aware)] + + class HostCreate(BaseModel): image: str | None = None template: uuid.UUID | None = Field( @@ -36,7 +40,7 @@ class HostCreate(BaseModel): "service handle. The sandbox receives a placeholder per entry, never the value." ), ) - expires_at: datetime | None = None + expires_at: ExpiresAt = None provider: str | None = Field( default=None, description="VM provider to provision on. Omit to use the service default.", @@ -92,22 +96,13 @@ def reject_reserved_env_keys(cls, env: dict[str, str]) -> dict[str, str]: environment.get_persist(env) return env - _validate_expires_at = field_validator("expires_at")(_expires_at_must_be_future_and_tz_aware) - - -class HostRenew(BaseModel): - # Omitted (or null) means "extend by LEASE_DEFAULT_TTL from now"; renewal - # never makes a host permanent — that is a create-time choice. - expires_at: datetime | None = None - - _validate_expires_at = field_validator("expires_at")(_expires_at_must_be_future_and_tz_aware) - class HostOut(BaseModel): model_config = ConfigDict(from_attributes=True) id: uuid.UUID name: str + service_account: str | None status: str provider: str image: str diff --git a/src/hosts/service.py b/src/hosts/service.py index fd79e78..49c1ca3 100644 --- a/src/hosts/service.py +++ b/src/hosts/service.py @@ -20,7 +20,7 @@ from host_secrets import catalog from host_secrets.exceptions import SecretsProxyNotConfiguredError from host_secrets.placeholder import Placeholder -from hosts.exceptions import HostStateError, ProvisioningFailedError +from hosts.exceptions import HostStateError, IdempotencyKeyConflictError, ProvisioningFailedError from hosts.models import Host, HostStatus, IdempotencyKey from networking.tailscale import ( DeviceDiscoveryTimeoutError, @@ -101,6 +101,7 @@ def _default_lease_expires_at(self) -> datetime: async def get_or_create_host( self, *, + service_account: str | None = None, env: dict[str, str], secrets: dict[str, dict[str, Any]] | None = None, image: str | None, @@ -123,7 +124,7 @@ async def get_or_create_host( raise UnknownProviderError(f"unknown provider {provider!r}; available: {available}") if idempotency_key: - existing = await self._lookup_idempotency_key(idempotency_key) + existing = await self._lookup_idempotency_key(idempotency_key, service_account) if existing: return existing @@ -136,10 +137,11 @@ async def get_or_create_host( customized = env or secrets or image or template or instance_type or disk_gb if not customized and self.settings.get_pool_targets().get(requested_provider): host = await self._try_claim_pool_host( - provider=requested_provider, expires_at=expires_at + service_account=service_account, provider=requested_provider, expires_at=expires_at ) if not host: host = await self.create_host( + service_account=service_account, env=env, secrets=secrets, image=image, @@ -158,14 +160,18 @@ async def get_or_create_host( host.claimed_at, ) await self._release_idempotency_loser(host) - winner = await self._lookup_idempotency_key(idempotency_key) + winner = await self._lookup_idempotency_key(idempotency_key, service_account) if not winner: raise HostStateError("idempotency race could not be resolved") from None return winner return host async def _try_claim_pool_host( - self, *, provider: str, expires_at: datetime | None | EllipsisType + self, + *, + service_account: str | None = None, + provider: str, + expires_at: datetime | None | EllipsisType, ) -> Host | None: # Pick a candidate, then atomically claim it with UPDATE ... WHERE # claimed_at IS NULL ... RETURNING. The WHERE predicate is the actual @@ -198,7 +204,12 @@ async def _try_claim_pool_host( update(Host) .where(Host.id == candidate_id) .where(Host.claimed_at.is_(None)) - .values(claimed_at=now, updated_at=now, expires_at=expires_at) + .values( + service_account=service_account, + claimed_at=now, + updated_at=now, + expires_at=expires_at, + ) .returning(Host) ) host = result.scalar_one_or_none() @@ -212,6 +223,7 @@ async def _try_claim_pool_host( async def create_host( self, *, + service_account: str | None = None, env: dict[str, str], secrets: dict[str, dict[str, Any]] | None = None, image: str | None, @@ -259,6 +271,7 @@ async def create_host( ) host = Host( id=uid, + service_account=service_account, env=env, secrets=secrets or {}, name=name, @@ -321,7 +334,7 @@ async def _resolve_template_image(self, *, template_id: uuid.UUID, provider: str template.last_used_at = utc_now() return template.image - async def _lookup_idempotency_key(self, key: str) -> Host | None: + async def _lookup_idempotency_key(self, key: str, service_account: str | None) -> Host | None: record = ( await self.session.execute(select(IdempotencyKey).where(IdempotencyKey.key == key)) ).scalar_one_or_none() @@ -329,10 +342,15 @@ async def _lookup_idempotency_key(self, key: str) -> Host | None: if not record: return - if record.expires_at > utc_now(): - host = await self.session.get(Host, record.host_id) - if host: - return host + host = ( + await self.session.get(Host, record.host_id) if record.expires_at > utc_now() else None + ) + if host: + if host.service_account != service_account: + raise IdempotencyKeyConflictError( + f"idempotency key {key} belongs to another service account" + ) + return host # Stale: expired, or the host vanished without the FK cascade firing. # GC in a dedicated session so we don't autoflush the caller's pending # state on `self.session`. @@ -374,6 +392,7 @@ async def _release_idempotency_loser(self, host: Host) -> None: now = utc_now() if fresh.claimed_at: fresh.claimed_at = None + fresh.service_account = None fresh.expires_at = now + timedelta(hours=self.settings.pool_host_max_age_hours) fresh.updated_at = now logger.info( diff --git a/src/hosts/tests/test_host_service_account.py b/src/hosts/tests/test_host_service_account.py new file mode 100644 index 0000000..73dd72c --- /dev/null +++ b/src/hosts/tests/test_host_service_account.py @@ -0,0 +1,85 @@ +from unittest.mock import AsyncMock + +import pytest +from httpx import AsyncClient + +from core.database import async_session_factory +from core.settings import get_settings +from hosts.models import HostStatus +from hosts.service import HostService + +ADMIN = {"Authorization": "Bearer service-token"} + + +@pytest.fixture(autouse=True) +def provision(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + provision = AsyncMock() + monkeypatch.setattr(HostService, "provision", provision) + return provision + + +async def issue(client: AsyncClient, name: str) -> dict[str, str]: + response = await client.post("/service-accounts", json={"name": name}, headers=ADMIN) + return {"Authorization": f"Bearer {response.json()['token']}"} + + +async def test_host_records_its_service_account(client: AsyncClient) -> None: + account = await issue(client, "account-a") + created = await client.post("/hosts", json={"service_account": "forged"}, headers=account) + assert created.status_code == 201 + assert created.json()["service_account"] == "account-a" + account_host = created.json()["id"] + + created = await client.post("/hosts", headers=ADMIN) + assert created.json()["service_account"] == "admin" + admin_host = created.json()["id"] + + listed = await client.get("/hosts", headers=account) + assert {host["id"]: host["service_account"] for host in listed.json()} == { + account_host: "account-a", + admin_host: "admin", + } + + assert (await client.delete("/service-accounts/account-a", headers=ADMIN)).status_code == 204 + fetched = await client.get(f"/hosts/{account_host}", headers=ADMIN) + assert fetched.json()["service_account"] == "account-a" + + +@pytest.mark.parametrize("name", ["admin", "account-a"]) +async def test_pool_claim_records_service_account_and_release_clears_it( + client: AsyncClient, monkeypatch: pytest.MonkeyPatch, name: str, provision: AsyncMock +) -> None: + monkeypatch.setattr(get_settings(), "pool_size", 1) + + async with async_session_factory() as session: + service = HostService(session) + pool_host = await service.create_host(env={}, image=None, pool_member=True) + assert pool_host.service_account is None + pool_host.status = HostStatus.ACTIVE.value + await session.commit() + pool_id = pool_host.id + + headers = ADMIN if name == "admin" else await issue(client, name) + claimed = await client.post("/hosts", headers=headers) + assert claimed.status_code == 201 + assert claimed.json()["id"] == str(pool_id) + assert claimed.json()["service_account"] == name + provision.assert_awaited_once() + + async with async_session_factory() as session: + await HostService(session)._release_idempotency_loser(pool_host) + + released = await client.get(f"/hosts/{pool_id}", headers=headers) + assert released.json()["service_account"] is None + + +async def test_idempotency_key_belongs_to_the_caller_that_used_it(client: AsyncClient) -> None: + account = await issue(client, "account-a") + key = {"Idempotency-Key": "shared-retry"} + first = await client.post("/hosts", headers={**account, **key}) + retry = await client.post("/hosts", headers={**account, **key}) + assert retry.json()["id"] == first.json()["id"] + + replay = await client.post("/hosts", headers={**ADMIN, **key}) + assert replay.status_code == 409 + assert replay.json()["error_code"] == "IDEMPOTENCY_KEY_CONFLICT" diff --git a/src/service_accounts/api.py b/src/service_accounts/api.py index d106989..d4b1147 100644 --- a/src/service_accounts/api.py +++ b/src/service_accounts/api.py @@ -7,7 +7,7 @@ from core.database import get_session from core.exceptions import ResourceNotFoundError from hosts.auth import require_admin_auth -from service_accounts.exceptions import ServiceAccountExistsError +from service_accounts.exceptions import ServiceAccountExistsError, ServiceAccountStateError from service_accounts.models import ServiceAccount SERVICE_ACCOUNT_NAME_PATTERN = r"^[a-z0-9][a-z0-9-]{0,63}$" @@ -45,6 +45,8 @@ async def create_service_account( @router.delete("/{name}", status_code=status.HTTP_204_NO_CONTENT) async def remove_service_account(name: ServiceAccountName, session: SessionDep) -> Response: if account := await session.get(ServiceAccount, name): + if account.name == ServiceAccount.ADMIN: + raise ServiceAccountStateError("the admin account has no token to revoke") await session.delete(account) await session.commit() return Response(status_code=status.HTTP_204_NO_CONTENT) diff --git a/src/service_accounts/exceptions.py b/src/service_accounts/exceptions.py index f87c41a..3a36e79 100644 --- a/src/service_accounts/exceptions.py +++ b/src/service_accounts/exceptions.py @@ -4,3 +4,13 @@ class ServiceAccountExistsError(AppException): status_code = 409 error_code = "SERVICE_ACCOUNT_EXISTS" + + +class ServiceAccountStateError(AppException): + status_code = 409 + error_code = "SERVICE_ACCOUNT_STATE" + + +class ServiceAccountTokenRejectedError(AppException): + status_code = 403 + error_code = "SERVICE_ACCOUNT_TOKEN_REJECTED" diff --git a/src/service_accounts/models.py b/src/service_accounts/models.py index bb5f4e4..1e1619b 100644 --- a/src/service_accounts/models.py +++ b/src/service_accounts/models.py @@ -1,18 +1,24 @@ import hashlib import secrets +from typing import ClassVar from sqlalchemy import String, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Mapped, mapped_column from core.database import Base +from service_accounts.exceptions import ServiceAccountTokenRejectedError class ServiceAccount(Base): __tablename__ = "service_accounts" + # The admin account is seeded by migration. It has no token: admin keys + # come from SERVICE_TOKENS. Its row reserves the name. + ADMIN: ClassVar[str] = "admin" + name: Mapped[str] = mapped_column(String(64), primary_key=True) - fingerprint: Mapped[str] = mapped_column(String(64), unique=True) + fingerprint: Mapped[str | None] = mapped_column(String(64), unique=True) def issue_token(self) -> str: token = f"drkb_{secrets.token_urlsafe(32)}" @@ -20,12 +26,12 @@ def issue_token(self) -> str: return token @classmethod - async def authenticate(cls, session: AsyncSession, token: str) -> bool: - return bool( - await session.scalar( - select(cls.name).where(cls.fingerprint == cls.get_fingerprint(token)) - ) - ) + async def authenticate(cls, session: AsyncSession, token: str) -> str: + if name := await session.scalar( + select(cls.name).where(cls.fingerprint == cls.get_fingerprint(token)) + ): + return name + raise ServiceAccountTokenRejectedError("service account token rejected") @staticmethod def get_fingerprint(token: str) -> str: diff --git a/src/service_accounts/tests/test_api.py b/src/service_accounts/tests/test_api.py index 0be8c1b..54a1e14 100644 --- a/src/service_accounts/tests/test_api.py +++ b/src/service_accounts/tests/test_api.py @@ -2,7 +2,6 @@ import pytest from httpx import AsyncClient -from sqlalchemy import select from api.app import app from core.database import Base, async_session_factory, engine @@ -22,7 +21,7 @@ async def test_create_stores_only_name_and_fingerprint(client: AsyncClient) -> N token = await create(client) async with async_session_factory() as session: - stored = (await session.scalars(select(ServiceAccount))).one() + stored = await session.get_one(ServiceAccount, "account-a") assert stored.name == "account-a" assert stored.fingerprint == ServiceAccount.get_fingerprint(token) assert set(ServiceAccount.__table__.columns.keys()) == {"name", "fingerprint"} @@ -120,6 +119,15 @@ async def test_removal_rejects_every_protected_route_without_restart(client: Asy ).status_code == 200 +async def test_admin_account_has_no_token_and_cannot_be_removed(client: AsyncClient) -> None: + response = await client.post("/service-accounts", json={"name": "admin"}, headers=ADMIN) + assert response.status_code == 409 + response = await client.delete("/service-accounts/admin", headers=ADMIN) + assert response.status_code == 409 + assert response.json()["error_code"] == "SERVICE_ACCOUNT_STATE" + assert (await client.get("/hosts", headers={"Authorization": "Bearer "})).status_code == 401 + + async def test_remove_unknown_name(client: AsyncClient) -> None: response = await client.delete("/service-accounts/unknown", headers=ADMIN) assert response.status_code == 404