feat: add ENABLE_WEBSOCKET_SUPPORT to force socket.io to ignore websocket upgrades

This commit is contained in:
Jun Siang Cheah 2024-09-08 12:00:36 +01:00
parent 1a6ce1d5d9
commit 698976add0
2 changed files with 15 additions and 1 deletions

View File

@ -2,9 +2,17 @@ import asyncio
import socketio import socketio
from open_webui.apps.webui.models.users import Users from open_webui.apps.webui.models.users import Users
from open_webui.config import ENABLE_WEBSOCKET_SUPPORT
from open_webui.utils.utils import decode_token from open_webui.utils.utils import decode_token
sio = socketio.AsyncServer(cors_allowed_origins=[], async_mode="asgi") sio = socketio.AsyncServer(
cors_allowed_origins=[],
async_mode="asgi",
transports=(
["polling", "websocket"] if ENABLE_WEBSOCKET_SUPPORT.value else ["polling"]
),
allow_upgrades=ENABLE_WEBSOCKET_SUPPORT.value,
)
app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io") app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io")
# Dictionary to maintain the user pool # Dictionary to maintain the user pool

View File

@ -810,6 +810,12 @@ ENABLE_MESSAGE_RATING = PersistentConfig(
os.environ.get("ENABLE_MESSAGE_RATING", "True").lower() == "true", os.environ.get("ENABLE_MESSAGE_RATING", "True").lower() == "true",
) )
ENABLE_WEBSOCKET_SUPPORT = PersistentConfig(
"ENABLE_WEBSOCKET_SUPPORT",
"ui.enable_websocket_support",
os.environ.get("ENABLE_WEBSOCKET_SUPPORT", "True").lower() == "true",
)
def validate_cors_origins(origins): def validate_cors_origins(origins):
for origin in origins: for origin in origins: