diff --git a/src/langbot/pkg/api/http/controller/groups/user.py b/src/langbot/pkg/api/http/controller/groups/user.py index 8fa38bf73..520d284e6 100644 --- a/src/langbot/pkg/api/http/controller/groups/user.py +++ b/src/langbot/pkg/api/http/controller/groups/user.py @@ -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: diff --git a/src/langbot/pkg/cloud/launch.py b/src/langbot/pkg/cloud/launch.py index ac5006870..3725d03d8 100644 --- a/src/langbot/pkg/cloud/launch.py +++ b/src/langbot/pkg/cloud/launch.py @@ -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: diff --git a/src/langbot/pkg/entity/persistence/cloud_directory.py b/src/langbot/pkg/entity/persistence/cloud_directory.py index 11e41cdfb..066873a14 100644 --- a/src/langbot/pkg/entity/persistence/cloud_directory.py +++ b/src/langbot/pkg/entity/persistence/cloud_directory.py @@ -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', + ), + ) diff --git a/src/langbot/pkg/persistence/alembic/versions/0016_space_launch_replay.py b/src/langbot/pkg/persistence/alembic/versions/0016_space_launch_replay.py new file mode 100644 index 000000000..33f44c19c --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0016_space_launch_replay.py @@ -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) diff --git a/src/langbot/pkg/persistence/tenant_uow.py b/src/langbot/pkg/persistence/tenant_uow.py index 27db2d64a..af472ac2d 100644 --- a/src/langbot/pkg/persistence/tenant_uow.py +++ b/src/langbot/pkg/persistence/tenant_uow.py @@ -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( diff --git a/tests/unit_tests/cloud/test_space_launch.py b/tests/unit_tests/cloud/test_space_launch.py index 14b6b7f17..07d0f80de 100644 --- a/tests/unit_tests/cloud/test_space_launch.py +++ b/tests/unit_tests/cloud/test_space_launch.py @@ -48,6 +48,7 @@ def _claims(*, now: int, jti: str | None = None, workspace_uuid: str = WORKSPACE 'payload': { 'account_uuid': ACCOUNT_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) - 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'): 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 with pytest.raises(SpaceLaunchError, match='verified Cloud mode'): 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)) diff --git a/web/src/app/auth/space/callback/page.tsx b/web/src/app/auth/space/callback/page.tsx index 9b3f9112a..cd87bdc56 100644 --- a/web/src/app/auth/space/callback/page.tsx +++ b/web/src/app/auth/space/callback/page.tsx @@ -29,6 +29,7 @@ type SpaceOAuthLoginResult = { token: string; user: string; workspace_uuid?: string; + return_path?: string; }; const pendingSpaceOAuthLogins = new Map< @@ -63,6 +64,10 @@ function SpaceOAuthCallbackContent() { const [searchParams] = useSearchParams(); const { t } = useTranslation(); const isMountedRef = useRef(true); + const directLaunchFragmentRef = useRef<{ + workspaceUuid: string | null; + launchAssertion: string | null; + } | null>(null); const [status, setStatus] = useState< 'loading' | 'confirm' | 'success' | 'error' @@ -106,7 +111,13 @@ function SpaceOAuthCallbackContent() { throw new Error('No Workspace is available for this Account'); } 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; } setStatus('success'); @@ -207,16 +218,29 @@ function SpaceOAuthCallbackContent() { const errorDescription = searchParams.get('error_description'); const mode = searchParams.get('mode'); const state = searchParams.get('state'); - const fragmentParams = new URLSearchParams( - window.location.hash.startsWith('#') - ? window.location.hash.slice(1) - : window.location.hash, - ); + if (directLaunchFragmentRef.current === null) { + const fragmentParams = new URLSearchParams( + window.location.hash.startsWith('#') + ? 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 = - fragmentParams.get('workspace_uuid') ?? searchParams.get('workspace_uuid'); + directLaunchFragmentRef.current.workspaceUuid ?? + searchParams.get('workspace_uuid'); const launchAssertion = - fragmentParams.get('launch_assertion') ?? - searchParams.get('launch_assertion'); + directLaunchFragmentRef.current.launchAssertion; if (error) { setStatus('error'); diff --git a/web/tests/unit/space-launch-callback.test.mjs b/web/tests/unit/space-launch-callback.test.mjs new file mode 100644 index 000000000..694cbcc6d --- /dev/null +++ b/web/tests/unit/space-launch-callback.test.mjs @@ -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 \}\)/); +});