From 314208c2406ecaf7cbe5c8cb14abba6cb939bfe2 Mon Sep 17 00:00:00 2001
From: Tom van der Lee <tom@vanderlee.io>
Date: Thu, 8 Jun 2023 23:06:05 +0200
Subject: Improved performace with v2

---
 setup.cfg        |  2 +-
 ttun/__init__.py |  5 +++++
 ttun/__main__.py |  6 +++++-
 ttun/client.py   | 47 ++++++++++++++++++++++++++++++++++++++++-------
 ttun/types.py    | 14 ++++++++++----
 5 files changed, 61 insertions(+), 13 deletions(-)

diff --git a/setup.cfg b/setup.cfg
index 8a52ffc..d58f78f 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,7 +1,7 @@
 [metadata]
 name = ttun
 author = Tom van der Lee
-description = Expose a localport via your selfhosted TTUN Server
+description = Expose a local port via your selfhosted TTUN Server
 long_description = file: README.rst
 license = BSD-3-Clause
 
diff --git a/ttun/__init__.py b/ttun/__init__.py
index e69de29..d6256a0 100644
--- a/ttun/__init__.py
+++ b/ttun/__init__.py
@@ -0,0 +1,5 @@
+from importlib.metadata import version
+
+
+__version__ = version('ttun')
+
diff --git a/ttun/__main__.py b/ttun/__main__.py
index b14d0e7..4e693fb 100644
--- a/ttun/__main__.py
+++ b/ttun/__main__.py
@@ -9,6 +9,8 @@ from asyncio.exceptions import TimeoutError
 from typing import Dict
 from typing import Tuple
 
+from websockets.exceptions import ConnectionClosedError
+
 from ttun.client import Client
 from ttun.inspect_server import Server
 from ttun.settings import SERVER_HOSTNAME
@@ -71,6 +73,7 @@ def main():
         headers=args.header,
     )
 
+
     try:
         loop = asyncio.get_running_loop()
     except RuntimeError:
@@ -96,7 +99,8 @@ def main():
 
     try:
         loop.run_until_complete(asyncio.wait(tasks, return_when=FIRST_EXCEPTION))
-    except (CancelledError, TimeoutError):
+    except (CancelledError, TimeoutError) as e:
+        print(e)
         for task in tasks:
             task.cancel()
         loop.close()
diff --git a/ttun/client.py b/ttun/client.py
index 1218283..a5c6e10 100644
--- a/ttun/client.py
+++ b/ttun/client.py
@@ -1,8 +1,11 @@
 import asyncio
 import json
+import logging
+from asyncio import create_task, get_running_loop
 from base64 import b64decode
 from base64 import b64encode
 from datetime import datetime
+from pprint import pformat
 from time import perf_counter
 from typing import Awaitable
 from typing import Callable
@@ -21,10 +24,12 @@ from websockets import WebSocketClientProtocol
 from websockets.exceptions import ConnectionClosed
 
 from ttun.pubsub import PubSub
-from ttun.types import Config
+from ttun.types import Config, Message, MessageType
 from ttun.types import RequestData
 from ttun.types import ResponseData
 
+from ttun import __version__
+
 
 class Client:
     def __init__(
@@ -36,6 +41,7 @@ class Client:
         https: bool = False,
         headers: List[Tuple[str, str]] = None,
     ):
+        self.version = __version__
         self.server = server
         self.subdomain = subdomain
 
@@ -47,10 +53,13 @@ class Client:
         self.headers = [] if headers is None else headers
 
     async def send(self, data: dict):
+        print('send {}'.format(pformat(data)))
         await self.connection.send(json.dumps(data))
 
     async def receive(self) -> dict:
-        return json.loads(await self.connection.recv())
+        data = json.loads(await self.connection.recv())
+        print('receive {}'.format(pformat(data)))
+        return data
 
     @staticmethod
     def loop(sleep: int = None):
@@ -69,8 +78,12 @@ class Client:
 
     async def connect(self) -> WebSocketClientProtocol:
         self.connection = await websockets.connect(f"{self.server}/tunnel/")
+        print(self.version)
 
-        await self.send({"subdomain": self.subdomain})
+        await self.send({
+            "subdomain": self.subdomain,
+            "version": self.version
+        })
 
         self.config = await self.receive()
 
@@ -81,20 +94,40 @@ class Client:
         return ClientSession(base_url=self.proxy_origin, cookie_jar=DummyCookieJar())
 
     async def handle_messages(self):
+        loop = get_running_loop()
         async with self.session() as session:
             while True:
                 try:
-                    request: RequestData = await self.receive()
+                    message: Message = await self.receive()
+
+                    try:
+                        if MessageType(message['type']) != MessageType.request:
+                            continue
+                    except ValueError:
+                        continue
+
+                    request: RequestData = message['payload']
 
                     request["headers"] = [
                         *request["headers"],
                         *self.headers,
                     ]
-                    await self.proxy_request(
+
+                    async def response_handler(
+                            response: ResponseData,
+                            identifier=message['identifier']
+                    ):
+                        await self.send(Message(
+                            type=MessageType.response.value,
+                            identifier=identifier,
+                            payload=response
+                        ))
+
+                    loop.create_task(self.proxy_request(
                         session=session,
                         request=request,
-                        on_response=lambda response: self.send(response),
-                    )
+                        on_response=response_handler,
+                    ))
                 except ConnectionClosed:
                     break
 
diff --git a/ttun/types.py b/ttun/types.py
index 3f4d9c9..59bfa79 100644
--- a/ttun/types.py
+++ b/ttun/types.py
@@ -1,10 +1,10 @@
+from enum import Enum
 from typing import Optional
 from typing import TypedDict
 
-
-class Message(TypedDict):
-    type: str
-
+class MessageType(Enum):
+    request = 'request'
+    response = 'response'
 
 class RequestData(TypedDict):
     method: str
@@ -19,5 +19,11 @@ class ResponseData(TypedDict):
     body: Optional[str]
 
 
+class Message(TypedDict):
+    type: MessageType
+    identifier: str
+    payload: RequestData | ResponseData
+
+
 class Config(TypedDict):
     url: str
-- 
cgit v1.2.3