657 lines
22 KiB
Python
657 lines
22 KiB
Python
"""Small in-memory HTTP/SOCKS proxy used by one browser generation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import select
|
|
import socket
|
|
import socketserver
|
|
import ssl
|
|
import threading
|
|
from collections.abc import Callable
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass
|
|
from typing import cast
|
|
from urllib.parse import urlsplit
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ProxyExit:
|
|
protocol: str
|
|
host: str
|
|
port: int
|
|
username: str = ""
|
|
credential: str = ""
|
|
|
|
|
|
class _ThreadingTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
|
|
allow_reuse_address = True
|
|
daemon_threads = True
|
|
|
|
|
|
class ProxyRegistry:
|
|
def __init__(self) -> None:
|
|
self._lock = threading.RLock()
|
|
self._proxies: dict[str, MemoryProxy] = {}
|
|
|
|
def configure(
|
|
self,
|
|
alias: str,
|
|
binding_version: int,
|
|
bind_host: str,
|
|
port: int,
|
|
exit: ProxyExit,
|
|
network_id: str = "",
|
|
) -> tuple[str, Callable[[], None]]:
|
|
with self._lock:
|
|
current = self._proxies.get(alias)
|
|
if current and current.matches(
|
|
binding_version, bind_host, port, exit, network_id
|
|
):
|
|
return current.url, lambda: self._remove_object(alias, current)
|
|
if current:
|
|
self._proxies.pop(alias, None)
|
|
current.close()
|
|
proxy = MemoryProxy(
|
|
alias, binding_version, bind_host, port, exit, network_id
|
|
)
|
|
self._proxies[alias] = proxy
|
|
return proxy.url, lambda: self._remove_object(alias, proxy)
|
|
|
|
def bind(
|
|
self,
|
|
alias: str,
|
|
binding_version: int,
|
|
proxy_url: str,
|
|
runtime_id: str,
|
|
network_id: str = "",
|
|
) -> bool:
|
|
with self._lock:
|
|
proxy = self._proxies.get(alias)
|
|
return bool(
|
|
proxy and proxy.bind(binding_version, proxy_url, runtime_id, network_id)
|
|
)
|
|
|
|
def ready(
|
|
self, alias: str, port: int, runtime_id: str, network_id: str = ""
|
|
) -> bool:
|
|
with self._lock:
|
|
proxy = self._proxies.get(alias)
|
|
return bool(proxy and proxy.ready(port, runtime_id, network_id))
|
|
|
|
def remove(
|
|
self, alias: str, binding_version: int, runtime_id: str, network_id: str = ""
|
|
) -> bool:
|
|
with self._lock:
|
|
proxy = self._proxies.get(alias)
|
|
if not proxy:
|
|
return True
|
|
if not proxy.identity_matches(binding_version, runtime_id, network_id):
|
|
return False
|
|
self._proxies.pop(alias, None)
|
|
proxy.close()
|
|
return True
|
|
|
|
def _remove_object(self, alias: str, proxy: MemoryProxy) -> None:
|
|
with self._lock:
|
|
if self._proxies.get(alias) is not proxy:
|
|
return
|
|
self._proxies.pop(alias, None)
|
|
proxy.close()
|
|
|
|
def close(self) -> None:
|
|
with self._lock:
|
|
proxies = list(self._proxies.values())
|
|
self._proxies.clear()
|
|
for proxy in proxies:
|
|
proxy.close()
|
|
|
|
|
|
class MemoryProxy:
|
|
def __init__(
|
|
self,
|
|
alias: str,
|
|
binding_version: int,
|
|
bind_host: str,
|
|
port: int,
|
|
exit: ProxyExit,
|
|
network_id: str,
|
|
) -> None:
|
|
self.alias = alias
|
|
self.binding_version = binding_version
|
|
self.bind_host = bind_host
|
|
self.exit = exit
|
|
self.network_id = network_id
|
|
self.runtime_id = ""
|
|
self._lock = threading.RLock()
|
|
self._tunnels: set[socket.socket] = set()
|
|
handler_type = type(
|
|
"ProxyHandler",
|
|
(_ProxyHandler,),
|
|
{"proxy": self},
|
|
)
|
|
self.server = _ThreadingTCPServer((bind_host, port), handler_type)
|
|
self.listener = self.server.socket
|
|
actual_port = self.listener.getsockname()[1]
|
|
self.url = f"http://{bind_host}:{actual_port}"
|
|
self._thread = threading.Thread(
|
|
target=self.server.serve_forever,
|
|
name=f"creatorhub-proxy-{alias}",
|
|
daemon=True,
|
|
)
|
|
self._thread.start()
|
|
|
|
def matches(
|
|
self,
|
|
binding_version: int,
|
|
bind_host: str,
|
|
port: int,
|
|
exit: ProxyExit,
|
|
network_id: str,
|
|
) -> bool:
|
|
actual_port = self.listener.getsockname()[1]
|
|
return (
|
|
self.binding_version == binding_version
|
|
and self.bind_host == bind_host
|
|
and (port == 0 or port == actual_port)
|
|
and self.exit == exit
|
|
and self.network_id == network_id
|
|
)
|
|
|
|
def bind(
|
|
self, binding_version: int, proxy_url: str, runtime_id: str, network_id: str
|
|
) -> bool:
|
|
with self._lock:
|
|
if (
|
|
binding_version != self.binding_version
|
|
or proxy_url != self.url
|
|
or not runtime_id
|
|
):
|
|
return False
|
|
self.runtime_id = runtime_id
|
|
if network_id:
|
|
self.network_id = network_id
|
|
return True
|
|
|
|
def ready(self, port: int, runtime_id: str, network_id: str) -> bool:
|
|
with self._lock:
|
|
return bool(
|
|
runtime_id
|
|
and self.runtime_id == runtime_id
|
|
and self.listener.getsockname()[1] == port
|
|
and (not network_id or self.network_id == network_id)
|
|
)
|
|
|
|
def identity_matches(
|
|
self, binding_version: int, runtime_id: str, network_id: str
|
|
) -> bool:
|
|
with self._lock:
|
|
return (
|
|
self.binding_version == binding_version
|
|
and bool(runtime_id)
|
|
and self.runtime_id == runtime_id
|
|
and (not network_id or self.network_id == network_id)
|
|
)
|
|
|
|
def close(self) -> None:
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
with self._lock:
|
|
tunnels = list(self._tunnels)
|
|
self._tunnels.clear()
|
|
for connection in tunnels:
|
|
with suppress(OSError):
|
|
connection.close()
|
|
|
|
def add_tunnel(self, connection: socket.socket) -> bool:
|
|
with self._lock:
|
|
if self.server.socket.fileno() < 0:
|
|
return False
|
|
self._tunnels.add(connection)
|
|
return True
|
|
|
|
def remove_tunnel(self, connection: socket.socket) -> None:
|
|
with self._lock:
|
|
self._tunnels.discard(connection)
|
|
|
|
def dial(self, target: str, timeout: float = 20.0) -> socket.socket:
|
|
exit = self.exit
|
|
if exit.protocol in ("http", "https"):
|
|
return _dial_http_proxy(exit, target, timeout)
|
|
if exit.protocol == "socks4":
|
|
return _dial_socks4(exit, target, timeout)
|
|
if exit.protocol == "socks5":
|
|
return _dial_socks5(exit, target, timeout)
|
|
raise OSError("unsupported proxy protocol")
|
|
|
|
def forward_http(
|
|
self,
|
|
client: socket.socket,
|
|
method: str,
|
|
target: str,
|
|
headers: list[tuple[str, str]],
|
|
body: bytes,
|
|
) -> None:
|
|
parsed = urlsplit(target)
|
|
if not parsed.hostname:
|
|
raise OSError("proxy request target is invalid")
|
|
target_host = parsed.hostname
|
|
target_port = parsed.port or (443 if parsed.scheme == "https" else 80)
|
|
origin_target = parsed.path or "/"
|
|
if parsed.query:
|
|
origin_target += "?" + parsed.query
|
|
if self.exit.protocol in ("http", "https"):
|
|
upstream = _open_host(
|
|
self.exit.host, self.exit.port, 20.0, self.exit.protocol == "https"
|
|
)
|
|
request_target = target
|
|
request_headers = _forward_headers(headers, len(body))
|
|
if self.exit.username:
|
|
token = base64.b64encode(
|
|
f"{self.exit.username}:{self.exit.credential}".encode()
|
|
).decode()
|
|
request_headers.append(("Proxy-Authorization", f"Basic {token}"))
|
|
else:
|
|
upstream = self.dial(f"{target_host}:{target_port}")
|
|
request_target = origin_target
|
|
request_headers = _forward_headers(headers, len(body))
|
|
try:
|
|
lines = [f"{method} {request_target} HTTP/1.1"]
|
|
lines.extend(f"{k}: {v}" for k, v in request_headers)
|
|
lines.append("Connection: close")
|
|
upstream.sendall(("\r\n".join(lines) + "\r\n\r\n").encode() + body)
|
|
_copy_until_close(upstream, client)
|
|
finally:
|
|
upstream.close()
|
|
|
|
def tunnel(self, client: socket.socket, target: str) -> None:
|
|
upstream = self.dial(target)
|
|
try:
|
|
client.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n")
|
|
if not self.add_tunnel(client):
|
|
return
|
|
_relay(client, upstream)
|
|
finally:
|
|
self.remove_tunnel(client)
|
|
with suppress(OSError):
|
|
upstream.close()
|
|
|
|
|
|
class _ProxyHandler(socketserver.BaseRequestHandler):
|
|
proxy: MemoryProxy
|
|
|
|
def handle(self) -> None:
|
|
client = self.request
|
|
client.settimeout(60.0)
|
|
try:
|
|
head, body = _read_request(client)
|
|
method, target, headers = _parse_request(head)
|
|
if method.upper() == "CONNECT":
|
|
self.proxy.tunnel(client, target)
|
|
else:
|
|
self.proxy.forward_http(client, method, target, headers, body)
|
|
except (OSError, ValueError):
|
|
with suppress(OSError):
|
|
client.sendall(
|
|
b"HTTP/1.1 502 Bad Gateway\r\nConnection: close\r\nContent-Length: 0\r\n\r\n"
|
|
)
|
|
|
|
|
|
def _read_request(connection: socket.socket) -> tuple[bytes, bytes]:
|
|
data = bytearray()
|
|
while b"\r\n\r\n" not in data:
|
|
chunk = connection.recv(65536)
|
|
if not chunk:
|
|
raise OSError("proxy request closed")
|
|
data.extend(chunk)
|
|
if len(data) > 256 * 1024:
|
|
raise OSError("proxy headers too large")
|
|
split = data.index(b"\r\n\r\n") + 4
|
|
head, initial_body = bytes(data[:split]), bytearray(data[split:])
|
|
fields = _header_fields(head)
|
|
lengths = fields.get("content-length", [])
|
|
transfer_encoding = fields.get("transfer-encoding", [])
|
|
if lengths and transfer_encoding:
|
|
raise OSError("proxy request has both content length and transfer encoding")
|
|
if transfer_encoding:
|
|
if (
|
|
len(transfer_encoding) != 1
|
|
or transfer_encoding[0].lower().strip() != "chunked"
|
|
):
|
|
raise OSError("unsupported proxy transfer encoding")
|
|
return head, _read_chunked_body(connection, initial_body)
|
|
content_length = 0
|
|
if lengths:
|
|
if len(lengths) != 1:
|
|
raise OSError("proxy content length is ambiguous")
|
|
try:
|
|
content_length = int(lengths[0])
|
|
except ValueError as exc:
|
|
raise OSError("invalid proxy content length") from exc
|
|
if content_length < 0 or content_length > 16 * 1024 * 1024:
|
|
raise OSError("proxy body too large")
|
|
while len(initial_body) < content_length:
|
|
chunk = connection.recv(min(65536, content_length - len(initial_body)))
|
|
if not chunk:
|
|
raise OSError("proxy body closed")
|
|
initial_body.extend(chunk)
|
|
return head, bytes(initial_body[:content_length])
|
|
|
|
|
|
def _header_fields(head: bytes) -> dict[str, list[str]]:
|
|
fields: dict[str, list[str]] = {}
|
|
for line in head.decode("iso-8859-1").split("\r\n")[1:]:
|
|
if not line:
|
|
continue
|
|
if ":" not in line:
|
|
raise OSError("invalid proxy header")
|
|
key, value = line.split(":", 1)
|
|
fields.setdefault(key.strip().lower(), []).append(value.strip())
|
|
return fields
|
|
|
|
|
|
def _read_line(connection: socket.socket, buffer: bytearray, limit: int) -> bytes:
|
|
while b"\r\n" not in buffer:
|
|
chunk = connection.recv(65536)
|
|
if not chunk:
|
|
raise OSError("proxy chunked body closed")
|
|
buffer.extend(chunk)
|
|
if len(buffer) > limit:
|
|
raise OSError("proxy chunk line too large")
|
|
end = buffer.index(b"\r\n")
|
|
line = bytes(buffer[:end])
|
|
del buffer[: end + 2]
|
|
return line
|
|
|
|
|
|
def _read_chunked_body(connection: socket.socket, buffer: bytearray) -> bytes:
|
|
body = bytearray()
|
|
while True:
|
|
line = _read_line(connection, buffer, 8192)
|
|
size_text = line.split(b";", 1)[0].strip()
|
|
try:
|
|
size = int(size_text, 16)
|
|
except ValueError as exc:
|
|
raise OSError("invalid proxy chunk size") from exc
|
|
if size < 0:
|
|
raise OSError("invalid proxy chunk size")
|
|
if size == 0:
|
|
while True:
|
|
trailer = _read_line(connection, buffer, 256 * 1024)
|
|
if not trailer:
|
|
return bytes(body)
|
|
if b":" not in trailer:
|
|
raise OSError("invalid proxy trailer")
|
|
if len(body) + size > 16 * 1024 * 1024:
|
|
raise OSError("proxy body too large")
|
|
while len(buffer) < size + 2:
|
|
chunk = connection.recv(65536)
|
|
if not chunk:
|
|
raise OSError("proxy chunked body closed")
|
|
buffer.extend(chunk)
|
|
body.extend(buffer[:size])
|
|
if buffer[size : size + 2] != b"\r\n":
|
|
raise OSError("invalid proxy chunk terminator")
|
|
del buffer[: size + 2]
|
|
|
|
|
|
def _forward_headers(
|
|
headers: list[tuple[str, str]], body_length: int
|
|
) -> list[tuple[str, str]]:
|
|
return [
|
|
(k, v)
|
|
for k, v in headers
|
|
if k.lower()
|
|
not in (
|
|
"proxy-authorization",
|
|
"connection",
|
|
"transfer-encoding",
|
|
"content-length",
|
|
)
|
|
] + [("Content-Length", str(body_length))]
|
|
|
|
|
|
def _parse_request(head: bytes) -> tuple[str, str, list[tuple[str, str]]]:
|
|
lines = head.decode("iso-8859-1").split("\r\n")
|
|
method, target, version = lines[0].split(" ", 2)
|
|
if version not in ("HTTP/1.0", "HTTP/1.1") or not target:
|
|
raise ValueError("invalid proxy request")
|
|
headers: list[tuple[str, str]] = []
|
|
for line in lines[1:]:
|
|
if not line:
|
|
continue
|
|
if ":" not in line:
|
|
raise ValueError("invalid proxy header")
|
|
key, value = line.split(":", 1)
|
|
headers.append((key.strip(), value.strip()))
|
|
if method.upper() == "CONNECT":
|
|
if ":" not in target:
|
|
raise ValueError("CONNECT target missing port")
|
|
return method, target, headers
|
|
return method, target, headers
|
|
|
|
|
|
def _open_host(
|
|
host: str, port: int, timeout: float, tls: bool = False
|
|
) -> socket.socket:
|
|
connection = socket.create_connection((host, port), timeout=timeout)
|
|
if tls:
|
|
context = ssl.create_default_context()
|
|
connection = context.wrap_socket(connection, server_hostname=host)
|
|
return connection
|
|
|
|
|
|
def _dial_http_proxy(exit: ProxyExit, target: str, timeout: float) -> socket.socket:
|
|
connection = _open_host(exit.host, exit.port, timeout, exit.protocol == "https")
|
|
try:
|
|
lines = [f"CONNECT {target} HTTP/1.1", f"Host: {target}"]
|
|
if exit.username:
|
|
token = base64.b64encode(
|
|
f"{exit.username}:{exit.credential}".encode()
|
|
).decode()
|
|
lines.append(f"Proxy-Authorization: Basic {token}")
|
|
connection.sendall(("\r\n".join(lines) + "\r\n\r\n").encode())
|
|
status, buffered = _read_connect_response(connection)
|
|
if status != 200:
|
|
raise OSError(f"upstream proxy returned {status}")
|
|
result = cast(
|
|
socket.socket,
|
|
_BufferedSocket(connection, buffered) if buffered else connection,
|
|
)
|
|
result.settimeout(None)
|
|
return result
|
|
except Exception:
|
|
connection.close()
|
|
raise
|
|
|
|
|
|
def _dial_socks4(exit: ProxyExit, target: str, timeout: float) -> socket.socket:
|
|
try:
|
|
host, port_text = target.rsplit(":", 1)
|
|
port = int(port_text)
|
|
except (ValueError, IndexError) as exc:
|
|
raise OSError("invalid SOCKS4 target") from exc
|
|
if not 1 <= port <= 65535:
|
|
raise OSError("invalid SOCKS4 target port")
|
|
connection = _open_host(exit.host, exit.port, timeout)
|
|
try:
|
|
ip = socket.inet_aton(host) if _is_ipv4(host) else b"\x00\x00\x00\x01"
|
|
payload = (
|
|
b"\x04\x01"
|
|
+ port.to_bytes(2, "big")
|
|
+ ip
|
|
+ exit.username.encode()
|
|
+ b"\x00"
|
|
)
|
|
if ip == b"\x00\x00\x00\x01":
|
|
payload += host.encode() + b"\x00"
|
|
connection.sendall(payload)
|
|
response = _recv_exact(connection, 8)
|
|
if response[1] != 90:
|
|
raise OSError("SOCKS4 proxy rejected connection")
|
|
connection.settimeout(None)
|
|
return connection
|
|
except Exception:
|
|
connection.close()
|
|
raise
|
|
|
|
|
|
def _dial_socks5(exit: ProxyExit, target: str, timeout: float) -> socket.socket:
|
|
try:
|
|
host, port_text = target.rsplit(":", 1)
|
|
port = int(port_text)
|
|
except (ValueError, IndexError) as exc:
|
|
raise OSError("invalid SOCKS5 target") from exc
|
|
if not 1 <= port <= 65535:
|
|
raise OSError("invalid SOCKS5 target port")
|
|
connection = _open_host(exit.host, exit.port, timeout)
|
|
try:
|
|
methods = b"\x05\x01\x02" if exit.username else b"\x05\x01\x00"
|
|
connection.sendall(methods)
|
|
version, selected = _recv_exact(connection, 2)
|
|
if version != 5 or selected == 255:
|
|
raise OSError("SOCKS5 authentication method rejected")
|
|
if selected == 2:
|
|
user, credential = exit.username.encode(), exit.credential.encode()
|
|
if len(user) > 255 or len(credential) > 255:
|
|
raise OSError("SOCKS5 credentials too long")
|
|
connection.sendall(
|
|
b"\x01"
|
|
+ bytes([len(user)])
|
|
+ user
|
|
+ bytes([len(credential)])
|
|
+ credential
|
|
)
|
|
if _recv_exact(connection, 2)[1] != 0:
|
|
raise OSError("SOCKS5 authentication rejected")
|
|
elif exit.username:
|
|
raise OSError("SOCKS5 proxy skipped required authentication")
|
|
if _is_ipv4(host):
|
|
address = b"\x01" + socket.inet_aton(host)
|
|
else:
|
|
try:
|
|
address = b"\x04" + socket.inet_pton(socket.AF_INET6, host)
|
|
except OSError:
|
|
encoded = host.encode()
|
|
if len(encoded) > 255:
|
|
raise OSError("SOCKS5 target too long") from None
|
|
address = b"\x03" + bytes([len(encoded)]) + encoded
|
|
connection.sendall(b"\x05\x01\x00" + address + port.to_bytes(2, "big"))
|
|
header = _recv_exact(connection, 4)
|
|
if header[0] != 5 or header[1] != 0:
|
|
raise OSError("SOCKS5 proxy rejected connection")
|
|
if header[3] == 1:
|
|
length = 4
|
|
elif header[3] == 3:
|
|
length = _recv_exact(connection, 1)[0]
|
|
elif header[3] == 4:
|
|
length = 16
|
|
else:
|
|
raise OSError("invalid SOCKS5 response")
|
|
_recv_exact(connection, length + 2)
|
|
connection.settimeout(None)
|
|
return connection
|
|
except Exception:
|
|
connection.close()
|
|
raise
|
|
|
|
|
|
def _read_status(connection: socket.socket) -> int:
|
|
return _read_connect_response(connection)[0]
|
|
|
|
|
|
def _read_connect_response(connection: socket.socket) -> tuple[int, bytes]:
|
|
# Consume the complete response header. A proxy may return the first
|
|
# tunnel bytes in the same read, so retain bytes after the header block.
|
|
data = bytearray()
|
|
while b"\r\n\r\n" not in data:
|
|
chunk = connection.recv(65536)
|
|
if not chunk:
|
|
raise OSError("proxy response closed before headers")
|
|
data.extend(chunk)
|
|
if len(data) > 256 * 1024:
|
|
raise OSError("proxy response headers too large")
|
|
split = data.index(b"\r\n\r\n") + 4
|
|
header = bytes(data[:split])
|
|
lines = header.decode("iso-8859-1").split("\r\n")
|
|
parts = lines[0].split(" ", 2)
|
|
if len(parts) < 2 or not parts[0].startswith("HTTP/"):
|
|
raise OSError("invalid proxy response status")
|
|
try:
|
|
status = int(parts[1])
|
|
except ValueError as exc:
|
|
raise OSError("invalid proxy response status") from exc
|
|
for line in lines[1:-2]:
|
|
if line and ":" not in line:
|
|
raise OSError("invalid proxy response header")
|
|
return status, bytes(data[split:])
|
|
|
|
|
|
class _BufferedSocket:
|
|
def __init__(self, connection: socket.socket, buffered: bytes) -> None:
|
|
self.connection = connection
|
|
self.buffered = bytearray(buffered)
|
|
|
|
def recv(self, size: int, flags: int = 0) -> bytes:
|
|
if self.buffered:
|
|
result = bytes(self.buffered[:size])
|
|
del self.buffered[:size]
|
|
return result
|
|
return (
|
|
self.connection.recv(size, flags) if flags else self.connection.recv(size)
|
|
)
|
|
|
|
def sendall(self, data: bytes) -> None:
|
|
self.connection.sendall(data)
|
|
|
|
def send(self, data: bytes, flags: int = 0) -> int:
|
|
return (
|
|
self.connection.send(data, flags) if flags else self.connection.send(data)
|
|
)
|
|
|
|
def settimeout(self, value: float | None) -> None:
|
|
self.connection.settimeout(value)
|
|
|
|
def close(self) -> None:
|
|
self.connection.close()
|
|
|
|
|
|
def _recv_exact(connection: socket.socket, size: int) -> bytes:
|
|
result = bytearray()
|
|
while len(result) < size:
|
|
chunk = connection.recv(size - len(result))
|
|
if not chunk:
|
|
raise OSError("proxy connection closed")
|
|
result.extend(chunk)
|
|
return bytes(result)
|
|
|
|
|
|
def _copy_until_close(source: socket.socket, target: socket.socket) -> None:
|
|
while True:
|
|
data = source.recv(65536)
|
|
if not data:
|
|
return
|
|
target.sendall(data)
|
|
|
|
|
|
def _relay(left: socket.socket, right: socket.socket) -> None:
|
|
sockets = [left, right]
|
|
while sockets:
|
|
readable, _, _ = select.select(sockets, [], [], 60.0)
|
|
if not readable:
|
|
return
|
|
for source in readable:
|
|
destination = right if source is left else left
|
|
data = source.recv(65536)
|
|
if not data:
|
|
return
|
|
destination.sendall(data)
|
|
|
|
|
|
def _is_ipv4(host: str) -> bool:
|
|
try:
|
|
socket.inet_aton(host)
|
|
return True
|
|
except OSError:
|
|
return False
|