1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
|
import asyncio
import logging
import os
from uuid import uuid4
from fastapi import WebSocket, WebSocketDisconnect
from starlette.authentication import requires
import conf
import ttun_server
from proxy.queue import ProxyQueue
from ttun_server.types import (
Config,
Message,
MessageType,
)
logger = logging.getLogger(__name__)
logger.setLevel('DEBUG')
async def assert_compatible_version(websocket: WebSocket, config: Config) -> None:
client_version = config.get('version', '1.0.0')
logger.debug('client_version %s', client_version)
if 'git' not in client_version and ttun_server.__version__ != 'development':
[client_major, *_] = [int(i) for i in client_version.split('.')[:3]]
[server_major, *_] = [int(i) for i in ttun_server.__version__.split('.')]
if client_major < server_major:
await websocket.close(4000, 'Your client is too old')
if client_major > server_major:
await websocket.close(4001, 'Your client is too new')
@requires(['authenticated', 'token'] if conf.ENABLE_AUTH else [])
async def tunnel(websocket: WebSocket) -> None:
request_tasks: dict[str, asyncio.Task] = {}
proxy_queues: dict[str, ProxyQueue] = {}
await websocket.accept()
config: Config = await websocket.receive_json()
await assert_compatible_version(websocket, config)
if 'subdomains' not in config:
config['subdomains'] = [config['subdomain']]
elif config['subdomains'] is None:
config['subdomains'] = [None]
for i, subdomain in enumerate(config['subdomains']):
if subdomain is None or await ProxyQueue.has_connection(subdomain):
config['subdomains'][i] = uuid4().hex
for subdomain in config['subdomains']:
proxy_queues[subdomain] = await ProxyQueue.create_for_identifier(subdomain)
hostname = os.environ.get('TUNNEL_DOMAIN')
protocol = 'https' if os.environ.get('SECURE', False) else 'http'
urls = [
f'{protocol}://{subdomain}.{hostname}'
for subdomain in config['subdomains']
]
await websocket.send_json({
'url': urls[0],
'urls': urls,
})
async def handle_requests(subdomain: str) -> None:
while request := await proxy_queues[subdomain].dequeue():
asyncio.create_task(websocket.send_json(request), name=request['identifier'])
for subdomain in config['subdomains']:
request_tasks[subdomain] = asyncio.create_task(
handle_requests(subdomain), name=subdomain
)
try:
while True:
data: Message = await websocket.receive_json()
try:
data['type'] = MessageType(data['type']).value
response_queue = await ProxyQueue.get_for_identifier(data['identifier'])
await response_queue.enqueue(data)
except AssertionError:
pass
except WebSocketDisconnect:
for proxy_queue in proxy_queues.values():
await proxy_queue.delete()
for request_task in request_tasks.values():
request_task.cancel()
|
