1 files changed, 3 insertions, 1 deletions

diff --git a/ttun_server/__init__.py b/ttun_server/__init__.py
index 81f8cd4..2f8fed0 100644
--- a/ttun_server/__init__.py
+++ b/ttun_server/__init__.py
@@ -4,7 +4,8 @@ import os
 from starlette.applications import Starlette
 from starlette.routing import Route, WebSocketRoute, Host, Router
 
-from ttun_server.endpoints import Proxy, Tunnel, Health
+from ttun_server.endpoints import Proxy, Health
+from .websockets import WebsocketProxy, Tunnel
 
 logging.basicConfig(level=getattr(logging, os.environ.get('LOG_LEVEL', 'INFO')))
 
@@ -18,6 +19,7 @@ server = Starlette(
     routes=[
         Host(os.environ['TUNNEL_DOMAIN'], base_router, 'base'),
         Route('/{path:path}', Proxy),
+        WebSocketRoute('/{path:path}', WebsocketProxy)
     ]
 )