xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-07-21 20:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): harden shared cloud runtime boundaries

Browse Source
This commit is contained in:
Junyan Qin
2026-07-20 01:47:42 +08:00
parent 41772920ef
commit a47bfe8167
121 changed files with 18152 additions and 5588 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit_tests/box/test_box_admission.py
+230
View File
@@ -0,0 +1,230 @@
from __future__ import annotations
import datetime as dt
import hashlib
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot_plugin.box.models import SandboxAdmissionPolicy
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.box.admission import (
BoxAdmissionError,
SandboxAdmissionController,
require_cloud_admission_policy,
)
from langbot.pkg.box.service import BoxService
from langbot.pkg.cloud.entitlements import (
EntitlementResolver,
EntitlementSnapshot,
EntitlementUnavailableError,
)
_UTC = dt.timezone.utc
_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=3,
entitlement_revision=7,
)
def _snapshot(
*,
revision: int = 7,
managed: bool = True,
sessions: int = 1,
expires_at: int = 2_000,
) -> EntitlementSnapshot:
return EntitlementSnapshot(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
entitlement_revision=revision,
status='active',
not_before=1,
expires_at=expires_at,
features={'managed_sandbox': managed},
limits={'managed_sandbox_sessions': sessions},
)
def _controller(snapshot: EntitlementSnapshot, *, now: float = 1_000.25):
provider = SimpleNamespace(get_workspace_entitlement=AsyncMock(return_value=snapshot))
resolver = EntitlementResolver('instance-a', provider)
client = SimpleNamespace(
upsert_sandbox_admission_grant=AsyncMock(
side_effect=lambda grant: {
'installed': True,
'workspace_uuid': grant.workspace_uuid,
'execution_generation': grant.execution_generation,
'entitlement_revision': grant.entitlement_revision,
'max_sessions': grant.max_sessions,
'max_managed_processes': grant.max_managed_processes,
}
),
revoke_sandbox_admission_grant=AsyncMock(
side_effect=lambda revocation: {
'revoked': True,
'workspace_uuid': revocation.workspace_uuid,
'entitlement_revision': revocation.entitlement_revision,
}
),
)
app = SimpleNamespace(entitlement_resolver=resolver, logger=Mock())
controller = SandboxAdmissionController(
app,
client,
policy=SandboxAdmissionPolicy(required=True, max_grant_ttl_sec=300),
wall_time=lambda: now,
)
return controller, client, provider
def test_cloud_admission_policy_requires_positive_workspace_quota():
with pytest.raises(BoxAdmissionError, match='workspace quota must be a positive integer'):
require_cloud_admission_policy(
{
'required': True,
'workspace_quota_mb': 0,
}
)
policy = require_cloud_admission_policy(
{
'required': True,
'workspace_quota_mb': 32,
}
)
assert policy.workspace_quota_mb == 32
@pytest.mark.asyncio
async def test_active_generic_entitlement_installs_short_lived_numeric_grant():
controller, client, provider = _controller(_snapshot())
grant = await controller.require(_CONTEXT)
assert grant.instance_uuid == _CONTEXT.instance_uuid
assert grant.workspace_uuid == _CONTEXT.workspace_uuid
assert grant.execution_generation == _CONTEXT.placement_generation
assert grant.entitlement_revision == 7
assert grant.max_sessions == 1
assert grant.max_managed_processes == 0
assert grant.expires_at == dt.datetime.fromtimestamp(1_300, tz=_UTC)
assert (grant.expires_at - dt.datetime.fromtimestamp(1_000.25, tz=_UTC)).total_seconds() < 300
provider.get_workspace_entitlement.assert_awaited_once_with('workspace-a')
client.upsert_sandbox_admission_grant.assert_awaited_once_with(grant)
client.revoke_sandbox_admission_grant.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
'snapshot',
[
_snapshot(managed=False),
_snapshot(sessions=0),
_snapshot(sessions=2),
],
)
async def test_non_eligible_entitlement_revokes_and_fails_closed(snapshot):
controller, client, _provider = _controller(snapshot)
with pytest.raises(EntitlementUnavailableError):
await controller.require(_CONTEXT)
client.upsert_sandbox_admission_grant.assert_not_awaited()
revocation = client.revoke_sandbox_admission_grant.await_args.args[0]
assert revocation.entitlement_revision == snapshot.entitlement_revision
@pytest.mark.asyncio
async def test_transient_entitlement_failure_does_not_tombstone_valid_revision():
controller, client, provider = _controller(_snapshot())
await controller.require(_CONTEXT)
provider.get_workspace_entitlement.side_effect = RuntimeError('control plane unavailable')
with pytest.raises(RuntimeError, match='control plane unavailable'):
await controller.require(_CONTEXT)
client.revoke_sandbox_admission_grant.assert_not_awaited()
provider.get_workspace_entitlement.side_effect = None
provider.get_workspace_entitlement.return_value = _snapshot()
recovered = await controller.require(_CONTEXT)
assert recovered.entitlement_revision == 7
@pytest.mark.asyncio
async def test_runtime_receipt_mismatch_is_revoked_and_never_admitted():
controller, client, _provider = _controller(_snapshot())
client.upsert_sandbox_admission_grant.return_value = {'installed': True, 'workspace_uuid': 'other'}
client.upsert_sandbox_admission_grant.side_effect = None
with pytest.raises(Exception, match='invalid sandbox admission receipt'):
await controller.require(_CONTEXT)
client.revoke_sandbox_admission_grant.assert_not_awaited()
@pytest.mark.asyncio
async def test_authoritative_cancelled_revision_is_revoked():
cancelled = _snapshot(revision=8).model_copy(update={'status': 'cancelled'})
controller, client, _provider = _controller(cancelled)
with pytest.raises(EntitlementUnavailableError, match='not active'):
await controller.require(_CONTEXT)
revocation = client.revoke_sandbox_admission_grant.await_args.args[0]
assert revocation.entitlement_revision == 8
@pytest.mark.asyncio
async def test_cloud_box_readiness_failure_aborts_service_initialization(tmp_path):
workspace_root = tmp_path / 'box' / 'workspaces'
workspace_root.mkdir(parents=True)
box_config = {
'enabled': True,
'backend': 'nsjail',
'runtime': {'endpoint': 'ws://box:5410'},
'local': {
'host_root': str(tmp_path / 'box'),
'default_workspace': str(workspace_root),
'allowed_mount_roots': [str(tmp_path / 'box')],
},
'admission': {
'required': True,
'logical_session_id': 'global',
'required_backend': 'nsjail',
'max_sessions': 1,
'max_managed_processes': 0,
'max_grant_ttl_sec': 300,
'workspace_quota_mb': 32,
},
}
client = SimpleNamespace(
initialize=AsyncMock(),
verify_shared_workspace=AsyncMock(
side_effect=lambda marker_name: {
'marker_name': marker_name,
'size': (workspace_root / marker_name).stat().st_size,
'sha256': hashlib.sha256((workspace_root / marker_name).read_bytes()).hexdigest(),
}
),
get_backend_info=AsyncMock(return_value={'name': 'docker', 'available': True}),
)
app = SimpleNamespace(
logger=Mock(),
deployment=SimpleNamespace(multi_workspace_enabled=True),
entitlement_resolver=Mock(),
workspace_service=SimpleNamespace(instance_uuid='instance-a'),
instance_config=SimpleNamespace(data={'box': box_config}),
)
service = BoxService(app, client=client)
with pytest.raises(Exception, match='nsjail isolation readiness failed'):
await service.initialize()
assert service.available is False
tests/unit_tests/box/test_box_service.py
+138 -13
View File
@@ -117,6 +117,9 @@ class _InProcessBoxRuntimeClient(BoxRuntimeClient):
async def init(self, config: dict) -> None:
self._runtime.init(config)
async def verify_shared_workspace(self, marker_name: str) -> dict:
return self._runtime.verify_shared_workspace(marker_name)
class FakeBackend(BaseSandboxBackend):
def __init__(self, logger: Mock, available: bool = True):
@@ -322,6 +325,43 @@ def test_separated_box_runtime_does_not_create_default_workspace_in_langbot(tmp_
assert not (host_root / 'default').exists()
@pytest.mark.asyncio
async def test_cloud_initialize_fails_when_core_and_runtime_volumes_are_separated(tmp_path):
logger = Mock()
core_root = tmp_path / 'core-box'
runtime_root = tmp_path / 'runtime-box'
(core_root / 'default').mkdir(parents=True)
runtime = BoxRuntime(logger=logger, backends=[FakeBackend(logger)], session_ttl_sec=300)
runtime.init(
{
'local': {
'host_root': str(runtime_root),
'default_workspace': 'default',
'allowed_mount_roots': [str(runtime_root)],
}
}
)
app = make_app(logger, host_root=str(core_root))
app.deployment = SimpleNamespace(multi_workspace_enabled=True)
app.instance_config.data['box'].update(
{
'backend': 'nsjail',
'admission': {'required': True, 'workspace_quota_mb': 32},
}
)
service = BoxService(
app,
client=_InProcessBoxRuntimeClient(logger, runtime),
)
with pytest.raises(BoxValidationError, match='shared durable Workspace volume'):
await service.initialize()
assert service.available is False
assert list((core_root / 'default').glob('.langbot-box-volume-probe-*')) == []
await runtime.shutdown()
def test_separated_box_runtime_allows_box_owned_missing_host_path(tmp_path):
logger = Mock()
runtime = BoxRuntime(logger=logger, backends=[FakeBackend(logger)], session_ttl_sec=300)
@@ -1888,7 +1928,7 @@ class TestInboundOutboundRoundTrip:
assert '/workspace/inbox/' in parameters['command']
return {
'ok': True,
'stdout': '["/workspace/inbox/42/image_1.png"]',
'stdout': '["/workspace/inbox/query-42/image_1.png"]',
'stderr': '',
}
@@ -1898,7 +1938,7 @@ class TestInboundOutboundRoundTrip:
assert len(descriptors) == 1
d = descriptors[0]
assert d['type'] == 'Image'
assert d['path'] == '/workspace/inbox/42/image_1.png'
assert d['path'] == '/workspace/inbox/query-42/image_1.png'
assert d['size'] == len(img_bytes)
@pytest.mark.asyncio
@@ -2011,7 +2051,7 @@ class TestAttachmentHostPath:
assert d['type'] == 'Image'
assert d['size'] == len(big)
# File actually landed on the host workspace.
host_file = os.path.join(ws, 'inbox', str(query.query_id), d['name'])
host_file = os.path.join(ws, 'inbox', str(query.query_uuid), d['name'])
assert os.path.isfile(host_file)
assert open(host_file, 'rb').read() == big
@@ -2023,7 +2063,7 @@ class TestAttachmentHostPath:
service, ws = self._service_with_workspace(tmp_path)
# Seed a stale file under the same query_id (simulates webchat id reuse).
stale_dir = os.path.join(ws, 'inbox', '42')
stale_dir = os.path.join(ws, 'inbox', 'query-42')
os.makedirs(stale_dir, exist_ok=True)
open(os.path.join(stale_dir, 'image_1.png'), 'wb').write(b'STALE-OLD-IMAGE')
@@ -2039,11 +2079,40 @@ class TestAttachmentHostPath:
# No leftover content from the stale image.
assert b'STALE-OLD-IMAGE' not in open(host_file, 'rb').read()
@pytest.mark.asyncio
async def test_inbound_host_replaces_query_symlink_without_touching_other_workspace(self, tmp_path):
import base64
import langbot_plugin.api.entities.builtin.platform.message as platform_message
service, ws = self._service_with_workspace(tmp_path)
other_workspace = tmp_path / 'other-workspace'
other_workspace.mkdir()
protected = other_workspace / 'protected.txt'
protected.write_bytes(b'workspace-b-secret')
inbox = os.path.join(ws, 'inbox')
os.makedirs(inbox, exist_ok=True)
os.symlink(other_workspace, os.path.join(inbox, 'query-42'))
query = make_query()
payload = b'workspace-a-input'
query.message_chain = platform_message.MessageChain(
[platform_message.File(name='input.bin', base64=base64.b64encode(payload).decode())]
)
service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path'))
descriptors = await service.materialize_inbound_attachments(query)
assert descriptors[0]['path'] == '/workspace/inbox/query-42/input.bin'
assert protected.read_bytes() == b'workspace-b-secret'
assert not os.path.islink(os.path.join(inbox, 'query-42'))
assert open(os.path.join(inbox, 'query-42', 'input.bin'), 'rb').read() == payload
@pytest.mark.asyncio
async def test_outbound_reads_host_and_clears(self, tmp_path):
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
outbox = os.path.join(ws, 'outbox', str(query.query_id))
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
os.makedirs(outbox, exist_ok=True)
# A large file that would be truncated on the exec/stdout path:
big_png = b'\x89PNG\r\n\x1a\n' + b'y' * (400 * 1024)
@@ -2063,13 +2132,69 @@ class TestAttachmentHostPath:
# Outbox cleared after collection.
assert os.listdir(outbox) == []
@pytest.mark.asyncio
async def test_outbound_host_never_follows_query_or_file_symlinks(self, tmp_path):
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
other_workspace = tmp_path / 'other-workspace'
other_workspace.mkdir()
secret = other_workspace / 'secret.txt'
secret.write_bytes(b'workspace-b-secret')
outbox_root = os.path.join(ws, 'outbox')
os.makedirs(outbox_root, exist_ok=True)
# A hostile query-directory replacement is rejected rather than read.
query_dir = os.path.join(outbox_root, str(query.query_uuid))
os.symlink(other_workspace, query_dir)
with pytest.raises(BoxValidationError, match='symbolic link'):
await service.collect_outbound_attachments(query)
assert secret.read_bytes() == b'workspace-b-secret'
os.unlink(query_dir)
os.makedirs(query_dir)
os.symlink(secret, os.path.join(query_dir, 'leak.txt'))
service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path'))
assert await service.collect_outbound_attachments(query) == []
assert secret.read_bytes() == b'workspace-b-secret'
@pytest.mark.asyncio
async def test_outbound_host_fails_closed_on_inode_bomb(self, tmp_path):
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
os.makedirs(outbox, exist_ok=True)
harmless_target = tmp_path / 'harmless-target'
harmless_target.write_bytes(b'x')
# Symlinks do not count toward the 20 returned files, so this proves
# traversal itself has a bounded entry budget.
for index in range(513):
os.symlink(harmless_target, os.path.join(outbox, f'entry-{index}'))
with pytest.raises(BoxValidationError, match='symbolic link'):
await service.collect_outbound_attachments(query)
assert harmless_target.read_bytes() == b'x'
def test_host_attachment_directories_use_query_uuid_not_process_local_id(self, tmp_path):
service, _ws = self._service_with_workspace(tmp_path)
first = make_query(query_id=7)
second = make_query(query_id=7)
object.__setattr__(first, 'query_uuid', 'replica-a-query')
object.__setattr__(second, 'query_uuid', 'replica-b-query')
first_path = service._host_query_dir(service.OUTBOX_SUBDIR, first)
second_path = service._host_query_dir(service.OUTBOX_SUBDIR, second)
assert first_path is not None and first_path.endswith('/outbox/replica-a-query')
assert second_path is not None and second_path.endswith('/outbox/replica-b-query')
assert first_path != second_path
@pytest.mark.asyncio
async def test_outbound_empty_clears_stale_host_dir(self, tmp_path):
# Reusing a query_id (counter resets on restart) must not re-send files
# a previous run left in the outbox: an empty collection still clears it.
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
outbox = os.path.join(ws, 'outbox', str(query.query_id))
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
os.makedirs(outbox, exist_ok=True)
# Stale file from a prior turn; the agent produced nothing this turn —
# but _read_outbox_host would still pick it up, so collection must drop
@@ -2118,16 +2243,16 @@ class TestAttachmentHostPath:
os.makedirs(os.path.join(outbox, '0'), exist_ok=True)
# Simulate a host delete that cannot remove the root-owned outbox.
import shutil as _shutil
from langbot.pkg.box import secure_fs
real_rmtree = _shutil.rmtree
real_purge = secure_fs.purge_subdirectory
def fake_rmtree(path, *a, **k):
if os.path.abspath(path) == os.path.abspath(outbox):
return # "permission denied" — silently leaves the dir
return real_rmtree(path, *a, **k)
def fake_purge(root, subdir):
if os.path.abspath(os.path.join(root, subdir)) == os.path.abspath(outbox):
raise PermissionError('root-owned')
return real_purge(root, subdir)
monkeypatch.setattr(_shutil, 'rmtree', fake_rmtree)
monkeypatch.setattr(secure_fs, 'purge_subdirectory', fake_purge)
service.build_spec = Mock()
service.client.execute = AsyncMock()