Merge pull request fastapi-admin#59 from PredaaA/feat/gateway-proxy-s…

…upport

Implement gateway proxy support

ballsdex/__main__.py

@@ -5,6 +5,8 @@
 import asyncio
 import logging
 import logging.handlers
+
+import yarl
 import discord
 import argparse
 
@@ -84,6 +86,61 @@ def print_welcome():
     print("")
 
 
+def patch_gateway(proxy_url: str):
+    """This monkeypatches discord.py in order to be able to use a custom gateway URL.
+
+    Parameters
+    ----------
+    proxy_url : str
+        The URL of the gateway proxy to use.
+    """
+
+    class ProductionHTTPClient(discord.http.HTTPClient):
+        async def get_gateway(self, **_):
+            return f"{proxy_url}?encoding=json&v=10"
+
+        async def get_bot_gateway(self, **_):
+            try:
+                data = await self.request(discord.http.Route("GET", "/gateway/bot"))
+            except discord.HTTPException as exc:
+                raise discord.GatewayNotFound() from exc
+            return data["shards"], f"{proxy_url}?encoding=json&v=10"
+
+    class ProductionDiscordWebSocket(discord.gateway.DiscordWebSocket):
+        def is_ratelimited(self):
+            return False
+
+        async def debug_send(self, data, /):
+            self._dispatch("socket_raw_send", data)
+            await self.socket.send_str(data)
+
+        async def send(self, data, /):
+            await self.socket.send_str(data)
+
+    class ProductionReconnectWebSocket(Exception):
+        def __init__(self, shard_id: int | None, *, resume: bool = False):
+            self.shard_id: int | None = shard_id
+            self.resume: bool = False
+            self.op: str = "IDENTIFY"
+
+    def is_ws_ratelimited(self):
+        return False
+
+    async def before_identify_hook(self, shard_id: int, *, initial: bool = False):
+        pass
+
+    discord.http.HTTPClient.get_gateway = ProductionHTTPClient.get_gateway
+    discord.http.HTTPClient.get_bot_gateway = ProductionHTTPClient.get_bot_gateway
+    discord.gateway.DiscordWebSocket._keep_alive = None
+    discord.gateway.DiscordWebSocket.is_ratelimited = ProductionDiscordWebSocket.is_ratelimited
+    discord.gateway.DiscordWebSocket.debug_send = ProductionDiscordWebSocket.debug_send
+    discord.gateway.DiscordWebSocket.send = ProductionDiscordWebSocket.send
+    discord.gateway.DiscordWebSocket.DEFAULT_GATEWAY = yarl.URL(proxy_url)
+    discord.gateway.ReconnectWebSocket.__init__ = ProductionReconnectWebSocket.__init__
+    BallsDexBot.is_ws_ratelimited = is_ws_ratelimited
+    BallsDexBot.before_identify_hook = before_identify_hook
+
+
 async def shutdown_handler(bot: BallsDexBot, signal_type: str | None = None):
     if signal_type:
         log.info(f"Received {signal_type}, stopping the bot...")
@@ -214,6 +271,10 @@ def main():
             time.sleep(1)
             sys.exit(0)
 
+        if settings.gateway_url is not None:
+            log.info("Using custom gateway URL: %s", settings.gateway_url)
+            patch_gateway(settings.gateway_url)
+
         prefix = settings.prefix
 
         try:

ballsdex/settings.py

@@ -19,6 +19,8 @@ class Settings:
     ----------
     bot_token: str
         Discord token for the bot to connect
+    gateway_url: str | None
+        The URL of the Discord gateway that this instance of the bot should connect to and use.
     prefix: str
         Prefix for text commands, mostly unused. Defaults to "b."
     collectible_name: str
@@ -46,6 +48,7 @@ class Settings:
     """
 
     bot_token: str = ""
+    gateway_url: str | None = None
     prefix: str = "b."
 
     collectible_name: str = "countryball"
@@ -77,6 +80,7 @@ def read_settings(path: "Path"):
     content = yaml.load(path.read_text(), yaml.Loader)
 
     settings.bot_token = content["discord-token"]
+    settings.gateway_url = content.get("gateway-url")
     settings.prefix = content["text-prefix"]
 
     settings.collectible_name = content["collectible-name"]

config.json.example

@@ -0,0 +1,21 @@
+{
+  "log_level": "info",
+  "token": "",
+  "intents": 33289,
+  "port": 5421,
+  "status": "online",
+  "backpressure": 250,
+  "externally_accessible_url": "ws://localhost:5421",
+  "cache": {
+    "channels": false,
+    "presences": false,
+    "emojis": false,
+    "current_member": false,
+    "members": false,
+    "roles": false,
+    "stage_instances": false,
+    "stickers": false,
+    "users": false,
+    "voice_states": false
+  }
+}

docker-compose.yml

@@ -99,6 +99,19 @@ services:
     volumes: 
       - cache-data:/data
 
+  # Uncomment to enable gateway proxy feature and
+  # add "gateway-url: ws://gateway_proxy:5421" to your config.yaml
+  # Also change tag corresponding to your platform if not x86-64
+  #
+  # gateway-proxy:
+  #   container_name: gateway_proxy
+  #   image: docker.io/gelbpunkt/gateway-proxy:x86-64
+  #   restart: always
+  #   volumes:
+  #     - ./gatewayproxy/config.json:/config.json
+  #   networks:
+  #     - internal
+
 volumes:
   database-data:
   cache-data: