Enhance error handling and add size limits for request and response bodies

This commit is contained in:
Abolfazl
2026-04-21 21:31:21 +03:30
parent d563c1516e
commit 607e226fff
2 changed files with 24 additions and 5 deletions
+10 -1
View File
@@ -1107,7 +1107,10 @@ class DomainFronter:
data = json.loads(text) data = json.loads(text)
except json.JSONDecodeError: except json.JSONDecodeError:
m = re.search(r'\{.*\}', text, re.DOTALL) 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: if not data:
raise RuntimeError(f"Bad batch response: {text[:200]}") 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).""" """Read one HTTP response. Keep-alive safe (no read-until-EOF)."""
raw = b"" raw = b""
while b"\r\n\r\n" not in raw: 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) chunk = await asyncio.wait_for(reader.read(8192), timeout=8)
if not chunk: if not chunk:
break break
@@ -1190,6 +1195,7 @@ class DomainFronter:
async def _read_chunked(self, reader, buf=b""): async def _read_chunked(self, reader, buf=b""):
"""Incrementally read chunked transfer-encoding.""" """Incrementally read chunked transfer-encoding."""
result = b"" result = b""
_MAX_BODY = 200 * 1024 * 1024 # 200 MB total body cap
while True: while True:
while b"\r\n" not in buf: while b"\r\n" not in buf:
data = await asyncio.wait_for(reader.read(8192), timeout=20) data = await asyncio.wait_for(reader.read(8192), timeout=20)
@@ -1209,6 +1215,9 @@ class DomainFronter:
break break
if size == 0: if size == 0:
break 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: while len(buf) < size + 2:
data = await asyncio.wait_for(reader.read(65536), timeout=20) data = await asyncio.wait_for(reader.read(65536), timeout=20)
+14 -4
View File
@@ -385,8 +385,14 @@ class ProxyServer:
# ── CONNECT (HTTPS tunnelling) ──────────────────────────────── # ── CONNECT (HTTPS tunnelling) ────────────────────────────────
async def _do_connect(self, target: str, reader, writer): async def _do_connect(self, target: str, reader, writer):
host, _, port = target.rpartition(":") host, _, port_str = target.rpartition(":")
port = int(port) if port else 443 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: if not host:
host, port = target, 443 host, port = target, 443
@@ -801,6 +807,8 @@ class ProxyServer:
for raw_line in header_block.split(b"\r\n"): for raw_line in header_block.split(b"\r\n"):
if raw_line.lower().startswith(b"content-length:"): if raw_line.lower().startswith(b"content-length:"):
length = int(raw_line.split(b":", 1)[1].strip()) 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) body = await reader.readexactly(length)
break break
@@ -1010,6 +1018,10 @@ class ProxyServer:
for raw_line in header_block.split(b"\r\n"): for raw_line in header_block.split(b"\r\n"):
if raw_line.lower().startswith(b"content-length:"): if raw_line.lower().startswith(b"content-length:"):
length = int(raw_line.split(b":", 1)[1].strip()) 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) body = await reader.readexactly(length)
break break
@@ -1082,8 +1094,6 @@ class ProxyServer:
to the target host and pipes raw HTTP through it. to the target host and pipes raw HTTP through it.
Much faster for rapid-fire requests (e.g., Telegram API). Much faster for rapid-fire requests (e.g., Telegram API).
""" """
import re as _re
# Parse target host:port from the raw HTTP request # Parse target host:port from the raw HTTP request
host = "" host = ""
port = 80 port = 80