FIXUP

synapse/module_api/__init__.py

@@ -24,10 +24,10 @@
     List,
     Optional,
     Tuple,
+    Union,
 )
 
 import jinja2
-from typing_extensions import TypedDict
 
 from twisted.internet import defer
 from twisted.web.resource import IResource
@@ -87,12 +87,21 @@
 logger = logging.getLogger(__name__)
 
 
-class UserIpAndAgent(TypedDict):
+class UserIpAndAgent:
+    """
+    A ip, user_agent pair used by a user to connect to this homeserver.
+    """
+
     ip: str
     user_agent: str
     # The time at which this user agent/ip was last seen.
     last_seen: int
 
+    def __init__(self, ip: str, user_agent: str, last_seen: int):
+        self.ip = ip
+        self.user_agent = user_agent
+        self.last_seen = last_seen
+
 
 class ModuleApi:
     """A proxy object that gets passed to various plugin modules so they
@@ -715,37 +724,45 @@ def read_templates(
             (td for td in (self.custom_template_dir, custom_template_directory) if td),
         )
 
-    def ts(self) -> int:
-        """Return a timestamp for the current time, in milliseconds since the epoch"""
-        return self._hs.get_clock().time_msec()
-
-    def is_mine(self, domain_specific_string: DomainSpecificString) -> bool:
+    def is_mine(self, id: Union[str, DomainSpecificString]) -> bool:
         """Checks whether an ID comes from this homeserver."""
-        return self._hs.is_mine(domain_specific_string)
-
-    def is_mine_id(self, string: str) -> bool:
-        """Checks whether an ID comes from this homeserver."""
-        return string.split(":", 1)[1] == self._server_name
+        if isinstance(id, DomainSpecificString):
+            return self._hs.is_mine(id)
+        else:
+            return self._hs.is_mine_id(id)
 
     async def get_user_ip_and_agents(
-        self, user: UserID, since_ts: Optional[int] = None
+        self, user_id: str, since_ts: Optional[float] = None
     ) -> List[UserIpAndAgent]:
         """
         Return the list of user IPs and agents for a user.
 
         Only useful for local users.
         """
-        raw_data = await self._store.get_user_ip_and_agents(user)
-        # Sanitize some of the data. We don't want to return tokens.
-        return [
-            {
-                "ip": str(data["ip"]),
-                "user_agent": str(data["user_agent"]),
-                "last_seen": int(data["last_seen"]),
-            }
-            for data in raw_data
-            if since_ts is None or int(data["last_seen"]) >= since_ts
-        ]
+        # Don't hit the db if this is not a local user.
+        is_mine = False
+        try:
+            # Let's be defensive against ill-formed strings.
+            if self.is_mine(user_id):
+                is_mine = True
+        except Exception:
+            pass
+        if is_mine:
+            raw_data = await self._store.get_user_ip_and_agents(
+                UserID.from_string(user_id)
+            )
+            # Sanitize some of the data. We don't want to return tokens.
+            return [
+                UserIpAndAgent(
+                    ip=str(data["ip"]),
+                    user_agent=str(data["user_agent"]),
+                    last_seen=int(data["last_seen"]),
+                )
+                for data in raw_data
+                if since_ts is None or int(data["last_seen"]) >= since_ts
+            ]
+        else:
+            return []
 
 
 class PublicRoomListManager:

tests/module_api/test_api.py

@@ -43,6 +43,7 @@ def prepare(self, reactor, clock, homeserver):
         self.module_api = homeserver.get_module_api()
         self.event_creation_handler = homeserver.get_event_creation_handler()
         self.sync_handler = homeserver.get_sync_handler()
+        self.auth_handler = homeserver.get_auth_handler()
 
     def make_homeserver(self, reactor, clock):
         return self.setup_test_homeserver(
@@ -92,15 +93,40 @@ def test_get_userinfo_by_id__no_user_found(self):
     def test_get_user_ip_and_agents(self):
         user_id = self.register_user("test_get_user_ip_and_agents_user", "1234")
         info = self.get_success(self.module_api.get_user_ip_and_agents(user_id))
-        self.assertIdentical(info, [])
+        self.assertEqual(info, [])
+
+        self.get_success(
+            self.store.insert_client_ip(
+                user_id, "access_token", "ip_1", "user_agent_1", None
+            )
+        )
+        self.get_success(
+            self.store.insert_client_ip(
+                user_id, "access_token", "ip_2", "user_agent_2", None
+            )
+        )
+        info = self.get_success(self.module_api.get_user_ip_and_agents(user_id))
+
+        self.assertEqual(len(info), 2)
+        ip_1_seen = False
+        ip_2_seen = False
+        for i in info:
+            if i.ip == "ip_1":
+                ip_1_seen = True
+                self.assertEqual(i.user_agent, "user_agent_1")
+            elif i.ip == "ip_2":
+                ip_2_seen = True
+                self.assertEqual(i.user_agent, "user_agent_2")
+        self.assertTrue(ip_1_seen)
+        self.assertTrue(ip_2_seen)
 
     def test_get_user_ip_and_agents__no_user_found(self):
         info = self.get_success(
             self.module_api.get_user_ip_and_agents(
-                "@test_get_user_ip_and_agents_user_nonexistent"
+                "@test_get_user_ip_and_agents_user_nonexistent:example.com"
             )
         )
-        self.assertIdentical(info, [])
+        self.assertEqual(info, [])
 
     def test_sending_events_into_room(self):
         """Tests that a module can send events into a room"""