Пул CF worker: закрытие при остановке, OPTI, сборка, доки (#1411)

docs/EN/BuildFromSource.md

@@ -55,11 +55,12 @@ tg-ws-proxy [--port PORT] [--host HOST] [--dc-ip DC:IP ...] [-v]
 | `--no-secure` | `false` | Use 80 port for CF-proxy and CF-worker connections |
 | `--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`) |
+| `--force-test-dc` | `false` | Route all traffic to Telegram test DCs [Learn more](./TestDc.md) |
 | `--buf-kb` | `256` | Buffer size in KB |
 | `--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-max-mb` | `5` | Maximum log file size in MB (afterwards overwrites) |
-| `--log-backups` | `0` | Number of log backups after overwrite |
+| `--log-backups` | `1` | Number of rotated log files to keep (minimum 1) |
 | `-v`, `--verbose` | disabled | Verbose logging (DEBUG) |
 
 **Examples:**

docs/EN/TrayConfig.md

@@ -18,13 +18,18 @@ The tray application stores data in:
   "verbose": false,
   "buf_kb": 256,
   "pool_size": 4,
-  "log_max_mb": 5.0,
+  "log_max_mb": 5,
   "check_updates": true,
   "cfproxy": true,
-  "cfproxy_user_domain": "",
-  "cfproxy_worker_domain": "",
+  "h2": true,
+  "cfproxy_user_domain_enabled": false,
+  "cfproxy_user_domain": [],
+  "cfproxy_worker_enabled": false,
+  "cfproxy_worker_domain": [],
   "force_test_dc": false,
-  "appearance": "auto"
+  "no_secure": false,
+  "appearance": "auto",
+  "language": "en"
 }
 ```
 

docs/RU/BuildFromSource.md

@@ -55,11 +55,12 @@ tg-ws-proxy [--port PORT] [--host HOST] [--dc-ip DC:IP ...] [-v]
 | `--no-secure` | `false` | Принудительно использует 80 порт для CF-proxy и CF-worker |
 | `--fake-tls-domain` | | Включить маскировку Fake TLS (ee-secret) с указанным SNI-доменом |
 | `--proxy-protocol` | выкл. | Принимать HAProxy PROXY protocol v1 (для работы за nginx/haproxy с `proxy_protocol on`) |
+| `--force-test-dc` | `false` | Направлять весь трафик на тестовые DC Telegram [Подробнее](./TestDc.md) |
 | `--buf-kb` | `256` | Размер буфера в КБ |
 | `--pool-size` | `4` | Количество заготовленных WS-соединений на каждый DC. `0` отключает прямой WS-маршрут DC→IP |
 | `--log-file` | выкл. | Путь к файлу, в который будут сохраняться логи |
 | `--log-max-mb` | `5` | Максимальный размер файла логов в МБ (после этого начинается перезапись) |
-| `--log-backups` | `0` | Количество сохранений логов после перезаписи |
+| `--log-backups` | `1` | Сколько старых файлов логов хранить при ротации (минимум 1) |
 | `-v`, `--verbose` | выкл. | Подробное логирование (DEBUG) |
 
 **Примеры:**

docs/RU/TrayConfig.md

@@ -18,13 +18,18 @@ Tray-приложение хранит данные в:
   "verbose": false,
   "buf_kb": 256,
   "pool_size": 4,
-  "log_max_mb": 5.0,
+  "log_max_mb": 5,
   "check_updates": true,
   "cfproxy": true,
-  "cfproxy_user_domain": "",
-  "cfproxy_worker_domain": "",
+  "h2": true,
+  "cfproxy_user_domain_enabled": false,
+  "cfproxy_user_domain": [],
+  "cfproxy_worker_enabled": false,
+  "cfproxy_worker_domain": [],
   "force_test_dc": false,
-  "appearance": "auto"
+  "no_secure": false,
+  "appearance": "auto",
+  "language": "ru"
 }
 ```
 

packaging/linux.spec

@@ -24,7 +24,7 @@ a = Analysis(
     [os.path.join(os.path.dirname(SPEC), os.pardir, 'linux.py')],
     pathex=[],
     binaries=appindicator_binaries,
-    datas=[(ctk_path, 'customtkinter/'), (_i18n_path, 'ui/i18n')],
+    datas=[(os.path.join(ctk_path, 'assets'), 'customtkinter/assets'), (_i18n_path, 'ui/i18n')],
     hiddenimports=[
         'httpx',
         'h2',

packaging/macos.spec

@@ -15,7 +15,7 @@ a = Analysis(
     [os.path.join(os.path.dirname(SPEC), os.pardir, 'macos.py')],
     pathex=[],
     binaries=[],
-    datas=[(ctk_path, 'customtkinter/'), (_i18n_path, 'ui/i18n')] + certifi_datas,
+    datas=[(os.path.join(ctk_path, 'assets'), 'customtkinter/assets'), (_i18n_path, 'ui/i18n')] + certifi_datas,
     hiddenimports=[
         'httpx',
         'h2',

packaging/windows.spec

@@ -18,7 +18,7 @@ a = Analysis(
     [os.path.join(os.path.dirname(SPEC), os.pardir, 'windows.py')],
     pathex=[],
     binaries=[],
-    datas=[(ctk_path, 'customtkinter/'), (_i18n_path, 'ui/i18n')] + certifi_datas,
+    datas=[(os.path.join(ctk_path, 'assets'), 'customtkinter/assets'), (_i18n_path, 'ui/i18n')] + certifi_datas,
     hiddenimports=[
         'httpx',
         'h2',

proxy/pool.py

@@ -5,7 +5,7 @@ import time
 
 from collections import deque
 from urllib.parse import urlencode
-from typing import Dict, List, Optional, Tuple, Set
+from typing import Dict, List, Optional, Tuple
 
 from .raw_websocket import RawWebSocket, WsHandshakeError
 from .stats import stats
@@ -237,7 +237,7 @@ class _CfWorkerPool:
 
     def __init__(self):
         self._idle: Dict[int, deque] = {}
-        self._refilling: Set[int] = set()
+        self._refilling: Dict[int, asyncio.Task] = {}
         self._exhausted_until: Dict[str, float] = {}
 
     async def get(self, dc: int, fallback_dst: str,
@@ -269,8 +269,7 @@ class _CfWorkerPool:
     def _schedule_refill(self, dc, fallback_dst, worker_domains):
         if dc in self._refilling:
             return
-        self._refilling.add(dc)
-        asyncio.create_task(self._refill(
+        self._refilling[dc] = asyncio.create_task(self._refill(
             dc, fallback_dst, list(worker_domains)))
 
     async def _refill(self, dc, fallback_dst, worker_domains):
@@ -287,11 +286,15 @@ class _CfWorkerPool:
                 if connected is None:
                     break
                 ws, worker_domain = connected
+                if self._refilling.get(dc) is not asyncio.current_task():
+                    await self._quiet_close(ws)
+                    return
                 bucket.append((ws, time.monotonic(), worker_domain))
             log.debug("CF worker pool refilled DC%d: %d ready",
                       dc, len(bucket))
         finally:
-            self._refilling.discard(dc)
+            if self._refilling.get(dc) is asyncio.current_task():
+                self._refilling.pop(dc, None)
 
     async def _connect_one(self, worker_domains, fallback_dst, dc):
         query = urlencode({
@@ -359,10 +362,25 @@ class _CfWorkerPool:
         log.info("CF worker pool warmup started for %d DC(s)", len(cf_fallbacks))
 
     def reset(self):
+        loop = asyncio.get_running_loop()
+        for task in self._refilling.values():
+            if not task.done() and task.get_loop() is loop:
+                task.cancel()
+        for bucket in self._idle.values():
+            for ws, _, _ in bucket:
+                try:
+                    ws.writer.close()
+                except Exception as exc:
+                    log.debug("CF worker pool close failed: %r", exc)
         self._idle.clear()
         self._refilling.clear()
         self._exhausted_until.clear()
 
+    async def close(self):
+        tasks = list(self._refilling.values())
+        self.reset()
+        await asyncio.gather(*tasks, return_exceptions=True)
+
 
 ws_pool = _WsPool()
 cf_worker_pool = _CfWorkerPool()

proxy/tg_ws_proxy.py

@@ -536,6 +536,7 @@ async def _run(stop_event: Optional[asyncio.Event] = None):
                 await _quiet_cancel(task)
         await server.wait_closed()
         await ws_pool.close()
+        await cf_worker_pool.close()
         if cf_h2_pool is not None:
             await cf_h2_pool.close()
             cf_h2_pool = None

proxy/utils.py

@@ -33,7 +33,8 @@ PROTO_PADDED_INTERMEDIATE_INT = 0xDDDDDDDD
 RESERVED_FIRST_BYTES = {0xEF}
 RESERVED_STARTS = {b'\x48\x45\x41\x44', b'\x50\x4F\x53\x54',
                     b'\x47\x45\x54\x20', b'\xee\xee\xee\xee',
-                    b'\xdd\xdd\xdd\xdd', b'\x16\x03\x01\x02'}
+                    b'\xdd\xdd\xdd\xdd', b'\x16\x03\x01\x02',
+                    b'\x4F\x50\x54\x49'}
 RESERVED_CONTINUE = b'\x00\x00\x00\x00'
 
 _GITHUB_IPS: Dict[str, str] = {

tests/test_pool.py

@@ -7,7 +7,7 @@ from types import SimpleNamespace
 from unittest import mock
 
 from proxy.config import proxy_config
-from proxy.pool import _WsPool
+from proxy.pool import _CfWorkerPool, _WsPool
 from proxy.raw_websocket import RawWebSocket, WsHandshakeError
 from proxy.utils import ws_domains
 
@@ -241,5 +241,31 @@ class WsPoolTest(unittest.IsolatedAsyncioTestCase):
         idle.writer.close.assert_called_once()
 
 
+class CfWorkerPoolTest(unittest.IsolatedAsyncioTestCase):
+    async def test_close_cancels_refill_and_closes_unclaimed_connections(self):
+        pool = _CfWorkerPool()
+        ready = _open_ws()
+        idle = _open_ws()
+        started = asyncio.Event()
+        pool._idle[2] = deque([(idle, time.monotonic(), 'a.example')])
+
+        async def connect(*args):
+            started.set()
+            try:
+                await asyncio.Future()
+            except asyncio.CancelledError:
+                return ready, 'b.example'
+
+        with mock.patch.object(proxy_config, 'pool_size', 1), \
+                mock.patch.object(pool, '_connect_one', side_effect=connect):
+            pool._schedule_refill(4, '192.0.2.4', ['b.example'])
+            await asyncio.wait_for(started.wait(), 1)
+            await pool.close()
+        self.assertFalse(pool._idle)
+        self.assertFalse(pool._refilling)
+        ready.writer.close.assert_called_once()
+        idle.writer.close.assert_called_once()
+
+
 if __name__ == '__main__':
     unittest.main()