Merge remote-tracking branch 'ShadowJonathan/master'

Author: tulir

mautrix/appservice/api/intent.py

@@ -92,7 +92,7 @@ async def wrapper(*args, __self=self, __method=method, **kwargs):
 
             setattr(self, method.__name__, wrapper)
 
-    def user(self, user_id: UserID, token: Optional[str] = None) -> 'IntentAPI':
+    def user(self, user_id: UserID, token: Optional[str] = None, base_url: Optional[str] = None) -> 'IntentAPI':
         """
         Get the intent API for a specific user.
         This is just a proxy to :meth:`AppServiceAPI.intent`.
@@ -103,15 +103,16 @@ def user(self, user_id: UserID, token: Optional[str] = None) -> 'IntentAPI':
         Args:
             user_id: The Matrix ID of the user whose intent API to get.
             token: The access token to use for the Matrix ID.
+            base_url: An optional URL to use for API requests.
 
         Returns:
             The IntentAPI for the given user.
         """
         if not self.bot:
-            return self.api.intent(user_id, token)
+            return self.api.intent(user_id, token, base_url)
         else:
             self.log.warning("Called IntentAPI#user() of child intent object.")
-            return self.bot.api.intent(user_id, token)
+            return self.bot.api.intent(user_id, token, base_url)
 
     # region User actions
 

mautrix/bridge/custom_puppet.py

@@ -11,6 +11,8 @@
 import hashlib
 import hmac
 import json
+import aiohttp
+import urllib.parse
 
 from aiohttp import ClientConnectionError
 
@@ -37,6 +39,27 @@ def __init__(self):
         super().__init__("You may only replace your puppet with your own Matrix account.")
 
 
+class CouldNotDetermineHomeServerURL(CustomPuppetError):
+    """
+    Will be raised when any are true:
+    - .well-known/matrix/client returns 200 with mangled JSON body
+    - .well-known's JSON key [""m.homeserver"]["base_url"] does not exist
+    - .well-known's JSON key [""m.homeserver"]["base_url"] is not a valid URL
+    - .well-known's supplied homeserver URL, or the base domain URL, errors when validating it's version endpoint
+
+    This is in accordance with: https://matrix.org/docs/spec/client_server/r0.6.1#id178
+    """
+
+    def __init__(self, domain: str):
+        super().__init__(f"Could not discover a valid homeserver URL from domain {domain}")
+
+
+class OnlyLoginLocalDomain(CustomPuppetError):
+    """Will be raised when CustomPuppetMixin.allow_external_custom_puppets is set to False"""
+    def __init__(self, domain: str):
+        super().__init__(f"You may only replace your puppet with an account from {domain}")
+
+
 class CustomPuppetMixin(ABC):
     """
     Mixin for the Puppet class to enable Matrix puppeting.
@@ -63,6 +86,7 @@ class CustomPuppetMixin(ABC):
     """
 
     sync_with_custom_puppets: bool = True
+    allow_external_custom_puppets: bool = False
     only_handle_own_synced_events: bool = True
     login_shared_secret: Optional[bytes] = None
     login_device_name: Optional[str] = None
@@ -78,6 +102,7 @@ class CustomPuppetMixin(ABC):
     default_mxid_intent: IntentAPI
     custom_mxid: Optional[UserID]
     access_token: Optional[str]
+    base_url: Optional[str]
     next_batch: Optional[SyncToken]
 
     intent: IntentAPI
@@ -99,9 +124,63 @@ def is_real_user(self) -> bool:
         return bool(self.custom_mxid and self.access_token)
 
     def _fresh_intent(self) -> IntentAPI:
-        return (self.az.intent.user(self.custom_mxid, self.access_token)
+        return (self.az.intent.user(self.custom_mxid, self.access_token, self.base_url)
                 if self.is_real_user else self.default_mxid_intent)
 
+    async def _discover_homeserver_endpoint(self, domain: str) -> str:
+        domain_is_valid = False
+
+        async def validate_versions_api(base_url: str) -> bool:
+
+            async with self.az.http_session.get(urllib.parse.urljoin(base_url, "_matrix/client/versions")) as response:
+                if response.status != 200:
+                    return False
+
+                try:
+                    obj = await response.json(content_type=None)
+                    if len(obj["versions"]) > 1:
+                        return True
+                except (KeyError, json.JSONDecodeError):
+                    return False
+
+        async def get_well_known_homeserver_base_url(probable_domain: str) -> Optional[str]:
+            async with self.az.http_session.get(f"https://{probable_domain}/.well-known/matrix/client") as response:
+                if response.status != 200:
+                    return None
+
+                try:
+                    obj = await response.json(content_type=None)
+                    return obj["m.homeserver"]["base_url"]
+                except (KeyError, json.JSONDecodeError) as e:
+                    raise CouldNotDetermineHomeServerURL(domain) from e
+
+        try:
+            if await validate_versions_api(f"https://{domain}"):
+                # Flag front domain as valid, but keep looking
+                domain_is_valid = True
+        except aiohttp.ClientError:
+            pass
+
+        try:
+            base_url = await get_well_known_homeserver_base_url(domain)
+
+            if base_url is None:
+                if domain_is_valid:
+                    # If we found a valid domain already, we just return that
+                    return f"https://{domain}"
+                else:
+                    raise CouldNotDetermineHomeServerURL(domain)
+
+            if await validate_versions_api(base_url):
+                return base_url
+            elif await validate_versions_api(base_url + "/"):
+                return base_url + "/"
+        except aiohttp.ClientError as e:
+            if domain_is_valid:
+                # Earlier we already found a valid domain, so we ignore the error and return the base domain instead
+                return f"https://{domain}"
+            raise CouldNotDetermineHomeServerURL(domain) from e
+
     @classmethod
     def can_auto_login(cls, mxid: UserID) -> bool:
         if not cls.login_shared_secret:
@@ -131,7 +210,8 @@ async def _login_with_shared_secret(cls, mxid: UserID) -> Optional[str]:
         data = await resp.json()
         return data["access_token"]
 
-    async def switch_mxid(self, access_token: Optional[str], mxid: Optional[UserID]) -> None:
+    async def switch_mxid(self, access_token: Optional[str], mxid: Optional[UserID],
+                          base_url: Optional[str] = None) -> None:
         """
         Switch to a real Matrix user or away from one.
 
@@ -140,15 +220,28 @@ async def switch_mxid(self, access_token: Optional[str], mxid: Optional[UserID])
                           the appservice-owned ID.
             mxid: The expected Matrix user ID of the custom account, or ``None`` when
                   ``access_token`` is None.
+            base_url: An optional base URL to direct API calls to. If ``None``, and ``mxid`` is not ``None``,
+                      and ``mxid`` ``server_part`` is the not the appservice domain, autodiscovery is tried.
         """
         if access_token == "auto":
             access_token = await self._login_with_shared_secret(mxid)
             if not access_token:
                 raise ValueError("Failed to log in with shared secret")
             self.log.debug(f"Logged in for {mxid} using shared secret")
+
+        if mxid is not None:
+            mxid_domain = self.az.intent.parse_user_id(mxid)[1]
+            if mxid_domain != self.az.domain:
+                if not self.allow_external_custom_puppets:
+                    raise OnlyLoginLocalDomain(self.az.domain)
+                elif base_url is None:
+                    # This can throw CouldNotDetermineHomeServerURL
+                    base_url = await self._discover_homeserver_endpoint(mxid_domain)
+
         prev_mxid = self.custom_mxid
         self.custom_mxid = mxid
         self.access_token = access_token
+        self.base_url = base_url
         self.intent = self._fresh_intent()
 
         await self.start()