procproxy 0.1.0__py3-none-any.whl

proxycore/tunnel/relay.py

@@ -0,0 +1,393 @@
+"""双向转发（事件驱动，不用 `time.sleep` 轮询）。
+
+实现：基于跨平台 `selectors` 的 reactor（Windows 上用 select）。
+
+语义要求：
+* 半关闭：一端 FIN 后把另一向 `shutdown(SHUT_WR)`，在对端也 FIN 或超时前保持另一向；
+* 背压：只在 `send()` 能全部写走时继续读，否则挂 写事件 等可写（不忙等）；
+* 计数：`bytes_up`（客户端→目标）、`bytes_down`（目标→客户端）；
+* 空闲与半关闭超时。
+"""
+
+from __future__ import annotations
+
+import contextlib
+import selectors
+import socket
+import time
+from dataclasses import dataclass, field
+
+CHUNK = 65536
+# 单连接单方向允许堆积的缓冲上限（背压阈值）
+MAX_BUFFER = 4 * 1024 * 1024
+
+REASON_OK = "ok"
+REASON_IDLE_TIMEOUT = "idle-timeout"
+REASON_HALF_CLOSE_TIMEOUT = "half-close-timeout"
+REASON_CLIENT_RESET = "client-reset"
+REASON_UPSTREAM_RESET = "upstream-reset"
+
+
+@dataclass
+class _Direction:
+    """一个方向的转发状态。"""
+
+    src: socket.socket
+    dst: socket.socket
+    is_up: bool
+    buffer: bytearray = field(default_factory=bytearray)
+    src_closed: bool = False
+    dst_shutdown: bool = False
+    done: bool = False
+    rst: bool = False
+    bytes_moved: int = 0
+
+    def flush(self) -> bool:
+        """尽力写走缓冲；返回 True 表示缓冲已空。"""
+        while self.buffer:
+            try:
+                sent = self.dst.send(self.buffer)
+            except BlockingIOError:
+                return False
+            except OSError:
+                raise
+            if sent <= 0:  # pragma: no cover - 对端已关闭
+                return False
+            del self.buffer[:sent]
+            self.bytes_moved += sent
+        return True
+
+
+@dataclass
+class _Watch:
+    """selector 的注册单位：按 socket 注册，不能按方向注册。
+
+    socketpair 是双工的：一个 socket 既是"某条方向的读端"，又是"另一条方向的写端"。
+    旧实现把 `EVENT_WRITE` 挂在 `direction.src` 上，而数据是往 `direction.dst` 写的，
+    于是"等对端可写"实际等的是读端，两个 socket 都几乎永远可写，`select()`
+    立刻返回，整个循环退化成 100% CPU 忙等，真正的背压从未被等待。
+    实测后果：1MB 回程在有背压时只送出去 128KB 就收工（见 `test_relay_1mb_*`）。
+    """
+
+    sock: socket.socket
+    read_from: _Direction  # 这个 socket 是它的 src
+    write_to: _Direction  # 这个 socket 是它的 dst
+
+
+def _finish(direction: _Direction) -> bool:
+    """缓冲区已空 → 补发 FIN 并宣告这一向结束；返回"是否已结束"。
+
+    FIN 必须跟在数据后面。 旧实现在读到 EOF 时先 `shutdown(SHUT_WR)` 再 flush，
+    结果 FIN 超车插到还没发出去的数据前面：上游立刻看到 EOF 提前收工，剩下的数据
+    在 Windows 上又会因"关闭后再 send"报 `10058` 直接丢掉，响应被截断。
+    """
+    if direction.buffer:
+        return False  # 还没写完，FIN 得等
+    if not direction.dst_shutdown:
+        _shutdown_write(direction.dst)
+        direction.dst_shutdown = True
+    return True
+
+
+def relay_bidirectional(
+    client: socket.socket,
+    upstream: socket.socket,
+    *,
+    idle_timeout: float = 300.0,
+    half_close_grace: float = 30.0,
+    response_grace: float = 2.0,
+    drain_grace: float = 2.0,
+    chunk: int = CHUNK,
+) -> tuple[int, int, str]:
+    """把 `client` 与 `upstream` 双向对接，返回 `(bytes_up, bytes_down, reason)`。
+
+    `response_grace`：客户端干净 FIN 之后，再给"响应"这么多时间。
+    这样既不会像"等双向都 FIN"那样在长连接上卡到 idle 超时，
+    也不会像"立刻退出"那样把还没到达的响应掐掉。
+
+    `drain_grace`：对端还有未被读走的残留字节时，最多再等这么久。
+    防止"发完就 FIN、不读响应"的客户端把连接挂到 `half_close_grace`。
+    """
+    directions = [
+        _Direction(src=client, dst=upstream, is_up=True),
+        _Direction(src=upstream, dst=client, is_up=False),
+    ]
+    for sock in (client, upstream):
+        with contextlib.suppress(OSError):
+            sock.setblocking(False)
+
+    up, down = directions[0], directions[1]
+    watch = [
+        _Watch(sock=client, read_from=up, write_to=down),
+        _Watch(sock=upstream, read_from=down, write_to=up),
+    ]
+
+    sel = selectors.DefaultSelector()
+    for entry in watch:
+        sel.register(entry.sock, _mask_for(entry), entry)
+
+    last_activity = time.monotonic()
+    half_closed_at: float | None = None
+    request_closed_at: float | None = None  # 客户端干净 FIN 的时刻
+    response_seen = False
+    upstream_quiet_since: float | None = None  # 上游恢复安静的时刻（响应已发完）
+    unread_since: float | None = None  # 首次发现"对端还有未读字节"的时刻
+    reason = REASON_OK
+
+    def request_still_pending() -> bool:
+        """请求方向是否还有"没送出去"的东西（缓冲没写空 / 源还没读到 EOF）。
+
+        为什么必须有这个判断：
+          `request_closed_at` 只记录"收到客户端 FIN"的时刻，而 FIN 不等于"数据已送达"
+          并发/高负载下，1MB 的尾部可能还堆在我们自己的发送缓冲里没 flush 出去。
+          如果这时候按 `response_grace` 直接收工，尾部就丢了，
+          表现是偶发的"数据短一截 / sha256 对不上"（这种 bug 最难查）。
+          所以：只要请求方向还有待发数据，任何"温和收工"都不许生效。
+        """
+        return bool(up.buffer) or not up.src_closed
+
+    def response_still_pending() -> bool:
+        """回程方向是否还有"没送到客户端"的东西。
+
+        与 `request_still_pending()` 同源：任何"温和收工"都必须等两个方向的缓冲
+        都写空。只检查请求方向的话，上游一 FIN 就会带着满缓冲的回程收工，同样是截断。
+        """
+        return bool(down.buffer)
+
+    try:
+        while True:
+            now = time.monotonic()
+            if now - last_activity > idle_timeout:
+                reason = REASON_IDLE_TIMEOUT
+                break
+            pending = request_still_pending() or response_still_pending()
+            if (
+                request_closed_at is not None
+                and now - request_closed_at > response_grace
+                and not response_seen
+                and not pending
+            ):
+                # 客户端已发完、请求也确实全部送达、且一直没响应 → 正常收工（不等 idle 超时）
+                break
+            if upstream_quiet_since is not None and now - upstream_quiet_since > response_grace and not pending:
+                # 响应已经发完并转发给客户端，上游也安静了 → 正常收工
+                break
+            if unread_since is not None and now - unread_since > drain_grace and not pending:
+                # 对端一直不读走残留字节（例如"发完就 FIN、不读响应"的客户端）
+                # → 不能无限等，按 drain_grace 收工，避免把连接挂到 half_close_grace。
+                break
+            if half_closed_at is not None and now - half_closed_at > half_close_grace:
+                reason = REASON_HALF_CLOSE_TIMEOUT
+                break
+
+            # 先把积压数据推出去（非阻塞，能写多少写多少）
+            # 注意变量名别叫 entry，那是 `_Watch`（上面注册用的），mypy 会串型。
+            for direction in directions:
+                if not direction.buffer:
+                    continue
+                try:
+                    direction.flush()
+                except OSError:
+                    reason = REASON_UPSTREAM_RESET if direction.is_up else REASON_CLIENT_RESET
+                    return _result(directions, reason)
+                if direction.src_closed and _finish(direction):
+                    direction.done = True
+                    if all(d.done for d in directions):
+                        return _result(directions, REASON_OK)
+
+            # 等"要写的那个 socket"可写，不是等读端（见 `_Watch`）
+            _sync_masks(sel, watch)
+
+            events = sel.select(timeout=0.25)
+            if not events:
+                continue
+
+            alive = True
+            for key, event_mask in events:
+                entry = key.data
+                assert isinstance(entry, _Watch)
+                if event_mask & selectors.EVENT_WRITE:
+                    direction = entry.write_to
+                    try:
+                        drained = direction.flush()
+                    except OSError:
+                        reason = REASON_UPSTREAM_RESET if direction.is_up else REASON_CLIENT_RESET
+                        return _result(directions, reason)
+                    if drained:
+                        if (
+                            not direction.is_up
+                            and request_closed_at is not None
+                            and response_seen
+                        ):
+                            # 回程缓冲写空 → 开始计时"上游安静"
+                            upstream_quiet_since = time.monotonic()
+                        if direction.src_closed and _finish(direction):
+                            direction.done = True
+                    if all(d.done for d in directions):
+                        return _result(directions, REASON_OK)
+                direction = entry.read_from
+                if event_mask & selectors.EVENT_READ and not direction.src_closed:
+                    try:
+                        data = direction.src.recv(chunk)
+                    except BlockingIOError:
+                        continue
+                    except OSError:
+                        # 连接被重置（RST）而不是干净 FIN
+                        direction.src_closed = True
+                        direction.rst = True
+                        data = b""
+                    if data:
+                        last_activity = time.monotonic()
+                        if not direction.is_up:
+                            response_seen = True  # 已经看到响应
+                            upstream_quiet_since = None
+                        if _append(direction, data):
+                            try:
+                                flushed = direction.flush()
+                            except OSError:
+                                reason = REASON_UPSTREAM_RESET if direction.is_up else REASON_CLIENT_RESET
+                                return _result(directions, reason)
+                            if flushed and not direction.is_up and request_closed_at is not None:
+                                # 响应已全部写回客户端 → 开始计时"上游安静"
+                                upstream_quiet_since = time.monotonic()
+                    else:
+                        # 读到 EOF：把 FIN 传给对端，然后按"这一向已结束"处理
+                        direction.src_closed = True
+                        # 顺序不能反：先 flush 再 FIN。缓冲区没写空时 `_finish()`
+                        # 会拒绝发 FIN，等写事件把缓冲推空后再补，否则 FIN 超车，
+                        # 上游提前收工，剩下的数据在 Windows 上还会因"关闭后再 send"
+                        # 报 10058 丢干净（实测 1MB 只回来 128KB）。
+                        try:
+                            direction.flush()
+                        except OSError:
+                            reason = REASON_UPSTREAM_RESET if direction.is_up else REASON_CLIENT_RESET
+                            return _result(directions, reason)
+                        direction.done = _finish(direction)
+
+                        if direction.rst:
+                            # RST：没有 FIN 语义，等另一半自然结束或超时
+                            if half_closed_at is None:
+                                half_closed_at = time.monotonic()
+                            continue
+
+                        if all(d.done for d in directions):
+                            return _result(directions, REASON_OK)
+
+                        if direction.is_up:
+                            # 客户端干净 FIN 且请求已全部转发：
+                            # 给响应留 `response_grace`，而不是傻等对端也 FIN。
+                            if _peer_has_unread(direction.dst) or request_still_pending():
+                                # 关键：对端（客户端）写端虽然收到 FIN，但内核里可能还有
+                                # 未读走的字节（并发 sendall 的最后一段）；
+                                # 或者是我们自己还有没 flush 出去的待发数据。
+                                # 这两种情况下收工都会丢尾部（表现为 1MB sha256 不一致）。
+                                # 但也不能无限等（有客户端"发完就 FIN、根本不读"），
+                                # 所以记一个起始时刻，由 drain_grace 兜底。
+                                if unread_since is None:
+                                    unread_since = time.monotonic()
+                                if half_closed_at is None:
+                                    half_closed_at = time.monotonic()
+                                continue
+                            unread_since = None
+                            if request_closed_at is None:
+                                request_closed_at = time.monotonic()
+                        else:
+                            # 上游 FIN：响应到此为止；
+                            # 只有"客户端此前已 FIN + 没有残留未读 + 请求方向确实发完了"才能立刻收工
+                            if (
+                                request_closed_at is not None
+                                and not _peer_has_unread(direction.dst)
+                                and not request_still_pending()
+                                and not response_still_pending()  # 回程缓冲没写空就不许收工
+                            ):
+                                return _result(directions, REASON_OK)
+                            if half_closed_at is None:
+                                half_closed_at = time.monotonic()
+            if not alive:  # pragma: no cover - 防御性分支
+                break
+    finally:
+        sel.close()
+
+    return _result(directions, reason)
+
+
+def _append(direction: _Direction, data: bytes) -> bool:
+    """写入待发缓冲；返回 True 表示需要尝试 flush。"""
+    if len(direction.buffer) + len(data) > MAX_BUFFER:
+        # 背压：超过上限时只保留到上限（多余丢弃会破坏 TCP 语义），
+        # 因此这里改为阻塞式写走，保证不丢数据（fd 已是非阻塞，用循环 + 短暂等待）。
+        direction.buffer.extend(data)
+        _drain_blocking(direction)
+        return False
+    direction.buffer.extend(data)
+    return True
+
+
+def _drain_blocking(direction: _Direction) -> None:
+    """缓冲超限时的兜底：带超时地写走（避免无界内存增长）。"""
+    deadline = time.monotonic() + 30.0
+    sel = selectors.DefaultSelector()
+    try:
+        sel.register(direction.dst, selectors.EVENT_WRITE)
+        while direction.buffer and time.monotonic() < deadline:
+            if direction.flush():
+                return
+            sel.select(timeout=0.5)
+    finally:
+        sel.close()
+
+
+def _mask_for(entry: _Watch) -> int:
+    """这个 socket 现在要等什么：读它（某条方向的 src）还是写它（另一条方向的 dst）。"""
+    mask = 0
+    if not entry.read_from.src_closed:
+        mask |= selectors.EVENT_READ
+    if entry.write_to.buffer:  # 有积压才等可写，否则就是忙等
+        mask |= selectors.EVENT_WRITE
+    return mask
+
+
+def _sync_masks(sel: selectors.BaseSelector, watch: list[_Watch]) -> None:
+    for entry in watch:
+        mask = _mask_for(entry)
+        try:
+            if mask:
+                sel.modify(entry.sock, mask, entry)
+            else:
+                sel.unregister(entry.sock)
+        except (KeyError, ValueError):  # pragma: no cover - 已注销
+            if mask:
+                with contextlib.suppress(KeyError, ValueError):
+                    sel.register(entry.sock, mask, entry)
+
+
+def _peer_has_unread(sock: socket.socket) -> bool:
+    """对端 socket 里是否还有未被我们读走的字节。
+
+    用途：决定"收到 FIN 后能不能立刻收工"。
+    只要内核缓冲里还有数据，收工就会把它丢掉（并发写入场景下表现为偶发截断）。
+    非阻塞 peek 失败时返回 False（保守放行，由超时兜底）。
+    """
+    try:
+        peek = getattr(sock, "recv", None)
+        if peek is None:  # pragma: no cover - 理论不可达
+            return False
+        data = sock.recv(1, socket.MSG_PEEK)
+        return bool(data)
+    except BlockingIOError:
+        return False
+    except OSError:
+        return False
+
+
+def _result(directions: list[_Direction], reason: str) -> tuple[int, int, str]:
+    up = next((d.bytes_moved for d in directions if d.is_up), 0)
+    down = next((d.bytes_moved for d in directions if not d.is_up), 0)
+    return up, down, reason
+
+
+def _shutdown_write(sock: socket.socket) -> None:
+    # 对端已 RST 时 shutdown 会报 ENOTCONN/WSAENOTCONN，这类错误没有补救动作；
+    # 统一忽略（之前的分支两个方向都是 pass，等价于吞掉所有 OSError）。
+    with contextlib.suppress(OSError):
+        sock.shutdown(socket.SHUT_WR)