2 files changed, 11 insertions, 1 deletions

diff --git a/ttun_server/__init__.py b/ttun_server/__init__.py
index 4af6ca0..f5d6b45 100644
--- a/ttun_server/__init__.py
+++ b/ttun_server/__init__.py
@@ -2,8 +2,10 @@ import logging
 import os
 
 from fastapi import FastAPI
+from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.routing import Host, Route, WebSocketRoute, Mount
 
+from authentication.backend import BearerTokenAuthBackend
 from authentication.endpoints import authentication, oauth
 from proxy import app as proxy_app
 from ttun_server.endpoints import health
@@ -16,12 +18,17 @@ app = FastAPI(
     routes=[
         Mount('/auth/', app=authentication),
         Mount('/oauth/', app=oauth),
-        WebSocketRoute('/tunnel/', endpoint=tunnel),
+        WebSocketRoute(
+            '/tunnel/',
+            endpoint=tunnel,
+        ),
         Route('/health/', endpoint=health),
         Host(f'{{subdomain}}.{os.environ['TUNNEL_DOMAIN']}', app=proxy_app)
     ]
 )
 
+app.add_middleware(AuthenticationMiddleware, backend=BearerTokenAuthBackend())
+
 try:
     from ._version import version
     __version__ = version
diff --git a/ttun_server/websockets.py b/ttun_server/websockets.py
index c791625..88d6cac 100644
--- a/ttun_server/websockets.py
+++ b/ttun_server/websockets.py
@@ -4,7 +4,9 @@ import os
 from uuid import uuid4
 
 from fastapi import WebSocket, WebSocketDisconnect
+from starlette.authentication import requires
 
+import conf
 import ttun_server
 from proxy.queue import ProxyQueue
 from ttun_server.types import (
@@ -31,6 +33,7 @@ async def assert_compatible_version(websocket: WebSocket, config: Config) -> Non
             await websocket.close(4001, 'Your client is too new')
 
 
+@requires(['authenticated', 'token'] if conf.ENABLE_AUTH else [])
 async def tunnel(websocket: WebSocket) -> None:
     request_tasks: dict[str, asyncio.Task] = {}
     proxy_queues: dict[str, ProxyQueue] = {}