Browse Source

Add config flag to toggle using proxy for media downloads

Nick Barrett 2 years ago
parent
commit
9129a5d16c
3 changed files with 33 additions and 23 deletions
  1. 1 0
      mautrix_instagram/config.py
  2. 2 0
      mautrix_instagram/example-config.yaml
  3. 30 23
      mautrix_instagram/portal.py

+ 1 - 0
mautrix_instagram/config.py

@@ -103,6 +103,7 @@ class Config(BaseBridgeConfig):
         copy("bridge.command_prefix")
 
         copy("bridge.get_proxy_api_url")
+        copy("bridge.use_proxy_for_media")
 
         copy_dict("bridge.permissions")
 

+ 2 - 0
mautrix_instagram/example-config.yaml

@@ -212,6 +212,8 @@ bridge:
 
     # URL to call to retrieve a proxy URL from (defaults to the http_proxy environment variable).
     get_proxy_api_url: null
+    # Whether to use proxy for downloading media from Instagram.
+    use_proxy_for_media: true
 
     # End-to-bridge encryption support options.
     #

+ 30 - 23
mautrix_instagram/portal.py

@@ -28,7 +28,7 @@ import re
 import sqlite3
 import time
 
-from aiohttp import ClientSession
+from aiohttp import ClientResponse, ClientSession
 from yarl import URL
 import asyncpg
 import magic
@@ -870,28 +870,35 @@ class Portal(DBPortal, BasePortal):
     async def _download_instagram_file(
         self, source: u.User, url: str
     ) -> tuple[Optional[bytes], str]:
-        async with ClientSession() as session:
-            async with session.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")
-                self.log.debug(f"Downloading file with length {length}: {url}")
-                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 handle_resp(resp: ClientResponse) -> tuple[Optional[bytes], str]:
+            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")
+            self.log.debug(f"Downloading file with length {length}: {url}")
+            data = await resp.read()
+            if not data:
+                return None, ""
+            mimetype = resp.headers["Content-Type"] or magic.from_buffer(data, mime=True)
+            return data, mimetype
+
+        if self.config["bridge.use_proxy_for_media"]:
+            async with source.client.raw_http_get(url) as resp:
+                return await handle_resp(resp)
+        else:
+            async with ClientSession() as session:
+                async with session.get(url) as resp:
+                    return await handle_resp(resp)
 
     async def _reupload_instagram_file(
         self,