feat: 添加对保留能力命名空间的检查,防止暴露注册

This commit is contained in:
whatevertogo
2026-03-12 22:56:27 +08:00
parent db9cb169e4
commit 61eecd2024
3 changed files with 84 additions and 22 deletions

View File

@@ -12,6 +12,7 @@ from ..protocol.descriptors import CapabilityDescriptor
CallHandler = Callable[[str, dict[str, Any], object], Awaitable[dict[str, Any]]]
StreamHandler = Callable[[str, dict[str, Any], object], AsyncIterator[dict[str, Any]]]
FinalizeHandler = Callable[[list[dict[str, Any]]], dict[str, Any]]
RESERVED_CAPABILITY_PREFIXES = ("handler.", "system.", "internal.")
@dataclass(slots=True)
@@ -53,6 +54,8 @@ class CapabilityRouter:
finalize: FinalizeHandler | None = None,
exposed: bool = True,
) -> None:
if exposed and descriptor.name.startswith(RESERVED_CAPABILITY_PREFIXES):
raise ValueError(f"保留 capability 命名空间仅供框架内部使用:{descriptor.name}")
self._registrations[descriptor.name] = _CapabilityRegistration(
descriptor=descriptor,
call_handler=call_handler,

View File

@@ -216,33 +216,19 @@ class Peer:
self.remote_handlers = message.handlers
self.remote_metadata = message.metadata
if self._initialize_handler is None:
error = AstrBotError.protocol_error("对端不接受 initialize")
await self._send(
ResultMessage(
id=message.id,
kind="initialize_result",
success=False,
error=ErrorPayload.model_validate(error.to_payload()),
)
await self._reject_initialize(
message,
AstrBotError.protocol_error("对端不接受 initialize"),
)
self._unusable = True
self._remote_initialized.set()
return
if message.protocol_version != self.protocol_version:
error = AstrBotError.protocol_version_mismatch(
f"服务端支持协议版本 {self.protocol_version},客户端请求版本 {message.protocol_version}"
await self._reject_initialize(
message,
AstrBotError.protocol_version_mismatch(
f"服务端支持协议版本 {self.protocol_version},客户端请求版本 {message.protocol_version}"
),
)
await self._send(
ResultMessage(
id=message.id,
kind="initialize_result",
success=False,
error=ErrorPayload.model_validate(error.to_payload()),
)
)
self._unusable = True
self._remote_initialized.set()
return
output = await self._initialize_handler(message)
@@ -338,6 +324,19 @@ class Peer:
)
)
async def _reject_initialize(self, message: InitializeMessage, error: AstrBotError) -> None:
await self._send(
ResultMessage(
id=message.id,
kind="initialize_result",
success=False,
error=ErrorPayload.model_validate(error.to_payload()),
)
)
self._unusable = True
self._remote_initialized.set()
await self.stop()
async def _send_cancelled_termination(self, message: InvokeMessage) -> None:
error = AstrBotError.cancelled()
await self._send_error_result(message, error)

View File

@@ -189,3 +189,63 @@ class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase):
self.assertEqual(result["text"], "Echo: ws")
await plugin.stop()
await core.stop()
async def test_initialize_failure_closes_receiver_connection(self) -> None:
core = Peer(
transport=self.left,
peer_info=PeerInfo(name="core", role="core", version="v4"),
protocol_version="1.0",
)
core.set_initialize_handler(
lambda _message: asyncio.sleep(
0,
result=InitializeOutput(
peer=PeerInfo(name="core", role="core", version="v4"),
capabilities=[],
metadata={},
),
)
)
plugin = Peer(
transport=self.right,
peer_info=PeerInfo(name="plugin", role="plugin", version="v4"),
protocol_version="2.0",
)
await core.start()
await plugin.start()
with self.assertRaises(AstrBotError) as raised:
await plugin.initialize([])
self.assertEqual(raised.exception.code, "protocol_version_mismatch")
await asyncio.wait_for(core.wait_closed(), timeout=1.0)
self.assertTrue(core._closed)
await plugin.stop()
class CapabilityRouterContractTest(unittest.TestCase):
def test_reserved_capability_namespaces_are_rejected_for_exposed_registrations(self) -> None:
router = CapabilityRouter()
for name in ("handler.demo", "system.health", "internal.trace"):
with self.assertRaises(ValueError) as raised:
router.register(
CapabilityDescriptor(
name=name,
description="reserved",
)
)
self.assertIn(name, str(raised.exception))
def test_reserved_capability_namespaces_remain_available_for_hidden_internal_registrations(self) -> None:
router = CapabilityRouter()
router.register(
CapabilityDescriptor(
name="system.health",
description="internal only",
),
exposed=False,
)
self.assertNotIn("system.health", [item.name for item in router.descriptors()])