WsPool refactoring, TCP fallback backoff, removed direct connections outside of a pool

This commit is contained in:
Flowseal committed 2026-10-04 21:53:26 +03:00
1 parent 48954ff9ab
commit 50c980e9ea
13 files changed
+563 -294

No files matched your search

+1 -1
View File
@@ -56,7 +56,7 @@ tg-ws-proxy [--port PORT] [--host HOST] [--dc-ip DC:IP ...] [-v]
| `--fake-tls-domain` | | Enable Fake TLS masquerading (ee-secret) with specified SNI domain | | `--fake-tls-domain` | | Enable Fake TLS masquerading (ee-secret) with specified SNI domain |
| `--proxy-protocol` | disabled | Accept HAProxy PROXY protocol v1 (for use behind nginx/haproxy with `proxy_protocol on`) | | `--proxy-protocol` | disabled | Accept HAProxy PROXY protocol v1 (for use behind nginx/haproxy with `proxy_protocol on`) |
| `--buf-kb` | `256` | Buffer size in KB | | `--buf-kb` | `256` | Buffer size in KB |
| `--pool-size` | `4` | Number of pre-allocated connections per DC | | `--pool-size` | `4` | Number of ready WS connections per DC. `0` disables the direct DC→IP WS route |
| `--log-file` | disabled | Path to file for saving logs | | `--log-file` | disabled | Path to file for saving logs |
| `--log-max-mb` | `5` | Maximum log file size in MB (afterwards overwrites) | | `--log-max-mb` | `5` | Maximum log file size in MB (afterwards overwrites) |
| `--log-backups` | `0` | Number of log backups after overwrite | | `--log-backups` | `0` | Number of log backups after overwrite |
+1 -1
View File
@@ -56,7 +56,7 @@ tg-ws-proxy [--port PORT] [--host HOST] [--dc-ip DC:IP ...] [-v]
| `--fake-tls-domain` | | Включить маскировку Fake TLS (ee-secret) с указанным SNI-доменом | | `--fake-tls-domain` | | Включить маскировку Fake TLS (ee-secret) с указанным SNI-доменом |
| `--proxy-protocol` | выкл. | Принимать HAProxy PROXY protocol v1 (для работы за nginx/haproxy с `proxy_protocol on`) | | `--proxy-protocol` | выкл. | Принимать HAProxy PROXY protocol v1 (для работы за nginx/haproxy с `proxy_protocol on`) |
| `--buf-kb` | `256` | Размер буфера в КБ | | `--buf-kb` | `256` | Размер буфера в КБ |
| `--pool-size` | `4` | Количество заготовленных соединений на каждый DC | | `--pool-size` | `4` | Количество заготовленных WS-соединений на каждый DC. `0` отключает прямой WS-маршрут DC→IP |
| `--log-file` | выкл. | Путь к файлу, в который будут сохраняться логи | | `--log-file` | выкл. | Путь к файлу, в который будут сохраняться логи |
| `--log-max-mb` | `5` | Максимальный размер файла логов в МБ (после этого начинается перезапись) | | `--log-max-mb` | `5` | Максимальный размер файла логов в МБ (после этого начинается перезапись) |
| `--log-backups` | `0` | Количество сохранений логов после перезаписи | | `--log-backups` | `0` | Количество сохранений логов после перезаписи |
+50 -14
View File
@@ -3,7 +3,7 @@ import logging
import struct import struct
import time import time
from typing import List, Optional from typing import Dict, List, Optional, Set, Tuple
from urllib.parse import urlencode from urllib.parse import urlencode
from .utils import * from .utils import *
@@ -20,6 +20,18 @@ from .cf_h2 import bridge_h2
log = logging.getLogger('tg-mtproto-proxy') log = logging.getLogger('tg-mtproto-proxy')
_st_I_le = struct.Struct('<I') _st_I_le = struct.Struct('<I')
TCP_BACKOFF_INITIAL = 30.0
TCP_BACKOFF_MAX = 3600.0
_tcp_failures: Dict[Tuple[str, int], int] = {}
_tcp_retry_after: Dict[Tuple[str, int], float] = {}
_tcp_connecting: Set[Tuple[str, int]] = set()
def reset_tcp_backoff() -> None:
_tcp_failures.clear()
_tcp_retry_after.clear()
_tcp_connecting.clear()
ZERO_64 = b'\x00' * 64 ZERO_64 = b'\x00' * 64
@@ -171,8 +183,6 @@ async def do_fallback(reader, writer, relay_init, label,
if ok: if ok:
return True return True
elif method == 'tcp' and fallback_dst: elif method == 'tcp' and fallback_dst:
log.info("[%s] DC%d%s -> TCP fallback to %s:443",
label, dc, media_tag, fallback_dst)
ok = await _tcp_fallback( ok = await _tcp_fallback(
reader, writer, fallback_dst, 443, reader, writer, fallback_dst, 443,
relay_init, label, ctx) relay_init, label, ctx)
@@ -271,19 +281,45 @@ async def _cfproxy_fallback(reader, writer, relay_init, label,
async def _tcp_fallback(reader, writer, dst, port, relay_init, label, ctx: CryptoCtx): async def _tcp_fallback(reader, writer, dst, port, relay_init, label, ctx: CryptoCtx):
try: key = (dst, port)
rr, rw = await asyncio.wait_for( if key in _tcp_connecting or time.monotonic() < _tcp_retry_after.get(key, 0):
asyncio.open_connection(dst, port), timeout=10) log.debug("[%s] TCP fallback to %s:%d skipped (backoff or connecting)",
except Exception as exc: label, dst, port)
log.warning("[%s] TCP fallback to %s:%d failed: %s",
label, dst, port, repr(exc))
return False return False
stats.connections_tcp_fallback += 1 _tcp_connecting.add(key)
rw.write(relay_init) rw = None
await rw.drain() try:
await _bridge_tcp_reencrypt(reader, writer, rr, rw, label, ctx) try:
return True log.info("[%s] TCP fallback to %s:%d", label, dst, port)
rr, rw = await asyncio.wait_for(
asyncio.open_connection(dst, port), timeout=10)
rw.write(relay_init)
await rw.drain()
except Exception as exc:
failures = _tcp_failures.get(key, 0) + 1
_tcp_failures[key] = failures
delay = min(TCP_BACKOFF_INITIAL * 2 ** min(failures - 1, 7),
TCP_BACKOFF_MAX)
_tcp_retry_after[key] = time.monotonic() + delay
log.warning("[%s] TCP fallback to %s:%d failed: %r; retry in %.0fs",
label, dst, port, exc, delay)
return False
finally:
_tcp_connecting.discard(key)
_tcp_failures.pop(key, None)
_tcp_retry_after.pop(key, None)
stats.connections_tcp_fallback += 1
await _bridge_tcp_reencrypt(reader, writer, rr, rw, label, ctx)
return True
finally:
if rw is not None:
rw.close()
try:
await rw.wait_closed()
except (OSError, ConnectionError):
pass
async def bridge_ws_reencrypt(reader, writer, ws: RawWebSocket, label, async def bridge_ws_reencrypt(reader, writer, ws: RawWebSocket, label,
+99 -86
View File
@@ -10,12 +10,11 @@ from typing import Dict, List, Optional, Tuple, Set
from .raw_websocket import RawWebSocket, WsHandshakeError from .raw_websocket import RawWebSocket, WsHandshakeError
from .stats import stats from .stats import stats
from .config import proxy_config from .config import proxy_config
from .utils import ws_domains, DC_DEFAULT_IPS from .utils import ws_domains, DC_DEFAULT_IPS, WS_PATH, WS_PATH_TEST
log = logging.getLogger('tg-mtproto-proxy') log = logging.getLogger('tg-mtproto-proxy')
# TODO: domains handling is broken: wrong is_media flag causes tcp_reset after handshake,
# but initial connection is still established no matter what is is_media flag is set to
class _WsPool: class _WsPool:
WS_POOL_MAX_AGE = 120.0 WS_POOL_MAX_AGE = 120.0
WS_POOL_CHECK_INTERVAL = 5.0 WS_POOL_CHECK_INTERVAL = 5.0
@@ -23,17 +22,20 @@ class _WsPool:
REFILL_BACKOFF_MAX = 3600.0 REFILL_BACKOFF_MAX = 3600.0
def __init__(self): def __init__(self):
self._idle: Dict[Tuple[int, bool], deque] = {} self._idle: Dict[Tuple[int, bool, bool], deque] = {}
self._refilling: Set[Tuple[int, bool]] = set() self._refilling: Dict[Tuple[int, bool, bool], asyncio.Task] = {}
self._rotating: Dict[Tuple[int, bool], asyncio.Task] = {} self._rotating: Dict[Tuple[int, bool, bool], asyncio.Task] = {}
self._refill_failures: Dict[Tuple[int, bool], int] = {} self._refill_failures: Dict[Tuple[int, bool, bool], int] = {}
self._refill_after: Dict[Tuple[int, bool], float] = {} self._refill_after: Dict[Tuple[int, bool, bool], float] = {}
self.try_fronting_first = True self.try_fronting_first = False
async def get(self, dc: int, is_media: bool, async def get(self, dc: int, is_media: bool, *, is_test_dc: bool = False
target_ip: str, domains: List[str]
) -> Optional[RawWebSocket]: ) -> Optional[RawWebSocket]:
key = (dc, is_media) target_ip = proxy_config.dc_redirects.get(dc)
if not target_ip or proxy_config.pool_size <= 0:
return None
key = (dc, is_media, is_test_dc)
domains = ws_domains(dc, is_media)
now = time.monotonic() now = time.monotonic()
bucket = self._idle.get(key) bucket = self._idle.get(key)
@@ -43,14 +45,12 @@ class _WsPool:
while bucket: while bucket:
ws, created = bucket.popleft() ws, created = bucket.popleft()
age = now - created age = now - created
if (age > self.WS_POOL_MAX_AGE or ws._closed if self._is_stale(ws, created, now):
or ws.writer.transport.is_closing()):
asyncio.create_task(self._quiet_close(ws)) asyncio.create_task(self._quiet_close(ws))
continue continue
stats.pool_hits += 1 stats.pool_hits += 1
log.debug("WS pool hit DC%d%s (age=%.1fs, left=%d)", log.debug("WS pool hit DC%d%s%s (age=%.1fs, left=%d)",
dc, 'm' if is_media else '', age, len(bucket)) dc, 't' if is_test_dc else '', 'm' if is_media else '', age, len(bucket))
self.report_success(dc, is_media)
self._schedule_refill(key, target_ip, domains) self._schedule_refill(key, target_ip, domains)
return ws return ws
@@ -58,20 +58,25 @@ class _WsPool:
self._schedule_refill(key, target_ip, domains) self._schedule_refill(key, target_ip, domains)
return None return None
def _is_stale(self, ws, created, now):
return (now - created >= self.WS_POOL_MAX_AGE or ws._closed
or ws.writer.transport.is_closing() or ws.reader.at_eof()
or ws.reader.exception() is not None)
def _schedule_refill(self, key, target_ip, domains): def _schedule_refill(self, key, target_ip, domains):
if proxy_config.pool_size <= 0:
return
self._schedule_rotation(key, target_ip, domains)
if (key in self._refilling if (key in self._refilling
or time.monotonic() < self._refill_after.get(key, 0)): or time.monotonic() < self._refill_after.get(key, 0)):
return return
self._refilling.add(key) self._refilling[key] = asyncio.create_task(
asyncio.create_task(self._refill(key, target_ip, domains)) self._refill(key, target_ip, domains))
def report_success(self, dc: int, is_media: bool) -> None:
key = (dc, is_media)
self._refill_failures.pop(key, None)
self._refill_after.pop(key, None)
async def _refill(self, key, target_ip, domains): async def _refill(self, key, target_ip, domains):
dc, is_media = key dc, is_media, is_test_dc = key
tasks = []
adopted = set()
try: try:
bucket = self._idle.setdefault(key, deque()) bucket = self._idle.setdefault(key, deque())
needed = proxy_config.pool_size - len(bucket) needed = proxy_config.pool_size - len(bucket)
@@ -79,19 +84,23 @@ class _WsPool:
return return
connected = 0 connected = 0
tasks = [asyncio.create_task( tasks = [asyncio.create_task(
self._connect_one(target_ip, domains)) self._connect_one(target_ip, domains,
WS_PATH_TEST if is_test_dc else WS_PATH))
for _ in range(needed)] for _ in range(needed)]
for t in tasks: for t in asyncio.as_completed(tasks):
try: try:
ws = await t ws = await t
if ws: if ws:
if self._refilling.get(key) is not asyncio.current_task():
return
bucket.append((ws, time.monotonic())) bucket.append((ws, time.monotonic()))
adopted.add(ws)
connected += 1 connected += 1
self._schedule_rotation(key, target_ip, domains) except Exception as exc:
except Exception: log.debug("WS pool connect failed: %r", exc)
pass
if connected: if connected:
self.report_success(dc, is_media) self._refill_failures.pop(key, None)
self._refill_after.pop(key, None)
else: else:
failures = self._refill_failures.get(key, 0) + 1 failures = self._refill_failures.get(key, 0) + 1
self._refill_failures[key] = failures self._refill_failures[key] = failures
@@ -102,12 +111,20 @@ class _WsPool:
) )
self._refill_after[key] = time.monotonic() + delay self._refill_after[key] = time.monotonic() + delay
log.info( log.info(
"WS pool refill failed for DC%d%s, retry in %.0fs", "WS pool refill failed for DC%d%s%s, retry in %.0fs",
dc, 'm' if is_media else '', delay) dc, 't' if is_test_dc else '', 'm' if is_media else '', delay)
log.debug("WS pool refilled DC%d%s: %d ready", log.debug("WS pool refilled DC%d%s%s: %d ready",
dc, 'm' if is_media else '', len(bucket)) dc, 't' if is_test_dc else '', 'm' if is_media else '', len(bucket))
finally: finally:
self._refilling.discard(key) for task in tasks:
if not task.done():
task.cancel()
results = await asyncio.gather(*tasks, return_exceptions=True)
for ws in results:
if ws is not None and not isinstance(ws, BaseException) and ws not in adopted:
await self._quiet_close(ws)
if self._refilling.get(key) is asyncio.current_task():
self._refilling.pop(key, None)
def _schedule_rotation(self, key, target_ip, domains): def _schedule_rotation(self, key, target_ip, domains):
if key in self._rotating: if key in self._rotating:
@@ -116,28 +133,16 @@ class _WsPool:
self._rotate(key, target_ip, domains)) self._rotate(key, target_ip, domains))
async def _rotate(self, key, target_ip, domains): async def _rotate(self, key, target_ip, domains):
dc, is_media = key dc, is_media, is_test_dc = key
try: try:
while True: while proxy_config.pool_size > 0:
bucket = self._idle.get(key) bucket = self._idle.setdefault(key, deque())
if not bucket:
return
expires_at = min(
created + self.WS_POOL_MAX_AGE
for _, created in bucket)
await asyncio.sleep(min(
self.WS_POOL_CHECK_INTERVAL,
max(0, expires_at - time.monotonic())))
now = time.monotonic() now = time.monotonic()
expired = [] expired = []
ready = deque() ready = deque()
while bucket: while bucket:
ws, created = bucket.popleft() ws, created = bucket.popleft()
if (now - created >= self.WS_POOL_MAX_AGE if self._is_stale(ws, created, now):
or ws._closed
or ws.writer.transport.is_closing()):
expired.append(ws) expired.append(ws)
else: else:
ready.append((ws, created)) ready.append((ws, created))
@@ -147,49 +152,45 @@ class _WsPool:
for ws in expired: for ws in expired:
asyncio.create_task(self._quiet_close(ws)) asyncio.create_task(self._quiet_close(ws))
log.debug( log.debug(
"WS pool rotated DC%d%s: %d stale, %d ready", "WS pool rotated DC%d%s%s: %d stale, %d ready",
dc, 'm' if is_media else '', len(expired), len(bucket)) dc, 't' if is_test_dc else '', 'm' if is_media else '',
len(expired), len(bucket))
if len(bucket) < proxy_config.pool_size: if len(bucket) < proxy_config.pool_size:
self._schedule_refill(key, target_ip, domains) self._schedule_refill(key, target_ip, domains)
wake_at = now + self.WS_POOL_CHECK_INTERVAL
if bucket:
wake_at = min(wake_at, min(
created + self.WS_POOL_MAX_AGE for _, created in bucket))
if key not in self._refilling and self._refill_after.get(key, 0) > now:
wake_at = min(wake_at, self._refill_after[key])
await asyncio.sleep(max(0, wake_at - time.monotonic()))
finally: finally:
if self._rotating.get(key) is asyncio.current_task(): if self._rotating.get(key) is asyncio.current_task():
self._rotating.pop(key, None) self._rotating.pop(key, None)
async def _connect_one(self, target_ip, domains) -> Optional[RawWebSocket]: async def _connect_one(self, target_ip, domains, path=WS_PATH) -> Optional[RawWebSocket]:
for domain in domains: for domain in domains:
if self.try_fronting_first: modes = (True, False) if self.try_fronting_first else (False, True)
ws = await self._connect_fronted(target_ip, domain) for fronted in modes:
if ws: try:
return ws ws = await RawWebSocket.connect(
try: target_ip, domain, timeout=7 if fronted else 8, path=path,
ws = await RawWebSocket.connect( sni="sprinthost.ru" if fronted else None)
target_ip, domain, timeout=8) except Exception as exc:
self.try_fronting_first = False stats.ws_errors += 1
return ws log.debug("WS pool connect %s%s via %s (fronting=%s): %r",
except (asyncio.TimeoutError, ConnectionResetError): domain, path, target_ip, fronted, exc)
if self.try_fronting_first:
return None
return await self._connect_fronted(target_ip, domain)
except WsHandshakeError as exc:
if exc.is_redirect:
continue continue
return None self.try_fronting_first = fronted
except Exception: if fronted:
return None stats.connections_fronting += 1
log.debug("WS pool connected %s%s via %s (fronting=%s)",
domain, path, target_ip, fronted)
return ws
return None return None
async def _connect_fronted(self, target_ip, domain) -> Optional[RawWebSocket]:
try:
ws = await RawWebSocket.connect(
target_ip, domain, timeout=7, sni="sprinthost.ru")
except Exception:
return None
stats.connections_fronting += 1
self.try_fronting_first = True
return ws
async def _quiet_close(self, ws): async def _quiet_close(self, ws):
try: try:
await ws.close() await ws.close()
@@ -202,14 +203,21 @@ class _WsPool:
continue continue
for is_media in (False, True): for is_media in (False, True):
domains = ws_domains(dc, is_media) domains = ws_domains(dc, is_media)
self._schedule_refill((dc, is_media), target_ip, domains) key = (dc, is_media, proxy_config.force_test_dc)
self._schedule_refill(key, target_ip, domains)
log.info("WS pool warmup started for %d DC(s)", len(proxy_config.dc_redirects)) log.info("WS pool warmup started for %d DC(s)", len(proxy_config.dc_redirects))
def reset(self): def reset(self):
loop = asyncio.get_running_loop() loop = asyncio.get_running_loop()
for task in self._rotating.values(): for task in list(self._rotating.values()) + list(self._refilling.values()):
if not task.done() and task.get_loop() is loop: if not task.done() and task.get_loop() is loop:
task.cancel() task.cancel()
for bucket in self._idle.values():
for ws, _ in bucket:
try:
ws.writer.close()
except Exception as exc:
log.debug("WS pool close failed: %r", exc)
self._idle.clear() self._idle.clear()
self._refilling.clear() self._refilling.clear()
self._rotating.clear() self._rotating.clear()
@@ -217,6 +225,11 @@ class _WsPool:
self._refill_after.clear() self._refill_after.clear()
self.try_fronting_first = False self.try_fronting_first = False
async def close(self):
tasks = list(self._rotating.values()) + list(self._refilling.values())
self.reset()
await asyncio.gather(*tasks, return_exceptions=True)
class _CfWorkerPool: class _CfWorkerPool:
WS_POOL_MAX_AGE = 100.0 WS_POOL_MAX_AGE = 100.0
+42 -44
View File
@@ -104,26 +104,26 @@ class RawWebSocket:
timeout=min(timeout, 10), timeout=min(timeout, 10),
) )
set_sock_opts(writer.transport, proxy_config.buffer_size)
ws_key = base64.b64encode(os.urandom(16)).decode()
req = (
f'GET {path} HTTP/1.1\r\n'
f'Host: {domain}\r\n'
f'Upgrade: websocket\r\n'
f'Connection: Upgrade\r\n'
f'Sec-WebSocket-Key: {ws_key}\r\n'
f'Sec-WebSocket-Version: 13\r\n'
f'Sec-WebSocket-Protocol: binary\r\n'
f'\r\n'
)
writer.write(req.encode())
await writer.drain()
response_lines: list[str] = []
try: try:
set_sock_opts(writer.transport, proxy_config.buffer_size)
ws_key = base64.b64encode(os.urandom(16)).decode()
req = (
f'GET {path} HTTP/1.1\r\n'
f'Host: {domain}\r\n'
f'Upgrade: websocket\r\n'
f'Connection: Upgrade\r\n'
f'Sec-WebSocket-Key: {ws_key}\r\n'
f'Sec-WebSocket-Version: 13\r\n'
f'Sec-WebSocket-Protocol: binary\r\n'
f'\r\n'
)
writer.write(req.encode())
await writer.drain()
response_lines: list[str] = []
while True: while True:
line = await asyncio.wait_for(reader.readline(), line = await asyncio.wait_for(reader.readline(),
timeout=timeout) timeout=timeout)
@@ -131,34 +131,32 @@ class RawWebSocket:
break break
response_lines.append( response_lines.append(
line.decode('utf-8', errors='replace').strip()) line.decode('utf-8', errors='replace').strip())
except asyncio.TimeoutError:
if not response_lines:
raise WsHandshakeError(0, 'empty response')
first_line = response_lines[0]
parts = first_line.split(' ', 2)
try:
status_code = int(parts[1]) if len(parts) >= 2 else 0
except ValueError:
status_code = 0
if status_code == 101:
return RawWebSocket(reader, writer)
headers: dict[str, str] = {}
for hl in response_lines[1:]:
if ':' in hl:
k, v = hl.split(':', 1)
headers[k.strip().lower()] = v.strip()
raise WsHandshakeError(status_code, first_line, headers,
location=headers.get('location'))
except BaseException:
writer.close() writer.close()
raise raise
if not response_lines:
writer.close()
raise WsHandshakeError(0, 'empty response')
first_line = response_lines[0]
parts = first_line.split(' ', 2)
try:
status_code = int(parts[1]) if len(parts) >= 2 else 0
except ValueError:
status_code = 0
if status_code == 101:
return RawWebSocket(reader, writer)
headers: dict[str, str] = {}
for hl in response_lines[1:]:
if ':' in hl:
k, v = hl.split(':', 1)
headers[k.strip().lower()] = v.strip()
writer.close()
raise WsHandshakeError(status_code, first_line, headers,
location=headers.get('location'))
async def send(self, data: bytes): async def send(self, data: bytes):
if self._closed: if self._closed:
raise ConnectionError("WebSocket closed") raise ConnectionError("WebSocket closed")
+14 -130
View File
@@ -10,7 +10,7 @@ import argparse
import logging import logging
import socket as _socket import socket as _socket
from typing import Dict, Optional, Set, Tuple from typing import Optional, Set, Tuple
if __name__ == '__main__' and (__package__ is None or __package__ == ''): if __name__ == '__main__' and (__package__ is None or __package__ == ''):
@@ -22,8 +22,8 @@ if __name__ == '__main__' and (__package__ is None or __package__ == ''):
from .utils import * from .utils import *
from .stats import stats from .stats import stats
from .config import proxy_config, parse_dc_ip_list, start_cfproxy_domain_refresh, coerce_domain_list from .config import proxy_config, parse_dc_ip_list, start_cfproxy_domain_refresh, coerce_domain_list
from .bridge import MsgSplitter, CryptoCtx, do_fallback, bridge_ws_reencrypt from .bridge import MsgSplitter, CryptoCtx, do_fallback, bridge_ws_reencrypt, reset_tcp_backoff
from .raw_websocket import RawWebSocket, WsHandshakeError, set_sock_opts from .raw_websocket import set_sock_opts
from .fake_tls import proxy_to_masking_domain, verify_client_hello, build_server_hello, FakeTlsStream, TLS_RECORD_HANDSHAKE from .fake_tls import proxy_to_masking_domain, verify_client_hello, build_server_hello, FakeTlsStream, TLS_RECORD_HANDSHAKE
from .balancer import balancer from .balancer import balancer
from .pool import ws_pool, cf_worker_pool from .pool import ws_pool, cf_worker_pool
@@ -34,14 +34,8 @@ from ._aes import Cipher, algorithms, modes
log = logging.getLogger('tg-mtproto-proxy') log = logging.getLogger('tg-mtproto-proxy')
IP_FAIL_COOLDOWN = 3600.0
DC_FAIL_COOLDOWN = 60.0
WS_FAIL_TIMEOUT = 2.0
LISTENER_CHECK_INTERVAL = 5.0 LISTENER_CHECK_INTERVAL = 5.0
LISTENER_RESTART_DELAY = 1.0 LISTENER_RESTART_DELAY = 1.0
ws_blacklist: Set[str] = set()
dc_fail_until: Dict[str, float] = {}
ip_fail_until: Dict[str, float] = {}
cf_h2_pool: Optional[CfH2Pool] = None cf_h2_pool: Optional[CfH2Pool] = None
@@ -298,119 +292,11 @@ async def _handle_client(reader, writer, secret: bytes):
relay_init = _generate_relay_init(proto_tag, dc_idx) relay_init = _generate_relay_init(proto_tag, dc_idx)
ctx = _build_crypto_ctx(client_dec_prekey_iv, secret, relay_init) ctx = _build_crypto_ctx(client_dec_prekey_iv, secret, relay_init)
dc_key = f'{dc}{"t" if is_test_dc else ""}{"m" if is_media else ""}'
media_tag = " media" if is_media else "" media_tag = " media" if is_media else ""
now = time.monotonic() ws = await ws_pool.get(dc, is_media, is_test_dc=is_test_dc)
ws_path = WS_PATH_TEST if is_test_dc else WS_PATH
target = proxy_config.dc_redirects.get(dc)
is_any_cf_fallback = proxy_config.fallback_cfproxy or proxy_config.cfproxy_worker_domains
domains = ws_domains(dc, is_media)
ws = None
# Fallback if DC not in config, if WS blacklisted for this DC/is_media or if connect to ip is timed out
if (dc not in proxy_config.dc_redirects
or dc_key in ws_blacklist
or now < ip_fail_until.get(target, 0) and is_any_cf_fallback):
if dc not in proxy_config.dc_redirects:
log.info("[%s] DC%d not in config -> fallback",
label, dc)
elif dc_key in ws_blacklist:
log.info("[%s] DC%d%s WS blacklisted -> fallback",
label, dc, media_tag)
else:
# Try to get WS from pool first, might be accidental timeout
ws = await ws_pool.get(
dc, is_media, target, domains
) if not is_test_dc else None
if not ws:
log.info("[%s] DC%d%s WS connect to %s was timed out -> fallback",
label, dc, media_tag, target)
else:
log.info("[%s] DC%d%s WS connect to %s was timed out, but pool hit -> using WS",
label, dc, media_tag, target)
if not ws:
splitter = None
try:
splitter = MsgSplitter(relay_init, proto_int)
except Exception:
pass
ok = await do_fallback(
clt_reader, clt_writer, relay_init, label,
dc, is_test_dc, is_media, media_tag,
ctx, splitter=splitter, h2_pool=cf_h2_pool, proto_tag=proto_tag)
if not ok:
log.warning("[%s] DC%d%s no fallback available",
label, dc, media_tag)
return
ws_timeout = WS_FAIL_TIMEOUT if now < dc_fail_until.get(dc_key, 0) else 5.0
ws_failed_redirect = False
ws_timed_out = False
all_redirects = True
ws = ws or await ws_pool.get(
dc, is_media, target, domains
) if not is_test_dc else None
if ws:
log.info("[%s] DC%d%s -> pool hit via %s",
label, dc, media_tag, target)
else:
for domain in domains:
url = f'wss://{domain}{ws_path}'
log.info("[%s] DC%d%s -> %s via %s",
label, dc, media_tag, url, target)
try:
ws = await RawWebSocket.connect(target, domain,
timeout=ws_timeout,
path=ws_path)
all_redirects = False
break
except WsHandshakeError as exc:
stats.ws_errors += 1
if exc.is_redirect:
ws_failed_redirect = True
log.warning("[%s] DC%d%s got %d from %s -> %s",
label, dc, media_tag,
exc.status_code, domain,
exc.location or '?')
continue
else:
all_redirects = False
log.warning("[%s] DC%d%s WS handshake: %s",
label, dc, media_tag, exc.status_line)
except asyncio.TimeoutError:
stats.ws_errors += 1
ws_timed_out = True
log.warning("[%s] DC%d%s WS connect timed out via %s",
label, dc, media_tag, domain)
break
except Exception as exc:
stats.ws_errors += 1
all_redirects = False
log.warning("[%s] DC%d%s WS connect failed: %s",
label, dc, media_tag, repr(exc))
# WS failed -> fallback
if ws is None: if ws is None:
if ws_timed_out: log.info("[%s] DC%d%s WS pool unavailable -> fallback",
ip_fail_until[target] = now + IP_FAIL_COOLDOWN label, dc, media_tag)
log.info("[%s] DC%d%s WS connect to %s timed out, cooldown for %ds",
label, dc, media_tag, target, int(IP_FAIL_COOLDOWN))
if ws_failed_redirect and all_redirects:
ws_blacklist.add(dc_key)
log.warning("[%s] DC%d%s blacklisted for WS (all 302)",
label, dc, media_tag)
elif ws_failed_redirect:
dc_fail_until[dc_key] = now + DC_FAIL_COOLDOWN
else:
dc_fail_until[dc_key] = now + DC_FAIL_COOLDOWN
log.info("[%s] DC%d%s WS failed for %ds",
label, dc, media_tag, int(DC_FAIL_COOLDOWN))
splitter_fb = None splitter_fb = None
try: try:
splitter_fb = MsgSplitter(relay_init, proto_int) splitter_fb = MsgSplitter(relay_init, proto_int)
@@ -421,12 +307,12 @@ async def _handle_client(reader, writer, secret: bytes):
dc, is_test_dc, is_media, media_tag, dc, is_test_dc, is_media, media_tag,
ctx, splitter=splitter_fb, h2_pool=cf_h2_pool, proto_tag=proto_tag) ctx, splitter=splitter_fb, h2_pool=cf_h2_pool, proto_tag=proto_tag)
if ok: if ok:
log.info("[%s] DC%d%s fallback closed", log.info("[%s] DC%d%s fallback closed", label, dc, media_tag)
label, dc, media_tag) else:
log.warning("[%s] DC%d%s no fallback available", label, dc, media_tag)
return return
ip_fail_until.pop(target, None) log.info("[%s] DC%d%s -> WS pool hit", label, dc, media_tag)
ws_pool.report_success(dc, is_media)
stats.connections_ws += 1 stats.connections_ws += 1
splitter = None splitter = None
@@ -476,11 +362,9 @@ async def _run(stop_event: Optional[asyncio.Event] = None):
global _server_instance, _server_stop_event, cf_h2_pool global _server_instance, _server_stop_event, cf_h2_pool
_server_stop_event = stop_event _server_stop_event = stop_event
reset_tcp_backoff()
ws_pool.reset() ws_pool.reset()
cf_worker_pool.reset() cf_worker_pool.reset()
ws_blacklist.clear()
dc_fail_until.clear()
ip_fail_until.clear()
_client_tasks.clear() _client_tasks.clear()
cf_h2_pool = CfH2Pool() if proxy_config.h2_enabled else None cf_h2_pool = CfH2Pool() if proxy_config.h2_enabled else None
@@ -567,8 +451,7 @@ async def _run(stop_event: Optional[asyncio.Event] = None):
if now < next_stats: if now < next_stats:
continue continue
next_stats = now + 60 next_stats = now + 60
bl = ', '.join(f'DC{k}' for k in sorted(ws_blacklist)) or 'none' log.info("stats: %s", stats.summary())
log.info("stats: %s | ws_bl: %s", stats.summary(), bl)
if cf_h2_pool is not None: if cf_h2_pool is not None:
cf_h2_pool.log_stats() cf_h2_pool.log_stats()
except asyncio.CancelledError: except asyncio.CancelledError:
@@ -650,6 +533,7 @@ async def _run(stop_event: Optional[asyncio.Event] = None):
if task is not None: if task is not None:
await _quiet_cancel(task) await _quiet_cancel(task)
await server.wait_closed() await server.wait_closed()
await ws_pool.close()
if cf_h2_pool is not None: if cf_h2_pool is not None:
await cf_h2_pool.close() await cf_h2_pool.close()
cf_h2_pool = None cf_h2_pool = None
@@ -695,7 +579,7 @@ def main():
ap.add_argument('--buf-kb', type=int, default=256, metavar='KB', ap.add_argument('--buf-kb', type=int, default=256, metavar='KB',
help='Socket send/recv buffer size in KB (default 256)') help='Socket send/recv buffer size in KB (default 256)')
ap.add_argument('--pool-size', type=int, default=4, metavar='N', ap.add_argument('--pool-size', type=int, default=4, metavar='N',
help='WS connection pool size per DC (default 4, min 0)') help='WS connection pool size per DC (default 4, 0 disables direct WS)')
ap.add_argument('--cfproxy-domain', action='append', default=None, ap.add_argument('--cfproxy-domain', action='append', default=None,
metavar='DOMAIN', metavar='DOMAIN',
help='User defined Cloudflare-proxied domain for H2 media and WS fallback ' help='User defined Cloudflare-proxied domain for H2 media and WS fallback '
+2 -2
View File
@@ -58,11 +58,11 @@ WS_PATH = '/apiws'
WS_PATH_TEST = WS_PATH + '_test' WS_PATH_TEST = WS_PATH + '_test'
def ws_domains(dc: int, is_media) -> List[str]: def ws_domains(dc: int, is_media: bool) -> List[str]:
if dc == 203: if dc == 203:
dc = 2 dc = 2
if not is_media: if not is_media:
return [f'kws{dc}.web.telegram.org', f'kws{dc}-1.web.telegram.org'] return [f'kws{dc}.web.telegram.org']
return [f'kws{dc}-1.web.telegram.org', f'kws{dc}.web.telegram.org'] return [f'kws{dc}-1.web.telegram.org', f'kws{dc}.web.telegram.org']
+24 -3
View File
@@ -19,8 +19,6 @@ class H2RoutingTest(unittest.IsolatedAsyncioTestCase):
with patch.object(proxy_config, 'dc_redirects', {2: '149.154.167.51'}), \ with patch.object(proxy_config, 'dc_redirects', {2: '149.154.167.51'}), \
patch.object(proxy_config, 'force_test_dc', False), \ patch.object(proxy_config, 'force_test_dc', False), \
patch.object(tg_ws_proxy, 'cf_h2_pool', pool), \ patch.object(tg_ws_proxy, 'cf_h2_pool', pool), \
patch.object(tg_ws_proxy, 'ws_blacklist', set()), \
patch.object(tg_ws_proxy, 'ip_fail_until', {}), \
patch.object(tg_ws_proxy, 'set_sock_opts'), \ patch.object(tg_ws_proxy, 'set_sock_opts'), \
patch.object(tg_ws_proxy, '_read_client_init', AsyncMock( patch.object(tg_ws_proxy, '_read_client_init', AsyncMock(
return_value=(b'', reader, writer, 'test'))), \ return_value=(b'', reader, writer, 'test'))), \
@@ -28,7 +26,6 @@ class H2RoutingTest(unittest.IsolatedAsyncioTestCase):
2, True, PROTO_TAG_INTERMEDIATE, b'')), \ 2, True, PROTO_TAG_INTERMEDIATE, b'')), \
patch.object(tg_ws_proxy, '_build_crypto_ctx'), \ patch.object(tg_ws_proxy, '_build_crypto_ctx'), \
patch.object(tg_ws_proxy.ws_pool, 'get', AsyncMock(return_value=ws)), \ patch.object(tg_ws_proxy.ws_pool, 'get', AsyncMock(return_value=ws)), \
patch.object(tg_ws_proxy.ws_pool, 'report_success'), \
patch.object(tg_ws_proxy, 'bridge_ws_reencrypt', AsyncMock()) as direct, \ patch.object(tg_ws_proxy, 'bridge_ws_reencrypt', AsyncMock()) as direct, \
patch.object(tg_ws_proxy, 'do_fallback', AsyncMock()) as fallback: patch.object(tg_ws_proxy, 'do_fallback', AsyncMock()) as fallback:
await tg_ws_proxy._handle_client(reader, writer, b'') await tg_ws_proxy._handle_client(reader, writer, b'')
@@ -36,6 +33,30 @@ class H2RoutingTest(unittest.IsolatedAsyncioTestCase):
fallback.assert_not_awaited() fallback.assert_not_awaited()
pool.open.assert_not_awaited() pool.open.assert_not_awaited()
async def test_pool_miss_goes_straight_to_fallback_including_test_dcs(self):
for dc, forced, expected_test in [(2, False, False),
(10002, False, True),
(2, True, True)]:
with self.subTest(dc=dc, forced=forced):
reader = Mock()
writer = Mock(wait_closed=AsyncMock())
writer.get_extra_info.return_value = None
with patch.object(proxy_config, 'force_test_dc', forced), \
patch.object(tg_ws_proxy, 'set_sock_opts'), \
patch.object(tg_ws_proxy, '_read_client_init', AsyncMock(
return_value=(b'', reader, writer, 'test'))), \
patch.object(tg_ws_proxy, '_try_handshake', return_value=(
dc, False, PROTO_TAG_INTERMEDIATE, b'')), \
patch.object(tg_ws_proxy, '_build_crypto_ctx'), \
patch.object(tg_ws_proxy.ws_pool, 'get', AsyncMock(return_value=None)) as get, \
patch('proxy.raw_websocket.RawWebSocket.connect', AsyncMock()) as connect, \
patch.object(tg_ws_proxy, 'do_fallback', AsyncMock(return_value=True)) as fallback:
await tg_ws_proxy._handle_client(reader, writer, b'')
get.assert_awaited_once_with(2, False, is_test_dc=expected_test)
fallback.assert_awaited_once()
self.assertEqual(fallback.call_args.args[4:7], (2, expected_test, False))
connect.assert_not_awaited()
async def test_cf_fallback_order_and_route_guards(self): async def test_cf_fallback_order_and_route_guards(self):
cases = [ cases = [
# media, test DC, CF enabled, Worker succeeds, H2 opens, expected # media, test DC, CF enabled, Worker succeeds, H2 opens, expected
+185 -11
View File
@@ -1,3 +1,4 @@
import asyncio
import time import time
import unittest import unittest
@@ -7,6 +8,8 @@ from unittest import mock
from proxy.config import proxy_config from proxy.config import proxy_config
from proxy.pool import _WsPool from proxy.pool import _WsPool
from proxy.raw_websocket import RawWebSocket, WsHandshakeError
from proxy.utils import ws_domains
class _StopRotation(Exception): class _StopRotation(Exception):
@@ -15,8 +18,9 @@ class _StopRotation(Exception):
def _open_ws(): def _open_ws():
transport = SimpleNamespace(is_closing=lambda: False) transport = SimpleNamespace(is_closing=lambda: False)
writer = SimpleNamespace(transport=transport) writer = mock.Mock(transport=transport, drain=mock.AsyncMock(),
return SimpleNamespace(_closed=False, writer=writer) wait_closed=mock.AsyncMock())
return RawWebSocket(asyncio.StreamReader(), writer)
class WsPoolRotationTest(unittest.IsolatedAsyncioTestCase): class WsPoolRotationTest(unittest.IsolatedAsyncioTestCase):
@@ -35,23 +39,18 @@ class WsPoolRotationTest(unittest.IsolatedAsyncioTestCase):
with mock.patch.object(pool, '_schedule_refill') as refill: with mock.patch.object(pool, '_schedule_refill') as refill:
await pool.warmup() await pool.warmup()
keys = [call.args[0] for call in refill.call_args_list] keys = [call.args[0] for call in refill.call_args_list]
self.assertIn((2, False), keys) self.assertIn((2, False, False), keys)
self.assertEqual((2, True) in keys, media) self.assertEqual((2, True, False) in keys, media)
async def test_refills_partially_populated_bucket(self): async def test_refills_partially_populated_bucket(self):
pool = _WsPool() pool = _WsPool()
key = (2, False) key = (2, False, False)
pool._idle[key] = deque([ pool._idle[key] = deque([
(_open_ws(), time.monotonic()), (_open_ws(), time.monotonic()),
(_open_ws(), time.monotonic()), (_open_ws(), time.monotonic()),
]) ])
sleep_calls = 0
async def stop_after_one_iteration(_delay): async def stop_after_one_iteration(_delay):
nonlocal sleep_calls raise _StopRotation
sleep_calls += 1
if sleep_calls > 1:
raise _StopRotation
with mock.patch.object(proxy_config, 'pool_size', 4): with mock.patch.object(proxy_config, 'pool_size', 4):
with mock.patch( with mock.patch(
@@ -67,5 +66,180 @@ class WsPoolRotationTest(unittest.IsolatedAsyncioTestCase):
key, '149.154.167.220', ['example.com']) key, '149.154.167.220', ['example.com'])
class WsPoolTest(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.pool = _WsPool()
self.key = (2, False, False)
self.pool.WS_POOL_CHECK_INTERVAL = .01
for name, value in [('dc_redirects', {2: '192.0.2.1'}),
('pool_size', 1), ('force_test_dc', False)]:
patcher = mock.patch.object(proxy_config, name, value)
patcher.start()
self.addCleanup(patcher.stop)
self.addAsyncCleanup(self.pool.close)
async def wait_for(self, predicate):
async def wait():
while not predicate():
await asyncio.sleep(.001)
await asyncio.wait_for(wait(), 1)
async def test_regular_sessions_never_use_media_domain(self):
for dc in (2, 4):
with self.subTest(dc=dc), mock.patch(
'proxy.pool.RawWebSocket.connect',
side_effect=WsHandshakeError(302, 'redirect')) as connect:
self.assertIsNone(await self.pool._connect_one('192.0.2.1', ws_domains(dc, False)))
self.assertEqual([call.args[1] for call in connect.call_args_list],
[f'kws{dc}.web.telegram.org'] * 2)
async def test_media_tries_regular_domain_after_both_modes_fail(self):
for error in (asyncio.TimeoutError(), ConnectionResetError(),
WsHandshakeError(302, 'redirect'), WsHandshakeError(500, 'error')):
with self.subTest(error=type(error)), mock.patch(
'proxy.pool.RawWebSocket.connect',
side_effect=[error, error, _open_ws()]) as connect:
self.assertIsNotNone(await self.pool._connect_one('192.0.2.1', ws_domains(2, True)))
self.assertEqual([call.args[1] for call in connect.call_args_list],
['kws2-1.web.telegram.org'] * 2 + ['kws2.web.telegram.org'])
async def test_concurrent_preference_change_does_not_skip_fronting(self):
ws = _open_ws()
async def connect(*args, **kwargs):
if kwargs['sni'] is None:
self.pool.try_fronting_first = True
raise asyncio.TimeoutError()
return ws
with mock.patch('proxy.pool.RawWebSocket.connect', side_effect=connect) as dial:
self.assertIs(await self.pool._connect_one('192.0.2.1', ws_domains(2, False)), ws)
self.assertEqual(dial.await_count, 2)
async def test_miss_returns_without_waiting_and_starts_only_one_refill(self):
started = asyncio.Event()
async def connect(*args):
started.set()
await asyncio.Future()
with mock.patch.object(self.pool, '_connect_one', side_effect=connect) as dial:
for _ in range(10):
self.assertIsNone(await self.pool.get(2, False))
await asyncio.wait_for(started.wait(), 1)
self.assertEqual(dial.await_count, 1)
await self.pool.close()
async def test_empty_pool_recovers_in_background_after_backoff(self):
ws = _open_ws()
self.pool.REFILL_BACKOFF_INITIAL = .02
with mock.patch.object(self.pool, '_connect_one',
side_effect=[None, ws]) as dial:
self.assertIsNone(await self.pool.get(2, False))
await self.wait_for(lambda: self.key in self.pool._refill_after)
self.assertEqual(dial.await_count, 1)
# No new client requests are needed to resume refilling.
await self.wait_for(lambda: bool(self.pool._idle[self.key]))
self.assertEqual(dial.await_count, 2)
self.assertNotIn(self.key, self.pool._refill_failures)
await self.pool.close()
async def test_ready_connection_is_available_before_slower_attempt(self):
ready = _open_ws()
slow_started = asyncio.Event()
calls = 0
async def connect(*args):
nonlocal calls
calls += 1
if calls == 1:
slow_started.set()
await asyncio.Future()
return ready
with mock.patch.object(proxy_config, 'pool_size', 2), \
mock.patch.object(self.pool, '_connect_one', side_effect=connect):
await self.pool.get(2, False)
await asyncio.wait_for(slow_started.wait(), 1)
await self.wait_for(lambda: bool(self.pool._idle[self.key]))
self.assertIs(await self.pool.get(2, False), ready)
await self.pool.close()
async def test_stale_connections_are_skipped(self):
for state in ('expired', 'eof', 'exception', 'closed', 'closing'):
with self.subTest(state=state):
stale, good = _open_ws(), _open_ws()
created = time.monotonic()
if state == 'expired':
created -= self.pool.WS_POOL_MAX_AGE
elif state == 'eof':
stale.reader.feed_eof()
elif state == 'exception':
stale.reader.set_exception(ConnectionResetError())
elif state == 'closed':
stale._closed = True
else:
stale.writer.transport.is_closing = lambda: True
self.pool._idle[self.key] = deque([(stale, created), (good, time.monotonic())])
with mock.patch.object(self.pool, '_schedule_refill'):
self.assertIs(await self.pool.get(2, False), good)
await asyncio.sleep(0)
async def test_hit_does_not_reset_failed_refill_backoff(self):
self.pool._idle[self.key] = deque([(_open_ws(), time.monotonic())])
self.pool._refill_failures[self.key] = 3
self.pool._refill_after[self.key] = time.monotonic() + 60
with mock.patch.object(self.pool, '_schedule_rotation'):
self.assertIsNotNone(await self.pool.get(2, False))
self.assertEqual(self.pool._refill_failures[self.key], 3)
self.assertFalse(self.pool._refilling)
async def test_zero_size_and_unconfigured_dc_do_not_connect(self):
with mock.patch.object(self.pool, '_connect_one') as dial:
self.assertIsNone(await self.pool.get(4, False))
with mock.patch.object(proxy_config, 'pool_size', 0):
await self.pool.warmup()
self.assertIsNone(await self.pool.get(2, False))
self.assertFalse(self.pool._rotating)
self.assertFalse(self.pool._refilling)
dial.assert_not_called()
async def test_test_dc_has_separate_pool_and_path(self):
prod = _open_ws()
self.pool._idle[self.key] = deque([(prod, time.monotonic())])
with mock.patch.object(self.pool, '_connect_one', return_value=_open_ws()) as dial:
self.assertIsNone(await self.pool.get(2, False, is_test_dc=True))
test_key = (2, False, True)
await self.wait_for(lambda: bool(self.pool._idle.get(test_key)))
dial.assert_awaited_once_with('192.0.2.1', ['kws2.web.telegram.org'], '/apiws_test')
self.assertIs(self.pool._idle[self.key][0][0], prod)
await self.pool.close()
async def test_close_cancels_refill_and_closes_unclaimed_connections(self):
ready = _open_ws()
idle = _open_ws()
started = asyncio.Event()
self.pool._idle[self.key] = deque([(idle, time.monotonic())])
async def connect(*args):
started.set()
try:
await asyncio.Future()
except asyncio.CancelledError:
# Completion raced with shutdown; it must not leak into the pool.
return ready
with mock.patch.object(proxy_config, 'pool_size', 2), \
mock.patch.object(self.pool, '_connect_one', side_effect=connect):
self.pool._schedule_refill(self.key, '192.0.2.1', ws_domains(2, False))
await asyncio.wait_for(started.wait(), 1)
await self.pool.close()
self.assertFalse(self.pool._idle)
self.assertFalse(self.pool._rotating)
self.assertFalse(self.pool._refilling)
ready.writer.close.assert_called_once()
idle.writer.close.assert_called_once()
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()
+34
View File
@@ -1,5 +1,6 @@
import asyncio import asyncio
import unittest import unittest
from unittest.mock import AsyncMock, Mock, patch
from proxy.raw_websocket import RawWebSocket, WsHandshakeError, _xor_mask from proxy.raw_websocket import RawWebSocket, WsHandshakeError, _xor_mask
@@ -126,5 +127,38 @@ class HandshakeErrorTest(unittest.TestCase):
self.assertFalse(WsHandshakeError(code, '').is_redirect) self.assertFalse(WsHandshakeError(code, '').is_redirect)
class HandshakeCleanupTest(unittest.IsolatedAsyncioTestCase):
async def test_cancelling_pending_upgrade_closes_socket(self):
reader = asyncio.StreamReader()
writer = Mock(drain=AsyncMock())
started = asyncio.Event()
async def drain():
started.set()
writer.drain.side_effect = drain
with patch('proxy.raw_websocket.asyncio.open_connection',
AsyncMock(return_value=(reader, writer))), \
patch('proxy.raw_websocket.set_sock_opts'):
task = asyncio.create_task(RawWebSocket.connect('192.0.2.1', 'example.org'))
await asyncio.wait_for(started.wait(), 1)
task.cancel()
with self.assertRaises(asyncio.CancelledError):
await task
writer.close.assert_called_once()
async def test_failed_upgrade_closes_socket(self):
for error in (ConnectionResetError(), asyncio.TimeoutError()):
with self.subTest(error=type(error)):
reader = Mock(readline=AsyncMock(side_effect=error))
writer = Mock(drain=AsyncMock())
with patch('proxy.raw_websocket.asyncio.open_connection',
AsyncMock(return_value=(reader, writer))), \
patch('proxy.raw_websocket.set_sock_opts'):
with self.assertRaises(type(error)):
await RawWebSocket.connect('192.0.2.1', 'example.org')
writer.close.assert_called_once()
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()
+109
View File
@@ -0,0 +1,109 @@
import asyncio
import time
import unittest
from unittest.mock import AsyncMock, Mock, patch
from proxy import bridge
class TcpFallbackTest(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
bridge.reset_tcp_backoff()
self.addCleanup(bridge.reset_tcp_backoff)
self.key = ('192.0.2.1', 443)
self.remote = Mock(drain=AsyncMock(), wait_closed=AsyncMock())
async def fallback(self, dst='192.0.2.1', port=443):
return await bridge._tcp_fallback(
Mock(), Mock(), dst, port, b'init', 'test', Mock())
async def test_failure_backoff_doubles_caps_and_skips_new_connections(self):
with patch.object(bridge.asyncio, 'open_connection',
AsyncMock(side_effect=asyncio.TimeoutError())) as connect:
for delay in (30, 60, 120, 240, 480, 960, 1920, 3600, 3600):
bridge._tcp_retry_after[self.key] = 0
start = time.monotonic()
self.assertFalse(await self.fallback())
remaining = bridge._tcp_retry_after[self.key] - start
self.assertGreaterEqual(remaining, delay - .001)
self.assertLess(remaining, delay + 1)
count = connect.await_count
self.assertFalse(await self.fallback())
self.assertEqual(connect.await_count, count)
self.assertFalse(bridge._tcp_connecting)
async def test_other_ips_and_ports_are_not_blocked(self):
bridge._tcp_retry_after[self.key] = time.monotonic() + 60
with patch.object(bridge.asyncio, 'open_connection', AsyncMock(
return_value=(Mock(), self.remote))) as connect, \
patch.object(bridge, '_bridge_tcp_reencrypt', AsyncMock()):
self.assertTrue(await self.fallback('192.0.2.2'))
self.assertTrue(await self.fallback(port=80))
self.assertEqual(connect.await_count, 2)
async def test_success_clears_backoff_and_allows_following_sessions(self):
bridge._tcp_failures[self.key] = 5
bridge._tcp_retry_after[self.key] = 0
with patch.object(bridge.asyncio, 'open_connection', AsyncMock(
return_value=(Mock(), self.remote))) as connect, \
patch.object(bridge, '_bridge_tcp_reencrypt', AsyncMock()):
self.assertTrue(await self.fallback())
self.assertNotIn(self.key, bridge._tcp_failures)
self.assertNotIn(self.key, bridge._tcp_retry_after)
self.assertTrue(await self.fallback())
self.assertEqual(connect.await_count, 2)
async def test_concurrent_requests_only_start_one_connection_attempt(self):
started = asyncio.Event()
release = asyncio.Event()
async def connect(*args):
started.set()
await release.wait()
raise ConnectionRefusedError()
with patch.object(bridge.asyncio, 'open_connection', side_effect=connect) as dial:
first = asyncio.create_task(self.fallback())
try:
await asyncio.wait_for(started.wait(), 1)
results = await asyncio.gather(*(self.fallback() for _ in range(8)))
self.assertEqual(results, [False] * 8)
self.assertEqual(dial.await_count, 1)
release.set()
self.assertFalse(await first)
finally:
first.cancel()
await asyncio.gather(first, return_exceptions=True)
async def test_init_failure_closes_socket_and_enters_backoff(self):
self.remote.drain.side_effect = ConnectionResetError()
with patch.object(bridge.asyncio, 'open_connection', AsyncMock(
return_value=(Mock(), self.remote))), \
patch.object(bridge, '_bridge_tcp_reencrypt', AsyncMock()) as forward:
self.assertFalse(await self.fallback())
self.remote.close.assert_called_once()
forward.assert_not_awaited()
self.assertIn(self.key, bridge._tcp_retry_after)
async def test_cancellation_releases_attempt_without_marking_ip_blocked(self):
started = asyncio.Event()
async def drain():
started.set()
await asyncio.Future()
self.remote.drain.side_effect = drain
with patch.object(bridge.asyncio, 'open_connection', AsyncMock(
return_value=(Mock(), self.remote))):
task = asyncio.create_task(self.fallback())
await asyncio.wait_for(started.wait(), 1)
task.cancel()
with self.assertRaises(asyncio.CancelledError):
await task
self.remote.close.assert_called_once()
self.assertFalse(bridge._tcp_connecting)
self.assertFalse(bridge._tcp_failures)
if __name__ == '__main__':
unittest.main()
+1 -1
View File
@@ -49,7 +49,7 @@
"tip.verbose": "When enabled, more details are written to the log file — useful for troubleshooting", "tip.verbose": "When enabled, more details are written to the log file — useful for troubleshooting",
"tip.no_secure": "When enabled, the proxy will use port 80 for CF-proxy and CF-worker connections", "tip.no_secure": "When enabled, the proxy will use port 80 for CF-proxy and CF-worker connections",
"tip.buf_kb": "Receive/send buffer size in kilobytes.\nA larger value allocates more memory per socket", "tip.buf_kb": "Receive/send buffer size in kilobytes.\nA larger value allocates more memory per socket",
"tip.pool": "How many parallel WebSocket sessions per data center can be kept open.\nIncreasing may help under high load", "tip.pool": "How many ready WebSocket connections to keep per data center.\n0 disables the direct DC->IP WebSocket route",
"tip.log_mb": "Maximum log file size; the file is overwritten when the limit is reached", "tip.log_mb": "Maximum log file size; the file is overwritten when the limit is reached",
"tip.autostart": "Launch TG WS Proxy on Windows login. If you move the app to another folder, autostart will reset", "tip.autostart": "Launch TG WS Proxy on Windows login. If you move the app to another folder, autostart will reset",
"tip.check_updates": "Check for updates on startup", "tip.check_updates": "Check for updates on startup",
+1 -1
View File
@@ -49,7 +49,7 @@
"tip.verbose": "Если включено, в файл логов пишется больше подробностей — необходимо при поиске неполадок", "tip.verbose": "Если включено, в файл логов пишется больше подробностей — необходимо при поиске неполадок",
"tip.no_secure": "Если включено, прокси будет использовать 80 порт для CF-прокси и CF-worker подключений", "tip.no_secure": "Если включено, прокси будет использовать 80 порт для CF-прокси и CF-worker подключений",
"tip.buf_kb": "Размер буфера приёма/передачи в килобайтах.\nБольше значение — больше выделение памяти на сокет", "tip.buf_kb": "Размер буфера приёма/передачи в килобайтах.\nБольше значение — больше выделение памяти на сокет",
"tip.pool": "Сколько параллельных WebSocket-сессий к одному датацентру можно держать.\nУвеличение может помочь при высокой нагрузке", "tip.pool": "Сколько готовых WebSocket-соединений держать для каждого датацентра.\n0 отключает прямой DC->IP WebSocket-маршрут",
"tip.log_mb": "Максимальный размер файла лога; при достижении лимита файл перезаписывается", "tip.log_mb": "Максимальный размер файла лога; при достижении лимита файл перезаписывается",
"tip.autostart": "Запускать TG WS Proxy при входе в Windows. Если вы переместите программу в другую папку, автозапуск сбросится", "tip.autostart": "Запускать TG WS Proxy при входе в Windows. Если вы переместите программу в другую папку, автозапуск сбросится",
"tip.check_updates": "При запуске проверять наличие обновлений", "tip.check_updates": "При запуске проверять наличие обновлений",