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
94
95
96
97
|
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") 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()
|
