diff options
Diffstat (limited to 'ttun_server/endpoints.py')
| -rw-r--r-- | ttun_server/endpoints.py | 22 |
1 files changed, 20 insertions, 2 deletions
diff --git a/ttun_server/endpoints.py b/ttun_server/endpoints.py index 6728c31..3e263da 100644 --- a/ttun_server/endpoints.py +++ b/ttun_server/endpoints.py | |||
| @@ -1,6 +1,7 @@ | |||
| 1 | import asyncio | 1 | import asyncio |
| 2 | import logging | 2 | import logging |
| 3 | import os | 3 | import os |
| 4 | from asyncio import create_task | ||
| 4 | from base64 import b64decode, b64encode | 5 | from base64 import b64decode, b64encode |
| 5 | from typing import Optional, Any | 6 | from typing import Optional, Any |
| 6 | from uuid import uuid4 | 7 | from uuid import uuid4 |
| @@ -11,6 +12,7 @@ from starlette.responses import Response | |||
| 11 | from starlette.types import Scope, Receive, Send | 12 | from starlette.types import Scope, Receive, Send |
| 12 | from starlette.websockets import WebSocket | 13 | from starlette.websockets import WebSocket |
| 13 | 14 | ||
| 15 | import ttun_server | ||
| 14 | from ttun_server.proxy_queue import ProxyQueue | 16 | from ttun_server.proxy_queue import ProxyQueue |
| 15 | from ttun_server.types import RequestData, Config, Message, MessageType | 17 | from ttun_server.types import RequestData, Config, Message, MessageType |
| 16 | 18 | ||
| @@ -40,9 +42,10 @@ class Proxy(HTTPEndpoint): | |||
| 40 | 42 | ||
| 41 | request_queue = await ProxyQueue.get_for_identifier(subdomain) | 43 | request_queue = await ProxyQueue.get_for_identifier(subdomain) |
| 42 | 44 | ||
| 45 | logger.debug('PROXY %s%s ', subdomain, request.url) | ||
| 43 | await request_queue.enqueue( | 46 | await request_queue.enqueue( |
| 44 | Message( | 47 | Message( |
| 45 | type=MessageType.request, | 48 | type=MessageType.request.value, |
| 46 | identifier=identifier, | 49 | identifier=identifier, |
| 47 | payload= | 50 | payload= |
| 48 | RequestData( | 51 | RequestData( |
| @@ -85,16 +88,31 @@ class Tunnel(WebSocketEndpoint): | |||
| 85 | 88 | ||
| 86 | async def handle_requests(self, websocket: WebSocket): | 89 | async def handle_requests(self, websocket: WebSocket): |
| 87 | while request := await self.proxy_queue.dequeue(): | 90 | while request := await self.proxy_queue.dequeue(): |
| 88 | await websocket.send_json(request) | 91 | create_task(websocket.send_json(request)) |
| 89 | 92 | ||
| 90 | async def on_connect(self, websocket: WebSocket) -> None: | 93 | async def on_connect(self, websocket: WebSocket) -> None: |
| 91 | await websocket.accept() | 94 | await websocket.accept() |
| 92 | self.config = await websocket.receive_json() | 95 | self.config = await websocket.receive_json() |
| 93 | 96 | ||
| 97 | client_version = self.config.get('version', '1.0.0') | ||
| 98 | logger.debug('client_version %s', client_version) | ||
| 99 | |||
| 100 | if 'git' not in client_version and ttun_server.__version__ != 'development': | ||
| 101 | [client_major, *_] = [int(i) for i in client_version.split('.')[:3]] | ||
| 102 | [server_major, *_] = [int(i) for i in ttun_server.__version__.split('.')] | ||
| 103 | |||
| 104 | if client_major < server_major: | ||
| 105 | await websocket.close(4000, 'Your client is too old') | ||
| 106 | |||
| 107 | if client_major > server_major: | ||
| 108 | await websocket.close(4001, 'Your client is too new') | ||
| 109 | |||
| 110 | |||
| 94 | if self.config['subdomain'] is None \ | 111 | if self.config['subdomain'] is None \ |
| 95 | or await ProxyQueue.has_connection(self.config['subdomain']): | 112 | or await ProxyQueue.has_connection(self.config['subdomain']): |
| 96 | self.config['subdomain'] = uuid4().hex | 113 | self.config['subdomain'] = uuid4().hex |
| 97 | 114 | ||
| 115 | |||
| 98 | self.proxy_queue = await ProxyQueue.create_for_identifier(self.config['subdomain']) | 116 | self.proxy_queue = await ProxyQueue.create_for_identifier(self.config['subdomain']) |
| 99 | 117 | ||
| 100 | hostname = os.environ.get("TUNNEL_DOMAIN") | 118 | hostname = os.environ.get("TUNNEL_DOMAIN") |
