lange-python 0.3.17__tar.gz → 0.3.19__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {lange_python-0.3.17 → lange_python-0.3.19}/PKG-INFO +5 -1
- {lange_python-0.3.17 → lange_python-0.3.19}/README.md +4 -0
- lange_python-0.3.19/lange/tunnel/_client.py +685 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/pyproject.toml +1 -1
- lange_python-0.3.17/lange/tunnel/_client.py +0 -477
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/__init__.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/__main__.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/_util/__init__.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/_util/_base_client.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/_util/_key_handling.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/__init__.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/build/__init__.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/build/_command.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/build/_discovery.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/build/_docker.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/build/_poetry.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/build/_types.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/code/__init__.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/code/_stats.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/code/audit/__init__.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/code/audit/_command.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/code/audit/_discovery.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/code/audit/_runner.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/code/audit/_types.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/distribution/__init__.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/cli/distribution/_command.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/distribution/__init__.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/distribution/_client.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/distribution/_update_macos.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/distribution/_util.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/tunnel/__init__.py +0 -0
- {lange_python-0.3.17 → lange_python-0.3.19}/lange/tunnel/_util.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: lange-python
|
|
3
|
-
Version: 0.3.
|
|
3
|
+
Version: 0.3.19
|
|
4
4
|
Summary: A bundeld set of tools, clients for the lange-suite of tools and more.
|
|
5
5
|
Author: contact@robertlange.me
|
|
6
6
|
Requires-Python: >=3.10
|
|
@@ -89,3 +89,7 @@ tunnel.stop()
|
|
|
89
89
|
Tunnel clients also expose a `status` property with one of:
|
|
90
90
|
`"unauthenticated"`, `"off"`, `"pending"`, `"connected"`, or `"failed"`.
|
|
91
91
|
|
|
92
|
+
The worker connects to `/api/v1/tunnels/worker`, consumes the
|
|
93
|
+
`worker_connected` welcome payload, forwards HTTP bodies as base64-encoded
|
|
94
|
+
bytes, and supports websocket relay sessions through the same control socket.
|
|
95
|
+
|
|
@@ -70,3 +70,7 @@ tunnel.stop()
|
|
|
70
70
|
|
|
71
71
|
Tunnel clients also expose a `status` property with one of:
|
|
72
72
|
`"unauthenticated"`, `"off"`, `"pending"`, `"connected"`, or `"failed"`.
|
|
73
|
+
|
|
74
|
+
The worker connects to `/api/v1/tunnels/worker`, consumes the
|
|
75
|
+
`worker_connected` welcome payload, forwards HTTP bodies as base64-encoded
|
|
76
|
+
bytes, and supports websocket relay sessions through the same control socket.
|
|
@@ -0,0 +1,685 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import base64
|
|
5
|
+
import json
|
|
6
|
+
import logging
|
|
7
|
+
import ssl
|
|
8
|
+
import threading
|
|
9
|
+
from typing import Any, Optional
|
|
10
|
+
from urllib.parse import urljoin, urlparse, urlunparse
|
|
11
|
+
|
|
12
|
+
import httpx
|
|
13
|
+
import websockets
|
|
14
|
+
from httpx import Timeout
|
|
15
|
+
|
|
16
|
+
from .._util import BaseLangeLabsClient
|
|
17
|
+
from .._util._base_client import _UNSET_API_KEY
|
|
18
|
+
from ._util import _filter_hop_by_hop_headers
|
|
19
|
+
|
|
20
|
+
logger = logging.getLogger("lange.tunnel")
|
|
21
|
+
HEARTBEAT_INTERVAL_SECONDS = 10.0
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class Tunnel(BaseLangeLabsClient):
|
|
25
|
+
"""Thread-based tunnel worker client."""
|
|
26
|
+
|
|
27
|
+
def __init__(
|
|
28
|
+
self,
|
|
29
|
+
host: str = "wss://tunnel.lange-labs.com",
|
|
30
|
+
api_key: str | None = None,
|
|
31
|
+
target: str = "http://localhost:80",
|
|
32
|
+
verify_ssl: bool = True,
|
|
33
|
+
max_retries: int = 5,
|
|
34
|
+
retry_delay: float = 5.0,
|
|
35
|
+
open_timeout: float = 20.0,
|
|
36
|
+
daemon: bool = True,
|
|
37
|
+
) -> None:
|
|
38
|
+
"""Initialize one tunnel worker client."""
|
|
39
|
+
super().__init__(api_key=api_key, daemon=daemon, host=host)
|
|
40
|
+
self.target = target.rstrip("/")
|
|
41
|
+
self.verify_ssl = verify_ssl
|
|
42
|
+
self.max_retries = max_retries
|
|
43
|
+
self.retry_delay = retry_delay
|
|
44
|
+
self.open_timeout = open_timeout
|
|
45
|
+
|
|
46
|
+
self._max_retry_delay = 60.0
|
|
47
|
+
self._stop_event = threading.Event()
|
|
48
|
+
self._reconnect_event = threading.Event()
|
|
49
|
+
self._connected = False
|
|
50
|
+
self._remote_address: Optional[str] = None
|
|
51
|
+
self._remote_address_roundrobin: Optional[str] = None
|
|
52
|
+
self._worker_index = -1
|
|
53
|
+
self._pool_size = 0
|
|
54
|
+
self._reconnect_count = 0
|
|
55
|
+
self._worker_id: Optional[str] = None
|
|
56
|
+
self._instance_id: Optional[str] = None
|
|
57
|
+
self._tunnel_id: Optional[str] = None
|
|
58
|
+
self._tunnel_name: Optional[str] = None
|
|
59
|
+
self._max_workers: int = -1
|
|
60
|
+
self._can_consume_tunnels = False
|
|
61
|
+
self._lock = threading.Lock()
|
|
62
|
+
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
|
63
|
+
self._active_ws: Any = None
|
|
64
|
+
self._ws_sessions: dict[str, Any] = {}
|
|
65
|
+
self._ws_reader_tasks: dict[str, asyncio.Task[None]] = {}
|
|
66
|
+
|
|
67
|
+
@property
|
|
68
|
+
def connected(self) -> bool:
|
|
69
|
+
"""Get current connection state."""
|
|
70
|
+
with self._lock:
|
|
71
|
+
return self._connected
|
|
72
|
+
|
|
73
|
+
@property
|
|
74
|
+
def status(self) -> str:
|
|
75
|
+
"""Get the current tunnel lifecycle status."""
|
|
76
|
+
with self._lock:
|
|
77
|
+
return self._status
|
|
78
|
+
|
|
79
|
+
@property
|
|
80
|
+
def remote_address(self) -> Optional[str]:
|
|
81
|
+
"""Get worker-specific public address when provided by the server."""
|
|
82
|
+
with self._lock:
|
|
83
|
+
return self._remote_address
|
|
84
|
+
|
|
85
|
+
@property
|
|
86
|
+
def remote_address_roundrobin(self) -> Optional[str]:
|
|
87
|
+
"""Get round-robin public address when provided by the server."""
|
|
88
|
+
with self._lock:
|
|
89
|
+
return self._remote_address_roundrobin
|
|
90
|
+
|
|
91
|
+
@property
|
|
92
|
+
def worker_index(self) -> int:
|
|
93
|
+
"""Get current worker index from the latest welcome payload."""
|
|
94
|
+
with self._lock:
|
|
95
|
+
return self._worker_index
|
|
96
|
+
|
|
97
|
+
@property
|
|
98
|
+
def pool_size(self) -> int:
|
|
99
|
+
"""Get current worker pool size from the latest welcome payload."""
|
|
100
|
+
with self._lock:
|
|
101
|
+
return self._pool_size
|
|
102
|
+
|
|
103
|
+
@property
|
|
104
|
+
def reconnect_count(self) -> int:
|
|
105
|
+
"""Get number of reconnect attempts after the last successful connection."""
|
|
106
|
+
with self._lock:
|
|
107
|
+
return self._reconnect_count
|
|
108
|
+
|
|
109
|
+
@property
|
|
110
|
+
def worker_id(self) -> Optional[str]:
|
|
111
|
+
"""Get the worker identifier assigned by the server."""
|
|
112
|
+
with self._lock:
|
|
113
|
+
return self._worker_id
|
|
114
|
+
|
|
115
|
+
@property
|
|
116
|
+
def instance_id(self) -> Optional[str]:
|
|
117
|
+
"""Get the API instance identifier from the welcome payload."""
|
|
118
|
+
with self._lock:
|
|
119
|
+
return self._instance_id
|
|
120
|
+
|
|
121
|
+
@property
|
|
122
|
+
def tunnel_name(self) -> Optional[str]:
|
|
123
|
+
"""Get the tunnel name from the welcome payload."""
|
|
124
|
+
with self._lock:
|
|
125
|
+
return self._tunnel_name
|
|
126
|
+
|
|
127
|
+
@property
|
|
128
|
+
def max_workers(self) -> int:
|
|
129
|
+
"""Get the configured tunnel worker limit from the welcome payload."""
|
|
130
|
+
with self._lock:
|
|
131
|
+
return self._max_workers
|
|
132
|
+
|
|
133
|
+
@property
|
|
134
|
+
def can_consume_tunnels(self) -> bool:
|
|
135
|
+
"""Get the worker tunnel-consumption capability from the welcome payload."""
|
|
136
|
+
with self._lock:
|
|
137
|
+
return self._can_consume_tunnels
|
|
138
|
+
|
|
139
|
+
def run(self) -> None:
|
|
140
|
+
"""Start the worker loop inside the thread."""
|
|
141
|
+
asyncio.run(self._run_async())
|
|
142
|
+
|
|
143
|
+
def stop(self) -> None:
|
|
144
|
+
"""Request a graceful shutdown."""
|
|
145
|
+
self._stop_event.set()
|
|
146
|
+
|
|
147
|
+
def reconnect(self) -> None:
|
|
148
|
+
"""Force a full reconnect cycle."""
|
|
149
|
+
self._request_reconnect()
|
|
150
|
+
|
|
151
|
+
def reload(self, api_key: str | None | object = _UNSET_API_KEY) -> None:
|
|
152
|
+
"""Reload tunnel authentication and restart the connection flow when running."""
|
|
153
|
+
super().reload(api_key=api_key)
|
|
154
|
+
self._set_connected(False)
|
|
155
|
+
|
|
156
|
+
if self.is_alive() or self._loop is not None:
|
|
157
|
+
self._request_reconnect()
|
|
158
|
+
return
|
|
159
|
+
|
|
160
|
+
self._set_status(self._status_for_auth())
|
|
161
|
+
|
|
162
|
+
async def _run_async(self) -> None:
|
|
163
|
+
"""Run the worker connection loop with reconnect backoff."""
|
|
164
|
+
tunnel_url = self._build_tunnel_url()
|
|
165
|
+
ssl_context = self._build_ssl_context(tunnel_url)
|
|
166
|
+
|
|
167
|
+
with self._lock:
|
|
168
|
+
self._loop = asyncio.get_running_loop()
|
|
169
|
+
|
|
170
|
+
current_delay = self.retry_delay
|
|
171
|
+
|
|
172
|
+
while not self._stop_event.is_set():
|
|
173
|
+
if self.api_key is None:
|
|
174
|
+
self._set_connected(False)
|
|
175
|
+
self._set_status("unauthenticated")
|
|
176
|
+
await self._wait_for_reconnect_signal()
|
|
177
|
+
continue
|
|
178
|
+
|
|
179
|
+
self._set_status("pending")
|
|
180
|
+
headers = self._build_connection_headers()
|
|
181
|
+
|
|
182
|
+
try:
|
|
183
|
+
logger.info("Connecting to %s", tunnel_url)
|
|
184
|
+
async with websockets.connect(
|
|
185
|
+
tunnel_url,
|
|
186
|
+
additional_headers=headers,
|
|
187
|
+
ssl=ssl_context,
|
|
188
|
+
proxy=None,
|
|
189
|
+
open_timeout=self.open_timeout,
|
|
190
|
+
) as ws:
|
|
191
|
+
with self._lock:
|
|
192
|
+
self._active_ws = ws
|
|
193
|
+
|
|
194
|
+
current_delay = self.retry_delay
|
|
195
|
+
await self._consume_welcome(ws)
|
|
196
|
+
|
|
197
|
+
with self._lock:
|
|
198
|
+
self._reconnect_count = 0
|
|
199
|
+
|
|
200
|
+
send_lock = asyncio.Lock()
|
|
201
|
+
heartbeat_task = asyncio.create_task(
|
|
202
|
+
self._send_heartbeats(ws, send_lock)
|
|
203
|
+
)
|
|
204
|
+
try:
|
|
205
|
+
await self._handle_messages(ws, send_lock)
|
|
206
|
+
finally:
|
|
207
|
+
heartbeat_task.cancel()
|
|
208
|
+
await asyncio.gather(heartbeat_task, return_exceptions=True)
|
|
209
|
+
await self._close_all_websocket_sessions()
|
|
210
|
+
with self._lock:
|
|
211
|
+
self._active_ws = None
|
|
212
|
+
self._set_connected(False)
|
|
213
|
+
except websockets.exceptions.ConnectionClosed as exc:
|
|
214
|
+
self._set_connected(False)
|
|
215
|
+
self._set_status("failed")
|
|
216
|
+
logger.warning("Connection closed: %s", exc)
|
|
217
|
+
except Exception as exc: # pragma: no cover - network error path
|
|
218
|
+
self._set_connected(False)
|
|
219
|
+
self._set_status("failed")
|
|
220
|
+
logger.error("Connection error: %s", exc)
|
|
221
|
+
finally:
|
|
222
|
+
with self._lock:
|
|
223
|
+
self._active_ws = None
|
|
224
|
+
|
|
225
|
+
if self._stop_event.is_set():
|
|
226
|
+
break
|
|
227
|
+
|
|
228
|
+
if self._reconnect_event.is_set():
|
|
229
|
+
self._set_connected(False)
|
|
230
|
+
with self._lock:
|
|
231
|
+
self._reconnect_count = 0
|
|
232
|
+
self._reconnect_event.clear()
|
|
233
|
+
current_delay = self.retry_delay
|
|
234
|
+
self._set_status(self._status_for_auth())
|
|
235
|
+
logger.info("Manual reconnect requested. Reconnecting now.")
|
|
236
|
+
continue
|
|
237
|
+
|
|
238
|
+
with self._lock:
|
|
239
|
+
self._reconnect_count += 1
|
|
240
|
+
attempt = self._reconnect_count
|
|
241
|
+
|
|
242
|
+
if self.max_retries > 0 and attempt > self.max_retries:
|
|
243
|
+
self._set_status("failed")
|
|
244
|
+
logger.error(
|
|
245
|
+
"Max reconnection attempts (%s) reached. Giving up.",
|
|
246
|
+
self.max_retries,
|
|
247
|
+
)
|
|
248
|
+
break
|
|
249
|
+
|
|
250
|
+
logger.info("Reconnecting in %.1fs (attempt %s)", current_delay, attempt)
|
|
251
|
+
|
|
252
|
+
wait_time = 0.0
|
|
253
|
+
while (
|
|
254
|
+
wait_time < current_delay
|
|
255
|
+
and not self._stop_event.is_set()
|
|
256
|
+
and not self._reconnect_event.is_set()
|
|
257
|
+
):
|
|
258
|
+
await asyncio.sleep(0.5)
|
|
259
|
+
wait_time += 0.5
|
|
260
|
+
|
|
261
|
+
if self._reconnect_event.is_set():
|
|
262
|
+
continue
|
|
263
|
+
|
|
264
|
+
current_delay = min(current_delay * 2, self._max_retry_delay)
|
|
265
|
+
|
|
266
|
+
with self._lock:
|
|
267
|
+
self._active_ws = None
|
|
268
|
+
self._loop = None
|
|
269
|
+
if self._status in {"connected", "pending"}:
|
|
270
|
+
self._status = self._status_for_auth()
|
|
271
|
+
|
|
272
|
+
logger.info("Disconnected from server.")
|
|
273
|
+
|
|
274
|
+
async def _consume_welcome(self, ws: Any) -> None:
|
|
275
|
+
"""Read and process the initial welcome message when present."""
|
|
276
|
+
try:
|
|
277
|
+
welcome_message = await ws.recv()
|
|
278
|
+
welcome = json.loads(welcome_message)
|
|
279
|
+
except Exception:
|
|
280
|
+
self._set_connected(True)
|
|
281
|
+
return
|
|
282
|
+
|
|
283
|
+
if isinstance(welcome, dict) and welcome.get("type") == "worker_connected":
|
|
284
|
+
self._set_connected(
|
|
285
|
+
True,
|
|
286
|
+
worker_id=str(welcome.get("worker_id") or welcome.get("id") or ""),
|
|
287
|
+
worker_index=int(welcome.get("worker_index", -1)),
|
|
288
|
+
instance_id=str(welcome.get("instance_id") or "") or None,
|
|
289
|
+
tunnel_id=str(welcome.get("tunnel_id") or "") or None,
|
|
290
|
+
tunnel_name=str(welcome.get("tunnel_name") or "") or None,
|
|
291
|
+
max_workers=int(welcome.get("max_workers", -1)),
|
|
292
|
+
can_consume_tunnels=bool(
|
|
293
|
+
welcome.get("can_consume_tunnels", False)
|
|
294
|
+
),
|
|
295
|
+
)
|
|
296
|
+
return
|
|
297
|
+
|
|
298
|
+
self._set_connected(True)
|
|
299
|
+
|
|
300
|
+
async def _handle_messages(
|
|
301
|
+
self, ws: Any, send_lock: asyncio.Lock
|
|
302
|
+
) -> None:
|
|
303
|
+
"""Handle request and websocket relay messages from the server."""
|
|
304
|
+
async with httpx.AsyncClient(timeout=Timeout(timeout=15.0)) as http_client:
|
|
305
|
+
while not self._stop_event.is_set() and not self._reconnect_event.is_set():
|
|
306
|
+
try:
|
|
307
|
+
message = await asyncio.wait_for(ws.recv(), timeout=1.0)
|
|
308
|
+
except asyncio.TimeoutError:
|
|
309
|
+
continue
|
|
310
|
+
except websockets.exceptions.ConnectionClosed:
|
|
311
|
+
break
|
|
312
|
+
|
|
313
|
+
try:
|
|
314
|
+
payload = json.loads(
|
|
315
|
+
message if isinstance(message, str) else message.decode("utf-8")
|
|
316
|
+
)
|
|
317
|
+
await self._handle_message(ws, http_client, payload, send_lock)
|
|
318
|
+
except Exception as exc:
|
|
319
|
+
logger.error("Error handling message: %s", exc)
|
|
320
|
+
|
|
321
|
+
async def _handle_message(
|
|
322
|
+
self,
|
|
323
|
+
server_ws: Any,
|
|
324
|
+
http_client: Any,
|
|
325
|
+
message: dict[str, Any],
|
|
326
|
+
send_lock: asyncio.Lock,
|
|
327
|
+
) -> None:
|
|
328
|
+
"""Dispatch one incoming server message to the matching handler."""
|
|
329
|
+
message_type = str(message.get("type", "")).strip().lower()
|
|
330
|
+
if message_type == "heartbeat_ack":
|
|
331
|
+
return
|
|
332
|
+
|
|
333
|
+
if message_type == "request":
|
|
334
|
+
response = await self._forward_request(http_client, message)
|
|
335
|
+
async with send_lock:
|
|
336
|
+
await server_ws.send(json.dumps(response))
|
|
337
|
+
return
|
|
338
|
+
|
|
339
|
+
if message_type == "ws_open":
|
|
340
|
+
await self._open_websocket_session(server_ws, message, send_lock)
|
|
341
|
+
return
|
|
342
|
+
|
|
343
|
+
if message_type == "ws_data":
|
|
344
|
+
await self._forward_websocket_data(message)
|
|
345
|
+
return
|
|
346
|
+
|
|
347
|
+
if message_type == "ws_close":
|
|
348
|
+
await self._close_websocket_session(message)
|
|
349
|
+
|
|
350
|
+
async def _forward_request(
|
|
351
|
+
self, client: Any, request: dict[str, Any]
|
|
352
|
+
) -> dict[str, Any]:
|
|
353
|
+
"""Forward one proxied HTTP request to the configured local target."""
|
|
354
|
+
request_id = str(request.get("id", ""))
|
|
355
|
+
method = str(request.get("method", "GET"))
|
|
356
|
+
path = str(request.get("path", "/"))
|
|
357
|
+
headers = request.get("headers", {})
|
|
358
|
+
body_b64 = str(request.get("body", ""))
|
|
359
|
+
|
|
360
|
+
body = base64.b64decode(body_b64) if body_b64 else None
|
|
361
|
+
url = urljoin(f"{self.target}/", path.lstrip("/"))
|
|
362
|
+
filtered_headers = _filter_hop_by_hop_headers(headers)
|
|
363
|
+
|
|
364
|
+
try:
|
|
365
|
+
response = await client.request(
|
|
366
|
+
method=method,
|
|
367
|
+
url=url,
|
|
368
|
+
headers=filtered_headers,
|
|
369
|
+
content=body,
|
|
370
|
+
follow_redirects=True,
|
|
371
|
+
)
|
|
372
|
+
return {
|
|
373
|
+
"type": "response",
|
|
374
|
+
"tunnel_id": str(request.get("tunnel_id") or self._tunnel_id or ""),
|
|
375
|
+
"id": request_id,
|
|
376
|
+
"status": response.status_code,
|
|
377
|
+
"headers": dict(response.headers.items()),
|
|
378
|
+
"body": base64.b64encode(response.content).decode("utf-8"),
|
|
379
|
+
"bodyEncoding": "base64",
|
|
380
|
+
}
|
|
381
|
+
except Exception as exc:
|
|
382
|
+
logger.error("Error forwarding request: %s", exc)
|
|
383
|
+
error_body = json.dumps({"error": str(exc)}).encode("utf-8")
|
|
384
|
+
return {
|
|
385
|
+
"type": "response",
|
|
386
|
+
"tunnel_id": str(request.get("tunnel_id") or self._tunnel_id or ""),
|
|
387
|
+
"id": request_id,
|
|
388
|
+
"status": 502,
|
|
389
|
+
"headers": {"Content-Type": "application/json"},
|
|
390
|
+
"body": base64.b64encode(error_body).decode("utf-8"),
|
|
391
|
+
"bodyEncoding": "base64",
|
|
392
|
+
"error": str(exc),
|
|
393
|
+
}
|
|
394
|
+
|
|
395
|
+
async def _open_websocket_session(
|
|
396
|
+
self,
|
|
397
|
+
server_ws: Any,
|
|
398
|
+
message: dict[str, Any],
|
|
399
|
+
send_lock: asyncio.Lock,
|
|
400
|
+
) -> None:
|
|
401
|
+
"""Open one websocket connection against the local target service."""
|
|
402
|
+
session_id = str(message.get("id", "")).strip()
|
|
403
|
+
if not session_id:
|
|
404
|
+
return
|
|
405
|
+
|
|
406
|
+
target_url = self._build_target_websocket_url(
|
|
407
|
+
path=str(message.get("path", "/")),
|
|
408
|
+
query_string=str(message.get("query_string", "") or ""),
|
|
409
|
+
)
|
|
410
|
+
headers = _filter_hop_by_hop_headers(message.get("headers", {}))
|
|
411
|
+
|
|
412
|
+
try:
|
|
413
|
+
local_ws = await websockets.connect(
|
|
414
|
+
target_url,
|
|
415
|
+
additional_headers=headers,
|
|
416
|
+
ssl=self._build_ssl_context(target_url),
|
|
417
|
+
proxy=None,
|
|
418
|
+
open_timeout=self.open_timeout,
|
|
419
|
+
)
|
|
420
|
+
except Exception as exc:
|
|
421
|
+
async with send_lock:
|
|
422
|
+
await server_ws.send(
|
|
423
|
+
json.dumps(
|
|
424
|
+
{
|
|
425
|
+
"type": "ws_error",
|
|
426
|
+
"tunnel_id": str(
|
|
427
|
+
message.get("tunnel_id") or self._tunnel_id or ""
|
|
428
|
+
),
|
|
429
|
+
"id": session_id,
|
|
430
|
+
"error": str(exc),
|
|
431
|
+
}
|
|
432
|
+
)
|
|
433
|
+
)
|
|
434
|
+
return
|
|
435
|
+
|
|
436
|
+
self._ws_sessions[session_id] = local_ws
|
|
437
|
+
self._ws_reader_tasks[session_id] = asyncio.create_task(
|
|
438
|
+
self._relay_local_websocket(server_ws, session_id, local_ws, send_lock)
|
|
439
|
+
)
|
|
440
|
+
async with send_lock:
|
|
441
|
+
await server_ws.send(
|
|
442
|
+
json.dumps(
|
|
443
|
+
{
|
|
444
|
+
"type": "ws_opened",
|
|
445
|
+
"tunnel_id": str(
|
|
446
|
+
message.get("tunnel_id") or self._tunnel_id or ""
|
|
447
|
+
),
|
|
448
|
+
"id": session_id,
|
|
449
|
+
}
|
|
450
|
+
)
|
|
451
|
+
)
|
|
452
|
+
|
|
453
|
+
async def _forward_websocket_data(self, message: dict[str, Any]) -> None:
|
|
454
|
+
"""Forward one websocket frame from the server to the local target."""
|
|
455
|
+
session_id = str(message.get("id", "")).strip()
|
|
456
|
+
local_ws = self._ws_sessions.get(session_id)
|
|
457
|
+
if local_ws is None:
|
|
458
|
+
return
|
|
459
|
+
|
|
460
|
+
is_binary = bool(message.get("isBinary", False))
|
|
461
|
+
data = str(message.get("data", ""))
|
|
462
|
+
if is_binary:
|
|
463
|
+
await local_ws.send(base64.b64decode(data.encode("ascii")))
|
|
464
|
+
return
|
|
465
|
+
|
|
466
|
+
await local_ws.send(data)
|
|
467
|
+
|
|
468
|
+
async def _close_websocket_session(self, message: dict[str, Any]) -> None:
|
|
469
|
+
"""Close one local websocket session at the server's request."""
|
|
470
|
+
session_id = str(message.get("id", "")).strip()
|
|
471
|
+
await self._cleanup_websocket_session(
|
|
472
|
+
session_id=session_id,
|
|
473
|
+
close_socket=True,
|
|
474
|
+
code=int(message.get("code") or 1000),
|
|
475
|
+
reason=str(message.get("reason") or ""),
|
|
476
|
+
)
|
|
477
|
+
|
|
478
|
+
async def _relay_local_websocket(
|
|
479
|
+
self,
|
|
480
|
+
server_ws: Any,
|
|
481
|
+
session_id: str,
|
|
482
|
+
local_ws: Any,
|
|
483
|
+
send_lock: asyncio.Lock,
|
|
484
|
+
) -> None:
|
|
485
|
+
"""Relay frames from one local websocket back to the server."""
|
|
486
|
+
try:
|
|
487
|
+
while True:
|
|
488
|
+
payload = await local_ws.recv()
|
|
489
|
+
if isinstance(payload, bytes):
|
|
490
|
+
message = {
|
|
491
|
+
"type": "ws_data",
|
|
492
|
+
"tunnel_id": self._tunnel_id or "",
|
|
493
|
+
"id": session_id,
|
|
494
|
+
"data": base64.b64encode(payload).decode("ascii"),
|
|
495
|
+
"isBinary": True,
|
|
496
|
+
}
|
|
497
|
+
else:
|
|
498
|
+
message = {
|
|
499
|
+
"type": "ws_data",
|
|
500
|
+
"tunnel_id": self._tunnel_id or "",
|
|
501
|
+
"id": session_id,
|
|
502
|
+
"data": payload,
|
|
503
|
+
"isBinary": False,
|
|
504
|
+
}
|
|
505
|
+
async with send_lock:
|
|
506
|
+
await server_ws.send(json.dumps(message))
|
|
507
|
+
except websockets.exceptions.ConnectionClosed as exc:
|
|
508
|
+
async with send_lock:
|
|
509
|
+
await server_ws.send(
|
|
510
|
+
json.dumps(
|
|
511
|
+
{
|
|
512
|
+
"type": "ws_close",
|
|
513
|
+
"tunnel_id": self._tunnel_id or "",
|
|
514
|
+
"id": session_id,
|
|
515
|
+
"code": exc.code,
|
|
516
|
+
"reason": exc.reason or None,
|
|
517
|
+
}
|
|
518
|
+
)
|
|
519
|
+
)
|
|
520
|
+
except Exception as exc:
|
|
521
|
+
async with send_lock:
|
|
522
|
+
await server_ws.send(
|
|
523
|
+
json.dumps(
|
|
524
|
+
{
|
|
525
|
+
"type": "ws_error",
|
|
526
|
+
"tunnel_id": self._tunnel_id or "",
|
|
527
|
+
"id": session_id,
|
|
528
|
+
"error": str(exc),
|
|
529
|
+
}
|
|
530
|
+
)
|
|
531
|
+
)
|
|
532
|
+
finally:
|
|
533
|
+
await self._cleanup_websocket_session(session_id=session_id)
|
|
534
|
+
|
|
535
|
+
async def _cleanup_websocket_session(
|
|
536
|
+
self,
|
|
537
|
+
session_id: str,
|
|
538
|
+
close_socket: bool = False,
|
|
539
|
+
code: int = 1000,
|
|
540
|
+
reason: str = "",
|
|
541
|
+
) -> None:
|
|
542
|
+
"""Remove one local websocket session and optionally close it."""
|
|
543
|
+
local_ws = self._ws_sessions.pop(session_id, None)
|
|
544
|
+
task = self._ws_reader_tasks.pop(session_id, None)
|
|
545
|
+
|
|
546
|
+
if close_socket and local_ws is not None:
|
|
547
|
+
try:
|
|
548
|
+
await local_ws.close(code=code, reason=reason)
|
|
549
|
+
except Exception:
|
|
550
|
+
pass
|
|
551
|
+
|
|
552
|
+
current_task = asyncio.current_task()
|
|
553
|
+
if task is not None and task is not current_task:
|
|
554
|
+
task.cancel()
|
|
555
|
+
await asyncio.gather(task, return_exceptions=True)
|
|
556
|
+
|
|
557
|
+
async def _close_all_websocket_sessions(self) -> None:
|
|
558
|
+
"""Close all tracked local websocket sessions."""
|
|
559
|
+
for session_id in list(self._ws_sessions):
|
|
560
|
+
await self._cleanup_websocket_session(session_id, close_socket=True)
|
|
561
|
+
|
|
562
|
+
async def _send_heartbeats(self, server_ws: Any, send_lock: asyncio.Lock) -> None:
|
|
563
|
+
"""Send periodic heartbeat messages while the worker socket is active."""
|
|
564
|
+
while not self._stop_event.is_set() and not self._reconnect_event.is_set():
|
|
565
|
+
await asyncio.sleep(HEARTBEAT_INTERVAL_SECONDS)
|
|
566
|
+
if self._tunnel_id is None or self._worker_id is None:
|
|
567
|
+
continue
|
|
568
|
+
|
|
569
|
+
heartbeat = {
|
|
570
|
+
"type": "heartbeat",
|
|
571
|
+
"tunnel_id": self._tunnel_id,
|
|
572
|
+
"id": self._worker_id,
|
|
573
|
+
"worker_id": self._worker_id,
|
|
574
|
+
}
|
|
575
|
+
try:
|
|
576
|
+
async with send_lock:
|
|
577
|
+
await server_ws.send(json.dumps(heartbeat))
|
|
578
|
+
except websockets.exceptions.ConnectionClosed:
|
|
579
|
+
return
|
|
580
|
+
|
|
581
|
+
def _build_tunnel_url(self) -> str:
|
|
582
|
+
"""Build the worker websocket URL."""
|
|
583
|
+
return f"{self.host}/api/v1/tunnels/worker"
|
|
584
|
+
|
|
585
|
+
def _build_target_websocket_url(self, path: str, query_string: str = "") -> str:
|
|
586
|
+
"""Build the local target websocket URL for one proxied session."""
|
|
587
|
+
parsed_target = urlparse(self.target)
|
|
588
|
+
scheme = parsed_target.scheme
|
|
589
|
+
if scheme == "http":
|
|
590
|
+
scheme = "ws"
|
|
591
|
+
elif scheme == "https":
|
|
592
|
+
scheme = "wss"
|
|
593
|
+
|
|
594
|
+
joined_url = urlparse(urljoin(f"{self.target}/", path.lstrip("/")))
|
|
595
|
+
return urlunparse(
|
|
596
|
+
(
|
|
597
|
+
scheme,
|
|
598
|
+
joined_url.netloc or parsed_target.netloc,
|
|
599
|
+
joined_url.path,
|
|
600
|
+
"",
|
|
601
|
+
query_string,
|
|
602
|
+
"",
|
|
603
|
+
)
|
|
604
|
+
)
|
|
605
|
+
|
|
606
|
+
def _set_connected(
|
|
607
|
+
self,
|
|
608
|
+
value: bool,
|
|
609
|
+
remote_address: Optional[str] = None,
|
|
610
|
+
remote_address_roundrobin: Optional[str] = None,
|
|
611
|
+
worker_index: int = -1,
|
|
612
|
+
pool_size: int = 0,
|
|
613
|
+
worker_id: Optional[str] = None,
|
|
614
|
+
instance_id: Optional[str] = None,
|
|
615
|
+
tunnel_id: Optional[str] = None,
|
|
616
|
+
tunnel_name: Optional[str] = None,
|
|
617
|
+
max_workers: int = -1,
|
|
618
|
+
can_consume_tunnels: bool = False,
|
|
619
|
+
) -> None:
|
|
620
|
+
"""Update connection state atomically."""
|
|
621
|
+
with self._lock:
|
|
622
|
+
self._connected = value
|
|
623
|
+
self._remote_address = remote_address if value else None
|
|
624
|
+
self._remote_address_roundrobin = (
|
|
625
|
+
remote_address_roundrobin if value else None
|
|
626
|
+
)
|
|
627
|
+
self._worker_index = worker_index if value else -1
|
|
628
|
+
self._pool_size = pool_size if value else 0
|
|
629
|
+
self._worker_id = worker_id if value else None
|
|
630
|
+
self._instance_id = instance_id if value else None
|
|
631
|
+
self._tunnel_id = tunnel_id if value else None
|
|
632
|
+
self._tunnel_name = tunnel_name if value else None
|
|
633
|
+
self._max_workers = max_workers if value else -1
|
|
634
|
+
self._can_consume_tunnels = can_consume_tunnels if value else False
|
|
635
|
+
if value:
|
|
636
|
+
self._status = "connected"
|
|
637
|
+
elif self.api_key is None:
|
|
638
|
+
self._status = "unauthenticated"
|
|
639
|
+
|
|
640
|
+
def _set_status(self, status: str) -> None:
|
|
641
|
+
"""Store the current tunnel lifecycle status."""
|
|
642
|
+
with self._lock:
|
|
643
|
+
self._status = status
|
|
644
|
+
|
|
645
|
+
def _request_reconnect(self) -> None:
|
|
646
|
+
"""Trigger a reconnect cycle and close the active websocket when needed."""
|
|
647
|
+
self._set_connected(False)
|
|
648
|
+
with self._lock:
|
|
649
|
+
self._reconnect_count = 0
|
|
650
|
+
ws = self._active_ws
|
|
651
|
+
loop = self._loop
|
|
652
|
+
self._status = self._status_for_auth()
|
|
653
|
+
self._reconnect_event.set()
|
|
654
|
+
|
|
655
|
+
if ws is not None and loop is not None and loop.is_running():
|
|
656
|
+
try:
|
|
657
|
+
future = asyncio.run_coroutine_threadsafe(
|
|
658
|
+
ws.close(code=1012, reason="Client reconnect"),
|
|
659
|
+
loop,
|
|
660
|
+
)
|
|
661
|
+
future.result(timeout=2)
|
|
662
|
+
except Exception as exc: # pragma: no cover - best-effort close path
|
|
663
|
+
logger.debug("Failed to close websocket during reconnect: %s", exc)
|
|
664
|
+
|
|
665
|
+
async def _wait_for_reconnect_signal(self) -> None:
|
|
666
|
+
"""Wait until the client is reloaded, stopped, or resumed from an unauthenticated state."""
|
|
667
|
+
while not self._stop_event.is_set() and not self._reconnect_event.is_set():
|
|
668
|
+
await asyncio.sleep(0.5)
|
|
669
|
+
|
|
670
|
+
if self._reconnect_event.is_set():
|
|
671
|
+
self._reconnect_event.clear()
|
|
672
|
+
|
|
673
|
+
def _build_ssl_context(self, tunnel_url: str) -> Optional[ssl.SSLContext]:
|
|
674
|
+
"""Create an SSL context for secure websocket URLs."""
|
|
675
|
+
if not tunnel_url.startswith("wss://"):
|
|
676
|
+
return None
|
|
677
|
+
|
|
678
|
+
ssl_context = ssl.create_default_context()
|
|
679
|
+
ssl_context.set_alpn_protocols(["http/1.1"])
|
|
680
|
+
if self.verify_ssl:
|
|
681
|
+
return ssl_context
|
|
682
|
+
|
|
683
|
+
ssl_context.check_hostname = False
|
|
684
|
+
ssl_context.verify_mode = ssl.CERT_NONE
|
|
685
|
+
return ssl_context
|
|
@@ -1,477 +0,0 @@
|
|
|
1
|
-
from __future__ import annotations
|
|
2
|
-
|
|
3
|
-
import asyncio
|
|
4
|
-
import base64
|
|
5
|
-
import json
|
|
6
|
-
import logging
|
|
7
|
-
import ssl
|
|
8
|
-
import threading
|
|
9
|
-
from typing import Any, Optional
|
|
10
|
-
from urllib.parse import urljoin
|
|
11
|
-
|
|
12
|
-
import httpx
|
|
13
|
-
import websockets
|
|
14
|
-
from httpx import Timeout
|
|
15
|
-
from .._util import BaseLangeLabsClient
|
|
16
|
-
from .._util._base_client import _UNSET_API_KEY
|
|
17
|
-
from ._util import _filter_hop_by_hop_headers
|
|
18
|
-
|
|
19
|
-
logger = logging.getLogger("lange.tunnel")
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
class Tunnel(BaseLangeLabsClient):
|
|
23
|
-
"""
|
|
24
|
-
Thread-based tunnel worker client.
|
|
25
|
-
|
|
26
|
-
Connects to the tunnel runtime WebSocket endpoint and forwards incoming
|
|
27
|
-
proxy messages to a local HTTP target.
|
|
28
|
-
|
|
29
|
-
:param host: Base service URL, e.g. ``wss://example.com``.
|
|
30
|
-
:param api_key: Bearer token used for worker authentication.
|
|
31
|
-
:param target: Local HTTP target URL to forward tunnel traffic to.
|
|
32
|
-
:param verify_ssl: Whether to verify TLS certificates for ``wss://`` hosts.
|
|
33
|
-
:param max_retries: Maximum reconnect attempts, ``0`` for infinite retries.
|
|
34
|
-
:param retry_delay: Initial reconnect delay in seconds.
|
|
35
|
-
:param open_timeout: Timeout in seconds for the WebSocket opening handshake.
|
|
36
|
-
:param daemon: Whether the worker thread is daemonized.
|
|
37
|
-
:raises ValueError: If API key is provided but empty.
|
|
38
|
-
"""
|
|
39
|
-
|
|
40
|
-
def __init__(
|
|
41
|
-
self,
|
|
42
|
-
host: str = "wss://tunnel.lange-labs.com",
|
|
43
|
-
api_key: str | None = None,
|
|
44
|
-
target: str = "http://localhost:80",
|
|
45
|
-
verify_ssl: bool = True,
|
|
46
|
-
max_retries: int = 5,
|
|
47
|
-
retry_delay: float = 5.0,
|
|
48
|
-
open_timeout: float = 20.0,
|
|
49
|
-
daemon: bool = True,
|
|
50
|
-
) -> None:
|
|
51
|
-
super().__init__(api_key=api_key, daemon=daemon, host=host)
|
|
52
|
-
self.target = target.rstrip("/")
|
|
53
|
-
self.verify_ssl = verify_ssl
|
|
54
|
-
self.max_retries = max_retries
|
|
55
|
-
self.retry_delay = retry_delay
|
|
56
|
-
self.open_timeout = open_timeout
|
|
57
|
-
|
|
58
|
-
self._max_retry_delay = 60.0
|
|
59
|
-
self._stop_event = threading.Event()
|
|
60
|
-
self._reconnect_event = threading.Event()
|
|
61
|
-
self._connected = False
|
|
62
|
-
self._remote_address: Optional[str] = None
|
|
63
|
-
self._remote_address_roundrobin: Optional[str] = None
|
|
64
|
-
self._worker_index = -1
|
|
65
|
-
self._pool_size = 0
|
|
66
|
-
self._reconnect_count = 0
|
|
67
|
-
self._lock = threading.Lock()
|
|
68
|
-
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
|
69
|
-
self._active_ws: Any = None
|
|
70
|
-
|
|
71
|
-
@property
|
|
72
|
-
def connected(self) -> bool:
|
|
73
|
-
"""
|
|
74
|
-
Get current connection state.
|
|
75
|
-
|
|
76
|
-
:returns: ``True`` when a worker socket is currently active.
|
|
77
|
-
"""
|
|
78
|
-
with self._lock:
|
|
79
|
-
return self._connected
|
|
80
|
-
|
|
81
|
-
@property
|
|
82
|
-
def status(self) -> str:
|
|
83
|
-
"""
|
|
84
|
-
Get the current tunnel lifecycle status.
|
|
85
|
-
|
|
86
|
-
:returns: Shared client status string.
|
|
87
|
-
"""
|
|
88
|
-
with self._lock:
|
|
89
|
-
return self._status
|
|
90
|
-
|
|
91
|
-
@property
|
|
92
|
-
def remote_address(self) -> Optional[str]:
|
|
93
|
-
"""
|
|
94
|
-
Get worker-specific public address when provided by the server.
|
|
95
|
-
|
|
96
|
-
:returns: Worker URL or ``None`` when unavailable.
|
|
97
|
-
"""
|
|
98
|
-
with self._lock:
|
|
99
|
-
return self._remote_address
|
|
100
|
-
|
|
101
|
-
@property
|
|
102
|
-
def remote_address_roundrobin(self) -> Optional[str]:
|
|
103
|
-
"""
|
|
104
|
-
Get round-robin public address when provided by the server.
|
|
105
|
-
|
|
106
|
-
:returns: Pool URL or ``None`` when unavailable.
|
|
107
|
-
"""
|
|
108
|
-
with self._lock:
|
|
109
|
-
return self._remote_address_roundrobin
|
|
110
|
-
|
|
111
|
-
@property
|
|
112
|
-
def worker_index(self) -> int:
|
|
113
|
-
"""
|
|
114
|
-
Get current worker index from the latest welcome payload.
|
|
115
|
-
|
|
116
|
-
:returns: Worker index, ``-1`` when disconnected.
|
|
117
|
-
"""
|
|
118
|
-
with self._lock:
|
|
119
|
-
return self._worker_index
|
|
120
|
-
|
|
121
|
-
@property
|
|
122
|
-
def pool_size(self) -> int:
|
|
123
|
-
"""
|
|
124
|
-
Get current worker pool size from the latest welcome payload.
|
|
125
|
-
|
|
126
|
-
:returns: Worker pool size, ``0`` when disconnected.
|
|
127
|
-
"""
|
|
128
|
-
with self._lock:
|
|
129
|
-
return self._pool_size
|
|
130
|
-
|
|
131
|
-
@property
|
|
132
|
-
def reconnect_count(self) -> int:
|
|
133
|
-
"""
|
|
134
|
-
Get number of reconnect attempts after the last successful connection.
|
|
135
|
-
|
|
136
|
-
:returns: Reconnect counter.
|
|
137
|
-
"""
|
|
138
|
-
with self._lock:
|
|
139
|
-
return self._reconnect_count
|
|
140
|
-
|
|
141
|
-
def run(self) -> None:
|
|
142
|
-
"""
|
|
143
|
-
Start the worker loop inside the thread.
|
|
144
|
-
|
|
145
|
-
:returns: ``None``.
|
|
146
|
-
"""
|
|
147
|
-
asyncio.run(self._run_async())
|
|
148
|
-
|
|
149
|
-
def stop(self) -> None:
|
|
150
|
-
"""
|
|
151
|
-
Request a graceful shutdown.
|
|
152
|
-
|
|
153
|
-
:returns: ``None``.
|
|
154
|
-
"""
|
|
155
|
-
self._stop_event.set()
|
|
156
|
-
|
|
157
|
-
def reconnect(self) -> None:
|
|
158
|
-
"""
|
|
159
|
-
Force a full reconnect cycle.
|
|
160
|
-
|
|
161
|
-
:returns: ``None``.
|
|
162
|
-
"""
|
|
163
|
-
self._request_reconnect()
|
|
164
|
-
|
|
165
|
-
def reload(self, api_key: str | None | object = _UNSET_API_KEY) -> None:
|
|
166
|
-
"""
|
|
167
|
-
Reload tunnel authentication and restart the connection flow when running.
|
|
168
|
-
|
|
169
|
-
:param api_key: Optional replacement API key. Pass ``None`` explicitly to clear authentication.
|
|
170
|
-
:returns: ``None``.
|
|
171
|
-
"""
|
|
172
|
-
super().reload(api_key=api_key)
|
|
173
|
-
self._set_connected(False)
|
|
174
|
-
|
|
175
|
-
if self.is_alive() or self._loop is not None:
|
|
176
|
-
self._request_reconnect()
|
|
177
|
-
return
|
|
178
|
-
|
|
179
|
-
self._set_status(self._status_for_auth())
|
|
180
|
-
|
|
181
|
-
async def _run_async(self) -> None:
|
|
182
|
-
"""
|
|
183
|
-
Run the worker connection loop with reconnect backoff.
|
|
184
|
-
|
|
185
|
-
:returns: ``None``.
|
|
186
|
-
"""
|
|
187
|
-
tunnel_url = self._build_tunnel_url()
|
|
188
|
-
ssl_context = self._build_ssl_context(tunnel_url)
|
|
189
|
-
|
|
190
|
-
with self._lock:
|
|
191
|
-
self._loop = asyncio.get_running_loop()
|
|
192
|
-
|
|
193
|
-
current_delay = self.retry_delay
|
|
194
|
-
|
|
195
|
-
while not self._stop_event.is_set():
|
|
196
|
-
if self.api_key is None:
|
|
197
|
-
self._set_connected(False)
|
|
198
|
-
self._set_status("unauthenticated")
|
|
199
|
-
await self._wait_for_reconnect_signal()
|
|
200
|
-
continue
|
|
201
|
-
|
|
202
|
-
self._set_status("pending")
|
|
203
|
-
headers = self._build_connection_headers()
|
|
204
|
-
|
|
205
|
-
try:
|
|
206
|
-
logger.info("Connecting to %s", tunnel_url)
|
|
207
|
-
async with websockets.connect(
|
|
208
|
-
tunnel_url,
|
|
209
|
-
additional_headers=headers,
|
|
210
|
-
ssl=ssl_context,
|
|
211
|
-
proxy=None,
|
|
212
|
-
open_timeout=self.open_timeout,
|
|
213
|
-
) as ws:
|
|
214
|
-
with self._lock:
|
|
215
|
-
self._active_ws = ws
|
|
216
|
-
|
|
217
|
-
current_delay = self.retry_delay
|
|
218
|
-
await self._consume_welcome(ws)
|
|
219
|
-
|
|
220
|
-
with self._lock:
|
|
221
|
-
self._reconnect_count = 0
|
|
222
|
-
|
|
223
|
-
try:
|
|
224
|
-
await self._handle_messages(ws)
|
|
225
|
-
finally:
|
|
226
|
-
with self._lock:
|
|
227
|
-
self._active_ws = None
|
|
228
|
-
self._set_connected(False)
|
|
229
|
-
except websockets.exceptions.ConnectionClosed as exc:
|
|
230
|
-
self._set_connected(False)
|
|
231
|
-
self._set_status("failed")
|
|
232
|
-
logger.warning("Connection closed: %s", exc)
|
|
233
|
-
except Exception as exc: # pragma: no cover - network error path
|
|
234
|
-
self._set_connected(False)
|
|
235
|
-
self._set_status("failed")
|
|
236
|
-
logger.error("Connection error: %s", exc)
|
|
237
|
-
finally:
|
|
238
|
-
with self._lock:
|
|
239
|
-
self._active_ws = None
|
|
240
|
-
|
|
241
|
-
if self._stop_event.is_set():
|
|
242
|
-
break
|
|
243
|
-
|
|
244
|
-
if self._reconnect_event.is_set():
|
|
245
|
-
self._set_connected(False)
|
|
246
|
-
with self._lock:
|
|
247
|
-
self._reconnect_count = 0
|
|
248
|
-
self._reconnect_event.clear()
|
|
249
|
-
current_delay = self.retry_delay
|
|
250
|
-
self._set_status(self._status_for_auth())
|
|
251
|
-
logger.info("Manual reconnect requested. Reconnecting now.")
|
|
252
|
-
continue
|
|
253
|
-
|
|
254
|
-
with self._lock:
|
|
255
|
-
self._reconnect_count += 1
|
|
256
|
-
attempt = self._reconnect_count
|
|
257
|
-
|
|
258
|
-
if self.max_retries > 0 and attempt > self.max_retries:
|
|
259
|
-
self._set_status("failed")
|
|
260
|
-
logger.error("Max reconnection attempts (%s) reached. Giving up.", self.max_retries)
|
|
261
|
-
break
|
|
262
|
-
|
|
263
|
-
logger.info("Reconnecting in %.1fs (attempt %s)", current_delay, attempt)
|
|
264
|
-
|
|
265
|
-
wait_time = 0.0
|
|
266
|
-
while (
|
|
267
|
-
wait_time < current_delay
|
|
268
|
-
and not self._stop_event.is_set()
|
|
269
|
-
and not self._reconnect_event.is_set()
|
|
270
|
-
):
|
|
271
|
-
await asyncio.sleep(0.5)
|
|
272
|
-
wait_time += 0.5
|
|
273
|
-
|
|
274
|
-
if self._reconnect_event.is_set():
|
|
275
|
-
continue
|
|
276
|
-
|
|
277
|
-
current_delay = min(current_delay * 2, self._max_retry_delay)
|
|
278
|
-
|
|
279
|
-
with self._lock:
|
|
280
|
-
self._active_ws = None
|
|
281
|
-
self._loop = None
|
|
282
|
-
if self._status in {"connected", "pending"}:
|
|
283
|
-
self._status = self._status_for_auth()
|
|
284
|
-
|
|
285
|
-
logger.info("Disconnected from server.")
|
|
286
|
-
|
|
287
|
-
async def _consume_welcome(self, ws: Any) -> None:
|
|
288
|
-
"""
|
|
289
|
-
Read and process the initial welcome message when present.
|
|
290
|
-
|
|
291
|
-
:param ws: Active websocket client.
|
|
292
|
-
:returns: ``None``.
|
|
293
|
-
"""
|
|
294
|
-
try:
|
|
295
|
-
welcome_message = await ws.recv()
|
|
296
|
-
welcome = json.loads(welcome_message)
|
|
297
|
-
except Exception:
|
|
298
|
-
self._set_connected(True)
|
|
299
|
-
return
|
|
300
|
-
|
|
301
|
-
if isinstance(welcome, dict) and welcome.get("type") == "welcome":
|
|
302
|
-
self._set_connected(
|
|
303
|
-
True,
|
|
304
|
-
remote_address=str(welcome.get("public_address") or "") or None,
|
|
305
|
-
remote_address_roundrobin=str(welcome.get("public_address_generic") or "") or None,
|
|
306
|
-
worker_index=int(welcome.get("worker_index", -1)),
|
|
307
|
-
pool_size=int(welcome.get("pool_size", 0)),
|
|
308
|
-
)
|
|
309
|
-
return
|
|
310
|
-
|
|
311
|
-
self._set_connected(True)
|
|
312
|
-
|
|
313
|
-
async def _handle_messages(self, ws: Any) -> None:
|
|
314
|
-
"""
|
|
315
|
-
Handle request/response proxy messages over the worker websocket.
|
|
316
|
-
|
|
317
|
-
:param ws: Active websocket client.
|
|
318
|
-
:returns: ``None``.
|
|
319
|
-
"""
|
|
320
|
-
async with httpx.AsyncClient(timeout=Timeout(timeout=15.0)) as http_client:
|
|
321
|
-
while not self._stop_event.is_set() and not self._reconnect_event.is_set():
|
|
322
|
-
try:
|
|
323
|
-
message = await asyncio.wait_for(ws.recv(), timeout=1.0)
|
|
324
|
-
request = json.loads(message)
|
|
325
|
-
response = await self._forward_request(http_client, request)
|
|
326
|
-
await ws.send(json.dumps(response))
|
|
327
|
-
except asyncio.TimeoutError:
|
|
328
|
-
continue
|
|
329
|
-
except websockets.exceptions.ConnectionClosed:
|
|
330
|
-
break
|
|
331
|
-
except Exception as exc:
|
|
332
|
-
logger.error("Error handling message: %s", exc)
|
|
333
|
-
|
|
334
|
-
async def _forward_request(self, client: Any, request: dict[str, Any]) -> dict[str, Any]:
|
|
335
|
-
"""
|
|
336
|
-
Forward a single proxied request to the configured local target.
|
|
337
|
-
|
|
338
|
-
:param client: HTTP client supporting ``request``.
|
|
339
|
-
:param request: Incoming tunnel request payload.
|
|
340
|
-
:returns: Tunnel response payload.
|
|
341
|
-
"""
|
|
342
|
-
request_id = str(request.get("id", ""))
|
|
343
|
-
method = str(request.get("method", "GET"))
|
|
344
|
-
path = str(request.get("path", "/"))
|
|
345
|
-
headers = request.get("headers", {})
|
|
346
|
-
body_b64 = str(request.get("body", ""))
|
|
347
|
-
|
|
348
|
-
body = base64.b64decode(body_b64) if body_b64 else None
|
|
349
|
-
|
|
350
|
-
url = urljoin(f"{self.target}/", path.lstrip("/"))
|
|
351
|
-
filtered_headers = _filter_hop_by_hop_headers(headers)
|
|
352
|
-
|
|
353
|
-
try:
|
|
354
|
-
response = await client.request(
|
|
355
|
-
method=method,
|
|
356
|
-
url=url,
|
|
357
|
-
headers=filtered_headers,
|
|
358
|
-
content=body,
|
|
359
|
-
follow_redirects=True,
|
|
360
|
-
)
|
|
361
|
-
return {
|
|
362
|
-
"id": request_id,
|
|
363
|
-
"status": response.status_code,
|
|
364
|
-
"headers": dict(response.headers.items()),
|
|
365
|
-
"body": base64.b64encode(response.content).decode("utf-8"),
|
|
366
|
-
}
|
|
367
|
-
except Exception as exc:
|
|
368
|
-
logger.error("Error forwarding request: %s", exc)
|
|
369
|
-
error_body = json.dumps({"error": str(exc)}).encode("utf-8")
|
|
370
|
-
return {
|
|
371
|
-
"id": request_id,
|
|
372
|
-
"status": 502,
|
|
373
|
-
"headers": {"Content-Type": "application/json"},
|
|
374
|
-
"body": base64.b64encode(error_body).decode("utf-8"),
|
|
375
|
-
}
|
|
376
|
-
|
|
377
|
-
def _build_tunnel_url(self) -> str:
|
|
378
|
-
"""
|
|
379
|
-
Build the worker websocket URL.
|
|
380
|
-
|
|
381
|
-
:returns: WebSocket endpoint URL.
|
|
382
|
-
"""
|
|
383
|
-
return f"{self.host}/api/tunnels/connection"
|
|
384
|
-
|
|
385
|
-
def _set_connected(
|
|
386
|
-
self,
|
|
387
|
-
value: bool,
|
|
388
|
-
remote_address: Optional[str] = None,
|
|
389
|
-
remote_address_roundrobin: Optional[str] = None,
|
|
390
|
-
worker_index: int = -1,
|
|
391
|
-
pool_size: int = 0,
|
|
392
|
-
) -> None:
|
|
393
|
-
"""
|
|
394
|
-
Update connection state atomically.
|
|
395
|
-
|
|
396
|
-
:param value: New connected state.
|
|
397
|
-
:param remote_address: Worker URL from welcome message.
|
|
398
|
-
:param remote_address_roundrobin: Pool URL from welcome message.
|
|
399
|
-
:param worker_index: Worker index from welcome message.
|
|
400
|
-
:param pool_size: Pool size from welcome message.
|
|
401
|
-
:returns: ``None``.
|
|
402
|
-
"""
|
|
403
|
-
with self._lock:
|
|
404
|
-
self._connected = value
|
|
405
|
-
self._remote_address = remote_address if value else None
|
|
406
|
-
self._remote_address_roundrobin = remote_address_roundrobin if value else None
|
|
407
|
-
self._worker_index = worker_index if value else -1
|
|
408
|
-
self._pool_size = pool_size if value else 0
|
|
409
|
-
if value:
|
|
410
|
-
self._status = "connected"
|
|
411
|
-
elif self.api_key is None:
|
|
412
|
-
self._status = "unauthenticated"
|
|
413
|
-
|
|
414
|
-
def _set_status(self, status: str) -> None:
|
|
415
|
-
"""
|
|
416
|
-
Store the current tunnel lifecycle status.
|
|
417
|
-
|
|
418
|
-
:param status: New lifecycle status.
|
|
419
|
-
:returns: ``None``.
|
|
420
|
-
"""
|
|
421
|
-
with self._lock:
|
|
422
|
-
self._status = status
|
|
423
|
-
|
|
424
|
-
def _request_reconnect(self) -> None:
|
|
425
|
-
"""
|
|
426
|
-
Trigger a reconnect cycle and close the active websocket when needed.
|
|
427
|
-
|
|
428
|
-
:returns: ``None``.
|
|
429
|
-
"""
|
|
430
|
-
self._set_connected(False)
|
|
431
|
-
with self._lock:
|
|
432
|
-
self._reconnect_count = 0
|
|
433
|
-
ws = self._active_ws
|
|
434
|
-
loop = self._loop
|
|
435
|
-
self._status = self._status_for_auth()
|
|
436
|
-
self._reconnect_event.set()
|
|
437
|
-
|
|
438
|
-
if ws is not None and loop is not None and loop.is_running():
|
|
439
|
-
try:
|
|
440
|
-
future = asyncio.run_coroutine_threadsafe(
|
|
441
|
-
ws.close(code=1012, reason="Client reconnect"),
|
|
442
|
-
loop,
|
|
443
|
-
)
|
|
444
|
-
future.result(timeout=2)
|
|
445
|
-
except Exception as exc: # pragma: no cover - best-effort close path
|
|
446
|
-
logger.debug("Failed to close websocket during reconnect: %s", exc)
|
|
447
|
-
|
|
448
|
-
async def _wait_for_reconnect_signal(self) -> None:
|
|
449
|
-
"""
|
|
450
|
-
Wait until the client is reloaded, stopped, or resumed from an unauthenticated state.
|
|
451
|
-
|
|
452
|
-
:returns: ``None``.
|
|
453
|
-
"""
|
|
454
|
-
while not self._stop_event.is_set() and not self._reconnect_event.is_set():
|
|
455
|
-
await asyncio.sleep(0.5)
|
|
456
|
-
|
|
457
|
-
if self._reconnect_event.is_set():
|
|
458
|
-
self._reconnect_event.clear()
|
|
459
|
-
|
|
460
|
-
def _build_ssl_context(self, tunnel_url: str) -> Optional[ssl.SSLContext]:
|
|
461
|
-
"""
|
|
462
|
-
Create SSL context for secure websocket URLs.
|
|
463
|
-
|
|
464
|
-
:param tunnel_url: Resolved websocket URL.
|
|
465
|
-
:returns: SSL context for ``wss`` or ``None`` for plain ``ws``.
|
|
466
|
-
"""
|
|
467
|
-
if not tunnel_url.startswith("wss://"):
|
|
468
|
-
return None
|
|
469
|
-
|
|
470
|
-
ssl_context = ssl.create_default_context()
|
|
471
|
-
ssl_context.set_alpn_protocols(["http/1.1"])
|
|
472
|
-
if self.verify_ssl:
|
|
473
|
-
return ssl_context
|
|
474
|
-
|
|
475
|
-
ssl_context.check_hostname = False
|
|
476
|
-
ssl_context.verify_mode = ssl.CERT_NONE
|
|
477
|
-
return ssl_context
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|