From 53a8f300859a50d9f99f1821c35bca999fced6d8 Mon Sep 17 00:00:00 2001
From: Tom van der Lee <tom@vanderlee.io>
Date: Thu, 8 Jun 2023 23:05:26 +0200
Subject: v2

---
 .github/workflows/docker-image.yml |  6 ++++--
 Dockerfile                         |  3 +++
 ttun_server/__init__.py            |  6 ++++++
 ttun_server/endpoints.py           | 22 ++++++++++++++++++++--
 ttun_server/redis.py               |  2 +-
 ttun_server/types.py               |  1 +
 6 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml
index 8945bfb..3f83351 100644
--- a/.github/workflows/docker-image.yml
+++ b/.github/workflows/docker-image.yml
@@ -16,7 +16,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Docker meta
         id: meta
-        uses: docker/metadata-action@v3
+        uses: docker/metadata-action@v4
         with:
           images: ghcr.io/tomvanderlee/ttun-server
           tags: |
@@ -31,9 +31,11 @@ jobs:
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
       - name: Build and push
-        uses: docker/build-push-action@v2
+        uses: docker/build-push-action@v4
         with:
           context: .
           push: ${{ github.event_name != 'pull_request' }}
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
+          build-args: |
+            DOCKER_METADATA_OUTPUT_VERSION
diff --git a/Dockerfile b/Dockerfile
index 66b4b70..fc4c26f 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -15,7 +15,10 @@ RUN pip install -r requirements.txt --root /buildroot
 FROM base
 
 COPY --from=build /buildroot /
+
 COPY . .
+ARG DOCKER_METADATA_OUTPUT_VERSION
+RUN echo "version='$DOCKER_METADATA_OUTPUT_VERSION'" > ttun_server/_version.py
 
 ENV TUNNEL_DOMAIN=
 ENV SECURE True
diff --git a/ttun_server/__init__.py b/ttun_server/__init__.py
index 17a8e7a..81f8cd4 100644
--- a/ttun_server/__init__.py
+++ b/ttun_server/__init__.py
@@ -20,3 +20,9 @@ server = Starlette(
         Route('/{path:path}', Proxy),
     ]
 )
+
+try:
+    from ._version import version
+    __version__ = version
+except ImportError:
+    __version__ = 'development'
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 @@
 import asyncio
 import logging
 import os
+from asyncio import create_task
 from base64 import b64decode, b64encode
 from typing import Optional, Any
 from uuid import uuid4
@@ -11,6 +12,7 @@ from starlette.responses import Response
 from starlette.types import Scope, Receive, Send
 from starlette.websockets import WebSocket
 
+import ttun_server
 from ttun_server.proxy_queue import ProxyQueue
 from ttun_server.types import RequestData, Config, Message, MessageType
 
@@ -40,9 +42,10 @@ class Proxy(HTTPEndpoint):
 
             request_queue = await ProxyQueue.get_for_identifier(subdomain)
 
+            logger.debug('PROXY %s%s ', subdomain, request.url)
             await request_queue.enqueue(
                 Message(
-                    type=MessageType.request,
+                    type=MessageType.request.value,
                     identifier=identifier,
                     payload=
                     RequestData(
@@ -85,16 +88,31 @@ class Tunnel(WebSocketEndpoint):
 
     async def handle_requests(self, websocket: WebSocket):
         while request := await self.proxy_queue.dequeue():
-            await websocket.send_json(request)
+            create_task(websocket.send_json(request))
 
     async def on_connect(self, websocket: WebSocket) -> None:
         await websocket.accept()
         self.config = await websocket.receive_json()
 
+        client_version = self.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')
+
+
         if self.config['subdomain'] is None \
                 or await ProxyQueue.has_connection(self.config['subdomain']):
             self.config['subdomain'] = uuid4().hex
 
+
         self.proxy_queue = await ProxyQueue.create_for_identifier(self.config['subdomain'])
 
         hostname = os.environ.get("TUNNEL_DOMAIN")
diff --git a/ttun_server/redis.py b/ttun_server/redis.py
index 344c107..3065dec 100644
--- a/ttun_server/redis.py
+++ b/ttun_server/redis.py
@@ -3,7 +3,7 @@ import os
 from aioredis import ConnectionPool, Redis
 
 
-class RedisConnectionPool():
+class RedisConnectionPool:
     instance: 'RedisConnectionPool' = None
 
     def __init__(self):
diff --git a/ttun_server/types.py b/ttun_server/types.py
index 2f1959f..8a4d929 100644
--- a/ttun_server/types.py
+++ b/ttun_server/types.py
@@ -10,6 +10,7 @@ class MessageType(Enum):
 
 class Config(TypedDict):
     subdomain: str
+    client_version: str
 
 
 class RequestData(TypedDict):
-- 
cgit v1.2.3