diff --git a/conf/consumers.py b/conf/consumers.py new file mode 100644 index 0000000..2e14406 --- /dev/null +++ b/conf/consumers.py @@ -0,0 +1,97 @@ +""" +WebSocket consumers for configuration updates +""" +import json +import logging +from channels.generic.websocket import AsyncWebsocketConsumer + +logger = logging.getLogger(__name__) + + +class ConfigConsumer(AsyncWebsocketConsumer): + """ + WebSocket consumer for real-time configuration updates + 当管理员修改配置后,通过 WebSocket 实时推送配置变化 + """ + + async def connect(self): + """处理 WebSocket 连接""" + self.user = self.scope["user"] + + # 只允许认证用户连接 + if not self.user.is_authenticated: + await self.close() + return + + # 使用全局配置组名,所有用户都能接收配置更新 + self.group_name = "config_updates" + + # 加入配置更新组 + await self.channel_layer.group_add( + self.group_name, + self.channel_name + ) + + await self.accept() + logger.info(f"Config WebSocket connected: user_id={self.user.id}, channel={self.channel_name}") + + async def disconnect(self, close_code): + """处理 WebSocket 断开连接""" + if hasattr(self, 'group_name'): + await self.channel_layer.group_discard( + self.group_name, + self.channel_name + ) + logger.info(f"Config WebSocket disconnected: user_id={self.user.id}, close_code={close_code}") + + async def receive(self, text_data): + """ + 接收客户端消息 + 客户端可以发送心跳包或配置更新请求 + """ + try: + data = json.loads(text_data) + message_type = data.get("type") + + if message_type == "ping": + # 响应心跳包 + await self.send(text_data=json.dumps({ + "type": "pong", + "timestamp": data.get("timestamp") + })) + elif message_type == "config_update": + # 处理配置更新请求 + key = data.get("key") + value = data.get("value") + if key and value is not None: + logger.info(f"User {self.user.id} requested config update: {key}={value}") + # 这里可以添加权限检查,只有管理员才能发送配置更新 + if self.user.is_superuser: + # 广播配置更新给所有连接的客户端 + await self.channel_layer.group_send( + self.group_name, + { + "type": "config_update", + "data": { + "type": "config_update", + "key": key, + "value": value + } + } + ) + except json.JSONDecodeError: + logger.error(f"Invalid JSON received from user {self.user.id}") + except Exception as e: + logger.error(f"Error handling message from user {self.user.id}: {str(e)}") + + async def config_update(self, event): + """ + 接收来自 channel layer 的配置更新消息并发送给客户端 + 这个方法名对应 group_send 中的 type 字段 + """ + try: + # 从 event 中提取数据并发送给客户端 + await self.send(text_data=json.dumps(event["data"])) + logger.debug(f"Sent config update to user {self.user.id}: {event['data']}") + except Exception as e: + logger.error(f"Error sending config update to user {self.user.id}: {str(e)}") diff --git a/conf/serializers.py b/conf/serializers.py index a428479..34a96a3 100644 --- a/conf/serializers.py +++ b/conf/serializers.py @@ -27,6 +27,7 @@ class CreateEditWebsiteConfigSerializer(serializers.Serializer): allow_register = serializers.BooleanField() submission_list_show_all = serializers.BooleanField() class_list = serializers.ListField(child=serializers.CharField(max_length=64)) + enable_maxkb = serializers.BooleanField() class JudgeServerSerializer(serializers.ModelSerializer): diff --git a/conf/views.py b/conf/views.py index ff24804..fa4b19e 100644 --- a/conf/views.py +++ b/conf/views.py @@ -24,6 +24,7 @@ from utils.api import APIView, CSRFExemptAPIView, validate_serializer from utils.cache import JsonDataLoader from utils.shortcuts import send_email, get_env from utils.xss_filter import XSSHtml +from utils.websocket import push_config_update from .models import JudgeServer from .serializers import ( CreateEditWebsiteConfigSerializer, @@ -107,6 +108,7 @@ class WebsiteConfigAPI(APIView): "allow_register", "submission_list_show_all", "class_list", + "enable_maxkb", ] } return self.success(ret) @@ -119,6 +121,10 @@ class WebsiteConfigAPI(APIView): with XSSHtml() as parser: v = parser.clean(v) setattr(SysOptions, k, v) + + # 推送配置更新到所有连接的客户端 + push_config_update(k, v) + return self.success() diff --git a/oj/routing.py b/oj/routing.py index 0041978..852da74 100644 --- a/oj/routing.py +++ b/oj/routing.py @@ -4,8 +4,10 @@ WebSocket URL Configuration for oj project. from django.urls import path from submission.consumers import SubmissionConsumer +from conf.consumers import ConfigConsumer websocket_urlpatterns = [ path("ws/submission/", SubmissionConsumer.as_asgi()), + path("ws/config/", ConfigConsumer.as_asgi()), ] diff --git a/options/options.py b/options/options.py index f37757a..6b4829a 100644 --- a/options/options.py +++ b/options/options.py @@ -104,6 +104,7 @@ class OptionKeys: judge_server_token = "judge_server_token" throttling = "throttling" languages = "languages" + enable_maxkb = "enable_maxkb" class OptionDefaultValue: @@ -119,6 +120,7 @@ class OptionDefaultValue: throttling = {"ip": {"capacity": 100, "fill_rate": 0.1, "default_capacity": 50}, "user": {"capacity": 20, "fill_rate": 0.03, "default_capacity": 10}} languages = languages + enable_maxkb = True class _SysOptionsMeta(type): @@ -283,6 +285,15 @@ class _SysOptionsMeta(type): def spj_language_names(cls): return [item["name"] for item in cls.languages if "spj" in item] + @my_property(ttl=DEFAULT_SHORT_TTL) + def enable_maxkb(cls): + return cls._get_option(OptionKeys.enable_maxkb) + + @enable_maxkb.setter + def enable_maxkb(cls, value): + cls._set_option(OptionKeys.enable_maxkb, value) + + def reset_languages(cls): cls.languages = languages diff --git a/utils/websocket.py b/utils/websocket.py index b486106..37ed3f3 100644 --- a/utils/websocket.py +++ b/utils/websocket.py @@ -73,3 +73,37 @@ def push_to_user(user_id: int, message_type: str, data: dict): except Exception as e: logger.error(f"Failed to push message to user {user_id}: error={str(e)}") + +def push_config_update(key: str, value): + """ + 推送配置更新到所有连接的客户端 + + Args: + key: 配置键名 + value: 配置值 + """ + channel_layer = get_channel_layer() + + if channel_layer is None: + logger.warning("Channel layer is not configured, cannot push config update") + return + + # 使用全局配置组名 + group_name = "config_updates" + + try: + # 向所有连接的客户端发送配置更新 + async_to_sync(channel_layer.group_send)( + group_name, + { + "type": "config_update", + "data": { + "type": "config_update", + "key": key, + "value": value + } + } + ) + logger.info(f"Pushed config update: {key}={value}") + except Exception as e: + logger.error(f"Failed to push config update: {key}={value}, error={str(e)}")