mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 01:40:15 +08:00
feat: 添加对保留能力命名空间的检查,防止暴露注册
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()])
|
||||
|
||||
Reference in New Issue
Block a user