mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
feat(cloud): add scoped support admin sessions
This commit is contained in:
dadachann
@@ -0,0 +1,453 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import datetime
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import sqlalchemy
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
|
||||
from quart import Quart
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
|
||||
from langbot.pkg.api.http.authz import Permission
|
||||
from langbot.pkg.api.http.context import PrincipalType, RequestContext
|
||||
from langbot.pkg.api.http.controller import group
|
||||
from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import WebSocketChatRouterGroup
|
||||
from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
|
||||
from langbot.pkg.cloud.launch import SpaceLaunchError, SpaceLaunchService
|
||||
from langbot.pkg.cloud.support_admin import SupportAdminSessionService
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.entity.persistence.support_admin import SupportAdminTemporarySession
|
||||
from langbot.pkg.entity.persistence.user import User
|
||||
from langbot.pkg.entity.persistence.workspace import (
|
||||
Workspace,
|
||||
WorkspaceExecutionState,
|
||||
WorkspaceMembership,
|
||||
)
|
||||
from langbot.pkg.workspace.service import WorkspaceService
|
||||
|
||||
|
||||
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
|
||||
|
||||
|
||||
INSTANCE_UUID = 'instance-support-admin'
|
||||
WORKSPACE_UUID = '10000000-0000-4000-8000-000000000001'
|
||||
OTHER_WORKSPACE_UUID = '10000000-0000-4000-8000-000000000002'
|
||||
ACTOR_ACCOUNT_UUID = '20000000-0000-4000-8000-000000000001'
|
||||
KEY_ID = 'support-admin-key-1'
|
||||
|
||||
|
||||
def _base64url(raw: bytes) -> str:
|
||||
return base64.urlsafe_b64encode(raw).rstrip(b'=').decode('ascii')
|
||||
|
||||
|
||||
def _sign(private_key: Ed25519PrivateKey, claims: dict, *, key_id: str = KEY_ID) -> str:
|
||||
header = {'alg': 'EdDSA', 'kid': key_id, 'typ': 'langbot-control-plane+jwt'}
|
||||
encoded_header = _base64url(json.dumps(header, separators=(',', ':')).encode('utf-8'))
|
||||
encoded_claims = _base64url(json.dumps(claims, separators=(',', ':')).encode('utf-8'))
|
||||
signing_input = f'{encoded_header}.{encoded_claims}'
|
||||
return f'{signing_input}.{_base64url(private_key.sign(signing_input.encode("ascii")))}'
|
||||
|
||||
|
||||
def _admin_claims(*, now: int, jti: str | None = None, workspace_uuid: str = WORKSPACE_UUID) -> dict:
|
||||
return {
|
||||
'iss': 'langbot-space',
|
||||
'aud': 'langbot-cloud-runtime',
|
||||
'sub': f'langbot-instance:{INSTANCE_UUID}',
|
||||
'jti': jti or str(uuid.uuid4()),
|
||||
'iat': now,
|
||||
'nbf': now - 5,
|
||||
'exp': now + 90,
|
||||
'instance_uuid': INSTANCE_UUID,
|
||||
'kind': 'workspace.support_admin_launch',
|
||||
'payload': {
|
||||
'workspace_uuid': workspace_uuid,
|
||||
'launch_mode': 'support_admin',
|
||||
'principal_type': 'support_admin',
|
||||
'actor_account_uuid': ACTOR_ACCOUNT_UUID,
|
||||
'effective_role': 'owner',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@group.group_class('support_admin_probe', '/api/v1/support-admin-probe')
|
||||
class SupportAdminProbeGroup(group.RouterGroup):
|
||||
async def initialize(self) -> None:
|
||||
@self.route('/user-token', auth_type=group.AuthType.USER_TOKEN, permission=Permission.WORKSPACE_VIEW)
|
||||
async def _(request_context: RequestContext) -> str:
|
||||
return self.success(data=_context_payload(request_context))
|
||||
|
||||
@self.route(
|
||||
'/member-operation',
|
||||
auth_type=group.AuthType.USER_TOKEN,
|
||||
permission=Permission.MEMBER_VIEW,
|
||||
)
|
||||
async def member_operation(request_context: RequestContext) -> str:
|
||||
return self.success(data=_context_payload(request_context))
|
||||
|
||||
@self.route(
|
||||
'/user-token-or-api-key',
|
||||
auth_type=group.AuthType.USER_TOKEN_OR_API_KEY,
|
||||
permission=Permission.WORKSPACE_VIEW,
|
||||
)
|
||||
async def _(request_context: RequestContext) -> str:
|
||||
return self.success(data=_context_payload(request_context))
|
||||
|
||||
|
||||
def _context_payload(request_context: RequestContext) -> dict:
|
||||
return {
|
||||
'principal_type': request_context.principal.principal_type.value,
|
||||
'actor_account_uuid': request_context.principal.actor_account_uuid,
|
||||
'account_uuid': request_context.principal.account_uuid,
|
||||
'role': request_context.workspace.role,
|
||||
'membership_uuid': request_context.workspace.membership_uuid,
|
||||
'permissions': sorted(request_context.workspace.permissions),
|
||||
}
|
||||
|
||||
|
||||
class _TenantUow:
|
||||
def __init__(self, engine):
|
||||
self._engine = engine
|
||||
self.session = None
|
||||
self._transaction = None
|
||||
|
||||
async def __aenter__(self):
|
||||
session_factory = async_sessionmaker(self._engine, expire_on_commit=False)
|
||||
self.session = session_factory()
|
||||
self._transaction = await self.session.begin()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, traceback):
|
||||
try:
|
||||
if exc_type is None:
|
||||
await self._transaction.commit()
|
||||
else:
|
||||
await self._transaction.rollback()
|
||||
finally:
|
||||
await self.session.close()
|
||||
|
||||
|
||||
class _TenantScope:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, traceback):
|
||||
return False
|
||||
|
||||
|
||||
class _PersistenceManager:
|
||||
def __init__(self, engine):
|
||||
self._engine = engine
|
||||
self.mode = SimpleNamespace(value='oss_compat')
|
||||
|
||||
def get_db_engine(self):
|
||||
return self._engine
|
||||
|
||||
def tenant_uow(self, workspace_uuid: str):
|
||||
del workspace_uuid
|
||||
return _TenantUow(self._engine)
|
||||
|
||||
def tenant_scope(self, workspace_uuid: str):
|
||||
del workspace_uuid
|
||||
return _TenantScope()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def support_admin_api(tmp_path):
|
||||
private_key = Ed25519PrivateKey.generate()
|
||||
public_key = private_key.public_key().public_bytes(
|
||||
encoding=serialization.Encoding.Raw,
|
||||
format=serialization.PublicFormat.Raw,
|
||||
)
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "support-admin.db"}')
|
||||
async with engine.begin() as connection:
|
||||
await connection.run_sync(
|
||||
Base.metadata.create_all,
|
||||
tables=[
|
||||
User.__table__,
|
||||
Workspace.__table__,
|
||||
WorkspaceExecutionState.__table__,
|
||||
WorkspaceMembership.__table__,
|
||||
SupportAdminTemporarySession.__table__,
|
||||
],
|
||||
)
|
||||
for workspace_uuid, slug in (
|
||||
(WORKSPACE_UUID, 'support-admin-a'),
|
||||
(OTHER_WORKSPACE_UUID, 'support-admin-b'),
|
||||
):
|
||||
await connection.execute(
|
||||
sqlalchemy.insert(Workspace).values(
|
||||
uuid=workspace_uuid,
|
||||
instance_uuid=INSTANCE_UUID,
|
||||
name=slug,
|
||||
slug=slug,
|
||||
type='team',
|
||||
status='active',
|
||||
source='cloud_projection',
|
||||
projection_revision=1,
|
||||
)
|
||||
)
|
||||
await connection.execute(
|
||||
sqlalchemy.insert(WorkspaceExecutionState).values(
|
||||
workspace_uuid=workspace_uuid,
|
||||
instance_uuid=INSTANCE_UUID,
|
||||
active_generation=1,
|
||||
state='active',
|
||||
write_fenced=False,
|
||||
source='cloud',
|
||||
desired_state_revision=1,
|
||||
)
|
||||
)
|
||||
|
||||
app = SimpleNamespace()
|
||||
app.persistence_mgr = _PersistenceManager(engine)
|
||||
app.instance_config = SimpleNamespace(
|
||||
data={
|
||||
'system': {
|
||||
'jwt': {'secret': 'support-admin-secret', 'expire': 3600},
|
||||
'websocket_retention': {},
|
||||
},
|
||||
'space': {
|
||||
'launch': {
|
||||
'control_plane_public_key': _base64url(public_key),
|
||||
}
|
||||
},
|
||||
'api': {'global_api_key': ''},
|
||||
}
|
||||
)
|
||||
app.logger = logging.getLogger('support-admin-test')
|
||||
app.deployment = SimpleNamespace(mode='cloud', multi_workspace_enabled=True, verification_key_id=KEY_ID)
|
||||
app.directory_projection_service = SimpleNamespace(require_ready=lambda: None)
|
||||
app.workspace_service = WorkspaceService(app, instance_uuid=INSTANCE_UUID)
|
||||
app.entitlement_resolver = SimpleNamespace(
|
||||
instance_uuid=INSTANCE_UUID,
|
||||
resolve=AsyncMock(return_value=SimpleNamespace(entitlement_revision=7)),
|
||||
)
|
||||
app.support_admin_session_service = SupportAdminSessionService(app)
|
||||
app.space_launch_service = SpaceLaunchService(app)
|
||||
app.user_service = SimpleNamespace()
|
||||
app.user_service.get_authenticated_account = AsyncMock(side_effect=AssertionError('normal account auth used'))
|
||||
app.user_service.verify_jwt_token = AsyncMock(side_effect=AssertionError('normal token verification used'))
|
||||
app.user_service.get_user_by_email = AsyncMock(side_effect=AssertionError('user lookup used'))
|
||||
app.apikey_service = SimpleNamespace()
|
||||
app.apikey_service.authenticate_api_key = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
instance_uuid=INSTANCE_UUID,
|
||||
workspace_uuid=OTHER_WORKSPACE_UUID,
|
||||
placement_generation=1,
|
||||
api_key_uuid='api-key',
|
||||
permissions=frozenset(permission.value for permission in Permission),
|
||||
)
|
||||
)
|
||||
|
||||
quart_app = Quart(__name__)
|
||||
await UserRouterGroup(app, quart_app).initialize()
|
||||
await SupportAdminProbeGroup(app, quart_app).initialize()
|
||||
|
||||
yield app, quart_app.test_client(), engine, private_key
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def _issue_support_token(app, private_key: Ed25519PrivateKey, *, jti: str | None = None) -> dict[str, str]:
|
||||
launch = await app.space_launch_service.consume_assertion(
|
||||
_sign(private_key, _admin_claims(now=int(time.time()), jti=jti)),
|
||||
expected_workspace_uuid=WORKSPACE_UUID,
|
||||
)
|
||||
return launch
|
||||
|
||||
|
||||
def _auth(token: str, workspace_uuid: str = WORKSPACE_UUID) -> dict[str, str]:
|
||||
return {'Authorization': f'Bearer {token}', 'X-Workspace-Id': workspace_uuid}
|
||||
|
||||
|
||||
async def test_support_admin_membership_only_routes_are_denied(support_admin_api):
|
||||
app, client, _engine, private_key = support_admin_api
|
||||
launch = await _issue_support_token(app, private_key)
|
||||
|
||||
response = await client.get(
|
||||
'/api/v1/support-admin-probe/member-operation',
|
||||
headers=_auth(launch['support_admin_token']),
|
||||
)
|
||||
|
||||
assert response.status_code == 403
|
||||
assert (await response.get_json())['code'] == 'permission_denied'
|
||||
|
||||
|
||||
async def test_support_admin_check_token_is_rejected(support_admin_api):
|
||||
app, client, _engine, private_key = support_admin_api
|
||||
launch = await _issue_support_token(app, private_key)
|
||||
|
||||
response = await client.get('/api/v1/user/check-token', headers=_auth(launch['support_admin_token']))
|
||||
|
||||
assert response.status_code == 401
|
||||
assert (await response.get_json())['code'] == 'invalid_authentication'
|
||||
|
||||
|
||||
async def test_support_admin_cross_workspace_denied_for_user_token_and_or_api_key(support_admin_api):
|
||||
app, client, _engine, private_key = support_admin_api
|
||||
launch = await _issue_support_token(app, private_key)
|
||||
|
||||
missing_selector = await client.get(
|
||||
'/api/v1/support-admin-probe/user-token',
|
||||
headers={'Authorization': f'Bearer {launch["support_admin_token"]}'},
|
||||
)
|
||||
user_response = await client.get(
|
||||
'/api/v1/support-admin-probe/user-token',
|
||||
headers=_auth(launch['support_admin_token'], OTHER_WORKSPACE_UUID),
|
||||
)
|
||||
either_response = await client.get(
|
||||
'/api/v1/support-admin-probe/user-token-or-api-key',
|
||||
headers={
|
||||
**_auth(launch['support_admin_token'], OTHER_WORKSPACE_UUID),
|
||||
'X-API-Key': 'valid-api-key',
|
||||
},
|
||||
)
|
||||
|
||||
assert missing_selector.status_code == 400
|
||||
assert user_response.status_code == 401
|
||||
assert either_response.status_code == 401
|
||||
app.apikey_service.authenticate_api_key.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_support_admin_request_context_has_actor_owner_and_no_membership(support_admin_api):
|
||||
app, client, engine, private_key = support_admin_api
|
||||
before_count = await _membership_count(engine)
|
||||
launch = await _issue_support_token(app, private_key)
|
||||
|
||||
response = await client.get(
|
||||
'/api/v1/support-admin-probe/user-token',
|
||||
headers=_auth(launch['support_admin_token']),
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = (await response.get_json())['data']
|
||||
permissions = set(data.pop('permissions'))
|
||||
assert Permission.WORKSPACE_VIEW.value in permissions
|
||||
assert Permission.RESOURCE_MANAGE.value in permissions
|
||||
assert not permissions.intersection(
|
||||
{
|
||||
Permission.OWNER_TRANSFER.value,
|
||||
Permission.MEMBER_VIEW.value,
|
||||
Permission.MEMBER_INVITE.value,
|
||||
Permission.MEMBER_UPDATE_ROLE.value,
|
||||
Permission.MEMBER_REMOVE.value,
|
||||
}
|
||||
)
|
||||
assert data == {
|
||||
'principal_type': PrincipalType.SUPPORT_ADMIN.value,
|
||||
'actor_account_uuid': ACTOR_ACCOUNT_UUID,
|
||||
'account_uuid': None,
|
||||
'role': 'owner',
|
||||
'membership_uuid': None,
|
||||
}
|
||||
assert await _membership_count(engine) == before_count
|
||||
|
||||
|
||||
async def test_support_admin_missing_workspace_is_controlled_launch_failure(support_admin_api):
|
||||
app, _client, engine, private_key = support_admin_api
|
||||
async with engine.begin() as connection:
|
||||
await connection.execute(
|
||||
sqlalchemy.delete(WorkspaceExecutionState).where(WorkspaceExecutionState.workspace_uuid == WORKSPACE_UUID)
|
||||
)
|
||||
|
||||
with pytest.raises(SpaceLaunchError, match='unavailable'):
|
||||
await _issue_support_token(app, private_key)
|
||||
|
||||
|
||||
async def test_support_admin_launch_replay_is_durable_across_service_instances(support_admin_api):
|
||||
app, _client, _engine, private_key = support_admin_api
|
||||
jti = str(uuid.uuid4())
|
||||
|
||||
await _issue_support_token(app, private_key, jti=jti)
|
||||
second_service = SpaceLaunchService(app)
|
||||
|
||||
with pytest.raises(SpaceLaunchError, match='already been consumed'):
|
||||
await second_service.consume_assertion(
|
||||
_sign(private_key, _admin_claims(now=int(time.time()), jti=jti)),
|
||||
expected_workspace_uuid=WORKSPACE_UUID,
|
||||
)
|
||||
|
||||
|
||||
async def test_support_admin_persisted_expiry_and_revocation_are_enforced(support_admin_api):
|
||||
app, client, engine, private_key = support_admin_api
|
||||
launch = await _issue_support_token(app, private_key)
|
||||
token = launch['support_admin_token']
|
||||
|
||||
async with engine.begin() as connection:
|
||||
await connection.execute(
|
||||
sqlalchemy.update(SupportAdminTemporarySession)
|
||||
.where(SupportAdminTemporarySession.grant_jti_hash == launch['grant_jti_hash'])
|
||||
.values(expires_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None) - datetime.timedelta(minutes=1))
|
||||
)
|
||||
expired = await client.get('/api/v1/support-admin-probe/user-token', headers=_auth(token))
|
||||
assert expired.status_code == 401
|
||||
|
||||
second = await _issue_support_token(app, private_key)
|
||||
async with engine.begin() as connection:
|
||||
await connection.execute(
|
||||
sqlalchemy.update(SupportAdminTemporarySession)
|
||||
.where(SupportAdminTemporarySession.grant_jti_hash == second['grant_jti_hash'])
|
||||
.values(revoked_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None))
|
||||
)
|
||||
revoked = await client.get('/api/v1/support-admin-probe/user-token', headers=_auth(second['support_admin_token']))
|
||||
assert revoked.status_code == 401
|
||||
|
||||
|
||||
async def test_support_admin_websocket_preserves_actor_and_revalidates(support_admin_api):
|
||||
app, _client, _engine, private_key = support_admin_api
|
||||
launch = await _issue_support_token(app, private_key)
|
||||
captured_contexts = []
|
||||
|
||||
class Adapter:
|
||||
async def handle_websocket_message(self, connection, data):
|
||||
del data
|
||||
captured_contexts.append(connection.execution_context)
|
||||
await connection.send_queue.put({'type': 'handled'})
|
||||
connection.is_active = False
|
||||
|
||||
app.pipeline_service = SimpleNamespace(get_pipeline=AsyncMock(return_value=SimpleNamespace(uuid='pipeline-1')))
|
||||
app.platform_mgr = SimpleNamespace(
|
||||
get_websocket_proxy_bot=AsyncMock(return_value=SimpleNamespace(adapter=Adapter()))
|
||||
)
|
||||
|
||||
quart_app = Quart(__name__)
|
||||
await WebSocketChatRouterGroup(app, quart_app).initialize()
|
||||
|
||||
async with quart_app.test_client().websocket('/api/v1/pipelines/pipeline-1/ws/connect') as websocket:
|
||||
await websocket.send(
|
||||
json.dumps(
|
||||
{
|
||||
'type': 'authenticate',
|
||||
'token': launch['support_admin_token'],
|
||||
'workspace_uuid': WORKSPACE_UUID,
|
||||
}
|
||||
)
|
||||
)
|
||||
connected = json.loads(await websocket.receive())
|
||||
assert connected['type'] == 'connected'
|
||||
await websocket.send(json.dumps({'type': 'message', 'message': [{'type': 'text', 'text': 'hi'}]}))
|
||||
handled = json.loads(await websocket.receive())
|
||||
assert handled['type'] == 'handled'
|
||||
|
||||
assert captured_contexts
|
||||
principal = captured_contexts[0].trigger_principal
|
||||
assert principal is not None
|
||||
assert principal.principal_type == PrincipalType.SUPPORT_ADMIN
|
||||
assert principal.actor_account_uuid == ACTOR_ACCOUNT_UUID
|
||||
|
||||
|
||||
async def _membership_count(engine) -> int:
|
||||
async with engine.connect() as connection:
|
||||
return int(
|
||||
await connection.scalar(
|
||||
sqlalchemy.select(sqlalchemy.func.count()).select_from(WorkspaceMembership),
|
||||
)
|
||||
or 0
|
||||
)
|
||||
Reference in New Issue
Block a user