mirror of
https://github.com/Flowseal/tg-ws-proxy.git
synced 2026-10-21 02:58:45 +00:00
WsPool refactoring, TCP fallback backoff, removed direct connections outside of a pool
This commit is contained in:
1 parent
48954ff9ab
commit
50c980e9ea
13 files changed
+563
-294
No files matched your search
@@ -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 |
|
||||||
|
|||||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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']
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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
@@ -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()
|
||||||
@@ -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()
|
||||||
@@ -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
@@ -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
@@ -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": "При запуске проверять наличие обновлений",
|
||||||
|
|||||||
Reference in new issue
Block a user