1392 lines
58 KiB
Python
1392 lines
58 KiB
Python
"""Single-file Python P2P channel: WebRTC first, encrypted relay fallback."""
|
||
|
||
import asyncio
|
||
import base64
|
||
import hashlib
|
||
import hmac
|
||
import inspect
|
||
import json
|
||
import logging
|
||
import math
|
||
import os
|
||
import random
|
||
import re
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
import zlib
|
||
from collections import deque
|
||
from contextlib import suppress
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from urllib.parse import parse_qsl, quote, unquote, urlencode, urlsplit, urlunsplit
|
||
|
||
import websockets
|
||
try:
|
||
from aiortc import (
|
||
RTCConfiguration,
|
||
RTCIceServer,
|
||
RTCPeerConnection,
|
||
RTCSessionDescription,
|
||
)
|
||
except ImportError: # 移动端(如Chaquopy)无aiortc轮子：自动降级为纯加密relay
|
||
RTCConfiguration = RTCIceServer = RTCPeerConnection = RTCSessionDescription = None
|
||
from cryptography.hazmat.primitives import hashes
|
||
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
|
||
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
|
||
from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
|
||
|
||
|
||
# ---- merged from auth.py ----
|
||
# 仅用于过滤不知道本项目的扫描和噪声，不是安全秘密。
|
||
# 即使该值公开，握手访问安全仍由不可猜测、不得泄露的 UUID 保证。
|
||
DEFAULT_ACCESS_KEY = "ljq-p2p-ws-default-access-key-v1-7f3c9a2e"
|
||
DEFAULT_STUN_URL = "stun:47.101.182.29:49157"
|
||
|
||
|
||
STANDARD_SIGNAL_URL = "ws://47.101.182.29:49157/ws"
|
||
|
||
|
||
def secure_signal_url(url: str) -> str:
|
||
"""信令统一走 47 VPS 明文 ws(无 TLS)。
|
||
|
||
传输层自带 X25519+HKDF 端到端加密, 信令只转发密文, TLS 不增加机密性,
|
||
反而多一跳 jump1 反代(且不转 UDP STUN)。2.33.4 曾把旧配对改写成
|
||
wss://p2p.gaagent.ai, 这里把这类已保存配对原房间迁回 47, 房间号不变。
|
||
"""
|
||
parts = urlsplit(url)
|
||
if parts.hostname == "p2p.gaagent.ai":
|
||
std = urlsplit(STANDARD_SIGNAL_URL)
|
||
return urlunsplit((std.scheme, std.netloc, parts.path or std.path, parts.query, parts.fragment))
|
||
return url
|
||
|
||
|
||
def signed_url(url: str, action: str, target: str, key: str = DEFAULT_ACCESS_KEY) -> str:
|
||
"""为一次请求添加短时有效、不可重放的 HMAC 查询参数。"""
|
||
ts = str(int(time.time()))
|
||
nonce = os.urandom(16).hex()
|
||
payload = f"{action}\n{target}\n{ts}\n{nonce}".encode()
|
||
sig = hmac.new(key.encode(), payload, hashlib.sha256).hexdigest()
|
||
parts = urlsplit(url)
|
||
query = dict(parse_qsl(parts.query, keep_blank_values=True))
|
||
query.update(ts=ts, nonce=nonce, sig=sig)
|
||
return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
|
||
|
||
def code_url(signal_url: str, endpoint: str) -> str:
|
||
"""把 ws(s)://host/[prefix/]ws 转成 http(s)://host/[prefix/]code/...。"""
|
||
parts = urlsplit(secure_signal_url(signal_url))
|
||
scheme = {"ws": "http", "wss": "https"}.get(parts.scheme)
|
||
if not scheme:
|
||
raise ValueError("signal_url must use ws:// or wss://")
|
||
parent = parts.path.rsplit("/", 1)[0]
|
||
path = f"{parent}/{endpoint.lstrip('/')}"
|
||
return urlunsplit((scheme, parts.netloc, path, "", ""))
|
||
|
||
# ---- merged from crypto.py ----
|
||
COMPRESS_MIN = 512
|
||
|
||
_FLAG_ZIP = 0x01
|
||
|
||
_FLAG_TEXT = 0x02
|
||
|
||
def _pub_raw(pk):
|
||
"""X25519 公钥原始 32 字节; 兼容 cryptography<40 (Android Chaquopy 无 public_bytes_raw)。"""
|
||
try:
|
||
return pk.public_bytes_raw()
|
||
except AttributeError:
|
||
from cryptography.hazmat.primitives import serialization as _ser
|
||
return pk.public_bytes(_ser.Encoding.Raw, _ser.PublicFormat.Raw)
|
||
|
||
def keypair():
|
||
"""生成一次性 X25519 密钥对，返回 (私钥, 公钥字节)。"""
|
||
sk = X25519PrivateKey.generate()
|
||
return sk, _pub_raw(sk.public_key())
|
||
|
||
def session_key(sk, peer_pub: bytes, psk: bytes = b"") -> bytes:
|
||
"""ECDH + HKDF 派生 32 字节会话密钥。
|
||
|
||
salt 取两个公钥排序后拼接，双方算出同一个值，且把密钥绑定到本次配对。
|
||
psk(共享秘密)也并入 salt：不知 psk 的信令方换公钥做中间人也解不开；空 psk 与旧版完全一致。
|
||
"""
|
||
mine = _pub_raw(sk.public_key())
|
||
salt = b"".join(sorted([mine, peer_pub])) + psk
|
||
shared = sk.exchange(X25519PublicKey.from_public_bytes(peer_pub))
|
||
return HKDF(
|
||
algorithm=hashes.SHA256(), length=32, salt=salt, info=b"p2p-ws/v1"
|
||
).derive(shared)
|
||
|
||
def seal(key: bytes, data, compress: bool = True) -> bytes:
|
||
"""把 str/bytes 压缩并加密成一帧。"""
|
||
flags = 0
|
||
if isinstance(data, str):
|
||
data = data.encode()
|
||
flags |= _FLAG_TEXT
|
||
if compress and len(data) >= COMPRESS_MIN:
|
||
packed = zlib.compress(data, 6)
|
||
if len(packed) < len(data):
|
||
data, flags = packed, flags | _FLAG_ZIP
|
||
nonce = os.urandom(12)
|
||
return nonce + AESGCM(key).encrypt(nonce, bytes([flags]) + data, None)
|
||
|
||
def open_(key: bytes, frame: bytes):
|
||
"""解密一帧，还原成原始 str 或 bytes；被篡改会抛异常。"""
|
||
plain = AESGCM(key).decrypt(frame[:12], frame[12:], None)
|
||
flags, data = plain[0], plain[1:]
|
||
if flags & _FLAG_ZIP:
|
||
data = zlib.decompress(data)
|
||
return data.decode() if flags & _FLAG_TEXT else data
|
||
|
||
# ---- merged from signal.py ----
|
||
log = logging.getLogger("p2p_ws.signal")
|
||
|
||
def _ws_connect_kw():
|
||
"""websockets.connect 的固定实参。
|
||
|
||
信令必须直连。websockets ≥14 的 connect() 默认读 HTTPS_PROXY/ALL_PROXY, 而 GA 会把给 LLM
|
||
配的翻墙代理写进这些变量(android_entry._apply_proxy) —— 于是短码换房(urlopen 已显式清代理,
|
||
见文末 _direct_opener)与信令 WS 分走两条路: 换码直连、信令走代理。代理多半不支持非标端口
|
||
49157, 且 socks5 还需 python_socks(未打包 → ImportError), 结果是配了代理的用户必然连不上,
|
||
报的还是与网络无关的错。这里显式声明不走代理, 与「两步必须同路」的不变量对齐。
|
||
proxy= 在 <14 不存在, 故按签名探测而非 try/except —— 后者会把真实的连接错误一并吞掉。"""
|
||
kw = {"max_size": 1 << 21}
|
||
try:
|
||
if "proxy" in inspect.signature(websockets.connect).parameters:
|
||
kw["proxy"] = None
|
||
except (TypeError, ValueError):
|
||
pass
|
||
return kw
|
||
|
||
|
||
_WS_CONNECT_KW = _ws_connect_kw()
|
||
|
||
|
||
class Signal:
|
||
def __init__(
|
||
self, url: str, room: str, retries: int = 6,
|
||
access_key: str = DEFAULT_ACCESS_KEY,
|
||
):
|
||
url = secure_signal_url(url)
|
||
separator = "&" if "?" in url else "?"
|
||
self.url = f"{url}{separator}room={room}"
|
||
self.room = room
|
||
self.access_key = access_key
|
||
self.retries = retries # 单次断线的最大重连尝试次数
|
||
self.on_data = None # 回调：dict -> None
|
||
self._ws = None
|
||
self._queues: dict[str, asyncio.Queue] = {}
|
||
self._ready = asyncio.Event()
|
||
self._closing = False
|
||
self._task = None
|
||
|
||
# ---------- 生命周期 ----------
|
||
|
||
def _signed_url(self):
|
||
return signed_url(self.url, "ws", self.room, self.access_key)
|
||
|
||
async def start(self):
|
||
# 半开连接(切网/熄屏)靠 ping 在 ~25s 内暴露，默认 20+20s 太慢
|
||
self._ws = await websockets.connect(
|
||
self._signed_url(), ping_interval=15, ping_timeout=10, **_WS_CONNECT_KW)
|
||
self._ready.set()
|
||
self._task = asyncio.create_task(self._loop())
|
||
return self
|
||
|
||
async def close(self):
|
||
self._closing = True
|
||
self._ready.clear()
|
||
if self._task:
|
||
self._task.cancel()
|
||
if self._ws:
|
||
# 失效会话不等对端关闭握手，先中止本地传输以释放房间。
|
||
self._ws.transport.abort()
|
||
|
||
@property
|
||
def alive(self) -> bool:
|
||
return self._ready.is_set() and not self._closing
|
||
|
||
# ---------- 收发 ----------
|
||
|
||
def queue(self, type_: str) -> asyncio.Queue:
|
||
return self._queues.setdefault(type_, asyncio.Queue())
|
||
|
||
async def expect(self, type_: str, timeout: float = 15):
|
||
"""等指定类型消息；连接关闭或服务端报错(如被顶号)时立即抛错，不傻等到超时。"""
|
||
want = asyncio.ensure_future(self.queue(type_).get())
|
||
fail = [asyncio.ensure_future(self.queue(t).get()) for t in ("closed", "error")]
|
||
try:
|
||
done, _ = await asyncio.wait(
|
||
[want, *fail], timeout=timeout, return_when=asyncio.FIRST_COMPLETED)
|
||
if want in done:
|
||
return want.result()
|
||
if not done:
|
||
raise asyncio.TimeoutError(f"no {type_!r} within {timeout}s")
|
||
raise ConnectionError(f"signal ended while waiting {type_!r}: {next(iter(done)).result()}")
|
||
finally:
|
||
for f in (want, *fail):
|
||
if not f.done():
|
||
f.cancel()
|
||
|
||
async def send(self, obj: dict, timeout: float = 15):
|
||
"""等连接就绪后发送；重连期间会阻塞而不是直接失败。"""
|
||
if self._closing:
|
||
raise ConnectionError("signal is closed")
|
||
await asyncio.wait_for(self._ready.wait(), timeout)
|
||
if self._closing:
|
||
raise ConnectionError("signal is closed")
|
||
await self._ws.send(json.dumps(obj))
|
||
|
||
async def _loop(self):
|
||
try:
|
||
async for raw in self._ws:
|
||
msg = json.loads(raw)
|
||
if msg.get("type") == "data" and self.on_data:
|
||
self.on_data(msg)
|
||
else:
|
||
self.queue(msg.get("type", "")).put_nowait(msg)
|
||
except Exception as exc:
|
||
log.debug("signal read stopped: %r", exc)
|
||
if not self._closing:
|
||
# 房间配对、ECDH 密钥与 relay 序号都属于本次 socket 会话；
|
||
# 信令原地重连会留下僵尸 P2PSocket，必须由上层完整重握手。
|
||
self._ready.clear()
|
||
self.queue("closed").put_nowait({"type": "closed"})
|
||
|
||
# ---- merged from socket.py ----
|
||
log = logging.getLogger("p2p_ws")
|
||
|
||
_CHUNK = 48 * 1024
|
||
# WebRTC DataChannel 单条消息有实现上限（aiortc 上超过 64KiB 会静默截断），
|
||
# 所以直连消息统一先整包压缩再分片，接收端重组后还原。
|
||
_DIRECT_CHUNK = 16 * 1024
|
||
_DIRECT_MAGIC = b"P2P1"
|
||
_DIRECT_HEAD = len(_DIRECT_MAGIC) + 1 + 6 + 2 + 2
|
||
|
||
def _pack_direct(data):
|
||
"""把一条消息封装成若干可安全通过 DataChannel 的帧。"""
|
||
raw = data.encode() if isinstance(data, str) else data
|
||
flags = _FLAG_TEXT if isinstance(data, str) else 0
|
||
# 保留小消息的旧版 wire format；魔数开头的 bytes 必须封装以免误判。
|
||
if len(raw) <= _DIRECT_CHUNK and not (
|
||
isinstance(data, bytes) and data.startswith(_DIRECT_MAGIC)):
|
||
yield data
|
||
return
|
||
if len(raw) >= COMPRESS_MIN:
|
||
packed = zlib.compress(raw, 6)
|
||
if len(packed) < len(raw):
|
||
raw, flags = packed, flags | _FLAG_ZIP
|
||
ident = os.urandom(6)
|
||
total = max(1, math.ceil(len(raw) / _DIRECT_CHUNK))
|
||
if total > 0xFFFF:
|
||
raise ValueError("message too large for direct mode")
|
||
for n in range(total):
|
||
head = (
|
||
_DIRECT_MAGIC + bytes([flags]) + ident
|
||
+ total.to_bytes(2, "big") + n.to_bytes(2, "big")
|
||
)
|
||
yield head + raw[n * _DIRECT_CHUNK : (n + 1) * _DIRECT_CHUNK]
|
||
|
||
class P2PSocket:
|
||
"""一条按 UUID 配对、用法近似 WebSocket 的异步消息通道。
|
||
|
||
用法::
|
||
|
||
ws = await P2PSocket.connect("ws://host:8765/ws", "same-uuid")
|
||
await ws.send("hello")
|
||
print(await ws.recv())
|
||
await ws.close()
|
||
|
||
``mode`` 为 ``"direct"`` 或 ``"relay"``；send/recv 的消息边界和
|
||
str/bytes 类型都会保留。
|
||
"""
|
||
|
||
def __init__(
|
||
self, signal_url, room, *, direct_timeout=15, peer_timeout=120, retries=6,
|
||
stun=True, access_key=DEFAULT_ACCESS_KEY,
|
||
):
|
||
self.signal = Signal(signal_url, room, retries, access_key)
|
||
self.direct_timeout = direct_timeout
|
||
# 等对端进房的时长；须 >= 配对码有效期(120s)，否则码没过期本端已离房
|
||
self.peer_timeout = peer_timeout
|
||
self.stun = stun
|
||
self.mode = "connecting"
|
||
self.closed = False
|
||
self._key = None
|
||
self._pc = None
|
||
self._dc = None
|
||
self._inbox = asyncio.Queue()
|
||
self._parts = {}
|
||
self._send_lock = asyncio.Lock()
|
||
self._relay_ready = asyncio.Event()
|
||
self._direct_open = asyncio.Event()
|
||
self._peer_watch = None
|
||
self._direct_task = None
|
||
self._close_task = None
|
||
self._upgrade_prepare = asyncio.Event()
|
||
self._upgrade_ack = asyncio.Event()
|
||
self._upgrade_commit = asyncio.Event()
|
||
self._direct_pending = []
|
||
|
||
@classmethod
|
||
async def connect(
|
||
cls, signal_url: str, room: str, *, direct_timeout=15, peer_timeout=120,
|
||
retries=6, stun=True, access_key=DEFAULT_ACCESS_KEY, psk=b"",
|
||
):
|
||
"""连接同一 UUID 的另一个客户端，直连失败自动回退中继。"""
|
||
self = cls(
|
||
signal_url, room, direct_timeout=direct_timeout,
|
||
peer_timeout=peer_timeout, retries=retries, stun=stun,
|
||
access_key=access_key,
|
||
)
|
||
self._psk = psk
|
||
try:
|
||
await self._connect()
|
||
except BaseException:
|
||
# 失败必须关闭信令连接，否则泄漏的连接会留在房间里
|
||
# 与本端下一次重连"自配对"，把双人房占满导致真对端永远进不来。
|
||
await self.close()
|
||
raise
|
||
return self
|
||
|
||
async def _connect(self):
|
||
await self.signal.start()
|
||
peer = await self.signal.expect("peer", self.peer_timeout)
|
||
|
||
# 一次性 ECDH：即使走中继，服务器也拿不到会话密钥。
|
||
private, public = keypair()
|
||
await self.signal.send(
|
||
{"type": "key", "key": base64.b64encode(public).decode()}
|
||
)
|
||
key_msg = await self.signal.expect("key")
|
||
self._key = session_key(private, base64.b64decode(key_msg["key"]), getattr(self, "_psk", b""))
|
||
self.signal.on_data = self._on_relay_data
|
||
self._relay_ready.set()
|
||
self.mode = "relay"
|
||
# 对端离开房间后本连接已无意义；若不关闭，本端会一直占着房间等待，
|
||
# 新 peer 加入时无人重做 key 握手，对方只能 15s 超时退出（死循环）。
|
||
self._peer_watch = asyncio.create_task(self._on_peer_left())
|
||
if RTCPeerConnection is not None and self.direct_timeout > 0:
|
||
self._direct_task = asyncio.create_task(
|
||
self._upgrade_direct(bool(peer["initiator"])))
|
||
return self
|
||
|
||
async def _upgrade_direct(self, initiator):
|
||
"""Relay 已可用时后台打洞，以三步屏障保序切到 DataChannel。"""
|
||
try:
|
||
await asyncio.wait_for(self._make_direct(initiator), self.direct_timeout)
|
||
await asyncio.wait_for(self._direct_open.wait(), self.direct_timeout)
|
||
if initiator:
|
||
async with self._send_lock:
|
||
await self._send_control(b"prepare")
|
||
await asyncio.wait_for(self._upgrade_ack.wait(), self.direct_timeout)
|
||
await self._send_control(b"commit")
|
||
self._commit_direct()
|
||
else:
|
||
await asyncio.wait_for(self._upgrade_prepare.wait(), self.direct_timeout)
|
||
async with self._send_lock:
|
||
await self._send_control(b"ack")
|
||
await asyncio.wait_for(self._upgrade_commit.wait(), self.direct_timeout)
|
||
self._commit_direct()
|
||
log.info("encrypted relay upgraded to direct DataChannel")
|
||
except Exception as exc:
|
||
log.info("background direct upgrade unavailable, keeping relay: %r", exc)
|
||
await self._drop_direct()
|
||
|
||
def _commit_direct(self):
|
||
self.mode = "direct"
|
||
for item in self._direct_pending:
|
||
self._inbox.put_nowait(item)
|
||
self._direct_pending.clear()
|
||
|
||
async def _on_peer_left(self):
|
||
"""对端离开或本端信令断开时，关闭连接让上层完整重握手。"""
|
||
waits = [asyncio.create_task(self.signal.queue(kind).get())
|
||
for kind in ("peer_left", "closed")]
|
||
try:
|
||
await asyncio.wait(waits, return_when=asyncio.FIRST_COMPLETED)
|
||
except asyncio.CancelledError:
|
||
return
|
||
finally:
|
||
for task in waits:
|
||
task.cancel()
|
||
if not self.closed:
|
||
log.info("P2P peer/signaling left; closing socket for a fresh handshake")
|
||
await self.close()
|
||
|
||
# ---------- WebSocket 风格接口 ----------
|
||
|
||
async def send(self, data):
|
||
"""发送一条 str 或 bytes 消息；并发 send 会按调用顺序串行化。"""
|
||
if self.closed:
|
||
raise ConnectionError("socket is closed")
|
||
if not isinstance(data, (str, bytes)):
|
||
raise TypeError("data must be str or bytes")
|
||
async with self._send_lock:
|
||
if self.mode == "direct" and self._dc and self._dc.readyState == "open":
|
||
for frame in _pack_direct(data):
|
||
self._dc.send(frame)
|
||
# 避免一次灌满发送缓冲导致底层丢帧。
|
||
while self._dc.bufferedAmount > 4 * 1024 * 1024:
|
||
await asyncio.sleep(0.01)
|
||
return
|
||
# 直连在运行中失效时，无需重新创建对象，立即降级。
|
||
self.mode = "relay"
|
||
await self._send_relay(data)
|
||
|
||
async def recv(self):
|
||
"""等待并返回下一条完整消息，保持原始 str/bytes 类型。"""
|
||
item = await self._inbox.get()
|
||
if isinstance(item, _Closed):
|
||
raise ConnectionError(item.reason)
|
||
return item
|
||
|
||
async def close(self):
|
||
if self._close_task is None:
|
||
self.closed = True
|
||
for task in (self._peer_watch, self._direct_task):
|
||
if task is not None and task is not asyncio.current_task():
|
||
task.cancel()
|
||
self._inbox.put_nowait(_Closed("socket is closed"))
|
||
self.mode = "closed"
|
||
self._close_task = asyncio.create_task(self._finish_close())
|
||
await asyncio.shield(self._close_task)
|
||
|
||
async def _finish_close(self):
|
||
await self.signal.close() # 先撤销信令，不能让 RTC 清理阻塞房间释放。
|
||
cleanup = asyncio.create_task(self._drop_direct())
|
||
_, pending = await asyncio.wait({cleanup}, timeout=1)
|
||
if pending:
|
||
cleanup.cancel() # wait_for 会等取消完成，遇到慢清理仍可能无界等待。
|
||
log.warning("direct cleanup exceeded 1s; reconnect may proceed")
|
||
|
||
async def __aenter__(self):
|
||
return self
|
||
|
||
async def __aexit__(self, *_):
|
||
await self.close()
|
||
|
||
# ---------- WebRTC 直连 ----------
|
||
|
||
async def _make_direct(self, initiator: bool):
|
||
if RTCPeerConnection is None:
|
||
raise RuntimeError("aiortc not installed; direct mode unavailable")
|
||
if self.stun:
|
||
if isinstance(self.stun, str):
|
||
stun_url = self.stun
|
||
else:
|
||
parts = urlsplit(self.signal.url)
|
||
host = parts.hostname
|
||
if not host:
|
||
raise ValueError("signal URL has no host for STUN")
|
||
if ":" in host: # IPv6 URI 必须保留方括号
|
||
host = f"[{host}]"
|
||
port = parts.port or (443 if parts.scheme == "wss" else 80)
|
||
stun_url = f"stun:{host}:{port}"
|
||
ice = [RTCIceServer(stun_url)]
|
||
else:
|
||
ice = []
|
||
pc = self._pc = RTCPeerConnection(RTCConfiguration(iceServers=ice))
|
||
|
||
@pc.on("connectionstatechange")
|
||
async def state_changed():
|
||
if pc.connectionState in ("failed", "closed", "disconnected"):
|
||
self._direct_open.clear()
|
||
if not self.closed and self.mode == "direct":
|
||
log.info("direct channel lost; switched to encrypted relay")
|
||
self.mode = "relay"
|
||
|
||
if initiator:
|
||
self._bind_datachannel(pc.createDataChannel("p2p-ws"))
|
||
await pc.setLocalDescription(await pc.createOffer())
|
||
await self.signal.send(
|
||
{
|
||
"type": "sdp",
|
||
"sdp_type": pc.localDescription.type,
|
||
"sdp": pc.localDescription.sdp,
|
||
}
|
||
)
|
||
answer = await self.signal.expect("sdp", self.direct_timeout)
|
||
await pc.setRemoteDescription(
|
||
RTCSessionDescription(answer["sdp"], answer["sdp_type"])
|
||
)
|
||
else:
|
||
channel_seen = asyncio.Event()
|
||
|
||
@pc.on("datachannel")
|
||
def datachannel(channel):
|
||
self._bind_datachannel(channel)
|
||
channel_seen.set()
|
||
|
||
offer = await self.signal.expect("sdp", self.direct_timeout)
|
||
await pc.setRemoteDescription(
|
||
RTCSessionDescription(offer["sdp"], offer["sdp_type"])
|
||
)
|
||
await pc.setLocalDescription(await pc.createAnswer())
|
||
await self.signal.send(
|
||
{
|
||
"type": "sdp",
|
||
"sdp_type": pc.localDescription.type,
|
||
"sdp": pc.localDescription.sdp,
|
||
}
|
||
)
|
||
await asyncio.wait_for(channel_seen.wait(), self.direct_timeout)
|
||
|
||
def _bind_datachannel(self, dc):
|
||
self._dc = dc
|
||
|
||
# 应答方拿到的通道可能已经是 open，此时不会再触发 open 事件
|
||
if dc.readyState == "open":
|
||
self._direct_open.set()
|
||
|
||
@dc.on("open")
|
||
def opened():
|
||
self._direct_open.set()
|
||
|
||
@dc.on("message")
|
||
def message(data):
|
||
self._recv_direct(data)
|
||
|
||
@dc.on("close")
|
||
def closed():
|
||
self._direct_open.clear()
|
||
if not self.closed and self.mode == "direct":
|
||
self.mode = "relay"
|
||
|
||
def _recv_direct(self, data):
|
||
"""重组直连分片；未封装的小消息按原样交付。"""
|
||
if not isinstance(data, bytes) or not data.startswith(_DIRECT_MAGIC):
|
||
self._inbox.put_nowait(data)
|
||
return
|
||
flags = data[4]
|
||
key = ("direct", data[5:11])
|
||
total = int.from_bytes(data[11:13], "big")
|
||
index = int.from_bytes(data[13:15], "big")
|
||
entry = self._parts.setdefault(
|
||
key, {"flags": flags, "total": total, "chunks": {}}
|
||
)
|
||
entry["chunks"][index] = data[_DIRECT_HEAD:]
|
||
if len(entry["chunks"]) < entry["total"]:
|
||
return
|
||
del self._parts[key]
|
||
try:
|
||
raw = b"".join(entry["chunks"][i] for i in range(entry["total"]))
|
||
if entry["flags"] & _FLAG_ZIP:
|
||
raw = zlib.decompress(raw)
|
||
item = raw.decode() if entry["flags"] & _FLAG_TEXT else raw
|
||
if self.mode == "direct":
|
||
self._inbox.put_nowait(item)
|
||
else:
|
||
self._direct_pending.append(item)
|
||
except Exception as exc:
|
||
log.warning("discarded invalid direct frame: %r", exc)
|
||
|
||
async def _drop_direct(self):
|
||
pc, self._pc, self._dc = self._pc, None, None
|
||
self._direct_open.clear()
|
||
if pc:
|
||
with suppress(Exception):
|
||
await pc.close()
|
||
|
||
# ---------- 加密中继 ----------
|
||
|
||
async def _send_control(self, command):
|
||
await self._send_relay(command, control=True)
|
||
|
||
async def _send_relay(self, data, control=False):
|
||
await self._relay_ready.wait()
|
||
raw = data.encode() if isinstance(data, str) else data
|
||
kind = 2 if control else int(isinstance(data, str))
|
||
total = max(1, math.ceil(len(raw) / _CHUNK))
|
||
msg_id = os.urandom(8).hex()
|
||
for n in range(total):
|
||
chunk = bytes([kind]) + raw[n * _CHUNK : (n + 1) * _CHUNK]
|
||
cipher = seal(self._key, chunk)
|
||
await self.signal.send(
|
||
{
|
||
"type": "data",
|
||
"id": msg_id,
|
||
"n": n,
|
||
"total": total,
|
||
"blob": base64.b64encode(cipher).decode(),
|
||
}
|
||
)
|
||
|
||
def _on_relay_data(self, msg):
|
||
try:
|
||
chunk = open_(self._key, base64.b64decode(msg["blob"]))
|
||
if not isinstance(chunk, bytes) or not chunk:
|
||
raise ValueError("invalid encrypted chunk")
|
||
ident = msg["id"]
|
||
entry = self._parts.setdefault(
|
||
ident, {"kind": chunk[0], "total": int(msg["total"]), "chunks": {}}
|
||
)
|
||
entry["chunks"][int(msg["n"])] = chunk[1:]
|
||
if len(entry["chunks"]) == entry["total"]:
|
||
raw = b"".join(entry["chunks"][i] for i in range(entry["total"]))
|
||
del self._parts[ident]
|
||
if entry["kind"] == 2:
|
||
if raw == b"prepare":
|
||
self._upgrade_prepare.set()
|
||
elif raw == b"ack":
|
||
self._upgrade_ack.set()
|
||
elif raw == b"commit":
|
||
self._upgrade_commit.set()
|
||
else:
|
||
self._inbox.put_nowait(raw.decode() if entry["kind"] else raw)
|
||
except Exception as exc:
|
||
# 篡改、错误密钥或畸形分片不会进入应用层。
|
||
log.warning("discarded invalid relay frame: %r", exc)
|
||
|
||
class _Closed:
|
||
def __init__(self, reason):
|
||
self.reason = reason
|
||
|
||
# ---- merged from invite.py ----
|
||
_CODE = re.compile(r"^[0-9]{9}$")
|
||
|
||
# 短码换房与随后的 WS 信令必须走同一条路: websockets 从不读环境代理, 而 urlopen 默认读
|
||
# (GA 给 LLM 配的翻墙代理会写进 HTTP_PROXY/http_proxy)。两步分道时, 换码走代理、信令直连,
|
||
# 代理不支持非标端口(如 49157)就只在换码这步炸, 报出与网络无关的 RemoteDisconnected。
|
||
# 信令服务器是直连服务, 这里显式清空代理与之对齐。
|
||
_direct_opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
|
||
|
||
|
||
async def _post_json(url: str, timeout: float = 15):
|
||
def request():
|
||
req = urllib.request.Request(url, data=b"", method="POST",
|
||
headers={"User-Agent": "GenericAgent-P2P/1.0"})
|
||
with _direct_opener.open(req, timeout=timeout) as response:
|
||
return json.load(response)
|
||
|
||
try:
|
||
return await asyncio.to_thread(request)
|
||
except urllib.error.HTTPError as exc:
|
||
reason = exc.read().decode(errors="replace").strip()
|
||
raise ConnectionError(f"short-code server rejected request: {exc.code} {reason}") from exc
|
||
except Exception as exc: # 网络层失败带上 URL 与异常类型: 原来只剩一句 RemoteDisconnected, 无从判断死在哪一步
|
||
raise ConnectionError(f"short-code server unreachable ({type(exc).__name__}: {exc}) url={url}") from exc
|
||
|
||
# ---- optional HTTP-over-P2P tunnel ----
|
||
_HTTP_REQ, _HTTP_RES = "p2p-http.req", "p2p-http.res"
|
||
_HTTP_HOP = {
|
||
"connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
|
||
"te", "trailer", "transfer-encoding", "upgrade", "content-length", "host",
|
||
}
|
||
|
||
|
||
def _http_modules():
|
||
"""The base P2P channel stays dependency-light; HTTP support is optional."""
|
||
try:
|
||
import aiohttp
|
||
from aiohttp import web
|
||
from multidict import CIMultiDict
|
||
except ImportError as exc:
|
||
raise RuntimeError(
|
||
'HTTP tunnelling requires aiohttp: pip install "p2p-ws[http]"'
|
||
) from exc
|
||
return aiohttp, web, CIMultiDict
|
||
|
||
|
||
def _safe_headers(items, drop=()):
|
||
blocked = _HTTP_HOP | {str(name).lower() for name in drop}
|
||
return [(str(k), str(v)) for k, v in items if str(k).lower() not in blocked]
|
||
|
||
|
||
def _decode_body(value, limit):
|
||
try:
|
||
body = base64.b64decode(value or "", validate=True)
|
||
except Exception as exc:
|
||
raise ValueError("invalid base64 HTTP body") from exc
|
||
if len(body) > limit:
|
||
raise ValueError("HTTP body exceeds configured limit")
|
||
return body
|
||
|
||
|
||
def _claim_receiver(socket, owner):
|
||
"""A tunnel owns recv(); two consumers would silently steal each other's frames."""
|
||
current = getattr(socket, "_http_tunnel_owner", None)
|
||
if current is not None and current is not owner:
|
||
raise RuntimeError("this P2PSocket already has an HTTP tunnel consumer")
|
||
socket._http_tunnel_owner = owner
|
||
|
||
|
||
def _release_receiver(socket, owner):
|
||
if getattr(socket, "_http_tunnel_owner", None) is owner:
|
||
delattr(socket, "_http_tunnel_owner")
|
||
|
||
|
||
class HTTPExporter:
|
||
"""Expose one fixed HTTP origin to the peer over a P2PSocket.
|
||
|
||
The peer supplies only a relative target. Redirects are not followed, so
|
||
this cannot be turned into a general SSRF proxy. This object exclusively
|
||
consumes ``socket.recv()`` until closed.
|
||
"""
|
||
|
||
def __init__(
|
||
self, socket, upstream, *, allow=("/",),
|
||
methods=("GET", "HEAD", "POST", "PUT", "PATCH", "DELETE"),
|
||
query=None, headers=None, drop_request_headers=(), timeout=30,
|
||
max_body=8 * 1024 * 1024, max_concurrency=16,
|
||
):
|
||
parts = urlsplit(str(upstream))
|
||
if parts.scheme not in ("http", "https") or not parts.hostname:
|
||
raise ValueError("upstream must be an absolute http(s) URL")
|
||
if parts.query or parts.fragment:
|
||
raise ValueError("upstream must not contain query or fragment")
|
||
self.socket, self.upstream = socket, str(upstream).rstrip("/")
|
||
self.allow = tuple(str(path) for path in allow)
|
||
if not self.allow or any(not path.startswith("/") for path in self.allow):
|
||
raise ValueError("allow must contain absolute path prefixes")
|
||
self.methods = {str(method).upper() for method in methods}
|
||
self.query = [(str(k), str(v)) for k, v in (query or {}).items()]
|
||
self.headers = [(str(k), str(v)) for k, v in (headers or {}).items()]
|
||
self.drop_request_headers = tuple(drop_request_headers)
|
||
self.timeout, self.max_body = float(timeout), int(max_body)
|
||
self._slots = asyncio.Semaphore(int(max_concurrency))
|
||
self._session = self._receiver = None
|
||
self._requests, self._closed = set(), asyncio.Event()
|
||
|
||
async def start(self):
|
||
if self._receiver:
|
||
return self
|
||
aiohttp, _, _ = _http_modules()
|
||
_claim_receiver(self.socket, self)
|
||
try:
|
||
self._session = aiohttp.ClientSession(
|
||
timeout=aiohttp.ClientTimeout(total=self.timeout)
|
||
)
|
||
self._receiver = asyncio.create_task(self._receive())
|
||
except Exception:
|
||
_release_receiver(self.socket, self)
|
||
raise
|
||
return self
|
||
|
||
async def _receive(self):
|
||
try:
|
||
while True:
|
||
raw = await self.socket.recv()
|
||
try:
|
||
message = json.loads(raw) if isinstance(raw, str) else None
|
||
except (TypeError, json.JSONDecodeError):
|
||
message = None
|
||
if not isinstance(message, dict) or message.get("type") != _HTTP_REQ:
|
||
log.warning("HTTPExporter ignored a non-HTTP message")
|
||
continue
|
||
task = asyncio.create_task(self._forward(message))
|
||
self._requests.add(task)
|
||
task.add_done_callback(self._requests.discard)
|
||
except asyncio.CancelledError:
|
||
pass
|
||
except Exception as exc:
|
||
log.debug("HTTP exporter stopped: %r", exc)
|
||
finally:
|
||
self._closed.set()
|
||
|
||
async def _forward(self, message):
|
||
request_id = str(message.get("id", ""))
|
||
status, response_headers, body = 502, [], b"upstream request failed"
|
||
try:
|
||
if not request_id:
|
||
raise ValueError("missing request id")
|
||
method = str(message.get("method", "")).upper()
|
||
if method not in self.methods:
|
||
status, body = 405, b"method not allowed"
|
||
else:
|
||
target = str(message.get("target", ""))
|
||
parts = urlsplit(target)
|
||
segments = parts.path.replace("\\", "/").split("/")
|
||
if (parts.scheme or parts.netloc or parts.fragment
|
||
or not parts.path.startswith("/") or ".." in segments):
|
||
status, body = 400, b"invalid relative HTTP target"
|
||
elif not any(parts.path.startswith(prefix) for prefix in self.allow):
|
||
status, body = 403, b"path not allowed"
|
||
else:
|
||
request_body = _decode_body(message.get("body", ""), self.max_body)
|
||
query = parse_qsl(parts.query, keep_blank_values=True) + self.query
|
||
url = self.upstream + parts.path
|
||
if query:
|
||
url += "?" + urlencode(query)
|
||
incoming = message.get("headers") or []
|
||
if not isinstance(incoming, list):
|
||
raise ValueError("headers must be a list")
|
||
request_headers = _safe_headers(incoming, self.drop_request_headers)
|
||
request_headers.extend(self.headers)
|
||
async with self._slots:
|
||
async with self._session.request(
|
||
method, url, headers=request_headers, data=request_body,
|
||
allow_redirects=False,
|
||
) as response:
|
||
if (response.content_length is not None
|
||
and response.content_length > self.max_body):
|
||
raise ValueError("HTTP response exceeds configured limit")
|
||
# StreamReader.read(n) means "at most n", not
|
||
# "read until EOF"; a single call often returns
|
||
# only aiohttp's current 64KiB buffer.
|
||
chunks, size = [], 0
|
||
async for chunk in response.content.iter_chunked(64 * 1024):
|
||
size += len(chunk)
|
||
if size > self.max_body:
|
||
raise ValueError("HTTP response exceeds configured limit")
|
||
chunks.append(chunk)
|
||
body = b"".join(chunks)
|
||
status = response.status
|
||
response_headers = _safe_headers(response.headers.items())
|
||
except ValueError as exc:
|
||
status, response_headers, body = 400, [], str(exc).encode()
|
||
except Exception as exc:
|
||
log.debug("HTTP upstream request failed: %r", exc)
|
||
response = {
|
||
"v": 1, "type": _HTTP_RES, "id": request_id,
|
||
"status": status, "headers": response_headers,
|
||
"body": base64.b64encode(body).decode("ascii"),
|
||
}
|
||
with suppress(Exception):
|
||
await self.socket.send(json.dumps(response, separators=(",", ":")))
|
||
|
||
async def close(self):
|
||
if self._receiver:
|
||
self._receiver.cancel()
|
||
with suppress(asyncio.CancelledError):
|
||
await self._receiver
|
||
self._receiver = None
|
||
for task in tuple(self._requests):
|
||
task.cancel()
|
||
if self._requests:
|
||
await asyncio.gather(*self._requests, return_exceptions=True)
|
||
self._requests.clear()
|
||
if self._session:
|
||
await self._session.close()
|
||
self._session = None
|
||
_release_receiver(self.socket, self)
|
||
self._closed.set()
|
||
|
||
async def wait_closed(self):
|
||
await self._closed.wait()
|
||
|
||
async def serve_forever(self):
|
||
await self.start()
|
||
await self.wait_closed()
|
||
|
||
async def __aenter__(self):
|
||
return await self.start()
|
||
|
||
async def __aexit__(self, *_):
|
||
await self.close()
|
||
|
||
|
||
class HTTPLocalProxy:
|
||
"""Expose the peer's exported HTTP origin as a local aiohttp server."""
|
||
|
||
def __init__(
|
||
self, socket, listen=("127.0.0.1", 0), *, cors="loopback",
|
||
timeout=30, max_body=8 * 1024 * 1024, allow_public=False,
|
||
):
|
||
host, port = listen
|
||
host = str(host)
|
||
if not allow_public and host not in ("127.0.0.1", "::1", "localhost"):
|
||
raise ValueError("HTTPLocalProxy is loopback-only by default")
|
||
self.socket, self.listen = socket, (host, int(port))
|
||
self.cors, self.timeout, self.max_body = cors, float(timeout), int(max_body)
|
||
self.allow_public = bool(allow_public)
|
||
self._app = self._runner = self._site = self._receiver = None
|
||
self._pending, self._closed = {}, asyncio.Event()
|
||
self._counter = 0
|
||
self.last_used = time.monotonic() # 供上层回收空闲代理
|
||
|
||
async def start(self):
|
||
if self._runner:
|
||
return self
|
||
_, web, _ = _http_modules()
|
||
_claim_receiver(self.socket, self)
|
||
try:
|
||
self._app = web.Application(client_max_size=self.max_body)
|
||
self._app.router.add_route("*", "/{tail:.*}", self._handle)
|
||
self._runner = web.AppRunner(self._app, access_log=None)
|
||
await self._runner.setup()
|
||
self._site = web.TCPSite(self._runner, *self.listen)
|
||
await self._site.start()
|
||
sockets = getattr(self._site, "_server", None).sockets or ()
|
||
if sockets:
|
||
address = sockets[0].getsockname()
|
||
self.url = f"http://{address[0] if ':' not in str(address[0]) else '[' + address[0] + ']'}:{address[1]}"
|
||
self.listen = (address[0], address[1])
|
||
self._receiver = asyncio.create_task(self._receive())
|
||
except Exception:
|
||
_release_receiver(self.socket, self)
|
||
if self._runner:
|
||
await self._runner.cleanup()
|
||
self._runner = self._site = None
|
||
raise
|
||
return self
|
||
|
||
def _cors_headers(self, request):
|
||
if self.cors is False or self.cors is None:
|
||
return {}
|
||
origin = request.headers.get("Origin")
|
||
if self.cors == "*":
|
||
return {"Access-Control-Allow-Origin": "*"}
|
||
if self.cors == "loopback":
|
||
allowed = not origin or origin.startswith(("http://127.0.0.1:", "http://localhost:", "http://[::1]:"))
|
||
return {"Access-Control-Allow-Origin": origin, "Vary": "Origin"} if allowed and origin else {}
|
||
if isinstance(self.cors, str):
|
||
return {"Access-Control-Allow-Origin": origin, "Vary": "Origin"} if origin == self.cors else {}
|
||
allowed = origin in self.cors if origin else False
|
||
return {"Access-Control-Allow-Origin": origin, "Vary": "Origin"} if allowed else {}
|
||
|
||
async def _handle(self, request):
|
||
self.last_used = time.monotonic()
|
||
_, web, _ = _http_modules()
|
||
cors_headers = self._cors_headers(request)
|
||
if request.method == "OPTIONS":
|
||
headers = {**cors_headers,
|
||
"Access-Control-Allow-Methods": "GET,HEAD,POST,PUT,PATCH,DELETE,OPTIONS",
|
||
"Access-Control-Allow-Headers": request.headers.get("Access-Control-Request-Headers", "Content-Type")}
|
||
return web.Response(status=204, headers=headers)
|
||
if request.headers.get("Content-Length"):
|
||
try:
|
||
if int(request.headers["Content-Length"]) > self.max_body:
|
||
return web.Response(status=413, text="request body too large", headers=cors_headers)
|
||
except ValueError:
|
||
return web.Response(status=400, text="invalid content length", headers=cors_headers)
|
||
try:
|
||
body = await request.read()
|
||
except Exception as exc:
|
||
return web.Response(status=400, text=str(exc), headers=cors_headers)
|
||
if len(body) > self.max_body:
|
||
return web.Response(status=413, text="request body too large", headers=cors_headers)
|
||
self._counter += 1
|
||
request_id = f"{id(self):x}-{self._counter:x}"
|
||
message = {
|
||
"v": 1, "type": _HTTP_REQ, "id": request_id,
|
||
"method": request.method, "target": request.raw_path,
|
||
"headers": _safe_headers(request.headers.items()),
|
||
"body": base64.b64encode(body).decode("ascii"),
|
||
}
|
||
future = asyncio.get_running_loop().create_future()
|
||
self._pending[request_id] = future
|
||
try:
|
||
await self.socket.send(json.dumps(message, separators=(",", ":")))
|
||
response = await asyncio.wait_for(future, self.timeout)
|
||
except asyncio.TimeoutError:
|
||
return web.Response(status=504, text="P2P HTTP request timed out", headers=cors_headers)
|
||
except (ConnectionError, asyncio.CancelledError):
|
||
return web.Response(status=502, text="P2P connection closed", headers=cors_headers)
|
||
finally:
|
||
self._pending.pop(request_id, None)
|
||
if not isinstance(response, dict):
|
||
return web.Response(status=502, text="invalid P2P HTTP response", headers=cors_headers)
|
||
try:
|
||
status = int(response["status"])
|
||
response_body = _decode_body(response.get("body", ""), self.max_body)
|
||
headers = dict(_safe_headers(response.get("headers") or []))
|
||
except (KeyError, TypeError, ValueError) as exc:
|
||
return web.Response(status=502, text=f"invalid P2P HTTP response: {exc}", headers=cors_headers)
|
||
headers.update(cors_headers)
|
||
return web.Response(status=status, body=response_body, headers=headers)
|
||
|
||
async def _receive(self):
|
||
try:
|
||
while True:
|
||
raw = await self.socket.recv()
|
||
try:
|
||
message = json.loads(raw) if isinstance(raw, str) else None
|
||
except (TypeError, json.JSONDecodeError):
|
||
message = None
|
||
if isinstance(message, dict) and message.get("type") == _HTTP_RES:
|
||
request_id = str(message.get("id", ""))
|
||
future = self._pending.get(request_id)
|
||
if future is not None and not future.done():
|
||
future.set_result(message)
|
||
except asyncio.CancelledError:
|
||
pass
|
||
except Exception as exc:
|
||
log.debug("HTTP local proxy stopped: %r", exc)
|
||
finally:
|
||
for future in tuple(self._pending.values()):
|
||
if not future.done():
|
||
future.set_exception(ConnectionError("P2P connection closed"))
|
||
self._closed.set()
|
||
|
||
async def close(self):
|
||
if self._receiver:
|
||
self._receiver.cancel()
|
||
with suppress(asyncio.CancelledError):
|
||
await self._receiver
|
||
self._receiver = None
|
||
for future in tuple(self._pending.values()):
|
||
if not future.done():
|
||
future.set_exception(ConnectionError("HTTP proxy closed"))
|
||
self._pending.clear()
|
||
if self._runner:
|
||
await self._runner.cleanup()
|
||
self._runner = self._site = self._app = None
|
||
_release_receiver(self.socket, self)
|
||
self._closed.set()
|
||
|
||
async def wait_closed(self):
|
||
await self._closed.wait()
|
||
|
||
async def serve_forever(self):
|
||
await self.start()
|
||
await self.wait_closed()
|
||
|
||
async def __aenter__(self):
|
||
return await self.start()
|
||
|
||
async def __aexit__(self, *_):
|
||
await self.close()
|
||
|
||
|
||
async def export_http(socket, upstream, **kwargs):
|
||
"""Create and start an :class:`HTTPExporter`."""
|
||
return await HTTPExporter(socket, upstream, **kwargs).start()
|
||
|
||
|
||
async def mount_http(socket, listen=("127.0.0.1", 0), **kwargs):
|
||
"""Create and start an :class:`HTTPLocalProxy`."""
|
||
return await HTTPLocalProxy(socket, listen, **kwargs).start()
|
||
|
||
|
||
ROOMS_FILE = Path.home() / ".p2p_ws" / "rooms.json"
|
||
|
||
|
||
def _load_rooms(path=ROOMS_FILE):
|
||
path = Path(path).expanduser()
|
||
try:
|
||
data = json.loads(path.read_text(encoding="utf-8"))
|
||
return data if isinstance(data, dict) else {}
|
||
except FileNotFoundError:
|
||
return {}
|
||
except (OSError, ValueError) as exc:
|
||
raise RuntimeError(f"cannot read saved rooms: {path}: {exc}") from exc
|
||
|
||
|
||
def save_room(signal_url: str, room: str, *, name: str = "default", path=ROOMS_FILE):
|
||
"""原子保存配对 UUID；访问密钥不会写入磁盘。"""
|
||
if not name or not isinstance(name, str):
|
||
raise ValueError("name must be a non-empty string")
|
||
path = Path(path).expanduser()
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
rooms = _load_rooms(path)
|
||
rooms[name] = {
|
||
"signal_url": secure_signal_url(signal_url),
|
||
"room": room,
|
||
"saved_at": int(time.time()),
|
||
}
|
||
temporary = path.with_name(path.name + ".tmp")
|
||
temporary.write_text(
|
||
json.dumps(rooms, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
|
||
)
|
||
if os.name != "nt":
|
||
temporary.chmod(0o600)
|
||
os.replace(temporary, path)
|
||
return room
|
||
|
||
|
||
def load_room(name: str = "default", *, path=ROOMS_FILE):
|
||
"""读取命名配对，返回包含 signal_url/room/saved_at 的字典。"""
|
||
try:
|
||
saved = _load_rooms(path)[name].copy()
|
||
if "signal_url" in saved:
|
||
saved["signal_url"] = secure_signal_url(saved["signal_url"])
|
||
return saved
|
||
except KeyError as exc:
|
||
raise KeyError(f"saved room not found: {name}") from exc
|
||
|
||
|
||
async def connect_saved(
|
||
name: str = "default", *, access_key: str = DEFAULT_ACCESS_KEY,
|
||
path=ROOMS_FILE, **kwargs
|
||
):
|
||
"""使用本机保存的 UUID 长期重连；另一端应使用相同的命名配对。"""
|
||
saved = load_room(name, path=path)
|
||
return await P2PSocket.connect(
|
||
saved["signal_url"], saved["room"], access_key=access_key, **kwargs
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class Invite:
|
||
"""短码创建结果；显示 ``code``，本端调用 ``connect()`` 等待对方。"""
|
||
|
||
code: str
|
||
room: str
|
||
expires_at: int
|
||
signal_url: str
|
||
access_key: str
|
||
name: str = "default"
|
||
rooms_file: object = ROOMS_FILE
|
||
|
||
async def connect(self, **kwargs):
|
||
return await P2PSocket.connect(
|
||
self.signal_url, self.room, access_key=self.access_key, **kwargs
|
||
)
|
||
|
||
async def create_code(
|
||
signal_url: str, *, access_key: str = DEFAULT_ACCESS_KEY,
|
||
name: str = "default", rooms_file=ROOMS_FILE,
|
||
) -> Invite:
|
||
"""创建两分钟短码并保存 UUID，短码过期后仍可用 connect_saved 重连。"""
|
||
url = signed_url(code_url(signal_url, "code/new"), "code:new", "", access_key)
|
||
data = await _post_json(url)
|
||
save_room(signal_url, data["room"], name=name, path=rooms_file)
|
||
return Invite(
|
||
code=data["code"], room=data["room"], expires_at=data["expires_at"],
|
||
signal_url=signal_url, access_key=access_key, name=name, rooms_file=rooms_file,
|
||
)
|
||
|
||
async def connect_code(
|
||
signal_url: str, code: str, *, access_key: str = DEFAULT_ACCESS_KEY,
|
||
name: str = "default", rooms_file=ROOMS_FILE, **kwargs
|
||
):
|
||
"""兑换一次性短码、保存所得 UUID 并立即连接创建者。"""
|
||
if not _CODE.fullmatch(str(code)):
|
||
raise ValueError("code must be exactly 9 digits")
|
||
code = str(code)
|
||
url = signed_url(code_url(signal_url, "code/use") + f"?code={code}",
|
||
"code:use", code, access_key)
|
||
data = await _post_json(url)
|
||
save_room(signal_url, data["room"], name=name, path=rooms_file)
|
||
return await P2PSocket.connect(
|
||
signal_url, data["room"], access_key=access_key, **kwargs
|
||
)
|
||
|
||
|
||
# ---- 多人组：同一 UUID 即全部权限；名单走大厅，数据走两两派生的普通 1v1 房间 ----
|
||
|
||
async def retry_forever(attempt, on_error=None, *, stable=30, window=60, budget=8):
|
||
"""永久重试 attempt()。节奏三层:
|
||
1) 本轮活过 stable 秒(正常会话被掉线/被顶号) -> 0.2~1s 秒级重连;
|
||
2) 秒断/连不上 -> 抖动指数退避 1,2,4..60s, 防"配对即断"空转;
|
||
3) 兜底: 任意 window 秒内最多 budget 次, 远低于信令服务器每 IP 限速。
|
||
抖动也避免服务端重启时大量客户端同一秒齐涌。"""
|
||
loop = asyncio.get_running_loop()
|
||
delay, recent = 0, deque()
|
||
while True:
|
||
now = loop.time()
|
||
while recent and now - recent[0] >= window:
|
||
recent.popleft()
|
||
if len(recent) >= budget:
|
||
await asyncio.sleep(recent[0] + window - now)
|
||
continue
|
||
recent.append(now)
|
||
try:
|
||
await attempt()
|
||
except asyncio.CancelledError:
|
||
raise
|
||
except Exception as exc:
|
||
if on_error:
|
||
on_error(exc)
|
||
delay = 0 if loop.time() - now >= stable else min(max(delay * 2, 1), 60)
|
||
await asyncio.sleep(random.uniform(0.2, 1) if delay == 0 else delay * random.uniform(0.5, 1.5))
|
||
|
||
|
||
def pair_room(room: str, caller: str, callee: str) -> str:
|
||
"""主叫到被叫的专属房间：双方独立算出同一值；有方向，A 呼 B 与 B 呼 A 互不干扰。"""
|
||
return hashlib.sha256(f"{room}\n{caller}\n{callee}".encode()).hexdigest()
|
||
|
||
|
||
class Group:
|
||
"""以成员名 me 加入 UUID room 的多人组。
|
||
|
||
大厅连接收在线名单(变化时推送，无轮询)并转发呼叫。两两连接按需建立：
|
||
open(peer) 经大厅呼叫对方，双方进入 pair_room(主叫, 被叫) 的普通 P2PSocket，
|
||
主叫交给 use(socket, peer)，被叫交给 serve(socket, peer)。主叫断线自动重呼；
|
||
close(peer) 停止，对方离组时暂停、回来后自动恢复。不 open 就不占连接。
|
||
同名重入会顶掉旧连接，因此每个进程须用唯一成员名。
|
||
"""
|
||
|
||
def __init__(self, signal_url, room, me, serve, use, *, access_key=DEFAULT_ACCESS_KEY, auto=()):
|
||
self.signal_url, self.room, self.me = secure_signal_url(signal_url), room, me
|
||
self.serve, self.use, self.access_key = serve, use, access_key
|
||
self.auto = tuple(auto) # 名字以这些前缀开头的成员一上线就自动 open
|
||
self.members: set[str] = set()
|
||
self._ws = None
|
||
self._wanted: set[str] = set()
|
||
self._out: dict[str, asyncio.Task] = {} # 本端主叫
|
||
self._in: dict[str, asyncio.Task] = {} # 对方呼入
|
||
|
||
async def run(self):
|
||
try:
|
||
await retry_forever(self._lobby)
|
||
finally:
|
||
for task in [*self._out.values(), *self._in.values()]:
|
||
task.cancel()
|
||
|
||
def open(self, peer):
|
||
self._wanted.add(peer)
|
||
self._sync(self.members)
|
||
|
||
def close(self, peer):
|
||
self._wanted.discard(peer)
|
||
if task := self._out.pop(peer, None):
|
||
task.cancel()
|
||
|
||
async def _lobby(self):
|
||
room = hashlib.sha256(f"{self.room}:lobby".encode()).hexdigest()
|
||
sep = "&" if "?" in self.signal_url else "?"
|
||
url = f"{self.signal_url}{sep}{urlencode({'room': room, 'member': self.me})}"
|
||
seq = 0 # 服务端广播在锁外发送可能乱序，只认更大的 seq
|
||
async with websockets.connect(
|
||
signed_url(url, "ws", room, self.access_key), ping_interval=15, ping_timeout=10,
|
||
**_WS_CONNECT_KW,
|
||
) as ws:
|
||
try:
|
||
async for raw in ws:
|
||
msg = json.loads(raw)
|
||
if msg.get("type") == "members" and msg["seq"] > seq:
|
||
seq = msg["seq"]
|
||
self._ws = ws # 收到名单才算在线：members 与 _ws 同时有效
|
||
self._sync(set(msg["members"]) - {self.me})
|
||
elif msg.get("type") == "call":
|
||
self._answer(msg["from"])
|
||
finally:
|
||
self._ws = None
|
||
|
||
def _sync(self, members):
|
||
self.members = members
|
||
for tasks in (self._out, self._in):
|
||
for peer in set(tasks) - members:
|
||
tasks.pop(peer).cancel()
|
||
wanted = (self._wanted & members) | {p for p in members if p.startswith(self.auto)}
|
||
for peer in wanted - set(self._out):
|
||
self._out[peer] = asyncio.create_task(retry_forever(lambda peer=peer: self._call(peer)))
|
||
|
||
async def _call(self, peer):
|
||
if self._ws is None:
|
||
raise ConnectionError("lobby offline")
|
||
await self._ws.send(json.dumps({"type": "call", "to": peer}))
|
||
await self._link(pair_room(self.room, self.me, peer), self.use, peer)
|
||
|
||
def _answer(self, peer):
|
||
"""每次呼叫只接一次；对方重呼时旧连接必已失效，直接替换。"""
|
||
if old := self._in.pop(peer, None):
|
||
old.cancel()
|
||
|
||
async def once():
|
||
with suppress(Exception):
|
||
await self._link(pair_room(self.room, peer, self.me), self.serve, peer)
|
||
if self._in.get(peer) is task:
|
||
del self._in[peer]
|
||
task = self._in[peer] = asyncio.create_task(once())
|
||
|
||
async def _link(self, room, handler, peer):
|
||
socket = await P2PSocket.connect(
|
||
self.signal_url, room, access_key=self.access_key,
|
||
direct_timeout=15, peer_timeout=15,
|
||
psk=hashlib.sha256(f"{self.room}:psk".encode()).digest(), # 组 UUID 绑进密钥, 防信令中间人
|
||
)
|
||
try:
|
||
await handler(socket, peer)
|
||
finally:
|
||
await socket.close()
|
||
|
||
|
||
class HubUnion:
|
||
"""customer 侧把组内所有 hub 合成一个 hub API；hub 本身不改。
|
||
|
||
作 Group(auto=("hub:",)) 的 use 回调，每台 hub 一条 socket。会话名对外为
|
||
name@hub(不同 PC 的 GA 常同名)；/api/peers 逐台带各自 psig 增量拉取再合并；
|
||
GET /api/hubs 列在线 hub；POST /api/sessions 按 body 的 "hub" 选机，缺省取名字最小者。
|
||
"""
|
||
|
||
def __init__(self, timeout=22):
|
||
self.hubs, self.rows, self.timeout = {}, {}, timeout # rows: hub -> (psig, rows)
|
||
self.lost = {} # 刚断线的 hub -> 时刻: 宽限期内保留它的会话行, 重连时界面不闪
|
||
self._pending, self._seq = {}, 0
|
||
|
||
async def use(self, socket, peer):
|
||
hub = peer.split(":", 1)[-1]
|
||
_claim_receiver(socket, self)
|
||
self.hubs[hub] = socket
|
||
self.lost.pop(hub, None)
|
||
try:
|
||
while True:
|
||
with suppress(TypeError, ValueError, AttributeError):
|
||
message = json.loads(await socket.recv())
|
||
future = self._pending.get(message.get("id"))
|
||
if future is not None and not future.done():
|
||
future.set_result(message)
|
||
finally:
|
||
if self.hubs.get(hub) is socket:
|
||
del self.hubs[hub]
|
||
self.lost[hub] = time.monotonic()
|
||
_release_receiver(socket, self)
|
||
|
||
def live(self, grace=20):
|
||
"""在线 hub + 宽限期内断线的 hub; 超期才清掉会话行。"""
|
||
for hub, at in list(self.lost.items()):
|
||
if time.monotonic() - at > grace:
|
||
del self.lost[hub]
|
||
self.rows.pop(hub, None)
|
||
return set(self.hubs) | set(self.lost)
|
||
|
||
async def _one(self, hub, method, target, body=b""):
|
||
self._seq += 1
|
||
rid = f"u{self._seq}"
|
||
future = self._pending[rid] = asyncio.get_running_loop().create_future()
|
||
try:
|
||
await self.hubs[hub].send(json.dumps({
|
||
"v": 1, "type": _HTTP_REQ, "id": rid, "method": method, "target": target,
|
||
"headers": [["Content-Type", "application/json"]] if body else [],
|
||
"body": base64.b64encode(body).decode("ascii")}, separators=(",", ":")))
|
||
res = await asyncio.wait_for(future, self.timeout)
|
||
finally:
|
||
self._pending.pop(rid, None)
|
||
return int(res.get("status") or 502), base64.b64decode(res.get("body") or "")
|
||
|
||
async def request(self, method, target, body=b""):
|
||
"""返回 (status, body bytes)。"""
|
||
parts = urlsplit(target)
|
||
path = parts.path.split("/")
|
||
if parts.path == "/api/hubs":
|
||
return 200, json.dumps(sorted(self.hubs)).encode()
|
||
if parts.path == "/api/peers":
|
||
return await self._peers(dict(parse_qsl(parts.query, keep_blank_values=True)).get("psig"))
|
||
hub = None
|
||
if len(path) > 3 and path[1] == "api" and "@" in unquote(path[2]): # webui 的 encodeURIComponent 把 @ 编成 %40
|
||
name, _, hub = unquote(path[2]).rpartition("@")
|
||
path[2] = quote(name, safe="")
|
||
target = urlunsplit(("", "", "/".join(path), parts.query, ""))
|
||
elif body:
|
||
with suppress(ValueError, AttributeError):
|
||
hub = json.loads(body).get("hub")
|
||
hub = hub if hub in self.hubs or hub in self.lost else min(self.hubs, default=None)
|
||
if hub not in self.hubs: # 指名的 hub 正在重连: 回 503, 不改投别台
|
||
return 503, b'{"error":"no hub online","code":503}'
|
||
status, data = await self._one(hub, method, target, body)
|
||
if parts.path == "/api/sessions" and status == 200:
|
||
with suppress(ValueError, TypeError):
|
||
data = json.loads(data)
|
||
data = json.dumps({**data, "name": f"{data['name']}@{hub}", "hub": hub}).encode()
|
||
return status, data
|
||
|
||
async def _peers(self, psig):
|
||
async def pull(hub):
|
||
target = "/api/peers?" + urlencode({"psig": self.rows.get(hub, ("",))[0]})
|
||
status, data = await self._one(hub, "GET", target)
|
||
reply = json.loads(data) if status == 200 else {}
|
||
if "rows" in reply and hub in self.hubs:
|
||
self.rows[hub] = (reply["psig"], [
|
||
{**row, "name": f"{row['name']}@{hub}", "hub": hub} for row in reply["rows"]])
|
||
|
||
hubs = sorted(self.live())
|
||
await asyncio.gather(*map(pull, sorted(self.hubs)), return_exceptions=True) # 超时/重连中沿用上次的行
|
||
rows = [row for hub in hubs for row in self.rows.get(hub, ("", []))[1]]
|
||
if psig is None:
|
||
return 200, json.dumps(rows).encode()
|
||
sig = hashlib.md5(json.dumps(rows, sort_keys=True).encode()).hexdigest()[:12]
|
||
return 200, json.dumps({"same": 1} if psig == sig else {"rows": rows, "psig": sig}).encode()
|
||
|
||
|
||
connect = P2PSocket.connect
|
||
|
||
__all__ = [
|
||
"DEFAULT_ACCESS_KEY", "ROOMS_FILE", "Invite", "P2PSocket", "Signal",
|
||
"HTTPExporter", "HTTPLocalProxy", "export_http", "mount_http",
|
||
"connect", "connect_code", "connect_saved", "create_code", "load_room",
|
||
"save_room", "signed_url", "Group", "HubUnion", "pair_room", "retry_forever",
|
||
]
|
||
|