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, 'token': token,
'user': account.user, 'user': account.user,
'workspace_uuid': access.workspace.uuid, 'workspace_uuid': access.workspace.uuid,
'return_path': launch.get('return_path', '/home'),
} }
) )
except SpaceLaunchError: except SpaceLaunchError:
+39 -6
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import asyncio import asyncio
import base64 import base64
import binascii import binascii
import datetime
import hashlib import hashlib
import heapq import heapq
import json import json
@@ -14,6 +15,10 @@ from collections.abc import Callable, Iterable
from cryptography.exceptions import InvalidSignature from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey 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: if typing.TYPE_CHECKING:
from ..core.app import Application from ..core.app import Application
@@ -125,12 +130,20 @@ class SpaceLaunchService:
raise SpaceLaunchError('Launch assertion payload must be a JSON object') raise SpaceLaunchError('Launch assertion payload must be a JSON object')
account_uuid = _required_string(payload, 'account_uuid') account_uuid = _required_string(payload, 'account_uuid')
workspace_uuid = _required_string(payload, 'workspace_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: if expected_workspace_uuid is not None and workspace_uuid != expected_workspace_uuid:
raise SpaceLaunchError('Launch assertion targets another Workspace') raise SpaceLaunchError('Launch assertion targets another Workspace')
await self._consume_jti(_required_string(claims, 'jti'), _required_int(claims, 'exp', minimum=1)) await self._consume_jti(_required_string(claims, 'jti'), _required_int(claims, 'exp', minimum=1))
return { return {
'account_uuid': account_uuid, 'account_uuid': account_uuid,
'workspace_uuid': workspace_uuid, 'workspace_uuid': workspace_uuid,
'return_path': return_path,
} }
def _verify_assertion(self, token: str) -> dict[str, typing.Any]: 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: async def _consume_jti(self, jti: str, expires_at: int) -> None:
digest = hashlib.sha256(jti.encode('utf-8')).hexdigest() digest = hashlib.sha256(jti.encode('utf-8')).hexdigest()
now = int(self._wall_time()) 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: async with self._replay_lock:
self._prune_consumed_jtis(now) self._prune_consumed_jtis(now)
if digest in self._consumed_jtis: if digest in self._consumed_jtis:
raise SpaceLaunchError('Launch assertion has already been consumed') raise SpaceLaunchError('Launch assertion has already been consumed')
if len(self._consumed_jtis) >= _CONSUMED_JTI_MAX_ENTRIES: 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') raise SpaceLaunchError('Launch assertion replay cache capacity reached')
self._consumed_jtis[digest] = expires_at self._consumed_jtis[digest] = expires_at
heapq.heappush( heapq.heappush(self._consumed_jti_expiry_heap, (expires_at, digest))
self._consumed_jti_expiry_heap,
(expires_at, digest),
)
def _prune_consumed_jtis(self, now: int) -> None: def _prune_consumed_jtis(self, now: int) -> None:
while self._consumed_jti_expiry_heap: while self._consumed_jti_expiry_heap:
@@ -67,3 +67,26 @@ class DirectoryProjectionInbox(Base):
name='ck_directory_projection_inbox_fingerprint', 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_TABLE_COLUMNS: dict[str, str] = {
'directory_projection_states': 'instance_uuid', 'directory_projection_states': 'instance_uuid',
'directory_projection_inbox': 'instance_uuid', 'directory_projection_inbox': 'instance_uuid',
'space_launch_assertion_consumptions': 'instance_uuid',
} }
DIRECTORY_PROJECTED_TENANT_TABLES = frozenset( DIRECTORY_PROJECTED_TENANT_TABLES = frozenset(
+18 -1
View File
@@ -48,6 +48,7 @@ def _claims(*, now: int, jti: str | None = None, workspace_uuid: str = WORKSPACE
'payload': { 'payload': {
'account_uuid': ACCOUNT_UUID, 'account_uuid': ACCOUNT_UUID,
'workspace_uuid': workspace_uuid, 'workspace_uuid': workspace_uuid,
'return_path': '/',
}, },
} }
@@ -81,7 +82,11 @@ async def test_consumes_valid_workspace_launch_assertion_once():
launch = await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID) launch = await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
assert launch == {'account_uuid': ACCOUNT_UUID, 'workspace_uuid': WORKSPACE_UUID} assert launch == {
'account_uuid': ACCOUNT_UUID,
'workspace_uuid': WORKSPACE_UUID,
'return_path': '/',
}
with pytest.raises(SpaceLaunchError, match='already been consumed'): with pytest.raises(SpaceLaunchError, match='already been consumed'):
await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID) await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
@@ -162,3 +167,15 @@ async def test_rejects_invalid_signature_and_non_cloud_mode():
oss_service.ap.deployment.multi_workspace_enabled = False oss_service.ap.deployment.multi_workspace_enabled = False
with pytest.raises(SpaceLaunchError, match='verified Cloud mode'): with pytest.raises(SpaceLaunchError, match='verified Cloud mode'):
await oss_service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID) await oss_service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
@pytest.mark.asyncio
async def test_rejects_unsafe_signed_return_path() -> None:
private_key = Ed25519PrivateKey.generate()
now = int(time.time())
service = _service(private_key, now=now)
claims = _claims(now=now)
claims['payload']['return_path'] = '//evil.example'
with pytest.raises(SpaceLaunchError, match='return path'):
await service.consume_assertion(_sign(private_key, claims))
+33 -9
View File
@@ -29,6 +29,7 @@ type SpaceOAuthLoginResult = {
token: string; token: string;
user: string; user: string;
workspace_uuid?: string; workspace_uuid?: string;
return_path?: string;
}; };
const pendingSpaceOAuthLogins = new Map< const pendingSpaceOAuthLogins = new Map<
@@ -63,6 +64,10 @@ function SpaceOAuthCallbackContent() {
const [searchParams] = useSearchParams(); const [searchParams] = useSearchParams();
const { t } = useTranslation(); const { t } = useTranslation();
const isMountedRef = useRef(true); const isMountedRef = useRef(true);
const directLaunchFragmentRef = useRef<{
workspaceUuid: string | null;
launchAssertion: string | null;
} | null>(null);
const [status, setStatus] = useState< const [status, setStatus] = useState<
'loading' | 'confirm' | 'success' | 'error' 'loading' | 'confirm' | 'success' | 'error'
@@ -106,7 +111,13 @@ function SpaceOAuthCallbackContent() {
throw new Error('No Workspace is available for this Account'); throw new Error('No Workspace is available for this Account');
} }
if (response.workspace_uuid) { if (response.workspace_uuid) {
navigate('/home', { replace: true }); const returnPath =
typeof response.return_path === 'string' &&
response.return_path.startsWith('/') &&
!response.return_path.startsWith('//')
? response.return_path
: '/home';
navigate(returnPath, { replace: true });
return; return;
} }
setStatus('success'); setStatus('success');
@@ -207,16 +218,29 @@ function SpaceOAuthCallbackContent() {
const errorDescription = searchParams.get('error_description'); const errorDescription = searchParams.get('error_description');
const mode = searchParams.get('mode'); const mode = searchParams.get('mode');
const state = searchParams.get('state'); const state = searchParams.get('state');
const fragmentParams = new URLSearchParams( if (directLaunchFragmentRef.current === null) {
window.location.hash.startsWith('#') const fragmentParams = new URLSearchParams(
? window.location.hash.slice(1) window.location.hash.startsWith('#')
: window.location.hash, ? window.location.hash.slice(1)
); : window.location.hash,
);
directLaunchFragmentRef.current = {
workspaceUuid: fragmentParams.get('workspace_uuid'),
launchAssertion: fragmentParams.get('launch_assertion'),
};
if (window.location.hash) {
window.history.replaceState(
null,
'',
`${window.location.pathname}${window.location.search}`,
);
}
}
const workspaceUuid = const workspaceUuid =
fragmentParams.get('workspace_uuid') ?? searchParams.get('workspace_uuid'); directLaunchFragmentRef.current.workspaceUuid ??
searchParams.get('workspace_uuid');
const launchAssertion = const launchAssertion =
fragmentParams.get('launch_assertion') ?? directLaunchFragmentRef.current.launchAssertion;
searchParams.get('launch_assertion');
if (error) { if (error) {
setStatus('error'); setStatus('error');
@@ -0,0 +1,24 @@
import assert from 'node:assert/strict';
import fs from 'node:fs';
import test from 'node:test';
const source = fs.readFileSync(
new URL('../../src/app/auth/space/callback/page.tsx', import.meta.url),
'utf8',
);
test('direct launch assertion is fragment-only and removed before exchange', () => {
assert.doesNotMatch(source, /searchParams\.get\(['"]launch_assertion['"]\)/);
const readIndex = source.indexOf("fragmentParams.get('launch_assertion')");
const clearIndex = source.indexOf('window.history.replaceState');
const exchangeIndex = source.indexOf('handleOAuthCallback(', clearIndex);
assert.ok(readIndex >= 0, 'fragment assertion read is missing');
assert.ok(clearIndex > readIndex, 'URL fragment is not cleared after copying the assertion');
assert.ok(exchangeIndex > clearIndex, 'assertion exchange starts before the fragment is cleared');
});
test('direct launch honors only a local signed return path', () => {
assert.match(source, /response\.return_path\.startsWith\(['"]\/['"]\)/);
assert.match(source, /!response\.return_path\.startsWith\(['"]\/\/['"]\)/);
assert.match(source, /navigate\(returnPath, \{ replace: true \}\)/);
});