From 607e226fffa7299f24c4fd3c79336b7c1bdddc17 Mon Sep 17 00:00:00 2001
From: Abolfazl <abolixes@gmail.com>
Date: Tue, 21 Apr 2026 21:31:21 +0330
Subject: [PATCH] Enhance error handling and add size limits for request and
 response bodies

---
 domain_fronter.py | 11 ++++++++++-
 proxy_server.py   | 18 ++++++++++++++----
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/domain_fronter.py b/domain_fronter.py
index 862c1ac..6f2bec3 100644
--- a/domain_fronter.py
+++ b/domain_fronter.py
@@ -1107,7 +1107,10 @@ class DomainFronter:
             data = json.loads(text)
         except json.JSONDecodeError:
             m = re.search(r'\{.*\}', text, re.DOTALL)
-            data = json.loads(m.group()) if m else None
+            try:
+                data = json.loads(m.group()) if m else None
+            except json.JSONDecodeError:
+                data = None
         if not data:
             raise RuntimeError(f"Bad batch response: {text[:200]}")
 
@@ -1131,6 +1134,8 @@ class DomainFronter:
         """Read one HTTP response. Keep-alive safe (no read-until-EOF)."""
         raw = b""
         while b"\r\n\r\n" not in raw:
+            if len(raw) > 65536:  # 64 KB header size limit
+                return 0, {}, b""
             chunk = await asyncio.wait_for(reader.read(8192), timeout=8)
             if not chunk:
                 break
@@ -1190,6 +1195,7 @@ class DomainFronter:
     async def _read_chunked(self, reader, buf=b""):
         """Incrementally read chunked transfer-encoding."""
         result = b""
+        _MAX_BODY = 200 * 1024 * 1024  # 200 MB total body cap
         while True:
             while b"\r\n" not in buf:
                 data = await asyncio.wait_for(reader.read(8192), timeout=20)
@@ -1209,6 +1215,9 @@ class DomainFronter:
                 break
             if size == 0:
                 break
+            if size > _MAX_BODY or len(result) + size > _MAX_BODY:
+                log.warning("Chunked body exceeds %d MB cap — truncating", _MAX_BODY // (1024 * 1024))
+                break
 
             while len(buf) < size + 2:
                 data = await asyncio.wait_for(reader.read(65536), timeout=20)
diff --git a/proxy_server.py b/proxy_server.py
index ea9117c..ae033df 100644
--- a/proxy_server.py
+++ b/proxy_server.py
@@ -385,8 +385,14 @@ class ProxyServer:
     # ── CONNECT (HTTPS tunnelling) ────────────────────────────────
 
     async def _do_connect(self, target: str, reader, writer):
-        host, _, port = target.rpartition(":")
-        port = int(port) if port else 443
+        host, _, port_str = target.rpartition(":")
+        try:
+            port = int(port_str) if port_str else 443
+        except ValueError:
+            log.warning("CONNECT invalid target: %r", target)
+            writer.write(b"HTTP/1.1 400 Bad Request\r\n\r\n")
+            await writer.drain()
+            return
         if not host:
             host, port = target, 443
 
@@ -801,6 +807,8 @@ class ProxyServer:
                 for raw_line in header_block.split(b"\r\n"):
                     if raw_line.lower().startswith(b"content-length:"):
                         length = int(raw_line.split(b":", 1)[1].strip())
+                        if length > 100 * 1024 * 1024:  # 100 MB cap
+                            raise ValueError(f"Request body too large: {length} bytes")
                         body = await reader.readexactly(length)
                         break
 
@@ -1010,6 +1018,10 @@ class ProxyServer:
         for raw_line in header_block.split(b"\r\n"):
             if raw_line.lower().startswith(b"content-length:"):
                 length = int(raw_line.split(b":", 1)[1].strip())
+                if length > 100 * 1024 * 1024:  # 100 MB cap
+                    writer.write(b"HTTP/1.1 413 Content Too Large\r\n\r\n")
+                    await writer.drain()
+                    return
                 body = await reader.readexactly(length)
                 break
 
@@ -1082,8 +1094,6 @@ class ProxyServer:
         to the target host and pipes raw HTTP through it.
         Much faster for rapid-fire requests (e.g., Telegram API).
         """
-        import re as _re
-
         # Parse target host:port from the raw HTTP request
         host = ""
         port = 80