socks5代理转发

[ColaSk]

Currently, we need to implement a function where other services request HTTP proxy services, and the proxy service performs upstream proxy selection (HTTP or SOCKS5). Our current tests have found that SOCKS5 type proxies are not supported. Is there currently any implementation method or development plan in the project

[ColaSk]

-- coding: utf-8 --

"""
SOCKS5 上游连接封装。

当代理池返回 socks5:// 端点时，通过 PySocks 建立到目标地址的 SOCKS5 隧道，
并封装为与 TcpServerConnection 兼容的接口，供 TcpUpstreamConnectionHandler 使用。
"""
import logging
import socket
import socks
from typing import List, Optional, Union

from proxy.common.types import HostPort, TcpOrTlsSocket
from proxy.common.constants import DEFAULT_BUFFER_SIZE, DEFAULT_MAX_SEND_SIZE

logger = logging.getLogger(name)

class Socks5UpstreamConnection:
"""对已建立的 SOCKS5 隧道 socket 的封装，接口与 TcpServerConnection 兼容。

用于经 SOCKS5 代理连到目标 (target_host, target_port) 后的上行连接，
供 proxy.core.base.TcpUpstreamConnectionHandler 的 self.upstream 使用。
"""

def __init__(self, host: str, port: int,
             username: str = None, password: str = None) -> None:
    """
    Args:
        conn: 已通过 SOCKS5 连接到目标地址的 socket。
        addr: 目标地址 (host, port)，用于日志等。
    """
    self._conn: TcpOrTlsSocket = None
    self.addr: HostPort = (host, port, username, password)
    self.closed: bool = False
    self.buffer: List[memoryview] = []
    self._num_buffer: int = 0

@property
def connection(self) -> TcpOrTlsSocket:
    """返回底层 socket，与 TcpServerConnection.connection 一致。"""
    return self._conn

def connect(
        self,
        addr: Optional[HostPort] = None,
        source_address: Optional[HostPort] = None,
) -> None:
    assert self._conn is None

    sock = socks.socksocket(socket.AF_INET, socket.SOCK_STREAM)
    sock.set_proxy(
        proxy_type=socks.SOCKS5,
        addr=self.addr[0],
        port=self.addr[1],
        username=self.addr[2] if len(self.addr) > 2 else None,  
        password=self.addr[3] if len(self.addr) > 3 else None,
    )
    sock.connect(source_address)
    # 与 proxy.py 事件循环一致：上游 socket 需为非阻塞，否则 recv/flush 会阻塞
    sock.setblocking(False)

    self._conn = sock
    self.closed = False

def send(self, data: Union[memoryview, bytes]) -> int:
    """发送数据。"""
    return self._conn.send(data)

def recv(
    self,
    buffer_size: int = DEFAULT_BUFFER_SIZE,
) -> Optional[memoryview]:
    """接收数据。返回 None 表示对端关闭。非阻塞下无数据时返回空 memoryview。"""
    try:
        data: bytes = self._conn.recv(buffer_size)
    except BlockingIOError:
        # 非阻塞 socket 无数据时，与 proxy.py 事件循环兼容
        return memoryview(b'')
    if len(data) == 0:
        return None
    logger.debug('received %d bytes from server', len(data))
    return memoryview(data)

def close(self) -> bool:
    """关闭连接。"""
    if not self.closed:
        self._conn.close()
        self.closed = True
    return self.closed

def has_buffer(self) -> bool:
    """是否还有待发送缓冲。"""
    return self._num_buffer != 0

def queue(self, mv: memoryview) -> None:
    """将数据加入发送缓冲。"""
    self.buffer.append(mv)
    self._num_buffer += 1

def flush(self, max_send_size: Optional[int] = None) -> int:
    """将缓冲数据写入 socket。"""
    if not self.has_buffer():
        return 0
    mv = self.buffer[0]
    max_send_size = max_send_size or DEFAULT_MAX_SEND_SIZE
    try:
        sent: int = self.send(mv[:max_send_size])
    except BlockingIOError:
        logger.warning('BlockingIOError when trying send to server')
        return 0
    if sent == len(mv):
        self.buffer.pop(0)
        self._num_buffer -= 1
    else:
        self.buffer[0] = mv[sent:]
    del mv
    logger.debug('flushed %d bytes to server', sent)
    return sent