diff --git a/src/a2a/server/tasks/inmemory_push_notification_config_store.py b/src/a2a/server/tasks/inmemory_push_notification_config_store.py index 19e35074a..b35f4294a 100644 --- a/src/a2a/server/tasks/inmemory_push_notification_config_store.py +++ b/src/a2a/server/tasks/inmemory_push_notification_config_store.py @@ -7,11 +7,20 @@ PushNotificationConfigStore, ) from a2a.types.a2a_pb2 import TaskPushNotificationConfig +from a2a.utils.errors import InvalidParamsError logger = logging.getLogger(__name__) +MAX_PUSH_NOTIFICATION_CONFIGS_PER_TASK = 50 +"""Maximum number of push notification configs a single task may have. + +Prevents a client from registering an unbounded number of webhooks for +one task (DoS / resource exhaustion). +""" + + class InMemoryPushNotificationConfigStore(PushNotificationConfigStore): """In-memory implementation of PushNotificationConfigStore interface. @@ -53,13 +62,33 @@ async def set_info( if not notification_config.id: notification_config.id = task_id + # Enforce the per-task cap. Updating an existing config does not + # count as a new entry, so deduplicate before the check. + existing_configs = owner_infos[task_id] + is_update = any( + config.id == notification_config.id + for config in existing_configs + ) + if ( + not is_update + and len(existing_configs) + >= MAX_PUSH_NOTIFICATION_CONFIGS_PER_TASK + ): + raise InvalidParamsError( + message=( + f'maximum of ' + f'{MAX_PUSH_NOTIFICATION_CONFIGS_PER_TASK} push ' + 'notification configs per task exceeded' + ) + ) + # Remove existing config with the same ID - for config in owner_infos[task_id]: + for config in existing_configs: if config.id == notification_config.id: - owner_infos[task_id].remove(config) + existing_configs.remove(config) break - owner_infos[task_id].append(notification_config) + existing_configs.append(notification_config) logger.debug( 'Push notification config for task %s with config id %s for owner %s saved/updated.', task_id, diff --git a/tests/server/tasks/test_inmemory_push_notifications.py b/tests/server/tasks/test_inmemory_push_notifications.py index f204e2181..a95f1cb7e 100644 --- a/tests/server/tasks/test_inmemory_push_notifications.py +++ b/tests/server/tasks/test_inmemory_push_notifications.py @@ -572,3 +572,81 @@ async def test_cross_user_dispatch_alice_registers_bob_triggers( if __name__ == '__main__': unittest.main() + + +class TestPushConfigPerTaskLimit(unittest.IsolatedAsyncioTestCase): + """set_info must not allow unbounded configs per task.""" + + def setUp(self) -> None: + self.config_store = InMemoryPushNotificationConfigStore() + + async def test_can_store_up_to_limit(self) -> None: + from a2a.server.tasks.inmemory_push_notification_config_store import ( + MAX_PUSH_NOTIFICATION_CONFIGS_PER_TASK, + ) + + limit = MAX_PUSH_NOTIFICATION_CONFIGS_PER_TASK + for i in range(limit): + await self.config_store.set_info( + 'task-limit', + _create_sample_push_config( + url=f'http://example.com/cb/{i}', config_id=f'cfg-{i}' + ), + MINIMAL_CALL_CONTEXT, + ) + retrieved = await self.config_store.get_info( + 'task-limit', MINIMAL_CALL_CONTEXT + ) + self.assertEqual(len(retrieved), limit) + + async def test_exceeding_limit_raises_invalid_params(self) -> None: + from a2a.server.tasks.inmemory_push_notification_config_store import ( + MAX_PUSH_NOTIFICATION_CONFIGS_PER_TASK, + ) + from a2a.utils.errors import InvalidParamsError + + limit = MAX_PUSH_NOTIFICATION_CONFIGS_PER_TASK + for i in range(limit): + await self.config_store.set_info( + 'task-limit', + _create_sample_push_config( + url=f'http://example.com/cb/{i}', config_id=f'cfg-{i}' + ), + MINIMAL_CALL_CONTEXT, + ) + with self.assertRaises(InvalidParamsError): + await self.config_store.set_info( + 'task-limit', + _create_sample_push_config( + url='http://example.com/cb/overflow', + config_id='cfg-overflow', + ), + MINIMAL_CALL_CONTEXT, + ) + + async def test_updating_existing_config_is_not_limited(self) -> None: + from a2a.server.tasks.inmemory_push_notification_config_store import ( + MAX_PUSH_NOTIFICATION_CONFIGS_PER_TASK, + ) + + limit = MAX_PUSH_NOTIFICATION_CONFIGS_PER_TASK + for i in range(limit): + await self.config_store.set_info( + 'task-limit', + _create_sample_push_config( + url=f'http://example.com/cb/{i}', config_id=f'cfg-{i}' + ), + MINIMAL_CALL_CONTEXT, + ) + # Overwriting an existing config id must still succeed beyond the cap. + await self.config_store.set_info( + 'task-limit', + _create_sample_push_config( + url='http://example.com/cb/updated', config_id='cfg-0' + ), + MINIMAL_CALL_CONTEXT, + ) + retrieved = await self.config_store.get_info( + 'task-limit', MINIMAL_CALL_CONTEXT + ) + self.assertEqual(len(retrieved), limit)