Add proxy util from IG/FB bridges (#134)

* Add proxy util from IG/FB bridges

* Add comment about blocking request

Author: Fizzadar

mautrix/util/proxy.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from typing import Awaitable, Callable, TypeVar
+import asyncio
+import json
+import logging
+import urllib.request
+
+from aiohttp import ClientConnectionError
+from yarl import URL
+
+from mautrix.util.logging import TraceLogger
+
+try:
+    from aiohttp_socks import ProxyConnectionError, ProxyError, ProxyTimeoutError
+except ImportError:
+
+    class ProxyError(Exception):
+        pass
+
+    ProxyConnectionError = ProxyTimeoutError = ProxyError
+
+RETRYABLE_PROXY_EXCEPTIONS = (
+    ProxyError,
+    ProxyTimeoutError,
+    ProxyConnectionError,
+    ClientConnectionError,
+    ConnectionError,
+    asyncio.TimeoutError,
+)
+
+
+class ProxyHandler:
+    current_proxy_url: str | None = None
+    log = logging.getLogger("mau.proxy")
+
+    def __init__(self, api_url: str | None) -> None:
+        self.api_url = api_url
+
+    def get_proxy_url_from_api(self, reason: str | None = None) -> str | None:
+        assert self.api_url is not None
+
+        api_url = str(URL(self.api_url).update_query({"reason": reason} if reason else {}))
+
+        # NOTE: using urllib.request to intentionally block the whole bridge until the proxy change applied
+        request = urllib.request.Request(api_url, method="GET")
+        self.log.debug("Requesting proxy from: %s", api_url)
+
+        try:
+            with urllib.request.urlopen(request) as f:
+                response = json.loads(f.read().decode())
+        except Exception:
+            self.log.exception("Failed to retrieve proxy from API")
+        else:
+            return response["proxy_url"]
+
+        return None
+
+    def update_proxy_url(self, reason: str | None = None) -> bool:
+        old_proxy = self.current_proxy_url
+        new_proxy = None
+
+        if self.api_url is not None:
+            new_proxy = self.get_proxy_url_from_api(reason)
+        else:
+            new_proxy = urllib.request.getproxies().get("http")
+
+        if old_proxy != new_proxy:
+            self.log.debug("Set new proxy URL: %s", new_proxy)
+            self.current_proxy_url = new_proxy
+            return True
+
+        self.log.debug("Got same proxy URL: %s", new_proxy)
+        return False
+
+    def get_proxy_url(self) -> str | None:
+        if not self.current_proxy_url:
+            self.update_proxy_url()
+
+        return self.current_proxy_url
+
+
+T = TypeVar("T")
+
+
+async def proxy_with_retry(
+    name: str,
+    func: Callable[[], Awaitable[T]],
+    logger: TraceLogger,
+    proxy_handler: ProxyHandler,
+    on_proxy_change: Callable[[], Awaitable[None]],
+    max_retries: int = 10,
+) -> T:
+    errors = 0
+
+    while True:
+        try:
+            return await func()
+        except RETRYABLE_PROXY_EXCEPTIONS as e:
+            errors += 1
+            if errors > max_retries:
+                raise
+            wait = min(errors * 10, 60)
+            logger.warning(
+                "%s while trying to %s, retrying in %d seconds",
+                e.__class__.__name__,
+                name,
+                wait,
+            )
+            if errors > 1 and proxy_handler.update_proxy_url(
+                f"{e.__class__.__name__} while trying to {name}"
+            ):
+                await on_proxy_change()