From 457f2edd2dca74509833c8fe59bd0b8cf112bd2a Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Thu, 5 Nov 2026 12:29:50 +0900 Subject: [PATCH] refac --- backend/open_webui/routers/tether.py | 41 ++++++++++++++++++++-------- backend/open_webui/utils/tether.py | 23 ++++++++++++++-- 2 files changed, 49 insertions(+), 15 deletions(-) diff --git a/backend/open_webui/routers/tether.py b/backend/open_webui/routers/tether.py index fb61645e96..2d7ee49c45 100644 --- a/backend/open_webui/routers/tether.py +++ b/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) diff --git a/backend/open_webui/utils/tether.py b/backend/open_webui/utils/tether.py index 297ce9c0eb..82d9da5211 100644 --- a/backend/open_webui/utils/tether.py +++ b/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,