234 lines
8.2 KiB
Python
234 lines
8.2 KiB
Python
"""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}
|