refactor: connection retrieval through route manager

Signed-off-by: Daniel Bluhm <dbluhm@pm.me>

Author: dbluhm

aries_cloudagent/multitenant/base.py

@@ -381,6 +381,19 @@ async def _get_wallet_by_key(self, recipient_key: str) -> Optional[WalletRecord]
         except (RouteNotFoundError):
             pass
 
+    async def get_profile_for_key(
+        self, context: InjectionContext, recipient_key: str
+    ) -> Optional[Profile]:
+        """Retrieve a wallet profile by recipient key."""
+        wallet = await self._get_wallet_by_key(recipient_key)
+        if not wallet:
+            return None
+
+        if wallet.requires_external_key:
+            raise WalletKeyMissingError()
+
+        return await self.get_wallet_profile(context, wallet)
+
     async def get_wallets_by_message(
         self, message_body, wire_format: BaseWireFormat = None
     ) -> List[WalletRecord]:

aries_cloudagent/multitenant/route_manager.py

@@ -4,6 +4,7 @@
 import logging
 from typing import List, Optional, Tuple
 
+from ..connections.models.conn_record import ConnRecord
 from ..core.profile import Profile
 from ..messaging.responder import BaseResponder
 from ..protocols.coordinate_mediation.v1_0.manager import MediationManager
@@ -14,6 +15,7 @@
 from ..protocols.routing.v1_0.manager import RoutingManager
 from ..protocols.routing.v1_0.models.route_record import RouteRecord
 from ..storage.error import StorageNotFoundError
+from .manager import MultitenantManager
 
 
 LOGGER = logging.getLogger(__name__)
@@ -103,3 +105,23 @@ async def routing_info(
             my_endpoint = mediation_record.endpoint
 
         return routing_keys, my_endpoint
+
+    async def connection_from_recipient_key(
+        self, profile: Profile, recipient_key: str
+    ) -> ConnRecord:
+        """Retrieve a connection by recipient key.
+
+        The recipient key is expected to be a local key owned by this agent.
+
+        Since the multi-tenant base wallet can receive and send keylist updates
+        for sub wallets, we check the sub wallet's connections before the base
+        wallet.
+        """
+        manager = MultitenantManager(self.root_profile)
+        profile_to_search = (
+            await manager.get_profile_for_key(profile.context, recipient_key) or profile
+        )
+
+        return await super().connection_from_recipient_key(
+            profile_to_search, recipient_key
+        )

aries_cloudagent/protocols/coordinate_mediation/v1_0/handlers/keylist_update_response_handler.py

@@ -1,11 +1,14 @@
 """Handler for keylist-update-response message."""
 
+from .....core.profile import Profile
 from .....messaging.base_handler import BaseHandler, HandlerException
 from .....messaging.request_context import RequestContext
 from .....messaging.responder import BaseResponder
-
-from ..messages.keylist_update_response import KeylistUpdateResponse
+from .....storage.error import StorageNotFoundError
+from .....wallet.error import WalletNotFoundError
 from ..manager import MediationManager
+from ..messages.keylist_update_response import KeylistUpdateResponse
+from ..route_manager import RouteManager
 
 
 class KeylistUpdateResponseHandler(BaseHandler):
@@ -25,6 +28,35 @@ async def handle(self, context: RequestContext, responder: BaseResponder):
         await mgr.store_update_results(
             context.connection_record.connection_id, context.message.updated
         )
-        await mgr.notify_keylist_updated(
-            context.connection_record.connection_id, context.message
+        await self.notify_keylist_updated(
+            context.profile, context.connection_record.connection_id, context.message
+        )
+
+    async def notify_keylist_updated(
+        self, profile: Profile, connection_id: str, response: KeylistUpdateResponse
+    ):
+        """Notify of keylist update response received."""
+        route_manager = profile.inject(RouteManager)
+        try:
+            key_to_connection = {
+                updated.recipient_key: await route_manager.connection_from_recipient_key(
+                    profile, updated.recipient_key
+                )
+                for updated in response.updated
+            }
+        except (StorageNotFoundError, WalletNotFoundError) as err:
+            raise HandlerException(
+                "Unknown recipient key received in keylist update response"
+            ) from err
+
+        await profile.notify(
+            MediationManager.KEYLIST_UPDATED_EVENT,
+            {
+                "connection_id": connection_id,
+                "thread_id": response._thread_id,
+                "updated": [update.serialize() for update in response.updated],
+                "mediated_connections": {
+                    key: conn.connection_id for key, conn in key_to_connection.items()
+                },
+            },
         )

aries_cloudagent/protocols/coordinate_mediation/v1_0/handlers/tests/test_keylist_update_response_handler.py

@@ -1,21 +1,29 @@
 """Test handler for keylist-update-response message."""
 
+from functools import partial
+from typing import AsyncGenerator
 import pytest
 from asynctest import TestCase as AsyncTestCase
 from asynctest import mock as async_mock
 
+
 from ......connections.models.conn_record import ConnRecord
+from ......core.event_bus import EventBus, MockEventBus
 from ......messaging.base_handler import HandlerException
 from ......messaging.request_context import RequestContext
 from ......messaging.responder import MockResponder
 from ...messages.inner.keylist_update_rule import KeylistUpdateRule
 from ...messages.inner.keylist_updated import KeylistUpdated
 from ...messages.keylist_update_response import KeylistUpdateResponse
 from ...manager import MediationManager
+from ...route_manager import RouteManager
+from ...tests.test_route_manager import MockRouteManager
 from ..keylist_update_response_handler import KeylistUpdateResponseHandler
 
 TEST_CONN_ID = "conn-id"
-TEST_VERKEY = "3Dn1SJNPaCXcvvJvSbsFWP2xaCjMom3can8CQNhWrTRx"
+TEST_THREAD_ID = "thread-id"
+TEST_VERKEY = "did:key:z6Mkgg342Ycpuk263R9d8Aq6MUaxPn1DDeHyGo38EefXmgDL"
+TEST_ROUTE_VERKEY = "did:key:z6MknxTj6Zj1VrDWc1ofaZtmCVv2zNXpD58Xup4ijDGoQhya"
 
 
 class TestKeylistUpdateResponseHandler(AsyncTestCase):
@@ -34,6 +42,10 @@ async def setUp(self):
         self.context.message = KeylistUpdateResponse(updated=self.updated)
         self.context.connection_ready = True
         self.context.connection_record = ConnRecord(connection_id=TEST_CONN_ID)
+        self.mock_event_bus = MockEventBus()
+        self.context.injector.bind_instance(EventBus, self.mock_event_bus)
+        self.route_manager = MockRouteManager()
+        self.context.injector.bind_instance(RouteManager, self.route_manager)
 
     async def test_handler_no_active_connection(self):
         handler, responder = KeylistUpdateResponseHandler(), MockResponder()
@@ -47,8 +59,86 @@ async def test_handler(self):
         with async_mock.patch.object(
             MediationManager, "store_update_results"
         ) as mock_store, async_mock.patch.object(
-            MediationManager, "notify_keylist_updated"
+            handler, "notify_keylist_updated"
         ) as mock_notify:
             await handler.handle(self.context, responder)
             mock_store.assert_called_once_with(TEST_CONN_ID, self.updated)
-            mock_notify.assert_called_once_with(TEST_CONN_ID, self.context.message)
+            mock_notify.assert_called_once_with(
+                self.context.profile, TEST_CONN_ID, self.context.message
+            )
+
+    async def test_notify_keylist_updated(self):
+        """test notify_keylist_updated."""
+        handler = KeylistUpdateResponseHandler()
+
+        async def _result_generator():
+            yield ConnRecord(connection_id="conn_id_1")
+            yield ConnRecord(connection_id="conn_id_2")
+
+        async def _retrieve_by_invitation_key(
+            generator: AsyncGenerator, *args, **kwargs
+        ):
+            return await generator.__anext__()
+
+        with async_mock.patch.object(
+            self.route_manager,
+            "connection_from_recipient_key",
+            partial(_retrieve_by_invitation_key, _result_generator()),
+        ):
+            response = KeylistUpdateResponse(
+                updated=[
+                    KeylistUpdated(
+                        recipient_key=TEST_ROUTE_VERKEY,
+                        action=KeylistUpdateRule.RULE_ADD,
+                        result=KeylistUpdated.RESULT_SUCCESS,
+                    ),
+                    KeylistUpdated(
+                        recipient_key=TEST_VERKEY,
+                        action=KeylistUpdateRule.RULE_REMOVE,
+                        result=KeylistUpdated.RESULT_SUCCESS,
+                    ),
+                ],
+            )
+
+            response.assign_thread_id(TEST_THREAD_ID)
+            await handler.notify_keylist_updated(
+                self.context.profile, TEST_CONN_ID, response
+            )
+        assert self.mock_event_bus.events
+        assert (
+            self.mock_event_bus.events[0][1].topic
+            == MediationManager.KEYLIST_UPDATED_EVENT
+        )
+        assert self.mock_event_bus.events[0][1].payload == {
+            "connection_id": TEST_CONN_ID,
+            "thread_id": TEST_THREAD_ID,
+            "updated": [result.serialize() for result in response.updated],
+            "mediated_connections": {
+                TEST_ROUTE_VERKEY: "conn_id_1",
+                TEST_VERKEY: "conn_id_2",
+            },
+        }
+
+    async def test_notify_keylist_updated_x_unknown_recip_key(self):
+        """test notify_keylist_updated."""
+        handler = KeylistUpdateResponseHandler()
+        response = KeylistUpdateResponse(
+            updated=[
+                KeylistUpdated(
+                    recipient_key=TEST_ROUTE_VERKEY,
+                    action=KeylistUpdateRule.RULE_ADD,
+                    result=KeylistUpdated.RESULT_SUCCESS,
+                ),
+                KeylistUpdated(
+                    recipient_key=TEST_VERKEY,
+                    action=KeylistUpdateRule.RULE_REMOVE,
+                    result=KeylistUpdated.RESULT_SUCCESS,
+                ),
+            ],
+        )
+
+        response.assign_thread_id(TEST_THREAD_ID)
+        with pytest.raises(HandlerException):
+            await handler.notify_keylist_updated(
+                self.context.profile, TEST_CONN_ID, response
+            )

aries_cloudagent/protocols/coordinate_mediation/v1_0/manager.py

@@ -3,17 +3,14 @@
 import logging
 from typing import Optional, Sequence, Tuple
 
-
 from ....core.error import BaseError
 from ....core.profile import Profile, ProfileSession
-from ....connections.models.conn_record import ConnRecord
 from ....storage.base import BaseStorage
 from ....storage.error import StorageNotFoundError
 from ....storage.record import StorageRecord
 from ....wallet.base import BaseWallet
 from ....wallet.did_info import DIDInfo
 from ....wallet.did_method import DIDMethod
-from ....wallet.error import WalletNotFoundError
 from ....wallet.key_type import KeyType
 from ...routing.v1_0.manager import RoutingManager
 from ...routing.v1_0.models.route_record import RouteRecord
@@ -602,48 +599,6 @@ async def store_update_results(
             for record_for_removal in to_remove:
                 await record_for_removal.delete_record(session)
 
-    async def _conn_id_from_recipient_key(
-        self, session: ProfileSession, wallet: BaseWallet, recipient_key: str
-    ) -> str:
-        try:
-            conn = await ConnRecord.retrieve_by_invitation_key(
-                session, invitation_key=normalize_from_did_key(recipient_key)
-            )
-        except StorageNotFoundError:
-            did_info = await wallet.get_local_did_for_verkey(
-                normalize_from_did_key(recipient_key)
-            )
-            conn = await ConnRecord.retrieve_by_did(session, my_did=did_info.did)
-        return conn.connection_id
-
-    async def notify_keylist_updated(
-        self, connection_id: str, response: KeylistUpdateResponse
-    ):
-        """Notify of keylist update response received."""
-        async with self._profile.session() as session:
-            wallet = session.inject(BaseWallet)
-            try:
-                routes = {
-                    updated.recipient_key: await self._conn_id_from_recipient_key(
-                        session, wallet, updated.recipient_key
-                    )
-                    for updated in response.updated
-                }
-            except (StorageNotFoundError, WalletNotFoundError) as err:
-                raise MediationManagerError(
-                    "Unknown recipient key received in keylist update response"
-                ) from err
-
-        await self._profile.notify(
-            self.KEYLIST_UPDATED_EVENT,
-            {
-                "connection_id": connection_id,
-                "thread_id": response._thread_id,
-                "updated": [update.serialize() for update in response.updated],
-                "mediated_connections": routes,
-            },
-        )
-
     async def get_my_keylist(
         self, connection_id: Optional[str] = None
     ) -> Sequence[RouteRecord]:

aries_cloudagent/protocols/coordinate_mediation/v1_0/route_manager.py

@@ -20,6 +20,7 @@
 from .manager import MediationManager
 from .messages.keylist_update import KeylistUpdate
 from .models.mediation_record import MediationRecord
+from .normalization import normalize_from_did_key
 
 
 LOGGER = logging.getLogger(__name__)
@@ -242,6 +243,27 @@ async def routing_info(
     ) -> Tuple[List[str], str]:
         """Retrieve routing keys."""
 
+    async def connection_from_recipient_key(
+        self, profile: Profile, recipient_key: str
+    ) -> ConnRecord:
+        """Retrieve connection for a recipient_key.
+
+        The recipient key is expected to be a local key owned by this agent.
+        """
+        async with profile.session() as session:
+            wallet = session.inject(BaseWallet)
+            try:
+                conn = await ConnRecord.retrieve_by_invitation_key(
+                    session, invitation_key=normalize_from_did_key(recipient_key)
+                )
+            except StorageNotFoundError:
+                did_info = await wallet.get_local_did_for_verkey(
+                    normalize_from_did_key(recipient_key)
+                )
+                conn = await ConnRecord.retrieve_by_did(session, my_did=did_info.did)
+
+            return conn
+
 
 class CoordinateMediationV1RouteManager(RouteManager):
     """Manage routes using Coordinate Mediation protocol."""