24024
af7c11d7f9
- 新增图像生成接口，支持试用、积分和自定义API Key模式 - 实现生成图片结果异步上传至MinIO存储，带重试机制 - 优化积分预扣除和异常退还逻辑，保障用户积分准确 - 添加获取生成历史记录接口，支持时间范围和分页 - 提供本地字典配置接口，支持模型、比例、提示模板和尺寸 - 实现图片批量上传接口，支持S3兼容对象存储 feat(admin): 增加管理员角色管理与权限分配接口 - 实现角色列表查询、角色创建、更新及删除功能 - 增加权限列表查询接口 - 实现用户角色分配接口，便于统一管理用户权限 - 增加系统字典增删查改接口，支持分类过滤和排序 - 权限控制全面覆盖管理接口，保证安全访问 feat(auth): 完善用户登录注册及权限相关接口与页面 - 实现手机号验证码发送及校验功能，保障注册安全 - 支持手机号注册、登录及退出接口，集成日志记录 - 增加修改密码功能，验证原密码后更新 - 提供动态导航菜单接口，基于权限展示不同菜单 - 实现管理界面路由及日志、角色、字典管理页面访问权限控制 - 添加系统日志查询接口，支持关键词和等级筛选 feat(app): 初始化Flask应用并配置蓝图与数据库 - 创建应用程序工厂，加载配置，初始化数据库和Redis客户端 - 注册认证、API及管理员蓝图，整合路由 - 根路由渲染主页模板 - 应用上下文中自动创建数据库表，保证运行环境准备完毕 feat(database): 提供数据库创建与迁移支持脚本 - 新增数据库创建脚本，支持自动检测是否已存在 - 添加数据库表初始化脚本，支持创建和删除所有表 - 实现RBAC权限初始化，包含基础权限和角色创建 - 新增字段手动修复脚本，添加用户API Key和积分字段 - 强制迁移脚本支持清理连接和修复表结构，初始化默认数据及角色分配 feat(config): 新增系统配置参数 - 配置数据库、Redis、Session和MinIO相关参数 - 添加AI接口地址及试用Key配置 - 集成阿里云短信服务配置及开发模式相关参数 feat(extensions): 初始化数据库、Redis和MinIO客户端 - 创建全局SQLAlchemy数据库实例和Redis客户端 - 配置基于boto3的MinIO兼容S3客户端 chore(logs): 添加示例系统日志文件 - 记录用户请求、验证码发送成功与失败的日志信息
263 lines
11 KiB
Python
263 lines
11 KiB
Python
"""WebSocket protocol versions 13 and 8."""
|
|
|
|
import asyncio
|
|
import random
|
|
import sys
|
|
from functools import partial
|
|
from typing import Final, Optional, Set, Union
|
|
|
|
from ..base_protocol import BaseProtocol
|
|
from ..client_exceptions import ClientConnectionResetError
|
|
from ..compression_utils import ZLibBackend, ZLibCompressor
|
|
from .helpers import (
|
|
MASK_LEN,
|
|
MSG_SIZE,
|
|
PACK_CLOSE_CODE,
|
|
PACK_LEN1,
|
|
PACK_LEN2,
|
|
PACK_LEN3,
|
|
PACK_RANDBITS,
|
|
websocket_mask,
|
|
)
|
|
from .models import WS_DEFLATE_TRAILING, WSMsgType
|
|
|
|
DEFAULT_LIMIT: Final[int] = 2**16
|
|
|
|
# WebSocket opcode boundary: opcodes 0-7 are data frames, 8-15 are control frames
|
|
# Control frames (ping, pong, close) are never compressed
|
|
WS_CONTROL_FRAME_OPCODE: Final[int] = 8
|
|
|
|
# For websockets, keeping latency low is extremely important as implementations
|
|
# generally expect to be able to send and receive messages quickly. We use a
|
|
# larger chunk size to reduce the number of executor calls and avoid task
|
|
# creation overhead, since both are significant sources of latency when chunks
|
|
# are small. A size of 16KiB was chosen as a balance between avoiding task
|
|
# overhead and not blocking the event loop too long with synchronous compression.
|
|
|
|
WEBSOCKET_MAX_SYNC_CHUNK_SIZE = 16 * 1024
|
|
|
|
|
|
class WebSocketWriter:
|
|
"""WebSocket writer.
|
|
|
|
The writer is responsible for sending messages to the client. It is
|
|
created by the protocol when a connection is established. The writer
|
|
should avoid implementing any application logic and should only be
|
|
concerned with the low-level details of the WebSocket protocol.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
protocol: BaseProtocol,
|
|
transport: asyncio.Transport,
|
|
*,
|
|
use_mask: bool = False,
|
|
limit: int = DEFAULT_LIMIT,
|
|
random: random.Random = random.Random(),
|
|
compress: int = 0,
|
|
notakeover: bool = False,
|
|
) -> None:
|
|
"""Initialize a WebSocket writer."""
|
|
self.protocol = protocol
|
|
self.transport = transport
|
|
self.use_mask = use_mask
|
|
self.get_random_bits = partial(random.getrandbits, 32)
|
|
self.compress = compress
|
|
self.notakeover = notakeover
|
|
self._closing = False
|
|
self._limit = limit
|
|
self._output_size = 0
|
|
self._compressobj: Optional[ZLibCompressor] = None
|
|
self._send_lock = asyncio.Lock()
|
|
self._background_tasks: Set[asyncio.Task[None]] = set()
|
|
|
|
async def send_frame(
|
|
self, message: bytes, opcode: int, compress: Optional[int] = None
|
|
) -> None:
|
|
"""Send a frame over the websocket with message as its payload."""
|
|
if self._closing and not (opcode & WSMsgType.CLOSE):
|
|
raise ClientConnectionResetError("Cannot write to closing transport")
|
|
|
|
if not (compress or self.compress) or opcode >= WS_CONTROL_FRAME_OPCODE:
|
|
# Non-compressed frames don't need lock or shield
|
|
self._write_websocket_frame(message, opcode, 0)
|
|
elif len(message) <= WEBSOCKET_MAX_SYNC_CHUNK_SIZE:
|
|
# Small compressed payloads - compress synchronously in event loop
|
|
# We need the lock even though sync compression has no await points.
|
|
# This prevents small frames from interleaving with large frames that
|
|
# compress in the executor, avoiding compressor state corruption.
|
|
async with self._send_lock:
|
|
self._send_compressed_frame_sync(message, opcode, compress)
|
|
else:
|
|
# Large compressed frames need shield to prevent corruption
|
|
# For large compressed frames, the entire compress+send
|
|
# operation must be atomic. If cancelled after compression but
|
|
# before send, the compressor state would be advanced but data
|
|
# not sent, corrupting subsequent frames.
|
|
# Create a task to shield from cancellation
|
|
# The lock is acquired inside the shielded task so the entire
|
|
# operation (lock + compress + send) completes atomically.
|
|
# Use eager_start on Python 3.12+ to avoid scheduling overhead
|
|
loop = asyncio.get_running_loop()
|
|
coro = self._send_compressed_frame_async_locked(message, opcode, compress)
|
|
if sys.version_info >= (3, 12):
|
|
send_task = asyncio.Task(coro, loop=loop, eager_start=True)
|
|
else:
|
|
send_task = loop.create_task(coro)
|
|
# Keep a strong reference to prevent garbage collection
|
|
self._background_tasks.add(send_task)
|
|
send_task.add_done_callback(self._background_tasks.discard)
|
|
await asyncio.shield(send_task)
|
|
|
|
# It is safe to return control to the event loop when using compression
|
|
# after this point as we have already sent or buffered all the data.
|
|
# Once we have written output_size up to the limit, we call the
|
|
# drain helper which waits for the transport to be ready to accept
|
|
# more data. This is a flow control mechanism to prevent the buffer
|
|
# from growing too large. The drain helper will return right away
|
|
# if the writer is not paused.
|
|
if self._output_size > self._limit:
|
|
self._output_size = 0
|
|
if self.protocol._paused:
|
|
await self.protocol._drain_helper()
|
|
|
|
def _write_websocket_frame(self, message: bytes, opcode: int, rsv: int) -> None:
|
|
"""
|
|
Write a websocket frame to the transport.
|
|
|
|
This method handles frame header construction, masking, and writing to transport.
|
|
It does not handle compression or flow control - those are the responsibility
|
|
of the caller.
|
|
"""
|
|
msg_length = len(message)
|
|
|
|
use_mask = self.use_mask
|
|
mask_bit = 0x80 if use_mask else 0
|
|
|
|
# Depending on the message length, the header is assembled differently.
|
|
# The first byte is reserved for the opcode and the RSV bits.
|
|
first_byte = 0x80 | rsv | opcode
|
|
if msg_length < 126:
|
|
header = PACK_LEN1(first_byte, msg_length | mask_bit)
|
|
header_len = 2
|
|
elif msg_length < 65536:
|
|
header = PACK_LEN2(first_byte, 126 | mask_bit, msg_length)
|
|
header_len = 4
|
|
else:
|
|
header = PACK_LEN3(first_byte, 127 | mask_bit, msg_length)
|
|
header_len = 10
|
|
|
|
if self.transport.is_closing():
|
|
raise ClientConnectionResetError("Cannot write to closing transport")
|
|
|
|
# https://datatracker.ietf.org/doc/html/rfc6455#section-5.3
|
|
# If we are using a mask, we need to generate it randomly
|
|
# and apply it to the message before sending it. A mask is
|
|
# a 32-bit value that is applied to the message using a
|
|
# bitwise XOR operation. It is used to prevent certain types
|
|
# of attacks on the websocket protocol. The mask is only used
|
|
# when aiohttp is acting as a client. Servers do not use a mask.
|
|
if use_mask:
|
|
mask = PACK_RANDBITS(self.get_random_bits())
|
|
message = bytearray(message)
|
|
websocket_mask(mask, message)
|
|
self.transport.write(header + mask + message)
|
|
self._output_size += MASK_LEN
|
|
elif msg_length > MSG_SIZE:
|
|
self.transport.write(header)
|
|
self.transport.write(message)
|
|
else:
|
|
self.transport.write(header + message)
|
|
|
|
self._output_size += header_len + msg_length
|
|
|
|
def _get_compressor(self, compress: Optional[int]) -> ZLibCompressor:
|
|
"""Get or create a compressor object for the given compression level."""
|
|
if compress:
|
|
# Do not set self._compress if compressing is for this frame
|
|
return ZLibCompressor(
|
|
level=ZLibBackend.Z_BEST_SPEED,
|
|
wbits=-compress,
|
|
max_sync_chunk_size=WEBSOCKET_MAX_SYNC_CHUNK_SIZE,
|
|
)
|
|
if not self._compressobj:
|
|
self._compressobj = ZLibCompressor(
|
|
level=ZLibBackend.Z_BEST_SPEED,
|
|
wbits=-self.compress,
|
|
max_sync_chunk_size=WEBSOCKET_MAX_SYNC_CHUNK_SIZE,
|
|
)
|
|
return self._compressobj
|
|
|
|
def _send_compressed_frame_sync(
|
|
self, message: bytes, opcode: int, compress: Optional[int]
|
|
) -> None:
|
|
"""
|
|
Synchronous send for small compressed frames.
|
|
|
|
This is used for small compressed payloads that compress synchronously in the event loop.
|
|
Since there are no await points, this is inherently cancellation-safe.
|
|
"""
|
|
# RSV are the reserved bits in the frame header. They are used to
|
|
# indicate that the frame is using an extension.
|
|
# https://datatracker.ietf.org/doc/html/rfc6455#section-5.2
|
|
compressobj = self._get_compressor(compress)
|
|
# (0x40) RSV1 is set for compressed frames
|
|
# https://datatracker.ietf.org/doc/html/rfc7692#section-7.2.3.1
|
|
self._write_websocket_frame(
|
|
(
|
|
compressobj.compress_sync(message)
|
|
+ compressobj.flush(
|
|
ZLibBackend.Z_FULL_FLUSH
|
|
if self.notakeover
|
|
else ZLibBackend.Z_SYNC_FLUSH
|
|
)
|
|
).removesuffix(WS_DEFLATE_TRAILING),
|
|
opcode,
|
|
0x40,
|
|
)
|
|
|
|
async def _send_compressed_frame_async_locked(
|
|
self, message: bytes, opcode: int, compress: Optional[int]
|
|
) -> None:
|
|
"""
|
|
Async send for large compressed frames with lock.
|
|
|
|
Acquires the lock and compresses large payloads asynchronously in
|
|
the executor. The lock is held for the entire operation to ensure
|
|
the compressor state is not corrupted by concurrent sends.
|
|
|
|
MUST be run shielded from cancellation. If cancelled after
|
|
compression but before sending, the compressor state would be
|
|
advanced but data not sent, corrupting subsequent frames.
|
|
"""
|
|
async with self._send_lock:
|
|
# RSV are the reserved bits in the frame header. They are used to
|
|
# indicate that the frame is using an extension.
|
|
# https://datatracker.ietf.org/doc/html/rfc6455#section-5.2
|
|
compressobj = self._get_compressor(compress)
|
|
# (0x40) RSV1 is set for compressed frames
|
|
# https://datatracker.ietf.org/doc/html/rfc7692#section-7.2.3.1
|
|
self._write_websocket_frame(
|
|
(
|
|
await compressobj.compress(message)
|
|
+ compressobj.flush(
|
|
ZLibBackend.Z_FULL_FLUSH
|
|
if self.notakeover
|
|
else ZLibBackend.Z_SYNC_FLUSH
|
|
)
|
|
).removesuffix(WS_DEFLATE_TRAILING),
|
|
opcode,
|
|
0x40,
|
|
)
|
|
|
|
async def close(self, code: int = 1000, message: Union[bytes, str] = b"") -> None:
|
|
"""Close the websocket, sending the specified code and message."""
|
|
if isinstance(message, str):
|
|
message = message.encode("utf-8")
|
|
try:
|
|
await self.send_frame(
|
|
PACK_CLOSE_CODE(code) + message, opcode=WSMsgType.CLOSE
|
|
)
|
|
finally:
|
|
self._closing = True
|