Added scaling support via redisv1.1.0-rc1

6 files changed, 206 insertions, 30 deletions

diff --git a/ttun_server/__init__.py b/ttun_server/__init__.py
index b8fd114..cf589cc 100644
--- a/ttun_server/__init__.py
+++ b/ttun_server/__init__.py
@@ -1,8 +1,13 @@
+import logging
+import os
+
 from starlette.applications import Starlette
 from starlette.routing import Route, WebSocketRoute
 
 from ttun_server.endpoints import Proxy, Tunnel
 
+logging.basicConfig(level=getattr(logging, os.environ.get('LOG_LEVEL', 'INFO')))
+
 server = Starlette(
     debug=True,
     routes=[
diff --git a/ttun_server/connections.py b/ttun_server/connections.py
deleted file mode 100644
index a8dabcf..0000000
--- a/ttun_server/connections.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from ttun_server.types import Connection
-
-connections: dict[str, Connection] = {}
diff --git a/ttun_server/endpoints.py b/ttun_server/endpoints.py
index d59cb7c..5b9e57f 100644
--- a/ttun_server/endpoints.py
+++ b/ttun_server/endpoints.py
@@ -1,4 +1,5 @@
 import asyncio
+import logging
 import os
 from asyncio import Queue
 from base64 import b64decode, b64encode
@@ -11,9 +12,10 @@ from starlette.responses import Response
 from starlette.types import Scope, Receive, Send
 from starlette.websockets import WebSocket
 
-from ttun_server.types import Connection, RequestData, Config, ResponseData
+from ttun_server.proxy_queue import ProxyQueue
+from ttun_server.types import RequestData, Config, ResponseData
 
-from ttun_server.connections import connections
+logger = logging.getLogger(__name__)
 
 
 class Proxy(HTTPEndpoint):
@@ -23,10 +25,10 @@ class Proxy(HTTPEndpoint):
         [subdomain, *_] = request.headers['host'].split('.')
         response = Response(content='Not Found', status_code=404)
 
-        if subdomain in connections:
-            connection = connections[subdomain]
+        try:
+            queue = await ProxyQueue.get_for_identifier(subdomain)
 
-            await connection['requests'].put(RequestData(
+            await queue.send_request(RequestData(
                 method=request.method,
                 path=str(request.url).replace(str(request.base_url), '/'),
                 headers=dict(request.headers),
@@ -34,12 +36,14 @@ class Proxy(HTTPEndpoint):
                 body=b64encode(await request.body()).decode()
             ))
 
-            _response = await connection['responses'].get()
+            _response = await queue.handle_response()
             response = Response(
                 status_code=_response['status'],
                 headers=_response['headers'],
                 content=b64decode(_response['body'].encode())
             )
+        except AssertionError:
+            pass
 
         await response(self.scope, self.receive, self.send)
 
@@ -52,16 +56,8 @@ class Tunnel(WebSocketEndpoint):
         self.request_task = None
         self.config: Optional[Config] = None
 
-    @property
-    def requests(self) -> Queue[RequestData]:
-        return connections[self.config['subdomain']]['requests']
-
-    @property
-    def responses(self) -> Queue[ResponseData]:
-        return connections[self.config['subdomain']]['responses']
-
     async def handle_requests(self, websocket: WebSocket):
-        while request := await self.requests.get():
+        while request := await self.proxy_queue.handle_request():
             await websocket.send_json(request)
 
     async def on_connect(self, websocket: WebSocket) -> None:
@@ -69,14 +65,10 @@ class Tunnel(WebSocketEndpoint):
         self.config = await websocket.receive_json()
 
         if self.config['subdomain'] is None \
-                or self.config['subdomain'] in connections:
+                or await ProxyQueue.has_connection(self.config['subdomain']):
             self.config['subdomain'] = uuid4().hex
 
-
-        connections[self.config['subdomain']] = Connection(
-            requests=Queue(),
-            responses=Queue(),
-        )
+        self.proxy_queue = await ProxyQueue.create_for_identifier(self.config['subdomain'])
 
         hostname = os.environ.get("TUNNEL_DOMAIN")
         protocol = "https" if os.environ.get("SECURE", False) else "http"
@@ -87,12 +79,11 @@ class Tunnel(WebSocketEndpoint):
 
         self.request_task = asyncio.create_task(self.handle_requests(websocket))
 
-    async def on_receive(self, websocket: WebSocket, data: Any) -> None:
-        await self.responses.put(data)
+    async def on_receive(self, websocket: WebSocket, data: Any):
+        await self.proxy_queue.send_response(data)
 
-    async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None:
-        if self.config is not None and self.config['subdomain'] in connections:
-            del connections[self.config['subdomain']]
+    async def on_disconnect(self, websocket: WebSocket, close_code: int):
+        await self.proxy_queue.delete()
 
         if self.request_task is not None:
             self.request_task.cancel()
diff --git a/ttun_server/proxy_queue.py b/ttun_server/proxy_queue.py
new file mode 100644
index 0000000..07e16e0
--- /dev/null
+++ b/ttun_server/proxy_queue.py
@@ -0,0 +1,163 @@
+import asyncio
+import json
+import logging
+import os
+from typing import Awaitable, Callable
+
+from ttun_server.redis import RedisConnectionPool
+from ttun_server.types import RequestData, ResponseData, MemoryConnection
+
+logger = logging.getLogger(__name__)
+
+class BaseProxyQueue:
+    def __init__(self, identifier: str):
+        self.identifier = identifier
+
+    @classmethod
+    async def create_for_identifier(cls, identifier: str) -> 'BaseProxyQueue':
+        raise NotImplementedError(f'Please implement create_for_identifier')
+
+    @classmethod
+    async def get_for_identifier(cls, identifier: str) -> 'BaseProxyQueue':
+        assert await cls.has_connection(identifier)
+        return cls(identifier)
+
+    @classmethod
+    async def has_connection(cls, identifier) -> bool:
+        raise NotImplementedError(f'Please implement has_connection')
+
+    async def send_request(self, request_data: RequestData):
+        raise NotImplementedError(f'Please implement send_request')
+
+    async def handle_request(self) -> RequestData:
+        raise NotImplementedError(f'Please implement handle_requests')
+
+    async def send_response(self, response_data: ResponseData):
+        raise NotImplementedError(f'Please implement send_request')
+
+    async def handle_response(self) -> ResponseData:
+        raise NotImplementedError(f'Please implement handle_response')
+
+    async def delete(self):
+        raise NotImplementedError(f'Please implement delete')
+
+class MemoryProxyQueue(BaseProxyQueue):
+    connections: dict[str, MemoryConnection] = {}
+
+    @classmethod
+    async def has_connection(cls, identifier) -> bool:
+        return identifier in cls.connections
+
+    @classmethod
+    async def create_for_identifier(cls, identifier: str) -> 'MemoryProxyQueue':
+        instance = cls(identifier)
+
+        cls.connections[identifier] = {
+            'requests': asyncio.Queue(),
+            'responses': asyncio.Queue(),
+        }
+
+        return instance
+
+    @property
+    def requests(self) -> asyncio.Queue[RequestData]:
+        return self.__class__.connections[self.identifier]['requests']
+
+    @property
+    def responses(self) -> asyncio.Queue[ResponseData]:
+        return self.__class__.connections[self.identifier]['responses']
+
+    async def send_request(self, request_data: RequestData):
+        await self.requests.put(request_data)
+
+    async def handle_request(self) -> RequestData:
+        return await self.requests.get()
+
+    async def send_response(self, response_data: ResponseData):
+        return await self.responses.put(response_data)