fix(cloud): make direct launch replay-safe

This commit is contained in:
dadachann
2026-07-30 21:02:38 +00:00
parent a5a26f81ee
commit 93dbd3541e
8 changed files with 196 additions and 16 deletions
@@ -410,6 +410,7 @@ class UserRouterGroup(group.RouterGroup):
'token': token,
'user': account.user,
'workspace_uuid': access.workspace.uuid,
'return_path': launch.get('return_path', '/home'),
}
)
except SpaceLaunchError:
+39 -6
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import asyncio
import base64
import binascii
import datetime
import hashlib
import heapq
import json
@@ -14,6 +15,10 @@ from collections.abc import Callable, Iterable
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
import sqlalchemy
from sqlalchemy.dialects.postgresql import insert as pg_insert
from ..entity.persistence.cloud_directory import SpaceLaunchAssertionConsumption
if typing.TYPE_CHECKING:
from ..core.app import Application
@@ -125,12 +130,20 @@ class SpaceLaunchService:
raise SpaceLaunchError('Launch assertion payload must be a JSON object')
account_uuid = _required_string(payload, 'account_uuid')
workspace_uuid = _required_string(payload, 'workspace_uuid')
return_path = _required_string(payload, 'return_path')
if (
not return_path.startswith('/')
or return_path.startswith('//')
or any(character in return_path for character in ('\\', '\r', '\n', '\t'))
):
raise SpaceLaunchError('Launch assertion return path is invalid')
if expected_workspace_uuid is not None and workspace_uuid != expected_workspace_uuid:
raise SpaceLaunchError('Launch assertion targets another Workspace')
await self._consume_jti(_required_string(claims, 'jti'), _required_int(claims, 'exp', minimum=1))
return {
'account_uuid': account_uuid,
'workspace_uuid': workspace_uuid,
'return_path': return_path,
}
def _verify_assertion(self, token: str) -> dict[str, typing.Any]:
@@ -214,19 +227,39 @@ class SpaceLaunchService:
async def _consume_jti(self, jti: str, expires_at: int) -> None:
digest = hashlib.sha256(jti.encode('utf-8')).hexdigest()
now = int(self._wall_time())
persistence_mgr = getattr(self.ap, 'persistence_mgr', None)
instance_uuid = str(self.ap.workspace_service.instance_uuid)
if persistence_mgr is not None:
expires_at_datetime = datetime.datetime.fromtimestamp(expires_at, tz=datetime.timezone.utc)
now_datetime = datetime.datetime.fromtimestamp(now, tz=datetime.timezone.utc)
async with persistence_mgr.directory_projection_uow(instance_uuid) as uow:
await uow.session.execute(
sqlalchemy.delete(SpaceLaunchAssertionConsumption).where(
SpaceLaunchAssertionConsumption.instance_uuid == instance_uuid,
SpaceLaunchAssertionConsumption.expires_at < now_datetime,
)
)
statement = (
pg_insert(SpaceLaunchAssertionConsumption)
.values(instance_uuid=instance_uuid, jti=digest, expires_at=expires_at_datetime)
.on_conflict_do_nothing(index_elements=['instance_uuid', 'jti'])
.returning(SpaceLaunchAssertionConsumption.jti)
)
result = await uow.session.execute(statement)
if result.scalar_one_or_none() is None:
raise SpaceLaunchError('Launch assertion has already been consumed')
return
# Lightweight unit-test and OSS compatibility fallback. Verified Cloud
# runtime always supplies the durable PostgreSQL persistence manager.
async with self._replay_lock:
self._prune_consumed_jtis(now)
if digest in self._consumed_jtis:
raise SpaceLaunchError('Launch assertion has already been consumed')
if len(self._consumed_jtis) >= _CONSUMED_JTI_MAX_ENTRIES:
# Evicting a still-valid digest would make a signed launch
# assertion replayable. Bound memory by failing closed instead.
raise SpaceLaunchError('Launch assertion replay cache capacity reached')
self._consumed_jtis[digest] = expires_at
heapq.heappush(
self._consumed_jti_expiry_heap,
(expires_at, digest),
)
heapq.heappush(self._consumed_jti_expiry_heap, (expires_at, digest))
def _prune_consumed_jtis(self, now: int) -> None:
while self._consumed_jti_expiry_heap:
@@ -67,3 +67,26 @@ class DirectoryProjectionInbox(Base):
name='ck_directory_projection_inbox_fingerprint',
),
)
class SpaceLaunchAssertionConsumption(Base):
"""Durable, instance-scoped replay ledger for signed Space launch assertions."""
__tablename__ = 'space_launch_assertion_consumptions'
instance_uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
jti = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
expires_at = sqlalchemy.Column(sqlalchemy.DateTime(timezone=True), nullable=False)
consumed_at = sqlalchemy.Column(
sqlalchemy.DateTime(timezone=True),
nullable=False,
server_default=sqlalchemy.func.now(),
)
__table_args__ = (
sqlalchemy.Index(
'ix_space_launch_assertion_consumptions_expiry',
'instance_uuid',
'expires_at',
),
)
@@ -0,0 +1,57 @@
"""add durable replay protection for signed Space launch assertions
Revision ID: 0016_space_launch_replay
Revises: 0015_cloud_core_collab
Create Date: 2026-07-31
"""
from __future__ import annotations
import sqlalchemy as sa
from alembic import op
revision = '0016_space_launch_replay'
down_revision = '0015_cloud_core_collab'
branch_labels = None
depends_on = None
_TABLE = 'space_launch_assertion_consumptions'
_POLICY = 'langbot_directory_projection'
_SETTING = "NULLIF(current_setting('langbot.directory_instance_uuid', true), '')"
def upgrade() -> None:
conn = op.get_bind()
if _TABLE not in set(sa.inspect(conn).get_table_names()):
op.create_table(
_TABLE,
sa.Column('instance_uuid', sa.String(255), nullable=False),
sa.Column('jti', sa.String(255), nullable=False),
sa.Column('expires_at', sa.DateTime(timezone=True), nullable=False),
sa.Column('consumed_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
sa.PrimaryKeyConstraint('instance_uuid', 'jti'),
)
op.create_index(
'ix_space_launch_assertion_consumptions_expiry',
_TABLE,
['instance_uuid', 'expires_at'],
unique=False,
)
if conn.dialect.name == 'postgresql':
table = conn.dialect.identifier_preparer.quote(_TABLE)
policy = conn.dialect.identifier_preparer.quote(_POLICY)
expression = f'instance_uuid::text = {_SETTING}'
op.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY'))
op.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY'))
op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}'))
op.execute(
sa.text(
f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC '
f'USING ({expression}) WITH CHECK ({expression})'
)
)
def downgrade() -> None:
if _TABLE in set(sa.inspect(op.get_bind()).get_table_names()):
op.drop_table(_TABLE)
@@ -75,6 +75,7 @@ TENANT_TABLE_COLUMNS: dict[str, str] = {
DIRECTORY_PROJECTION_TABLE_COLUMNS: dict[str, str] = {
'directory_projection_states': 'instance_uuid',
'directory_projection_inbox': 'instance_uuid',
'space_launch_assertion_consumptions': 'instance_uuid',
}
DIRECTORY_PROJECTED_TENANT_TABLES = frozenset(