Add nginx anti ddos by whitelisting validators ip

image_generation_subnet/base/miner.py

@@ -23,6 +23,8 @@
 import bittensor as bt
 
 from image_generation_subnet.base.neuron import BaseNeuron
+from image_generation_subnet.miner.constants import NGINX_CONF
+import os
 
 
 class BaseMinerNeuron(BaseNeuron):
@@ -42,7 +44,9 @@ def __init__(self, config=None):
             bt.logging.warning(
                 "You are allowing non-registered entities to send requests to your miner. This is a security risk."
             )
-
+        if self.config.miner.use_nginx:
+            self.config.axon.external_port = self.config.axon.port
+            self.config.axon.port = 8091
         # The axon handles request processing, allowing validators to send this miner requests.
         self.axon = bt.axon(wallet=self.wallet, config=self.config)
 
@@ -101,7 +105,6 @@ def run(self):
             f"Serving miner axon {self.axon} on network: {self.config.subtensor.chain_endpoint} with netuid: {self.config.netuid}"
         )
         self.axon.serve(netuid=self.config.netuid, subtensor=self.subtensor)
-
         # Start  starts the miner's axon, making it active on the network.
         self.axon.start()
 
@@ -198,3 +201,20 @@ def resync_metagraph(self):
 
         # Sync the metagraph.
         self.metagraph.sync(subtensor=self.subtensor)
+        if self.config.miner.use_nginx:
+            try:
+                axons = self.metagraph.axons
+                whitelist = []
+                for k in self.volume_per_validator.keys():
+                    whitelist.append(f"allow {axons[int(k)].ip};")
+                whitelist = "\n".join(whitelist)
+                nginx_conf = NGINX_CONF.replace(
+                    "{{external_axon_port}}", str(self.config.axon.external_port)
+                ).replace("{{internal_axon_port}}", str(self.config.axon.port))
+                with open("/etc/nginx/nginx.conf", "w") as f:
+                    f.write(nginx_conf)
+                os.system("nginx -s reload")
+                bt.logging.info("Nginx configuration updated.")
+            except Exception as e:
+                bt.logging.error(f"Error in updating nginx configuration: {e}")
+                traceback.print_exc()

image_generation_subnet/miner/__init__.py

@@ -1,9 +1,11 @@
 from .blacklist import check_limit, check_min_stake
 from .forward import set_info, generate
+from .constants import NGINX_CONF
 
 __all__ = [
     "check_limit",
     "check_min_stake",
     "set_info",
     "generate",
+    "NGINX_CONF",
 ]

image_generation_subnet/miner/constants.py

@@ -0,0 +1,27 @@
+NGINX_CONF = """
+worker_processes 1;
+events {
+    worker_connections 1024;
+}
+
+http {
+    limit_conn_zone $binary_remote_addr zone=addr:10m;
+    ignore_invalid_headers off;
+    client_max_body_size 0;
+    proxy_intercept_errors on;
+    server {
+        listen {{external_axon_port}};  # Port to listen on
+
+        # Whitelist IP addresses
+        # allow 123.45.67.89;  # Replace with your allowed IPs
+        # allow 98.76.54.32;   # You can add multiple allowed IPs
+        # deny all;            # Deny all other IPs
+
+        location / {
+            proxy_pass http://127.0.0.1:{{internal_axon_port}};
+            limit_conn addr 20;
+        }
+    }
+}
+
+"""

image_generation_subnet/utils/config.py

@@ -422,6 +422,19 @@ def add_args(cls, parser):
             default=4,
         )
 
+        parser.add_argument(
+            "--miner.use_nginx",
+            action="store_true",
+            help="If set, we will use nginx to serve requests.",
+            default=False,
+        )
+        parser.add_argument(
+            "--miner.nginx_port",
+            type=int,
+            help="The port to run nginx on.",
+            default=8091,
+        )
+
 
 def config(cls):
     """

neurons/miner/miner.py

@@ -185,17 +185,17 @@ async def priority(self, synapse: ImageGenerating) -> float:
                 )
                 start_time = time.time()
                 miner.total_request_in_interval = 0
-            try:
-                bt.logging.debug("Syncing metagraph")
-                miner.resync_metagraph()
-                bt.logging.debug("Synced metagraph")
-                miner.volume_per_validator = image_generation_subnet.utils.volume_setting.get_volume_per_validator(
-                    miner.metagraph,
-                    miner.config.miner.total_volume,
-                    miner.config.miner.size_preference_factor,
-                    miner.config.miner.min_stake,
-                    log=False,
-                )
-            except Exception as e:
-                print(e)
+                try:
+                    bt.logging.debug("Syncing metagraph")
+                    miner.resync_metagraph()
+                    bt.logging.debug("Synced metagraph")
+                    miner.volume_per_validator = image_generation_subnet.utils.volume_setting.get_volume_per_validator(
+                        miner.metagraph,
+                        miner.config.miner.total_volume,
+                        miner.config.miner.size_preference_factor,
+                        miner.config.miner.min_stake,
+                        log=False,
+                    )
+                except Exception as e:
+                    print(e)
             time.sleep(60)