mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-21 20:06:06 +00:00
Junyan Qin
467 lines
19 KiB
Python
467 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import datetime
|
|
import hashlib
|
|
import logging
|
|
import uuid
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
import sqlalchemy
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from langbot.pkg.api.http.authz import Permission, PermissionDeniedError
|
|
from langbot.pkg.api.http.context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext
|
|
from langbot.pkg.api.http.service.apikey import ApiKeyService
|
|
from langbot.pkg.entity.persistence.apikey import ApiKey
|
|
from langbot.pkg.entity.persistence.base import Base
|
|
from langbot.pkg.entity.persistence.user import User
|
|
from langbot.pkg.entity.persistence.workspace import (
|
|
Workspace,
|
|
WorkspaceExecutionSource,
|
|
WorkspaceExecutionState,
|
|
WorkspaceSource,
|
|
)
|
|
from langbot.pkg.workspace.policy import SingleWorkspacePolicy
|
|
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
|
|
from langbot.pkg.workspace.service import WorkspaceService
|
|
|
|
|
|
pytestmark = pytest.mark.asyncio
|
|
|
|
|
|
class _PersistenceManager:
|
|
def __init__(self, engine):
|
|
self.engine = engine
|
|
|
|
def get_db_engine(self):
|
|
return self.engine
|
|
|
|
async def execute_async(self, *args, **kwargs):
|
|
async with self.engine.connect() as connection:
|
|
result = await connection.execute(*args, **kwargs)
|
|
await connection.commit()
|
|
return result
|
|
|
|
@staticmethod
|
|
def serialize_model(model, row, masked_columns=()):
|
|
return {
|
|
column.name: (
|
|
getattr(row, column.name).isoformat()
|
|
if isinstance(getattr(row, column.name), datetime.datetime)
|
|
else getattr(row, column.name)
|
|
)
|
|
for column in model.__table__.columns
|
|
if column.name not in masked_columns
|
|
}
|
|
|
|
|
|
def _context(workspace_uuid: str, account_uuid: str, permissions: set[Permission]) -> RequestContext:
|
|
return RequestContext(
|
|
instance_uuid='api-key-instance',
|
|
placement_generation=1,
|
|
request_id=str(uuid.uuid4()),
|
|
auth_type='user-token',
|
|
principal=PrincipalContext(PrincipalType.ACCOUNT, account_uuid=account_uuid),
|
|
workspace=WorkspaceContext(
|
|
workspace_uuid=workspace_uuid,
|
|
membership_uuid=str(uuid.uuid4()),
|
|
role='owner',
|
|
permissions=frozenset(permission.value for permission in permissions),
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
async def api_key_context(tmp_path):
|
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "api-keys.db"}')
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
|
|
application = SimpleNamespace(
|
|
persistence_mgr=_PersistenceManager(engine),
|
|
instance_config=SimpleNamespace(data={'api': {'global_api_key': ''}}),
|
|
logger=logging.getLogger('api-key-test'),
|
|
)
|
|
application.workspace_service = WorkspaceService(application, instance_uuid='api-key-instance')
|
|
workspace = await application.workspace_service.ensure_singleton_workspace()
|
|
account_uuid = str(uuid.uuid4())
|
|
session_factory = async_sessionmaker(engine, expire_on_commit=False)
|
|
async with session_factory.begin() as session:
|
|
session.add(
|
|
User(
|
|
uuid=account_uuid,
|
|
user='owner@example.com',
|
|
normalized_email='owner@example.com',
|
|
password='hash',
|
|
account_type='local',
|
|
)
|
|
)
|
|
service = ApiKeyService(application)
|
|
context = _context(workspace.uuid, account_uuid, set(Permission))
|
|
yield application, service, context, engine
|
|
await engine.dispose()
|
|
|
|
|
|
async def test_secret_is_returned_once_and_only_hash_is_persisted(api_key_context):
|
|
_application, service, context, engine = api_key_context
|
|
|
|
created = await service.create_api_key(context, 'Automation', 'CI key')
|
|
secret = created['key']
|
|
assert secret.startswith('lbk_')
|
|
assert created['secret_available'] is True
|
|
assert 'key_hash' not in created
|
|
|
|
listed = await service.get_api_keys(context)
|
|
assert len(listed) == 1
|
|
assert 'key' not in listed[0]
|
|
assert 'key_hash' not in listed[0]
|
|
assert listed[0]['secret_available'] is False
|
|
|
|
async with engine.connect() as connection:
|
|
stored = await connection.scalar(sqlalchemy.select(ApiKey.key_hash))
|
|
assert stored == hashlib.sha256(secret.encode()).hexdigest()
|
|
assert secret not in stored
|
|
|
|
|
|
async def test_authentication_derives_workspace_scopes_and_updates_usage(api_key_context):
|
|
_application, service, context, engine = api_key_context
|
|
created = await service.create_api_key(
|
|
context,
|
|
'Read only',
|
|
scopes=[Permission.RESOURCE_VIEW.value],
|
|
)
|
|
|
|
identity = await service.authenticate_api_key(created['key'])
|
|
assert identity is not None
|
|
assert identity.workspace_uuid == context.workspace_uuid
|
|
assert identity.permissions == frozenset({Permission.RESOURCE_VIEW.value})
|
|
|
|
async with engine.connect() as connection:
|
|
last_used_at = await connection.scalar(sqlalchemy.select(ApiKey.last_used_at))
|
|
assert last_used_at is not None
|
|
|
|
|
|
async def test_revoked_expired_and_unknown_keys_fail_closed(api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
created = await service.create_api_key(context, 'Revocable')
|
|
await service.delete_api_key(context, created['id'])
|
|
assert await service.authenticate_api_key(created['key']) is None
|
|
assert await service.verify_api_key('') is False
|
|
assert await service.verify_api_key('plain-secret') is False
|
|
assert await service.verify_api_key('lbk_unknown') is False
|
|
|
|
expired_secret = 'lbk_expired'
|
|
await service.ap.persistence_mgr.execute_async(
|
|
sqlalchemy.insert(ApiKey).values(
|
|
workspace_uuid=context.workspace_uuid,
|
|
name='Expired',
|
|
key_hash=hashlib.sha256(expired_secret.encode()).hexdigest(),
|
|
scopes=[Permission.RESOURCE_VIEW.value],
|
|
status='active',
|
|
expires_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None) - datetime.timedelta(seconds=1),
|
|
)
|
|
)
|
|
assert await service.authenticate_api_key(expired_secret) is None
|
|
|
|
|
|
async def test_revoke_winning_last_used_update_race_fails_authentication(api_key_context):
|
|
application, service, context, _engine = api_key_context
|
|
created = await service.create_api_key(context, 'Racing revoke')
|
|
original_execute = application.persistence_mgr.execute_async
|
|
injected_revoke = False
|
|
|
|
async def execute_with_revoke(statement, *args, **kwargs):
|
|
nonlocal injected_revoke
|
|
if (
|
|
not injected_revoke
|
|
and isinstance(statement, sqlalchemy.sql.dml.Update)
|
|
and statement.table.name == ApiKey.__tablename__
|
|
):
|
|
injected_revoke = True
|
|
await original_execute(sqlalchemy.update(ApiKey).where(ApiKey.id == created['id']).values(status='revoked'))
|
|
return await original_execute(statement, *args, **kwargs)
|
|
|
|
application.persistence_mgr.execute_async = execute_with_revoke
|
|
|
|
assert await service.authenticate_api_key(created['key']) is None
|
|
assert injected_revoke is True
|
|
|
|
|
|
async def test_cross_workspace_crud_and_secret_guessing_are_isolated(api_key_context):
|
|
application, service, first_context, engine = api_key_context
|
|
second_workspace_uuid = str(uuid.uuid4())
|
|
async with async_sessionmaker(engine, expire_on_commit=False).begin() as session:
|
|
session.add(
|
|
Workspace(
|
|
uuid=second_workspace_uuid,
|
|
instance_uuid='api-key-instance',
|
|
name='Second',
|
|
slug='second',
|
|
source=WorkspaceSource.CLOUD_PROJECTION.value,
|
|
)
|
|
)
|
|
session.add(
|
|
WorkspaceExecutionState(
|
|
workspace_uuid=second_workspace_uuid,
|
|
instance_uuid='api-key-instance',
|
|
active_generation=3,
|
|
state='active',
|
|
write_fenced=False,
|
|
source=WorkspaceExecutionSource.CLOUD.value,
|
|
)
|
|
)
|
|
second_context = _context(second_workspace_uuid, first_context.account_uuid or '', set(Permission))
|
|
created = await service.create_api_key(first_context, 'First only')
|
|
|
|
assert await service.get_api_key(second_context, created['id']) is None
|
|
assert await service.get_api_keys(second_context) == []
|
|
identity = await service.authenticate_api_key(created['key'])
|
|
assert identity is not None
|
|
assert identity.workspace_uuid == first_context.workspace_uuid
|
|
assert identity.workspace_uuid != second_workspace_uuid
|
|
|
|
# Prove the explicit multi-Workspace policy does not change key-derived routing.
|
|
application.workspace_service.policy = SingleWorkspacePolicy(workspace_limit=10, multi_workspace_enabled=True)
|
|
identity = await service.authenticate_api_key(created['key'])
|
|
assert identity is not None
|
|
assert identity.workspace_uuid == first_context.workspace_uuid
|
|
|
|
|
|
async def test_global_config_key_is_oss_singleton_only(api_key_context):
|
|
application, service, _context_value, _engine = api_key_context
|
|
application.instance_config.data['api']['global_api_key'] = 'configured-secret'
|
|
|
|
identity = await service.authenticate_api_key('configured-secret')
|
|
assert identity is not None
|
|
assert identity.api_key_uuid == 'global-oss-api-key'
|
|
|
|
application.workspace_service.policy = SingleWorkspacePolicy(workspace_limit=10, multi_workspace_enabled=True)
|
|
assert await service.authenticate_api_key('configured-secret') is None
|
|
|
|
|
|
async def test_explicit_scopes_cannot_exceed_callers_workspace_permissions(api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
limited_context = _context(
|
|
context.workspace_uuid,
|
|
context.account_uuid or '',
|
|
{Permission.API_KEY_MANAGE, Permission.RESOURCE_VIEW},
|
|
)
|
|
|
|
created = await service.create_api_key(
|
|
limited_context,
|
|
'Read only',
|
|
scopes=[Permission.RESOURCE_VIEW.value],
|
|
)
|
|
identity = await service.authenticate_api_key(created['key'])
|
|
assert identity is not None
|
|
assert identity.permissions == frozenset({Permission.RESOURCE_VIEW.value})
|
|
|
|
with pytest.raises(PermissionDeniedError) as exc_info:
|
|
await service.create_api_key(
|
|
limited_context,
|
|
'Escalated',
|
|
scopes=[Permission.WORKSPACE_DELETE.value],
|
|
)
|
|
assert exc_info.value.permission == Permission.WORKSPACE_DELETE.value
|
|
|
|
|
|
# Preserve the pre-tenancy CRUD and verification regression matrix while
|
|
# exercising it through the new Workspace-bound API. The assertions reflect
|
|
# intentional security changes: secrets are returned once, deletion revokes,
|
|
# and missing Workspace resources are reported as not found.
|
|
class TestApiKeyServiceGetApiKeys:
|
|
async def test_get_api_keys_empty_list(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
|
|
assert await service.get_api_keys(context) == []
|
|
|
|
async def test_get_api_keys_returns_serialized_list(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
await service.create_api_key(context, 'Test Key 1', 'First test key')
|
|
await service.create_api_key(context, 'Test Key 2', 'Second test key')
|
|
|
|
result = await service.get_api_keys(context)
|
|
|
|
assert [item['name'] for item in result] == ['Test Key 1', 'Test Key 2']
|
|
assert [item['description'] for item in result] == ['First test key', 'Second test key']
|
|
assert all('key' not in item and 'key_hash' not in item for item in result)
|
|
|
|
|
|
class TestApiKeyServiceCreateApiKey:
|
|
async def test_create_api_key_generates_key_with_prefix(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
|
|
with pytest.MonkeyPatch.context() as monkeypatch:
|
|
monkeypatch.setattr(
|
|
'langbot.pkg.api.http.service.apikey.secrets.token_urlsafe', lambda _size: 'fixed-token'
|
|
)
|
|
result = await service.create_api_key(context, 'New Key', 'Test description')
|
|
|
|
assert result['key'] == 'lbk_fixed-token'
|
|
assert result['name'] == 'New Key'
|
|
assert result['description'] == 'Test description'
|
|
assert result['secret_available'] is True
|
|
|
|
async def test_create_api_key_without_description(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
|
|
result = await service.create_api_key(context, 'No Desc Key')
|
|
|
|
assert result['description'] == ''
|
|
|
|
|
|
class TestApiKeyServiceGetApiKey:
|
|
async def test_get_api_key_by_id_found(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
created = await service.create_api_key(context, 'Found Key', 'Found')
|
|
|
|
result = await service.get_api_key(context, created['id'])
|
|
|
|
assert result is not None
|
|
assert result['id'] == created['id']
|
|
assert result['name'] == 'Found Key'
|
|
assert 'key' not in result and 'key_hash' not in result
|
|
|
|
async def test_get_api_key_by_id_not_found(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
|
|
assert await service.get_api_key(context, 999) is None
|
|
|
|
async def test_get_api_key_by_id_zero(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
|
|
assert await service.get_api_key(context, 0) is None
|
|
|
|
|
|
class TestApiKeyServiceVerifyApiKey:
|
|
async def test_verify_api_key_valid(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
created = await service.create_api_key(context, 'Valid')
|
|
|
|
assert await service.verify_api_key(created['key']) is True
|
|
|
|
async def test_verify_api_key_invalid(self, api_key_context):
|
|
_application, service, _context, _engine = api_key_context
|
|
|
|
assert await service.verify_api_key('lbk_invalid_key') is False
|
|
|
|
async def test_verify_api_key_empty_string(self, api_key_context):
|
|
_application, service, _context, _engine = api_key_context
|
|
|
|
assert await service.verify_api_key('') is False
|
|
|
|
async def test_verify_api_key_unknown_key(self, api_key_context):
|
|
_application, service, _context, _engine = api_key_context
|
|
|
|
assert await service.verify_api_key('unknown_key') is False
|
|
|
|
async def test_verify_global_api_key_match(self, api_key_context):
|
|
application, service, context, _engine = api_key_context
|
|
application.instance_config.data['api']['global_api_key'] = 'my-global-secret'
|
|
|
|
identity = await service.authenticate_api_key('my-global-secret')
|
|
|
|
assert identity is not None
|
|
assert identity.workspace_uuid == context.workspace_uuid
|
|
assert identity.api_key_uuid == 'global-oss-api-key'
|
|
|
|
async def test_verify_global_api_key_no_prefix_required(self, api_key_context):
|
|
application, service, _context, _engine = api_key_context
|
|
application.instance_config.data['api']['global_api_key'] = 'plainsecret123'
|
|
|
|
assert await service.verify_api_key('plainsecret123') is True
|
|
|
|
async def test_verify_global_api_key_mismatch_falls_back_to_db(self, api_key_context):
|
|
application, service, context, _engine = api_key_context
|
|
application.instance_config.data['api']['global_api_key'] = 'my-global-secret'
|
|
created = await service.create_api_key(context, 'DB key')
|
|
|
|
identity = await service.authenticate_api_key(created['key'])
|
|
|
|
assert identity is not None
|
|
assert identity.api_key_uuid == created['uuid']
|
|
|
|
async def test_verify_empty_global_api_key_disabled(self, api_key_context):
|
|
application, service, _context, _engine = api_key_context
|
|
application.instance_config.data['api']['global_api_key'] = ''
|
|
|
|
assert await service.verify_api_key('') is False
|
|
assert await service.verify_api_key(' ') is False
|
|
|
|
async def test_verify_api_key_missing_global_config_key(self, api_key_context):
|
|
application, service, _context, _engine = api_key_context
|
|
application.instance_config.data = {'api': {}}
|
|
|
|
assert await service.verify_api_key('lbk_some_key') is False
|
|
|
|
|
|
class TestApiKeyServiceDeleteApiKey:
|
|
async def test_delete_api_key_by_id(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
created = await service.create_api_key(context, 'Delete me')
|
|
|
|
await service.delete_api_key(context, created['id'])
|
|
|
|
stored = await service.get_api_key(context, created['id'])
|
|
assert stored is not None
|
|
assert stored['status'] == 'revoked'
|
|
assert await service.verify_api_key(created['key']) is False
|
|
|
|
async def test_delete_api_key_nonexistent_id(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
|
|
with pytest.raises(WorkspaceNotFoundError, match='API key not found'):
|
|
await service.delete_api_key(context, 999)
|
|
|
|
|
|
class TestApiKeyServiceUpdateApiKey:
|
|
async def test_update_api_key_name_only(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
created = await service.create_api_key(context, 'Original', 'Description')
|
|
|
|
await service.update_api_key(context, created['id'], name='Updated Name')
|
|
|
|
stored = await service.get_api_key(context, created['id'])
|
|
assert stored is not None
|
|
assert stored['name'] == 'Updated Name'
|
|
assert stored['description'] == 'Description'
|
|
|
|
async def test_update_api_key_description_only(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
created = await service.create_api_key(context, 'Original', 'Description')
|
|
|
|
await service.update_api_key(context, created['id'], description='Updated description')
|
|
|
|
stored = await service.get_api_key(context, created['id'])
|
|
assert stored is not None
|
|
assert stored['name'] == 'Original'
|
|
assert stored['description'] == 'Updated description'
|
|
|
|
async def test_update_api_key_both_fields(self, api_key_context):
|
|
_application, service, context, _engine = api_key_context
|
|
created = await service.create_api_key(context, 'Original', 'Description')
|
|
|
|
await service.update_api_key(
|
|
context,
|
|
created['id'],
|
|
name='New Name',
|
|
description='New description',
|
|
)
|
|
|
|
stored = await service.get_api_key(context, created['id'])
|
|
assert stored is not None
|
|
assert stored['name'] == 'New Name'
|
|
assert stored['description'] == 'New description'
|
|
|
|
async def test_update_api_key_no_fields(self, api_key_context):
|
|
application, service, context, _engine = api_key_context
|
|
created = await service.create_api_key(context, 'Original')
|
|
original_execute = application.persistence_mgr.execute_async
|
|
application.persistence_mgr.execute_async = AsyncMock(wraps=original_execute)
|
|
|
|
await service.update_api_key(context, created['id'])
|
|
|
|
application.persistence_mgr.execute_async.assert_not_awaited()
|