refac

backend/open_webui/routers/tether.py

@@ -9,7 +9,7 @@ from uuid import UUID, uuid4
 from fastapi import APIRouter, WebSocket
 from open_webui.models.auth_sessions import AuthSessions
 from open_webui.utils.auth import decode_token, get_verified_user_by_token
-from open_webui.utils.tether import hub, encode, decode
+from open_webui.utils.tether import hub, encode, decode, get_services
 
 router = APIRouter()
 
@@ -21,7 +21,7 @@ async def connect(ws: WebSocket):
     watcher = None
     try:
         hello = json.loads(await asyncio.wait_for(ws.receive_text(), 10))
-        if hello.get('type') != 'hello' or hello.get('version') != 1:
+        if hello.get('type') != 'hello' or hello.get('version') not in (1, 2):
             raise ValueError('Unsupported Tether handshake')
         token = hello.get('token', '')
         user = await get_verified_user_by_token(token, getattr(ws.app.state, 'redis', None))
@@ -43,6 +43,7 @@ async def connect(ws: WebSocket):
             'user': user.id,
             'connection_id': connection_id,
             'name': name.strip(),
+            'services': get_services(hello.get('services', ['terminal'] if hello['version'] == 1 else None)),
         }
         connection = SimpleNamespace(ws=ws, presence=presence, routes={}, lock=asyncio.Lock())
 
@@ -58,6 +59,7 @@ async def connect(ws: WebSocket):
         hub.connections[presence['connection']] = connection
         key = hub.key(user.id, connection_id)
         record = json.dumps(presence)
+        presence_lock = asyncio.Lock()
         if hub.redis:
             await hub.redis.set(key, record, ex=60)
         await ws.send_json({'type': 'ready'})
@@ -69,16 +71,17 @@ async def connect(ws: WebSocket):
                 if not current or current.id != user.id:
                     await ws.close(code=4001)
                     return
-                if hub.redis:
-                    refreshed = await hub.redis.eval(
-                        "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('expire', KEYS[1], 60) else return 0 end",
-                        1,
-                        key,
-                        record,
-                    )
-                    if not refreshed:
-                        await ws.close(code=1000)
-                        return
+                async with presence_lock:
+                    if hub.redis:
+                        refreshed = await hub.redis.eval(
+                            "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('expire', KEYS[1], 60) else return 0 end",
+                            1,
+                            key,
+                            record,
+                        )
+                        if not refreshed:
+                            await ws.close(code=1000)
+                            return
 
         async def guarded_watch():
             try:
@@ -97,6 +100,20 @@ async def connect(ws: WebSocket):
                 await hub.returned(connection, header, data)
             elif message.get('text') is not None:
                 update = json.loads(message['text'])
+                if update.get('type') == 'services':
+                    services = get_services(update.get('services'))
+                    async with presence_lock:
+                        updated_record = json.dumps({**presence, 'services': services})
+                        if hub.redis:
+                            replaced = await hub.redis.eval(
+                                "if redis.call('get', KEYS[1]) == ARGV[1] then redis.call('set', KEYS[1], ARGV[2], 'EX', 60); return 1 else return 0 end",
+                                1, key, record, updated_record,
+                            )
+                            if not replaced:
+                                raise ValueError('Tether connection was replaced')
+                        presence['services'] = services
+                        record = updated_record
+                    continue
                 updated = decode_token(update.get('token', ''))
                 current = await get_verified_user_by_token(
                     update.get('token', ''), getattr(ws.app.state, 'redis', None)

backend/open_webui/utils/tether.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import json
+import re
 import struct
 import time
 from contextlib import suppress
@@ -21,6 +22,16 @@ WINDOW = 262144
 MAX_MESSAGE = 16 * 1024 * 1024
 
 
+def get_services(value):
+    if (
+        not isinstance(value, list)
+        or len(value) > 16
+        or any(not isinstance(name, str) or not re.fullmatch(r'[a-z][a-z0-9_]{0,31}', name) for name in value)
+    ):
+        raise ValueError('Invalid Tether services')
+    return list(dict.fromkeys(value))
+
+
 def encode(header, data=b''):
     raw = json.dumps(header, separators=(',', ':')).encode()
     if len(raw) > 4096 or len(data) > CHUNK:
@@ -222,13 +233,15 @@ class Hub:
         # SCAN may return a key more than once while connections are changing.
         return list({record['connection_id']: record for record in records if record['user'] == user_id}.values())
 
-    async def open(self, user_id, connection_id):
+    async def open(self, user_id, connection_id, service='terminal'):
         presence = await self.presence(user_id, connection_id)
         if not presence or presence['user'] != user_id or presence['connection_id'] != connection_id:
             raise HTTPException(503, 'Tether is offline.')
         row = await AuthSessions.get(presence['session'])
         if not row or row.user_id != user_id or row.refresh_token_id is None or row.expires_at <= time.time():
             raise HTTPException(403, 'Tether session is no longer available.')
+        if service not in presence.get('services', ['terminal']):
+            raise HTTPException(503, 'Tether service is unavailable.')
         if sum(s.presence['connection'] == presence['connection'] for s in self.streams.values()) >= 16:
             raise HTTPException(429, 'Tether has too many active operations.')
         stream = Stream(self, presence)
@@ -277,6 +290,8 @@ async def get_terminal_connections(user):
     user_id = user['id'] if isinstance(user, dict) else user.id
     connections = list(await Config.get('terminal_server.connections', []) or [])
     for presence in await hub.presences(user_id):
+        if 'terminal' not in presence.get('services', ['terminal']):
+            continue
         row = await AuthSessions.get(presence['session'])
         if not row or row.user_id != user_id or row.refresh_token_id is None or row.expires_at <= time.time():
             continue
@@ -410,10 +425,11 @@ class RelaySocket:
 
 
 class RelayClient:
-    def __init__(self, user_id):
+    def __init__(self, user_id, service='terminal'):
         if not user_id:
             raise HTTPException(403, 'A user is required for Tether.')
         self.user_id = user_id
+        self.service = get_services([service])[0]
         self.operations = []
 
     async def __aenter__(self):
@@ -445,7 +461,7 @@ class RelayClient:
         query = parsed.query
         if kwargs.get('params'):
             query += ('&' if query else '') + urlencode(kwargs['params'], doseq=True)
-        stream = await hub.open(self.user_id, parsed.hostname)
+        stream = await hub.open(self.user_id, parsed.hostname, self.service)
         body = kwargs.get('data')
         if 'json' in kwargs:
             body = json.dumps(kwargs['json']).encode()
@@ -457,6 +473,7 @@ class RelayClient:
         try:
             await stream.send(
                 'http-open' if method else 'ws-open',
+                service=self.service,
                 method=method,
                 path=parsed.path or '/',
                 query=query,