aboutsummaryrefslogtreecommitdiffstats
path: root/proxy/websockets.py
diff options
context:
space:
mode:
authorGravatar Tom van der Lee <tom@vanderlee.io>2026-09-04 17:09:24 +0200
committerGravatar Tom van der Lee <tom@vanderlee.io>2026-09-04 17:09:24 +0200
commit6b6f20a032eb31d747366c8bc270ca81731eaf6a (patch)
treeada48080495267bebde4ee1251bbbe77d629745f /proxy/websockets.py
parent2d22158dd95f76399f5716d556be60c77da04a46 (diff)
downloadserver-6b6f20a032eb31d747366c8bc270ca81731eaf6a.tar.gz
server-6b6f20a032eb31d747366c8bc270ca81731eaf6a.tar.bz2
server-6b6f20a032eb31d747366c8bc270ca81731eaf6a.zip
Fixed pre-commit issuesv3
Diffstat (limited to 'proxy/websockets.py')
-rw-r--r--proxy/websockets.py48
1 files changed, 31 insertions, 17 deletions
diff --git a/proxy/websockets.py b/proxy/websockets.py
index ec65e6d..a09e261 100644
--- a/proxy/websockets.py
+++ b/proxy/websockets.py
@@ -9,17 +9,22 @@ from uuid import uuid4
9from starlette.endpoints import WebSocketEndpoint 9from starlette.endpoints import WebSocketEndpoint
10from starlette.websockets import WebSocket 10from starlette.websockets import WebSocket
11 11
12from proxy.utils import get_path_with_query_string
13from proxy.queue import ProxyQueue 12from proxy.queue import ProxyQueue
14from ttun_server.types import WebsocketMessage, WebsocketMessageType, WebsocketConnectData, WebsocketMessageData, \ 13from proxy.utils import get_path_with_query_string
15 WebsocketDisconnectData 14from ttun_server.types import (
15 WebsocketConnectData,
16 WebsocketDisconnectData,
17 WebsocketMessage,
18 WebsocketMessageData,
19 WebsocketMessageType,
20)
16 21
17logger = logging.getLogger(__name__) 22logger = logging.getLogger(__name__)
18logger.setLevel('DEBUG') 23logger.setLevel("DEBUG")
19 24
20 25
21class WebsocketProxy(WebSocketEndpoint): 26class WebsocketProxy(WebSocketEndpoint):
22 encoding = 'json' 27 encoding = "json"
23 websocket_listen_task = None 28 websocket_listen_task = None
24 29
25 def __init__(self, *args, **kwargs): 30 def __init__(self, *args, **kwargs):
@@ -28,16 +33,20 @@ class WebsocketProxy(WebSocketEndpoint):
28 33
29 @asynccontextmanager 34 @asynccontextmanager
30 async def proxy(self, websocket: WebSocket, message: WebsocketMessage): 35 async def proxy(self, websocket: WebSocket, message: WebsocketMessage):
31 [subdomain, *_] = websocket.url.hostname.split('.') 36 [subdomain, *_] = websocket.url.hostname.split(".")
32 37
33 expect_ack = WebsocketMessageType(message['type']) == WebsocketMessageType.connect 38 expect_ack = (
39 WebsocketMessageType(message["type"]) == WebsocketMessageType.connect
40 )
34 41
35 try: 42 try:
36 request_queue = await ProxyQueue.get_for_identifier(subdomain) 43 request_queue = await ProxyQueue.get_for_identifier(subdomain)
37 await request_queue.enqueue(message) 44 await request_queue.enqueue(message)
38 45
39 if expect_ack: 46 if expect_ack:
40 response_queue = await ProxyQueue.create_for_identifier(message["identifier"]) 47 response_queue = await ProxyQueue.create_for_identifier(
48 message["identifier"]
49 )
41 yield await response_queue.dequeue() 50 yield await response_queue.dequeue()
42 await response_queue.delete() 51 await response_queue.delete()
43 else: 52 else:
@@ -51,7 +60,9 @@ class WebsocketProxy(WebSocketEndpoint):
51 while True: 60 while True:
52 message: WebsocketMessage = await response_queue.dequeue() 61 message: WebsocketMessage = await response_queue.dequeue()
53 logger.debug(message) 62 logger.debug(message)
54 await websocket.send_text(b64decode(message['payload']['body'].encode()).decode()) 63 await websocket.send_text(
64 b64decode(message["payload"]["body"].encode()).decode()
65 )
55 66
56 async def on_connect(self, websocket: WebSocket) -> None: 67 async def on_connect(self, websocket: WebSocket) -> None:
57 message = WebsocketMessage( 68 message = WebsocketMessage(
@@ -60,18 +71,21 @@ class WebsocketProxy(WebSocketEndpoint):
60 payload=WebsocketConnectData( 71 payload=WebsocketConnectData(
61 path=get_path_with_query_string(websocket), 72 path=get_path_with_query_string(websocket),
62 headers=[ 73 headers=[
63 (k.decode(), v.decode()) 74 (k.decode(), v.decode()) for k, v in websocket.scope["headers"]
64 for k, v
65 in websocket.scope['headers']
66 ], 75 ],
67 ) 76 ),
68 ) 77 )
69 78
70 async with self.proxy(websocket, message) as m: 79 async with self.proxy(websocket, message) as m:
71 if m is not None and WebsocketMessageType(m['type']) == WebsocketMessageType.ack: 80 if (
81 m is not None
82 and WebsocketMessageType(m["type"]) == WebsocketMessageType.ack
83 ):
72 await super().on_connect(websocket) 84 await super().on_connect(websocket)
73 85
74 self.websocket_listen_task = asyncio.create_task(self.listen_for_messages(websocket)) 86 self.websocket_listen_task = asyncio.create_task(
87 self.listen_for_messages(websocket)
88 )
75 89
76 def callback(*args, **kwargs): 90 def callback(*args, **kwargs):
77 self.websocket_listen_task = None 91 self.websocket_listen_task = None
@@ -92,7 +106,7 @@ class WebsocketProxy(WebSocketEndpoint):
92 identifier=self.id, 106 identifier=self.id,
93 payload=WebsocketMessageData( 107 payload=WebsocketMessageData(
94 body=b64encode(data_bytes).decode(), 108 body=b64encode(data_bytes).decode(),
95 ) 109 ),
96 ) 110 )
97 111
98 async with self.proxy(websocket, message): 112 async with self.proxy(websocket, message):
@@ -104,7 +118,7 @@ class WebsocketProxy(WebSocketEndpoint):
104 identifier=self.id, 118 identifier=self.id,
105 payload=WebsocketDisconnectData( 119 payload=WebsocketDisconnectData(
106 close_code=close_code, 120 close_code=close_code,
107 ) 121 ),
108 ) 122 )
109 123
110 async with self.proxy(websocket, message): 124 async with self.proxy(websocket, message):