Use proxy retry util for all API requests (#104)

* Pass `on_proxy_update` through to `AndroidAPI`

* Use `proxy_with_retry` in base HTTP methods in `AndroidAPI`

* Revert "Retry all provisioning related HTTP requests"

This reverts commit ff28865699784af8edb0e4da3d24b6180be71333, which
is no longer needed as the retry logic has been moved to the base
API class.

* Revert "Use proxy retry wrapper when getting mobileconfig"

This reverts commit ce47ad1a25c0f02ff3b5a8cd4108d12d03bef618.

* Set `min_wait_seconds=1` for HTTP retries on proxy error

* Add `AndroidAPI.proxy_with_retry`

* Wrap raw HTTP calls using `proxy_with_retry`

* Fix typo

Co-authored-by: Tulir Asokan <tulir@maunium.net>

---------

Co-authored-by: Tulir Asokan <tulir@maunium.net>

mauigpapi/http/base.py

@@ -15,7 +15,8 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 from __future__ import annotations
 
-from typing import Any, Type, TypeVar
+from typing import Any, Awaitable, Callable, Type, TypeVar
+from functools import partial
 import json
 import logging
 import time
@@ -25,7 +26,7 @@ from yarl import URL
 
 from mautrix.types import JSON, Serializable
 from mautrix.util.logging import TraceLogger
-from mautrix.util.proxy import ProxyHandler
+from mautrix.util.proxy import ProxyHandler, proxy_with_retry
 
 from ..errors import (
     IG2FACodeExpiredError,
@@ -80,14 +81,24 @@ class BaseAndroidAPI:
         state: AndroidState,
         log: TraceLogger | None = None,
         proxy_handler: ProxyHandler | None = None,
+        on_proxy_update: Callable[[], Awaitable[None]] | None = None,
     ) -> None:
         self.log = log or logging.getLogger("mauigpapi.http")
 
         self.proxy_handler = proxy_handler
+        self.on_proxy_update = on_proxy_update
         self.setup_http(cookie_jar=state.cookies.jar)
 
         self.state = state
 
+        self.proxy_with_retry = partial(
+            proxy_with_retry,
+            logger=self.log,
+            proxy_handler=self.proxy_handler,
+            on_proxy_change=self.on_proxy_update,
+            min_wait_seconds=1,  # we want to retry these pretty fast
+        )
+
     @staticmethod
     def sign(req: Any, filter_nulls: bool = False) -> dict[str, str]:
         if isinstance(req, Serializable):
@@ -178,7 +189,10 @@ class BaseAndroidAPI:
         if not raw:
             data = self.sign(data, filter_nulls=filter_nulls)
         url = self.url.with_path(path).with_query(query or {})
-        resp = await self.http.post(url=url, headers=headers, data=data)
+        resp = await self.proxy_with_retry(
+            f"AndroidAPI.std_http_post: {url}",
+            lambda: self.http.post(url=url, headers=headers, data=data),
+        )
         self.log.trace(f"{path} response: {await resp.text()}")
         if response_type is str or response_type is None:
             self._handle_response_headers(resp)
@@ -199,7 +213,11 @@ class BaseAndroidAPI:
     ) -> T:
         headers = {**self._headers, **headers} if headers else self._headers
         query = {k: v for k, v in (query or {}).items() if v is not None}
-        resp = await self.http.get(url=self.url.with_path(path).with_query(query), headers=headers)
+        url = self.url.with_path(path).with_query(query)
+        resp = await self.proxy_with_retry(
+            f"AndroidAPI.std_http_get: {url}",
+            lambda: self.http.get(url=url, headers=headers),
+        )
         self.log.trace(f"{path} response: {await resp.text()}")
         if response_type is None:
             self._handle_response_headers(resp)

mautrix_instagram/commands/auth.py

@@ -34,7 +34,6 @@ from mauigpapi.state import AndroidState
 from mauigpapi.types import BaseResponseUser
 from mautrix.bridge.commands import HelpSection, command_handler
 from mautrix.types import EventID
-from mautrix.util.proxy import proxy_with_retry
 
 from .. import user as u
 from .typehint import CommandEvent
@@ -50,14 +49,13 @@ async def get_login_state(user: u.User, seed: str) -> tuple[AndroidAPI, AndroidS
         state = AndroidState()
         seed = hmac.new(seed.encode("utf-8"), user.mxid.encode("utf-8"), hashlib.sha256).digest()
         state.device.generate(seed)
-        api = AndroidAPI(state, log=user.api_log, proxy_handler=user.proxy_handler)
-        await proxy_with_retry(
-            "get_login_state",
-            lambda: api.get_mobile_config(),
-            logger=user.log,
+        api = AndroidAPI(
+            state,
+            log=user.api_log,
             proxy_handler=user.proxy_handler,
-            on_proxy_change=user.on_proxy_update,
+            on_proxy_update=user.on_proxy_update,
         )
+        await api.get_mobile_config()
         user.command_status = {
             "action": "Login",
             "state": state,

mautrix_instagram/portal.py

@@ -870,26 +870,29 @@ class Portal(DBPortal, BasePortal):
     async def _download_instagram_file(
         self, source: u.User, url: str
     ) -> tuple[Optional[bytes], str]:
-        async with source.client.raw_http_get(url) as resp:
-            try:
-                length = int(resp.headers["Content-Length"])
-            except KeyError:
-                # TODO can the download be short-circuited if there's too much data?
-                self.log.warning(
-                    "Got file download response with no Content-Length header,"
-                    "reading data dangerously"
-                )
-                length = 0
-            if length > self.matrix.media_config.upload_size:
-                self.log.debug(
-                    f"{url} was too large ({length} > {self.matrix.media_config.upload_size})"
-                )
-                raise ValueError("Attachment not available: too large")
-            data = await resp.read()
-            if not data:
-                return None, ""
-            mimetype = resp.headers["Content-Type"] or magic.from_buffer(data, mime=True)
-            return data, mimetype
+        async def download():
+            async with source.client.raw_http_get(url) as resp:
+                try:
+                    length = int(resp.headers["Content-Length"])
+                except KeyError:
+                    # TODO can the download be short-circuited if there's too much data?
+                    self.log.warning(
+                        "Got file download response with no Content-Length header,"
+                        "reading data dangerously"
+                    )
+                    length = 0
+                if length > self.matrix.media_config.upload_size:
+                    self.log.debug(
+                        f"{url} was too large ({length} > {self.matrix.media_config.upload_size})"
+                    )
+                    raise ValueError("Attachment not available: too large")
+                data = await resp.read()
+                if not data:
+                    return None, ""
+                mimetype = resp.headers["Content-Type"] or magic.from_buffer(data, mime=True)
+                return data, mimetype
+
+        return await source.client.proxy_with_retry("Portal._download_instagram_file", download)
 
     async def _reupload_instagram_file(
         self,

mautrix_instagram/puppet.py

@@ -160,9 +160,12 @@ class Puppet(DBPuppet, BasePuppet):
             if info.has_anonymous_profile_picture:
                 mxc = ""
             else:
-                async with source.client.raw_http_get(info.profile_pic_url) as resp:
-                    content_type = resp.headers["Content-Type"]
-                    resp_data = await resp.read()
+                resp = await source.client.proxy_with_retry(
+                    "Puppet._update_avatar",
+                    lambda: source.client.raw_http_get(info.profile_pic_url),
+                )
+                content_type = resp.headers["Content-Type"]
+                resp_data = await resp.read()
                 mxc = await self.default_mxid_intent.upload_media(
                     data=resp_data,
                     mime_type=content_type,

mautrix_instagram/user.py

@@ -242,6 +242,7 @@ class User(DBUser, BaseUser):
             self.state,
             log=self.api_log,
             proxy_handler=self.proxy_handler,
+            on_proxy_update=self.on_proxy_update,
         )
 
         if not user:
@@ -441,7 +442,7 @@ class User(DBUser, BaseUser):
             self.log.exception("Failed to update own puppet info")
         try:
             if puppet.custom_mxid != self.mxid and puppet.can_auto_login(self.mxid):
-                self.log.info(f"Automatically enabling custom puppet")
+                self.log.info("Automatically enabling custom puppet")
                 await puppet.switch_mxid(access_token="auto", mxid=self.mxid)
         except Exception:
             self.log.exception("Failed to automatically enable custom puppet")

mautrix_instagram/web/provisioning_api.py

@@ -50,7 +50,6 @@ from mauigpapi.errors import (
 from mauigpapi.types import ChallengeStateResponse, FacebookLoginResponse, LoginResponse
 from mautrix.types import JSON, Serializable, UserID
 from mautrix.util.logging import TraceLogger
-from mautrix.util.proxy import proxy_with_retry
 
 from .. import user as u
 from ..commands.auth import get_login_state
@@ -278,13 +277,7 @@ class ProvisioningAPI:
         track(user, "$login_start", {"type": "instagram"})
         api, state = await get_login_state(user, self.device_seed)
         try:
-            resp = await proxy_with_retry(
-                "provisioning_api.login",
-                lambda: api.login(username, password),
-                logger=self.log,
-                proxy_handler=user.proxy_handler,
-                on_proxy_change=user.on_proxy_update,
-            )
+            resp = await api.login(username, password)
         except IGLoginTwoFactorRequiredError as e:
             return self._2fa_required(user, username, state, e)
         except IGChallengeError as e:
@@ -399,13 +392,7 @@ class ProvisioningAPI:
             username = state.login_2fa_username
         track(user, "$login_resend_2fa_sms")
         try:
-            resp = await proxy_with_retry(
-                "provisioning_api.send_two_factor_login_sms",
-                lambda: api.send_two_factor_login_sms(username, identifier=identifier),
-                logger=self.log,
-                proxy_handler=user.proxy_handler,
-                on_proxy_change=user.on_proxy_update,
-            )
+            resp = await api.send_two_factor_login_sms(username, identifier=identifier)
         except IGRateLimitError as e:
             track(user, "$login_resend_2fa_sms_fail", {"error": "ratelimit"})
             try:
@@ -464,14 +451,8 @@ class ProvisioningAPI:
             username = state.login_2fa_username
         track(user, "$login_submit_2fa")
         try:
-            resp = await proxy_with_retry(
-                "provisioning_api.two_factor_login",
-                lambda: api.two_factor_login(
-                    username, code=code, identifier=identifier, is_totp=is_totp
-                ),
-                logger=self.log,
-                proxy_handler=user.proxy_handler,
-                on_proxy_change=user.on_proxy_update,
+            resp = await api.two_factor_login(
+                username, code=code, identifier=identifier, is_totp=is_totp
             )
         except IGBad2FACodeError:
             self.log.debug("%s submitted an incorrect 2-factor auth code", user.mxid)
@@ -511,13 +492,7 @@ class ProvisioningAPI:
         self, user: u.User, state: AndroidState, api: AndroidAPI, err: IGChallengeError, after: str
     ) -> web.Response:
         try:
-            resp = await proxy_with_retry(
-                "provisioning_api.challenge_auto",
-                lambda: api.challenge_auto(reset=after == "2fa"),
-                logger=self.log,
-                proxy_handler=user.proxy_handler,
-                on_proxy_change=user.on_proxy_update,
-            )
+            resp = await api.challenge_auto(reset=after == "2fa")
         except Exception:
             self.log.exception("Challenge reset failed for %s", user.mxid)
             track(user, "$login_failed", {"error": "challenge-reset-fail", "after": after})
@@ -583,13 +558,7 @@ class ProvisioningAPI:
         state: AndroidState = user.command_status["state"]
         track(user, "$login_submit_challenge")
         try:
-            resp = await proxy_with_retry(
-                "provisioning_api.challenge_send_security_code",
-                lambda: api.challenge_send_security_code(code=code),
-                logger=self.log,
-                proxy_handler=user.proxy_handler,
-                on_proxy_change=user.on_proxy_update,
-            )
+            resp = await api.challenge_send_security_code(code=code)
         except IGChallengeWrongCodeError:
             self.log.debug("%s submitted an incorrect challenge code", user.mxid)
             track(user, "$login_failed", {"error": "incorrect-challenge-code"})
@@ -643,13 +612,7 @@ class ProvisioningAPI:
             else "<username not known>"
         )
         try:
-            resp = await proxy_with_retry(
-                "provisioning_api.finish_login",
-                lambda: api.current_user(),
-                logger=self.log,
-                proxy_handler=user.proxy_handler,
-                on_proxy_change=user.on_proxy_update,
-            )
+            resp = await api.current_user()
         except IGChallengeError as e:
             if isinstance(login_resp, ChallengeStateResponse):
                 track(user, "$login_failed", {"error": "repeat-challenge", "after": after})
@@ -741,13 +704,7 @@ class ProvisioningAPI:
         track(user, "$login_start", {"type": "facebook"})
         api, state = await get_login_state(user, self.device_seed)
         try:
-            resp: FacebookLoginResponse = await proxy_with_retry(
-                "provisioning_api.post_fb_login_token",
-                lambda: api.facebook_signup(fb_access_token),
-                logger=self.log,
-                proxy_handler=user.proxy_handler,
-                on_proxy_change=user.on_proxy_update,
-            )
+            resp = await api.facebook_signup(fb_access_token)
             if not resp.account_created:
                 return self._no_fb_account(user, resp=resp)
         except IGFBNoContactPointFoundError as e: