更新
This commit is contained in:
Your Name
@@ -0,0 +1,233 @@
|
||||
"""Loopback HTTP proxy whose outbound sockets bind to one local IPv4.
|
||||
|
||||
Playwright does not expose a ``local_address`` option. Accounts that select a
|
||||
specific server egress channel therefore use this tiny process-local proxy so
|
||||
their browser login/refresh traffic leaves through the same interface as IM
|
||||
HTTP and WebSocket traffic. The listener is loopback-only and does not rotate
|
||||
or retry public addresses.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import ipaddress
|
||||
import logging
|
||||
import socket
|
||||
import weakref
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
logger = logging.getLogger("rpa_engine.source_proxy")
|
||||
|
||||
_MAX_HEADER_BYTES = 64 * 1024
|
||||
_HEADER_TIMEOUT_SECONDS = 20.0
|
||||
|
||||
|
||||
class SourceBoundProxy:
|
||||
"""Minimal HTTP/HTTPS CONNECT proxy bound to a fixed source address."""
|
||||
|
||||
def __init__(self, source_ip: str):
|
||||
address = ipaddress.ip_address(str(source_ip or "").strip())
|
||||
if address.version != 4 or address.is_unspecified or address.is_multicast:
|
||||
raise ValueError(f"invalid IPv4 source address: {source_ip!r}")
|
||||
self.source_ip = str(address)
|
||||
self._server: asyncio.AbstractServer | None = None
|
||||
|
||||
@property
|
||||
def server_url(self) -> str:
|
||||
if self._server is None or not self._server.sockets:
|
||||
raise RuntimeError("source-bound proxy has not started")
|
||||
port = int(self._server.sockets[0].getsockname()[1])
|
||||
return f"http://127.0.0.1:{port}"
|
||||
|
||||
async def start(self) -> "SourceBoundProxy":
|
||||
if self._server is None:
|
||||
self._server = await asyncio.start_server(
|
||||
self._handle_client,
|
||||
host="127.0.0.1",
|
||||
port=0,
|
||||
family=socket.AF_INET,
|
||||
)
|
||||
logger.info(
|
||||
"source-bound browser proxy ready: %s -> source %s",
|
||||
self.server_url,
|
||||
self.source_ip,
|
||||
)
|
||||
return self
|
||||
|
||||
async def close(self) -> None:
|
||||
server = self._server
|
||||
self._server = None
|
||||
if server is not None:
|
||||
server.close()
|
||||
await server.wait_closed()
|
||||
|
||||
async def _open_upstream(
|
||||
self,
|
||||
host: str,
|
||||
port: int,
|
||||
) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
|
||||
return await asyncio.open_connection(
|
||||
host=host,
|
||||
port=port,
|
||||
family=socket.AF_INET,
|
||||
local_addr=(self.source_ip, 0),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _relay(
|
||||
source: asyncio.StreamReader,
|
||||
destination: asyncio.StreamWriter,
|
||||
) -> None:
|
||||
try:
|
||||
while True:
|
||||
chunk = await source.read(64 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
destination.write(chunk)
|
||||
await destination.drain()
|
||||
except (ConnectionError, asyncio.CancelledError):
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
destination.write_eof()
|
||||
except (AttributeError, OSError, RuntimeError):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
async def _bridge(
|
||||
cls,
|
||||
client_reader: asyncio.StreamReader,
|
||||
client_writer: asyncio.StreamWriter,
|
||||
upstream_reader: asyncio.StreamReader,
|
||||
upstream_writer: asyncio.StreamWriter,
|
||||
) -> None:
|
||||
tasks = (
|
||||
asyncio.create_task(cls._relay(client_reader, upstream_writer)),
|
||||
asyncio.create_task(cls._relay(upstream_reader, client_writer)),
|
||||
)
|
||||
try:
|
||||
await asyncio.gather(*tasks)
|
||||
finally:
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
@staticmethod
|
||||
def _parse_authority(authority: str, default_port: int) -> tuple[str, int]:
|
||||
parsed = urlsplit(f"//{authority}")
|
||||
host = str(parsed.hostname or "").strip()
|
||||
if not host:
|
||||
raise ValueError("proxy request is missing a host")
|
||||
return host, int(parsed.port or default_port)
|
||||
|
||||
async def _handle_client(
|
||||
self,
|
||||
client_reader: asyncio.StreamReader,
|
||||
client_writer: asyncio.StreamWriter,
|
||||
) -> None:
|
||||
upstream_writer: asyncio.StreamWriter | None = None
|
||||
try:
|
||||
header = await asyncio.wait_for(
|
||||
client_reader.readuntil(b"\r\n\r\n"),
|
||||
timeout=_HEADER_TIMEOUT_SECONDS,
|
||||
)
|
||||
if len(header) > _MAX_HEADER_BYTES:
|
||||
raise ValueError("proxy request headers are too large")
|
||||
lines = header.decode("latin-1").split("\r\n")
|
||||
request_line = lines[0].split(" ", 2)
|
||||
if len(request_line) != 3:
|
||||
raise ValueError("malformed proxy request line")
|
||||
method, target, version = request_line
|
||||
|
||||
if method.upper() == "CONNECT":
|
||||
host, port = self._parse_authority(target, 443)
|
||||
upstream_reader, upstream_writer = await self._open_upstream(host, port)
|
||||
client_writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
|
||||
await client_writer.drain()
|
||||
else:
|
||||
parsed = urlsplit(target)
|
||||
host_header = next(
|
||||
(
|
||||
line.partition(":")[2].strip()
|
||||
for line in lines[1:]
|
||||
if line.lower().startswith("host:")
|
||||
),
|
||||
"",
|
||||
)
|
||||
authority = parsed.netloc or host_header
|
||||
host, port = self._parse_authority(
|
||||
authority,
|
||||
443 if parsed.scheme.lower() == "https" else 80,
|
||||
)
|
||||
upstream_reader, upstream_writer = await self._open_upstream(host, port)
|
||||
origin_target = parsed.path or "/"
|
||||
if parsed.query:
|
||||
origin_target += f"?{parsed.query}"
|
||||
forwarded = [f"{method} {origin_target} {version}"]
|
||||
forwarded.extend(
|
||||
line for line in lines[1:]
|
||||
if line and not line.lower().startswith("proxy-connection:")
|
||||
)
|
||||
upstream_writer.write(("\r\n".join(forwarded) + "\r\n\r\n").encode("latin-1"))
|
||||
await upstream_writer.drain()
|
||||
|
||||
await self._bridge(
|
||||
client_reader,
|
||||
client_writer,
|
||||
upstream_reader,
|
||||
upstream_writer,
|
||||
)
|
||||
except asyncio.IncompleteReadError:
|
||||
pass
|
||||
except asyncio.CancelledError:
|
||||
# Event-loop shutdown may cancel an in-flight browser tunnel.
|
||||
# Closing both writers below is sufficient; do not leak a noisy
|
||||
# cancelled handler callback into the server log.
|
||||
pass
|
||||
except Exception as exc:
|
||||
logger.warning("source-bound browser proxy request failed: %s", exc)
|
||||
try:
|
||||
client_writer.write(
|
||||
b"HTTP/1.1 502 Bad Gateway\r\nConnection: close\r\n\r\n"
|
||||
)
|
||||
await client_writer.drain()
|
||||
except (ConnectionError, RuntimeError):
|
||||
pass
|
||||
finally:
|
||||
for writer in (upstream_writer, client_writer):
|
||||
if writer is None:
|
||||
continue
|
||||
try:
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
except (ConnectionError, RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class _LoopProxyState:
|
||||
def __init__(self) -> None:
|
||||
self.lock = asyncio.Lock()
|
||||
self.proxies: dict[str, SourceBoundProxy] = {}
|
||||
|
||||
|
||||
_loop_states: weakref.WeakKeyDictionary[
|
||||
asyncio.AbstractEventLoop, _LoopProxyState
|
||||
] = weakref.WeakKeyDictionary()
|
||||
|
||||
|
||||
async def playwright_proxy_for_source(source_ip: str) -> dict[str, str]:
|
||||
"""Return a Playwright proxy config fixed to ``source_ip``."""
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
state = _loop_states.get(loop)
|
||||
if state is None:
|
||||
state = _LoopProxyState()
|
||||
_loop_states[loop] = state
|
||||
normalized = str(ipaddress.ip_address(str(source_ip or "").strip()))
|
||||
async with state.lock:
|
||||
proxy = state.proxies.get(normalized)
|
||||
if proxy is None:
|
||||
proxy = await SourceBoundProxy(normalized).start()
|
||||
state.proxies[normalized] = proxy
|
||||
return {"server": proxy.server_url}
|
||||
Reference in New Issue
Block a user