CheeseAPI 2.0.8b7__tar.gz → 2.0.8b8__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,7 +1,7 @@
1
1
  from CheeseAPI.app import CheeseAPI, AppProxy
2
2
  from CheeseAPI.printer import Printer
3
3
  from CheeseAPI.request import Request, RequestProxy
4
- from CheeseAPI.response import Response, ResponseProxy, FileResponse
4
+ from CheeseAPI.response import Response, ResponseProxy, FileResponse, RedirectResponse
5
5
  from CheeseAPI.websocket import Websocket
6
6
  from CheeseAPI.file import File
7
7
  from CheeseAPI.route import Route, RouteProxy
@@ -1,4 +1,4 @@
1
- import os, pathlib, multiprocessing, ssl, socket, asyncio, concurrent.futures, inspect, importlib.util
1
+ import os, re, pathlib, multiprocessing, ssl, socket, asyncio, concurrent.futures, inspect, importlib.util
2
2
  from typing import Type, Literal, Callable, AsyncIterable, TYPE_CHECKING
3
3
 
4
4
  import signal, redis
@@ -44,8 +44,6 @@ class AppProxy:
44
44
 
45
45
  self.app._is_running = True
46
46
 
47
- self.app.signal.before_server_start.send()
48
-
49
47
  self.server_start()
50
48
 
51
49
  waiting_list = self.worker_start()
@@ -215,6 +213,7 @@ class AppProxy:
215
213
 
216
214
  processes = []
217
215
  if workers == 1:
216
+ ''' 单进程模式下工作进程复用当前进程,after_workers_start 在 worker 内、开始接收请求前触发 '''
218
217
  self.worker_running(True)
219
218
  else:
220
219
  processes: list[multiprocessing.Process] = []
@@ -223,10 +222,10 @@ class AppProxy:
223
222
  process.start()
224
223
  processes.append(process)
225
224
 
226
- self.after_workers_start(workers)
227
- self.app.signal.after_workers_start.send(kwargs = {
228
- 'workers': workers
229
- })
225
+ self.after_workers_start(workers)
226
+ self.app.signal.after_workers_start.send(kwargs = {
227
+ 'workers': workers
228
+ })
230
229
  return processes
231
230
 
232
231
  def before_workers_start(self, workers: int) -> int:
@@ -272,6 +271,12 @@ class AppProxy:
272
271
  'is_first': is_first
273
272
  })
274
273
 
274
+ if is_first and self.app.workers == 1:
275
+ self.after_workers_start(self.app.workers)
276
+ self.app.signal.after_workers_start.send(kwargs = {
277
+ 'workers': self.app.workers
278
+ })
279
+
275
280
  while True:
276
281
  client_socket, addr = await loop.sock_accept(self.server_socket)
277
282
  loop.create_task(self.client_socket_process(client_socket, addr))
@@ -289,6 +294,7 @@ class AppProxy:
289
294
  self.app.printer.app_error(e)
290
295
 
291
296
  async def client_socket_process(self, client_socket: socket.socket, addr: tuple[str, int]):
297
+ request = None
292
298
  try:
293
299
  loop = asyncio.get_event_loop()
294
300
 
@@ -298,6 +304,8 @@ class AppProxy:
298
304
  async for request, response in self.get_request(client_socket, addr):
299
305
  if not response:
300
306
  response = await self.get_response(request)
307
+ if inspect.isfunction(request.fn):
308
+ await self.attach_cors_headers(request, response)
301
309
 
302
310
  if not response._proxy:
303
311
  response = self.app.ResponseProxy_Class(self.app, response).response
@@ -313,7 +321,11 @@ class AppProxy:
313
321
  if client_socket._closed is False:
314
322
  client_socket.close()
315
323
  except Exception as e:
316
- self.app.printer.fn_error(e, request)
324
+ ''' 请求尚未创建时无法定位请求上下文,直接上报原始异常,而不是引用未赋值的 request '''
325
+ if request is None:
326
+ self.app.printer.app_error(e)
327
+ else:
328
+ self.app.printer.fn_error(e, request)
317
329
 
318
330
  if client_socket._closed is False:
319
331
  client_socket.close()
@@ -331,6 +343,25 @@ class AppProxy:
331
343
  'response': response
332
344
  })
333
345
 
346
+ async def attach_cors_headers(self, request: Request, response: Response):
347
+ ''' 简单请求:把 CORS 响应头并入路由响应;来源不在白名单内时不附加任何 CORS 头 '''
348
+ if not request.headers.get('origin'):
349
+ return
350
+
351
+ cors_response = self.get_request_cors(request).get_response(request)
352
+ if cors_response.status != 204:
353
+ return
354
+
355
+ for key, value in cors_response.headers.items():
356
+ response.headers.setdefault(key, value)
357
+
358
+ def get_request_cors(self, request: Request) -> CORS:
359
+ ''' 取请求命中路由的 CORS 配置,未配置时回退应用级 '''
360
+ route = self.app.route._proxy.get_route(request.method, request.path)
361
+ if route != 404 and route != 405 and route[0]['cors'] is not None:
362
+ return route[0]['cors']
363
+ return self.app.cors
364
+
334
365
  async def get_response(self, request: Request) -> Response:
335
366
  if inspect.isfunction(request.fn):
336
367
  try:
@@ -352,7 +383,7 @@ class AppProxy:
352
383
  async def get_request(self, client_socket: socket.socket, addr: tuple[str, int]) -> AsyncIterable[tuple[Request, Response | None]]:
353
384
  keep_alive_max_requests = 0
354
385
  request = None
355
- while keep_alive_max_requests < self.app.keep_alive_max_requests and client_socket._closed is False:
386
+ while client_socket._closed is False:
356
387
  if request:
357
388
  keep_alive_max_requests += 1
358
389
 
@@ -364,6 +395,9 @@ class AppProxy:
364
395
  if request._proxy.protocol is None or not self.app.keep_alive or (request._proxy.protocol == 'HTTP/1.0' and request.headers.get('connection') != 'keep-alive') or (request._proxy.protocol == 'HTTP/1.1' and request.headers.get('connection') == 'close'):
365
396
  break
366
397
 
398
+ if keep_alive_max_requests >= self.app.keep_alive_max_requests:
399
+ break
400
+
367
401
  client_socket, addr = await self.before_request(client_socket, addr)
368
402
  await self.app.signal.before_request.async_send(kwargs = {
369
403
  'client_socket': client_socket,
@@ -426,7 +460,9 @@ class AppProxy:
426
460
  relative_path = relative_path[1:]
427
461
 
428
462
  path = os.path.join(_path, relative_path)
429
- if not os.path.abspath(path).startswith(os.path.abspath(_path)):
463
+ static_root = os.path.abspath(_path)
464
+ target_path = os.path.abspath(path)
465
+ if target_path != static_root and target_path.startswith(f'{static_root}{os.sep}') is False:
430
466
  return Response(status = 403)
431
467
 
432
468
  if os.path.exists(path):
@@ -440,13 +476,33 @@ class AppProxy:
440
476
  return Response(status = 404)
441
477
 
442
478
  async def get_cors_response(self, request: Request) -> Response:
443
- route = self.app.route._proxy.get_route(request.method, request.path)
444
- if route != 405:
445
- if route[0]['cors']:
446
- return route[0]['cors'].get_response(request)
447
- else:
448
- return self.app.cors.get_response(request)
449
- return Response(status = 405)
479
+ ''' 预检请求(OPTIONS + Origin + Access-Control-Request-Method)返回 CORS 响应,其余 405 场景保持裸 405 '''
480
+ if request.method != 'OPTIONS' or not request.headers.get('origin') or not request.headers.get('access-control-request-method'):
481
+ return Response(status = 405)
482
+
483
+ cors = self.get_route_cors(request.path) or self.app.cors
484
+ return cors.get_response(request)
485
+
486
+ def get_route_cors(self, path: str) -> CORS | None:
487
+ ''' 按路径取路由级 CORS 配置(与请求方法无关),供预检请求使用 '''
488
+ routes = self.app.route.routes.get(path)
489
+ if routes is None:
490
+ for _path, dynamic_routes in self.app.route._proxy.dynamic_routes.items():
491
+ if re.match(_path, path) is not None:
492
+ routes = dynamic_routes
493
+ break
494
+
495
+ if not routes:
496
+ return None
497
+
498
+ if routes.get('OPTIONS') and routes['OPTIONS']['cors'] is not None:
499
+ return routes['OPTIONS']['cors']
500
+
501
+ for route in routes.values():
502
+ if route['cors'] is not None:
503
+ return route['cors']
504
+
505
+ return None
450
506
 
451
507
  async def before_request(self, client_socket: socket.socket, addr: tuple[str, int]) -> tuple[socket.socket, tuple]:
452
508
  return client_socket, addr
@@ -27,8 +27,9 @@ class CORS:
27
27
 
28
28
  if origin in self.allow_origins:
29
29
  headers['access-control-allow-origin'] = origin
30
- elif '*' in self.allow_origins and not self.allow_credentials:
31
- headers['access-control-allow-origin'] = '*'
30
+ elif '*' in self.allow_origins:
31
+ ''' 通配来源与凭证不能同用,携带凭证时回显具体 origin '''
32
+ headers['access-control-allow-origin'] = origin if self.allow_credentials else '*'
32
33
 
33
34
  headers['access-control-allow-methods'] = ', '.join(self.allow_methods)
34
35
 
@@ -29,6 +29,7 @@ class File:
29
29
  elif len(args) == 2:
30
30
  self._name = args[0]
31
31
  self._data = args[1]
32
+ self._data_in_file = False
32
33
 
33
34
  def save(self, path: str, update_path: bool = False, data_in_file: bool = False):
34
35
  '''
@@ -65,7 +66,7 @@ class File:
65
66
 
66
67
  @property
67
68
  def data(self) -> bytes:
68
- if self._data:
69
+ if self._data is not None:
69
70
  return self._data
70
71
  else:
71
72
  with open(self._path, 'rb') as f:
@@ -12,6 +12,9 @@ if TYPE_CHECKING:
12
12
 
13
13
  HTTP_STATUS_COLOR: tuple[str] = ('blue', 'green', 'cyan', 'yellow', 'red')
14
14
 
15
+ MAX_ERROR_LINES: int = 50
16
+ ''' 单条错误日志中堆栈的最大行数;超出部分会被截断,避免异常堆栈无限增长 '''
17
+
15
18
  class Printer:
16
19
  __slots__ = ('_app', '_progress_bar')
17
20
 
@@ -33,8 +36,19 @@ class Printer:
33
36
  def server_start(self):
34
37
  self.app.logger.print('START', f'CheeseAPI is running on {self.app.host}:{self.app.port}', f'CheeseAPI is running on <cyan>{self.app.host}:{self.app.port}</cyan>')
35
38
 
39
+ def _error(self) -> str:
40
+ '''
41
+ 格式化当前异常的堆栈
42
+
43
+ 限制最大行数,避免反复报错时堆栈越来越长把日志写爆
44
+ '''
45
+ lines = traceback.format_exc()[:-1].splitlines()
46
+ if len(lines) > MAX_ERROR_LINES:
47
+ lines = lines[:MAX_ERROR_LINES] + [f'... truncated {len(lines) - MAX_ERROR_LINES} lines ...']
48
+ return '\n'.join(lines).replace('\n', '\n ')
49
+
36
50
  def app_error(self, e: Exception):
37
- error = traceback.format_exc()[:-1].replace("\n", "\n ")
51
+ error = self._error()
38
52
  self.app.logger.error(f'An error occurred causing the server to stop:\n {error}', f'An error occurred causing the server to stop:\n {self.app.logger.encode(error)}')
39
53
 
40
54
  def server_stop(self):
@@ -44,11 +58,11 @@ class Printer:
44
58
  self.app.logger.print('STOP', 'CheeseAPI has stopped')
45
59
 
46
60
  def fn_error(self, e: Exception, request: 'Request'):
47
- error = traceback.format_exc()[:-1].replace("\n", "\n ")
61
+ error = self._error()
48
62
  self.app.logger.danger(f'An error occurred causing the {request.ip} visited {request.method} {request.full_path}:\n {error}', f'An error occurred causing the <cyan>{request.ip}</cyan> visited <cyan>{request.method} {self.app.logger.encode(request.full_path)}</cyan>:\n {self.app.logger.encode(error)}')
49
63
 
50
64
  def websocket_error(self, e: Exception, websocket: 'Websocket'):
51
- error = traceback.format_exc()[:-1].replace("\n", "\n ")
65
+ error = self._error()
52
66
  self.app.logger.danger(f'An error occurred causing the {websocket.request.ip} disconnected {websocket.request.method} {websocket.request.full_path}:\n {error}', f'An error occurred causing the <cyan>{websocket.request.ip}</cyan> disconnected <cyan>{websocket.request.method} {self.app.logger.encode(websocket.request.full_path)}</cyan>:\n {self.app.logger.encode(error)}')
53
67
 
54
68
  def response(self, request: 'Request', response: 'Response'):
@@ -65,11 +79,11 @@ class Printer:
65
79
  self.app.logger.print('WEBSOCKET', f'The {websocket.request.ip} disconnected {websocket.request.method} {websocket.request.full_path}', f'The <cyan>{websocket.request.ip}</cyan> disconnected <cyan>{websocket.request.method} {self.app.logger.encode(websocket.request.full_path)}</cyan>')
66
80
 
67
81
  def websocket_message_error(self, e: Exception, websocket: 'Websocket'):
68
- error = traceback.format_exc()[:-1].replace("\n", "\n ")
82
+ error = self._error()
69
83
  self.app.logger.danger(f'An error occurred causing the {websocket.request.ip} received a message to {websocket.request.method} {websocket.request.full_path}:\n {error}', f'An error occurred causing the <cyan>{websocket.request.ip}</cyan> receive a message to <cyan>{websocket.request.method} {self.app.logger.encode(websocket.request.full_path)}</cyan>:\n {self.app.logger.encode(error)}')
70
84
 
71
85
  def scheduler_error(self, e: Exception, task: 'Task'):
72
- error = traceback.format_exc()[:-1].replace("\n", "\n ")
86
+ error = self._error()
73
87
  self.app.logger.danger(f'An error occurred in the scheduled task {task.key} running:\n {error}', f'An error occurred in the scheduled task running <green>{self.app.logger.encode(task.key)}</green>:\n {self.app.logger.encode(error)}')
74
88
 
75
89
  @property
@@ -169,7 +169,12 @@ class RequestProxy:
169
169
  range_part = range_part.strip()
170
170
  if '-' in range_part:
171
171
  start, end = range_part.split('-', 1)
172
- self.request.ranges.append((int(start) if start else 0, int(end) if end else None))
172
+ if start:
173
+ self.request.ranges.append((int(start), int(end) if end else None))
174
+ elif end:
175
+ self.request.ranges.append((-int(end), None)) # 后缀 Range(`bytes=-50`):负起点表示从末尾倒数,末尾 50 字节
176
+ else:
177
+ self.request.ranges.append((0, None))
173
178
 
174
179
  if 'upgrade' in self.request.headers and self.request.headers['upgrade'] == 'websocket':
175
180
  self.request._method = 'WEBSOCKET'
@@ -246,17 +251,20 @@ class RequestProxy:
246
251
  if self.request.body is None:
247
252
  return
248
253
 
254
+ body = self.request.body
249
255
  content_type = self.request.headers.get('content-type')
250
- if content_type == 'text/plain' or content_type is None:
251
- self.request._body = self.request.body.decode()
252
- elif content_type == 'application/json':
253
- self.request._json = json.loads(self.request.body)
254
- elif content_type == 'application/x-www-form-urlencoded':
256
+ media_type = (content_type or '').split(';')[0].strip().lower() or None # 忽略 `; charset=utf-8` 等参数
257
+
258
+ if media_type == 'text/plain' or media_type is None:
259
+ self.request._body = body.decode()
260
+ elif media_type == 'application/json':
261
+ self.request._json = json.loads(body)
262
+ elif media_type == 'application/x-www-form-urlencoded':
255
263
  self.request._form = {
256
- key: value[0] for key, value in urllib.parse.parse_qs(self.request.body.decode()).items()
264
+ key: value[0] for key, value in urllib.parse.parse_qs(body.decode()).items()
257
265
  }
258
- elif content_type.startswith('multipart/form-data'):
259
- for part in self.request.body.split(f'--{content_type.split("boundary=")[1].strip()}'.encode()):
266
+ elif media_type == 'multipart/form-data':
267
+ for part in body.split(f'--{content_type.split("boundary=")[1].strip()}'.encode()):
260
268
  if part == b'' or part == b'--\r\n':
261
269
  continue
262
270
 
@@ -290,4 +298,4 @@ class RequestProxy:
290
298
  if self.request.headers.get('content-disposition'):
291
299
  match = re.search(r'filename="([^"]*)"', self.request.headers['content-disposition'])
292
300
  if match:
293
- self.request._file = File(match.group(1), self.request.body)
301
+ self.request._file = File(match.group(1), body) # 始终用原始 bytes,避免 text/plain 分支已解码成 str
@@ -1,4 +1,4 @@
1
- import socket, asyncio, datetime, json, zlib, gzip, os, mimetypes, uuid, math
1
+ import socket, asyncio, datetime, json, zlib, gzip, os, mimetypes, uuid
2
2
  from typing import TYPE_CHECKING, TypedDict, Literal, AsyncIterable
3
3
 
4
4
  import brotli, zstandard
@@ -127,7 +127,7 @@ class RedirectResponse(Response):
127
127
  headers = {}
128
128
  headers['location'] = location
129
129
 
130
- super().__init__(status, body, headers)
130
+ super().__init__(body, status, headers)
131
131
 
132
132
  class FileResponse(Response):
133
133
  __slots__ = ('file', 'preview', 'transmission_type', 'chunked_size')
@@ -169,21 +169,34 @@ class ResponseProxy:
169
169
  status, headers, data = await anext(gen)
170
170
 
171
171
  bytes = [f'HTTP/1.1 {status} {HTTP_STATUS[status]}']
172
- bytes.extend(f'{key}: {value}' for key, value in headers.items())
172
+ for key, value in headers.items():
173
+ if isinstance(value, list): # 某些响应头(如 set-cookie)需要写成多行
174
+ bytes.extend(f'{key}: {_value}' for _value in value)
175
+ else:
176
+ bytes.append(f'{key}: {value}')
173
177
  bytes.extend(['', ''])
174
178
  bytes = '\r\n'.join(bytes).encode()
175
179
  if not no_body:
176
- if self.response.headers.get('transfer-encoding') == 'chunked':
177
- bytes += hex(len(data)).encode() + b'\r\n' + data + b'\r\n'
180
+ if headers.get('transfer-encoding') == 'chunked':
181
+ bytes += f'{len(data):x}'.encode() + b'\r\n' + data + b'\r\n'
178
182
  else:
179
183
  bytes += data
180
184
  await loop.sock_sendall(client_socket, bytes)
181
185
 
182
186
  if not no_body:
183
- if self.response.headers.get('transfer-encoding') == 'chunked':
184
- async for _, _, data in gen:
185
- await loop.sock_sendall(client_socket, hex(len(data)).encode() + b'\r\n' + data + b'\r\n')
186
- await loop.sock_sendall(client_socket, b'0\r\n\r\n')
187
+ if headers.get('transfer-encoding') == 'chunked':
188
+ '''
189
+ 分块传输:无论生成器是否抛异常,都必须补发 0 长度块结束分块流,
190
+ 否则客户端会认为响应被截断
191
+ '''
192
+ try:
193
+ async for _, _, data in gen:
194
+ await loop.sock_sendall(client_socket, f'{len(data):x}'.encode() + b'\r\n' + data + b'\r\n')
195
+ finally:
196
+ try:
197
+ await loop.sock_sendall(client_socket, b'0\r\n\r\n')
198
+ except Exception:
199
+ ...
187
200
  elif self.request.ranges:
188
201
  async for _, _, data in gen:
189
202
  await loop.sock_sendall(client_socket, data)
@@ -192,19 +205,14 @@ class ResponseProxy:
192
205
  self.app.printer.response(self.request, self.response)
193
206
 
194
207
  async def get_status(self, status: int, headers: dict[str, str], body: dict | list | str | bytes | None) -> tuple[int, dict[str, str], dict | list | str | bytes | None]:
195
- if isinstance(self.response, FileResponse):
196
- if self.request.headers.get('range') is not None:
197
- max_range = -1
198
- for range in self.request.ranges:
199
- if range[1] is not None:
200
- max_range = max(max_range, range[1])
201
- max_range += 1
202
- if self.response.file._data is not None:
203
- if len(self.response.file.data) < max_range:
204
- status = 416
205
- else:
206
- if os.path.getsize(self.response.file.path) < max_range:
207
- status = 416
208
+ if isinstance(self.response, FileResponse) and self.request.ranges:
209
+ max_range = -1
210
+ for range in self.request.ranges:
211
+ if range[1] is not None:
212
+ max_range = max(max_range, range[1])
213
+ max_range += 1
214
+ size = len(self.response.file.data) if self.response.file._data is not None else os.path.getsize(self.response.file.path)
215
+ status = 416 if size < max_range else 206
208
216
 
209
217
  return status, headers, body
210
218
 
@@ -236,7 +244,7 @@ class ResponseProxy:
236
244
  headers['content-type'] = f'{mime_type}; charset=utf-8'
237
245
  headers['content-disposition'] = f'{"inline" if self.response.preview and mime_type in PREVIEWABLE_TYPES else "attachment"}; filename="{self.response.file.name}"'
238
246
 
239
- if isinstance(self.response.body, AsyncIterable):
247
+ if isinstance(body, AsyncIterable):
240
248
  if 'transfer-encoding' not in headers:
241
249
  headers['transfer-encoding'] = 'chunked'
242
250
 
@@ -265,6 +273,8 @@ class ResponseProxy:
265
273
  else:
266
274
  encoding_quality = True
267
275
  encoding_split[1] = float(encoding_split[1].split('=')[1])
276
+ if float(encoding_split[1]) <= 0: # q=0 表示该算法不可接受
277
+ continue
268
278
  encodings.append(encoding_split)
269
279
  encodings.sort(key = lambda x: float(x[1]), reverse = True)
270
280
  encodings = [encoding[0] for encoding in encodings]
@@ -304,51 +314,55 @@ class ResponseProxy:
304
314
  if cookie['http_only']:
305
315
  _cookie += '; HttpOnly'
306
316
  cookies.append(_cookie)
307
- headers['set-cookie'] = ', '.join(cookies)
317
+ headers['set-cookie'] = cookies # set-cookie 需要写成多行,此处保留列表,由 ResponseProxy.send 展开
308
318
 
309
319
  return status, headers, body
310
320
 
311
321
  async def get_body(self, status: int, headers: dict[str, str], body: dict | list | str | bytes | AsyncIterable | None) -> AsyncIterable[tuple[int, dict[str, str], bytes]]:
312
322
  if type(self.response) is FileResponse and self.request.ranges and status != 416:
313
- if self.response.file._data is None:
323
+ if self.response.file._data is not None:
324
+ size = len(self.response.file.data)
325
+ else:
326
+ size = os.path.getsize(self.response.file.path)
314
327
  handler = open(self.response.file.path, 'rb')
328
+
315
329
  if len(self.request.ranges) == 1:
330
+ start, end = self.request.ranges[0]
331
+ end = size - 1 if end is None else end # 区间为闭区间,末尾含在响应内
316
332
  if self.response.file._data is not None:
317
- size = len(self.response.file.data)
318
- data = self.response.file.data[self.request.ranges[0][0] or 0:self.request.ranges[0][1] or len(self.response.file.data) + 1]
333
+ data = self.response.file.data[start:end + 1]
319
334
  else:
320
- size = os.path.getsize(self.response.file.path)
321
- handler.seek(self.request.ranges[0][0] or 0)
322
- data = handler.read((self.request.ranges[0][1] or size) - (self.request.ranges[0][0]))
335
+ handler.seek(start)
336
+ data = handler.read(end + 1 - start)
323
337
  handler.close()
338
+
324
339
  headers['content-length'] = str(len(data))
325
- headers['content-range'] = f'bytes {self.request.ranges[0][0]}-{(self.request.ranges[0][1] or size) - 1}/{size}'
340
+ headers['content-range'] = f'bytes {start}-{end}/{size}'
326
341
  yield status, headers, data
327
342
  else:
328
343
  boundary = uuid.uuid4().hex
329
344
  content_type = headers['content-type']
330
345
  headers['content-type'] = f'multipart/byteranges; boundary={boundary}'
331
- if self.response.file._data is not None:
332
- size = len(self.response.file.data)
333
- else:
334
- size = os.path.getsize(self.response.file.path)
335
-
336
- content_length = 0
337
- for range in self.request.ranges:
338
- content_length += 2 + 32 + 2 + 14 + len(content_type) + 2 + 21 + (1 if range[0] == 0 else int(math.log10(range[0]))) + 1 + 1 + int(math.log10(range[1] or size)) + 1 + 1 + int(math.log10(size)) + 1 + 4 + (range[1] or size) - range[0] + 1
339
- headers['content-length'] = str(content_length)
340
346
 
341
- for range in self.request.ranges:
342
- data = [b'--', boundary.encode(), b'\r\n', b'content-type: ', content_type.encode(), b'\r\n', b'content-range: bytes ', str(range[0]).encode(), b'-', str(range[1] or size).encode(), b'/', str(size).encode() + b'\r\n\r\n']
347
+ parts = []
348
+ for start, end in self.request.ranges:
349
+ end = size - 1 if end is None else end
343
350
  if self.response.file._data is not None:
344
- data.append(self.response.file.data[range[0]:range[1] or size + 1])
351
+ data = self.response.file.data[start:end + 1]
345
352
  else:
346
- handler.seek(range[0])
347
- data.append(handler.read((range[1] or size) - (range[0])))
348
- yield status, headers, b''.join(data)
353
+ handler.seek(start)
354
+ data = handler.read(end + 1 - start)
355
+ parts.append(b''.join([b'--', boundary.encode(), b'\r\n', b'content-type: ', content_type.encode(), b'\r\n', b'content-range: bytes ', str(start).encode(), b'-', str(end).encode(), b'/', str(size).encode(), b'\r\n\r\n', data]))
356
+
349
357
  if self.response.file._data is None:
350
358
  handler.close()
351
- yield status, headers, b'--' + boundary.encode() + b'--'
359
+
360
+ trailer = b'--' + boundary.encode() + b'--'
361
+ headers['content-length'] = str(sum(len(part) for part in parts) + len(trailer))
362
+
363
+ for part in parts:
364
+ yield status, headers, part
365
+ yield status, headers, trailer
352
366
  else:
353
367
  if isinstance(body, AsyncIterable):
354
368
  data = await anext(body)
@@ -386,11 +400,11 @@ class ResponseProxy:
386
400
 
387
401
  async def get_encode_body(self, status: int, headers: dict[str, str], body: bytes) -> tuple[int, dict[str, str], bytes]:
388
402
  content_length = headers.get('content-length')
389
- if content_length and int(content_length) < self.app.compress_min_length:
403
+ if self.response.compress is None and content_length and int(content_length) < self.app.compress_min_length:
390
404
  if 'content-encoding' in headers:
391
405
  del headers['content-encoding']
392
406
 
393
- if 'content-encoding' in headers and 'content-length' in headers and (int(headers['content-length']) > self.app.compress_min_length or self.response.compress is not None):
407
+ if 'content-encoding' in headers and 'content-length' in headers and (self.response.compress is not None or int(headers['content-length']) >= self.app.compress_min_length):
394
408
  compress_level = self.response.compress_level if self.response.compress_level is not None else self.app.compress_level
395
409
  if headers['content-encoding'] == 'gzip':
396
410
  body = gzip.compress(body, compress_level)
@@ -1,5 +1,6 @@
1
1
  import inspect, os
2
2
  import datetime, uuid, threading, multiprocessing, asyncio, time, json
3
+ from queue import Empty
3
4
  from typing import Callable, Literal, TYPE_CHECKING
4
5
 
5
6
  import redis, redis.exceptions
@@ -13,22 +14,25 @@ class Task:
13
14
  @classmethod
14
15
  def from_dict(cls, data: dict[str, any], _scheduler_proxy) -> 'Task':
15
16
  instance = cls.__new__(cls)
17
+
16
18
  for key, value in data.items():
17
19
  if key == '_queue':
18
- if value:
19
- value = multiprocessing.get_context('spawn').Queue()
20
- value.put(None)
21
- else:
22
- value = multiprocessing.get_context('spawn').Queue()
20
+ value = multiprocessing.get_context('spawn').Queue()
21
+ elif key == '_stop_event':
22
+ value = threading.Event()
23
23
  elif key == 'first_run_timer':
24
24
  value = datetime.datetime.fromtimestamp(value) if value else None
25
25
  elif key == '_last_run_timer':
26
26
  value = datetime.datetime.fromtimestamp(value) if value else None
27
27
  setattr(instance, key, value)
28
+
29
+ if '_running_remote' not in data:
30
+ instance._running_remote = not data.get('_queue', True) # 旧数据只有 `_queue`(非空代表未运行)
31
+
28
32
  setattr(instance, '_scheduler_proxy', _scheduler_proxy)
29
33
  return instance
30
34
 
31
- __slots__ = ('fn', 'interval_time', 'first_run_timer', 'expected_run_num', '_key', 'run_type', 'args', 'kwargs', 'auto_remove', '_last_run_timer', '_last_run_time', '_run_num', '_handler', '_queue', '_scheduler_proxy', 'timeout')
35
+ __slots__ = ('fn', 'interval_time', 'first_run_timer', 'expected_run_num', '_key', 'run_type', 'args', 'kwargs', 'auto_remove', '_last_run_timer', '_last_run_time', '_run_num', '_handler', '_queue', '_stop_event', '_running_remote', '_scheduler_proxy', 'timeout')
32
36
 
33
37
  def __init__(self, fn: Callable, interval_time: float, *, first_run_timer: datetime.datetime | float | None = None, expected_run_num: int | None = None, key: str | None = None, run_type: Literal['THREAD', 'PROCESS', 'ASYNC'] = 'THREAD', args: tuple = (), kwargs: dict = {}, auto_remove: bool = False, timeout: float | None = None, _scheduler_proxy: 'SchedulerProxy'):
34
38
  '''
@@ -67,23 +71,66 @@ class Task:
67
71
  self._last_run_time: float | None = None
68
72
  self._run_num: int = 0
69
73
  self._handler: threading.Thread | multiprocessing.Process | asyncio.Task | None = None
74
+
70
75
  self._queue = multiprocessing.get_context('spawn').Queue()
76
+ ''' 进程任务的停止信号队列(父进程放入、子进程取出) '''
77
+ self._stop_event = threading.Event()
78
+ ''' 线程 / 协程任务的停止信号(同进程内使用) '''
79
+ self._running_remote: bool = False
80
+ ''' 任务是否运行在其它进程(由 `sync_server` 反序列化而来的任务) '''
81
+
82
+ '''
83
+ 停止信号的收发
84
+
85
+ 不使用 `multiprocessing.Queue.qsize()`(macOS 未实现,调用即 `NotImplementedError`):
86
+ 线程 / 协程任务用 `threading.Event`,进程任务用队列的 `empty()` / `get_nowait()` 判断与取出信号。
87
+ '''
88
+
89
+ def _stop(self):
90
+ ''' 发送停止信号 '''
91
+
92
+ if self.run_type == 'PROCESS':
93
+ self._queue.put(None)
94
+ else:
95
+ self._stop_event.set()
96
+
97
+ def _reset_stop(self):
98
+ ''' 清除停止信号(启动任务前调用) '''
99
+
100
+ if self.run_type == 'PROCESS':
101
+ while True:
102
+ try:
103
+ self._queue.get_nowait()
104
+ except Empty:
105
+ break
106
+ else:
107
+ self._stop_event.clear()
108
+
109
+ def _stops(self) -> bool:
110
+ ''' 是否收到了停止信号 '''
111
+
112
+ if self.run_type == 'PROCESS':
113
+ return not self._queue.empty()
114
+ return self._stop_event.is_set()
71
115
 
72
116
  def __getstate__(self) -> tuple[None, dict[str, any]]:
73
117
  state = {
74
118
  key: getattr(self, key) for key in self.__slots__
75
119
  }
76
120
  state['_handler'] = None
121
+ state['_stop_event'] = None # `threading.Event` 不能被 pickle,进程任务只用队列信号
77
122
  return None, state
78
123
 
79
124
  def _to_dict(self) -> dict[str, any]:
80
125
  data = {
81
126
  key: getattr(self, key) for key in self.__slots__
82
127
  }
83
- data['_queue'] = bool(data['_queue'].qsize())
128
+ data['_queue'] = None
129
+ data['_stop_event'] = None
130
+ data['_handler'] = None
131
+ data['_running_remote'] = self.is_running
84
132
  data['first_run_timer'] = self.first_run_timer.timestamp() if self.first_run_timer else None
85
133
  data['_last_run_timer'] = self._last_run_timer.timestamp() if self._last_run_timer else None
86
- data['_handler'] = None
87
134
  data['fn'] = None
88
135
  data['args'] = tuple()
89
136
  data['kwargs'] = {}
@@ -95,7 +142,7 @@ class Task:
95
142
 
96
143
  async def async_start(self):
97
144
  if not self._handler:
98
- self._handler = asyncio.create_task(self._scheduler_proxy.async_task_processing(self.key, self._queue, self.fn, *self.args, **self.kwargs))
145
+ self._handler = asyncio.create_task(self._scheduler_proxy.async_task_processing(self.key, self.fn, *self.args, **self.kwargs))
99
146
 
100
147
  @property
101
148
  def key(self) -> str:
@@ -105,7 +152,7 @@ class Task:
105
152
  def run_num_completed(self) -> bool:
106
153
  if self.expected_run_num is None:
107
154
  return False
108
- return self._run_num >= self.expected_run_num
155
+ return self.run_num >= self.expected_run_num
109
156
 
110
157
  @property
111
158
  def last_run_timer(self) -> datetime.datetime | None:
@@ -129,7 +176,11 @@ class Task:
129
176
  def is_running(self) -> bool:
130
177
  ''' 任务是否在运行中 '''
131
178
 
132
- return not self._queue.qsize()
179
+ if self._handler is None:
180
+ return self._running_remote
181
+ if isinstance(self._handler, asyncio.Task):
182
+ return not self._handler.done()
183
+ return self._handler.is_alive()
133
184
 
134
185
  class Scheduler:
135
186
  __slots__ = ('_proxy',)
@@ -310,7 +361,6 @@ class SchedulerProxy:
310
361
  coro.close()
311
362
 
312
363
  task = Task(fn, interval_time, first_run_timer = first_run_timer, expected_run_num = expected_run_num, key = key, run_type = run_type, args = args, kwargs = kwargs, auto_remove = auto_remove, timeout = timeout, _scheduler_proxy = self)
313
- task._queue.put(None)
314
364
 
315
365
  if task.key in self.get_tasks():
316
366
  raise KeyError(f'Task with key "{task.key}" already exists')
@@ -339,7 +389,6 @@ class SchedulerProxy:
339
389
  coro.close()
340
390
 
341
391
  task = Task(fn, interval_time = interval_time, first_run_timer = first_run_timer, expected_run_num = expected_run_num, key = key, run_type = 'ASYNC', args = args, kwargs = kwargs, auto_remove = auto_remove, timeout = timeout, _scheduler_proxy = self)
342
- task._queue.put(None)
343
392
 
344
393
  if task.key in await self.async_get_tasks():
345
394
  raise KeyError(f'Task with key "{task.key}" already exists')
@@ -356,16 +405,12 @@ class SchedulerProxy:
356
405
  return _fn
357
406
  return wrapper
358
407
 
359
- def task_processing(self, task: Task, queue: multiprocessing.Queue, fn, *args, **kwargs):
408
+ def task_processing(self, task: Task, fn, *args, **kwargs):
360
409
  try:
361
- queue.get()
362
-
363
- if static.scheduler_sync_servers:
364
- self.get_task(task.key)._queue.get()
365
410
  if task.first_run_timer:
366
411
  time.sleep(max(0, task.first_run_timer.timestamp() - time.time()))
367
412
 
368
- while not queue.qsize():
413
+ while not task._stops():
369
414
  now = time.time()
370
415
 
371
416
  try:
@@ -378,7 +423,7 @@ class SchedulerProxy:
378
423
  except Exception as e:
379
424
  self.app.printer.scheduler_error(e, task)
380
425
 
381
- if queue.qsize():
426
+ if task._stops():
382
427
  break
383
428
 
384
429
  task._last_run_time = time.time() - now
@@ -390,6 +435,9 @@ class SchedulerProxy:
390
435
  sync_server.hset('CheeseAPI_scheduler_tasks', task.key, json.dumps(task._to_dict()))
391
436
  sync_server.hpexpire('CheeseAPI_scheduler_tasks', int(task.timeout * 1000), task.key)
392
437
 
438
+ if task.run_num_completed:
439
+ break
440
+
393
441
  time.sleep(max(0, task.interval_time - time.time() + now))
394
442
  except (KeyboardInterrupt, SystemExit):
395
443
  ...
@@ -399,16 +447,16 @@ class SchedulerProxy:
399
447
  if _redis.hexists('CheeseAPI_scheduler_tasks', task.key):
400
448
  _redis.hpersist('CheeseAPI_scheduler_tasks', task.key)
401
449
 
402
- async def async_task_processing(self, key: str, queue: multiprocessing.Queue, fn, *args, **kwargs):
403
- queue.get()
450
+ if task.run_type == 'PROCESS':
451
+ task._reset_stop() # 取出残留的停止信号,保证下次启动时队列是干净的
404
452
 
453
+ async def async_task_processing(self, key: str, fn, *args, **kwargs):
405
454
  task = await self.async_get_task(key)
406
- if static.scheduler_sync_servers:
407
- task._queue.get()
455
+
408
456
  if task.first_run_timer:
409
457
  await asyncio.sleep(max(0, task.first_run_timer.timestamp() - time.time()))
410
458
 
411
- while not queue.qsize():
459
+ while not task._stops():
412
460
  now = time.time()
413
461
 
414
462
  try:
@@ -421,7 +469,7 @@ class SchedulerProxy:
421
469
  except Exception as e:
422
470
  self.app.printer.scheduler_error(e, task)
423
471
 
424
- if queue.qsize():
472
+ if task._stops():
425
473
  break
426
474
 
427
475
  task._last_run_time = time.time() - now
@@ -469,10 +517,11 @@ class SchedulerProxy:
469
517
  redis.Redis(connection_pool = static.scheduler_sync_servers[0]).publish('CheeseAPI_scheduler', json.dumps(['start', key]))
470
518
  return
471
519
 
520
+ task._reset_stop()
472
521
  if task.run_type == 'THREAD':
473
- task._handler = threading.Thread(target = self.task_processing, args = (task, task._queue, task.fn, *task.args), kwargs = task.kwargs, daemon = True)
522
+ task._handler = threading.Thread(target = self.task_processing, args = (task, task.fn, *task.args), kwargs = task.kwargs, daemon = True)
474
523
  elif task.run_type == 'PROCESS':
475
- task._handler = multiprocessing.get_context('spawn').Process(target = self.task_processing, args = (task, task._queue, task.fn, *task.args), kwargs = task.kwargs, daemon = True)
524
+ task._handler = multiprocessing.get_context('spawn').Process(target = self.task_processing, args = (task, task.fn, *task.args), kwargs = task.kwargs, daemon = True)
476
525
  task._handler.start()
477
526
 
478
527
  if task.auto_remove:
@@ -489,7 +538,8 @@ class SchedulerProxy:
489
538
  if not task and static.scheduler_sync_servers:
490
539
  await redis.asyncio.Redis(connection_pool = static.scheduler_sync_servers[1]).publish('CheeseAPI_scheduler', json.dumps(['start', key]))
491
540
  else:
492
- task._handler = asyncio.create_task(self.async_task_processing(key, task._queue, task.fn, *task.args, **task.kwargs))
541
+ task._reset_stop()
542
+ task._handler = asyncio.create_task(self.async_task_processing(key, task.fn, *task.args, **task.kwargs))
493
543
 
494
544
  def stop(self, key: str):
495
545
  task = self.get_task(key)
@@ -503,7 +553,7 @@ class SchedulerProxy:
503
553
  if static.scheduler_sync_servers:
504
554
  redis.Redis(connection_pool = static.scheduler_sync_servers[0]).publish('CheeseAPI_scheduler', json.dumps(['stop', key]))
505
555
  else:
506
- local_task._queue.put(None)
556
+ local_task._stop()
507
557
  if static.scheduler_sync_servers:
508
558
  redis.Redis(connection_pool = static.scheduler_sync_servers[0]).hset('CheeseAPI_scheduler_tasks', key, json.dumps(task._to_dict()))
509
559
 
@@ -519,7 +569,7 @@ class SchedulerProxy:
519
569
  if static.scheduler_sync_servers:
520
570
  redis.Redis(connection_pool = static.scheduler_sync_servers[0]).publish('CheeseAPI_scheduler', json.dumps(['remove', key]))
521
571
  else:
522
- local_task._queue.put(None)
572
+ local_task._stop()
523
573
  self._tasks.pop(key, None)
524
574
  if static.scheduler_sync_servers:
525
575
  redis.Redis(connection_pool = static.scheduler_sync_servers[0]).hdel('CheeseAPI_scheduler_tasks', key)
@@ -536,7 +586,7 @@ class SchedulerProxy:
536
586
  if static.scheduler_sync_servers:
537
587
  await redis.asyncio.Redis(connection_pool = static.scheduler_sync_servers[1]).publish('CheeseAPI_scheduler', json.dumps(['stop', key]))
538
588
  else:
539
- local_task._queue.put(None)
589
+ local_task._stop()
540
590
  if static.scheduler_sync_servers:
541
591
  await redis.asyncio.Redis(connection_pool = static.scheduler_sync_servers[1]).hset('CheeseAPI_scheduler_tasks', key, json.dumps(task._to_dict()))
542
592
 
@@ -550,7 +600,7 @@ class SchedulerProxy:
550
600
  if static.scheduler_sync_servers:
551
601
  await redis.asyncio.Redis(connection_pool = static.scheduler_sync_servers[1]).publish('CheeseAPI_scheduler', json.dumps(['remove', key]))
552
602
  else:
553
- local_task._queue.put(None)
603
+ local_task._stop()
554
604
  self._tasks.pop(key, None)
555
605
  if static.scheduler_sync_servers:
556
606
  await redis.asyncio.Redis(connection_pool = static.scheduler_sync_servers[1]).hdel('CheeseAPI_scheduler_tasks', key)
@@ -111,6 +111,37 @@ class Websocket:
111
111
  class WebsocketProxy:
112
112
  _sync_tasks: dict[str, asyncio.Task] = {}
113
113
 
114
+ @staticmethod
115
+ def _consume_exception(future: asyncio.Future):
116
+ ''' 消费投递结果的异常,避免出现 "exception was never retrieved" '''
117
+ if not future.cancelled():
118
+ future.exception()
119
+
120
+ @staticmethod
121
+ def _schedule(connector: Websocket, coroutine):
122
+ '''
123
+ 同步入口(`_static_send` / `_static_close`)不能 await,将协程投递到连接所属的事件循环执行
124
+
125
+ 既不能直接丢弃(会报 `coroutine ... was never awaited` 且消息实际发不出去),
126
+ 也不能在调用方的循环里执行(连接的 writer 属于创建它的那个循环)
127
+ '''
128
+ loop: asyncio.AbstractEventLoop | None = connector._proxy._loop
129
+
130
+ try:
131
+ current_loop = asyncio.get_running_loop()
132
+ except RuntimeError:
133
+ current_loop = None
134
+
135
+ if loop is None:
136
+ coroutine.close()
137
+ elif current_loop is loop:
138
+ loop.create_task(coroutine)
139
+ elif loop.is_running():
140
+ ''' 调用方在其他线程 / 非协程上下文,线程安全地投递回该循环 '''
141
+ asyncio.run_coroutine_threadsafe(coroutine, loop).add_done_callback(WebsocketProxy._consume_exception)
142
+ else:
143
+ coroutine.close()
144
+
114
145
  @staticmethod
115
146
  def _static_send(path: str, data: bytes | list | str | dict, *, websocket_key_or_keys: str | list[str] | None = None):
116
147
  if static.websocket_sync_servers is not None:
@@ -131,16 +162,16 @@ class WebsocketProxy:
131
162
  elif path in Websocket.connectors:
132
163
  if websocket_key_or_keys is None:
133
164
  for connector in Websocket.connectors[path]:
134
- connector.send(data)
165
+ WebsocketProxy._schedule(connector, connector.send(data))
135
166
  elif isinstance(websocket_key_or_keys, str):
136
167
  for connector in Websocket.connectors[path]:
137
168
  if connector.key == websocket_key_or_keys:
138
- connector.send(data)
169
+ WebsocketProxy._schedule(connector, connector.send(data))
139
170
  break
140
171
  elif isinstance(websocket_key_or_keys, list):
141
172
  for connector in Websocket.connectors[path]:
142
173
  if connector.key in websocket_key_or_keys:
143
- connector.send(data)
174
+ WebsocketProxy._schedule(connector, connector.send(data))
144
175
 
145
176
  @staticmethod
146
177
  async def async_send(path: str, data: bytes | list | str | dict, *, websocket_key_or_keys: str | list[str] | None = None):
@@ -182,16 +213,16 @@ class WebsocketProxy:
182
213
  elif path in Websocket.connectors:
183
214
  if websocket_key_or_keys is None:
184
215
  for connector in Websocket.connectors[path]:
185
- connector.close(code, message)
216
+ WebsocketProxy._schedule(connector, connector.close(code, message))
186
217
  elif isinstance(websocket_key_or_keys, str):
187
218
  for connector in Websocket.connectors[path]:
188
219
  if connector.key == websocket_key_or_keys:
189
- connector.close(code, message)
220
+ WebsocketProxy._schedule(connector, connector.close(code, message))
190
221
  break
191
222
  elif isinstance(websocket_key_or_keys, list):
192
223
  for connector in Websocket.connectors[path]:
193
224
  if connector.key in websocket_key_or_keys:
194
- connector.close(code, message)
225
+ WebsocketProxy._schedule(connector, connector.close(code, message))
195
226
 
196
227
  @staticmethod
197
228
  async def async_close(path: str, code: int = 1000, message: str = '', *, websocket_key_or_keys: str | list[str] | None = None):
@@ -218,16 +249,25 @@ class WebsocketProxy:
218
249
 
219
250
  self.reader: asyncio.StreamReader | None = None
220
251
  self.writer: asyncio.StreamWriter | None = None
252
+ self._loop: asyncio.AbstractEventLoop | None = None
253
+ self._is_disconnected: bool = False
221
254
 
222
255
  async def running(self) -> AsyncIterable[Response]:
223
256
  try:
224
257
  await self.connect()
225
258
  await self.message()
226
- await self.disconnect()
227
259
  except Exception as e:
228
260
  result = self.app.printer.websocket_error(e, self.websocket)
229
261
  if inspect.isawaitable(result):
230
262
  await result
263
+ finally:
264
+ ''' 无论连接、接收如何结束(含异常断连),都必须执行一次清理,否则心跳与推送任务不会停止 '''
265
+ try:
266
+ await self.disconnect()
267
+ except Exception as e:
268
+ result = self.app.printer.websocket_error(e, self.websocket)
269
+ if inspect.isawaitable(result):
270
+ await result
231
271
 
232
272
  async def get_response(self) -> Response:
233
273
  headers = {
@@ -251,6 +291,7 @@ class WebsocketProxy:
251
291
  WebsocketProxy._sync_tasks[self.websocket.request.path] = asyncio.create_task(self.sync_server_running())
252
292
 
253
293
  loop = asyncio.get_running_loop()
294
+ self._loop = loop
254
295
  self.reader = asyncio.StreamReader()
255
296
  protocol = asyncio.StreamReaderProtocol(self.reader)
256
297
  if isinstance(self.websocket.request._proxy.client_socket, ssl.SSLSocket):
@@ -346,13 +387,52 @@ class WebsocketProxy:
346
387
  await self.websocket.on_pong()
347
388
 
348
389
  async def disconnect(self):
390
+ '''
391
+ 清理连接
392
+
393
+ - 允许重复调用:接收、发送两侧可能同时发现断连,重复清理直接返回
394
+ - 兼容连接未完全建立的情况:reader / writer 可能为 None,也未加入 connectors
395
+ - 保证执行 on_disconnect():心跳任务与业务推送任务依赖它停止
396
+ - 用 finally 保证底层连接关闭,即使回调清理报错也要关闭
397
+ '''
398
+ if self._is_disconnected:
399
+ return
400
+ self._is_disconnected = True
401
+
349
402
  self.websocket._is_running = False
350
- Websocket.connectors[self.websocket.request.path].remove(self.websocket)
351
- self.app.printer.websocket_disconnect(self.websocket)
352
- await self.websocket.on_disconnect()
353
- if self.writer:
354
- self.writer.close()
355
- await self.writer.wait_closed()
403
+
404
+ connectors = Websocket.connectors.get(self.websocket.request.path)
405
+ if connectors is not None:
406
+ try:
407
+ connectors.remove(self.websocket)
408
+ except ValueError:
409
+ ...
410
+
411
+ try:
412
+ self.app.printer.websocket_disconnect(self.websocket)
413
+ except Exception:
414
+ ...
415
+
416
+ try:
417
+ await self.websocket.on_disconnect()
418
+ except Exception as e:
419
+ result = self.app.printer.websocket_error(e, self.websocket)
420
+ if inspect.isawaitable(result):
421
+ await result
422
+ finally:
423
+ writer, self.writer = self.writer, None
424
+ self.reader = None
425
+
426
+ if writer is not None:
427
+ try:
428
+ writer.close()
429
+ except Exception:
430
+ ...
431
+ try:
432
+ await writer.wait_closed()
433
+ except Exception:
434
+ ''' 对端已断开时 wait_closed 可能再报 OSError,属于预期错误,不能中断清理 '''
435
+ ...
356
436
 
357
437
  def encode(self, opcode: int, data: bytes) -> bytes:
358
438
  _bytes = bytearray()
@@ -406,23 +486,51 @@ class WebsocketProxy:
406
486
  except asyncio.TimeoutError:
407
487
  return opcode, full_payload
408
488
 
489
+ def _check_running(self):
490
+ ''' 发送前的连接校验,连接已断开或未建立时直接拒绝,避免继续向坏连接写入 '''
491
+ if self._is_disconnected or self.writer is None or not self.websocket._is_running:
492
+ raise ConnectionError('The websocket is disconnected, data will not be sent')
493
+
494
+ def _disconnected(self, e: Exception) -> ConnectionError:
495
+ ''' 标记连接断开并返回统一的异常,阻止后续继续向坏连接发送 '''
496
+ self.websocket._is_running = False
497
+ self._is_disconnected = True
498
+ return ConnectionError('The websocket is disconnected, data will not be sent')
499
+
409
500
  async def _instance_send(self, data: bytes | list | str | dict, **kwargs):
410
- if isinstance(data, str):
411
- self.writer.write(self.encode(0x1, data.encode()))
412
- elif isinstance(data, bytes):
413
- self.writer.write(self.encode(0x2, data))
414
- else:
415
- self.writer.write(self.encode(0x1, json.dumps(data).encode()))
416
- await self.writer.drain()
501
+ self._check_running()
502
+ try:
503
+ if isinstance(data, str):
504
+ self.writer.write(self.encode(0x1, data.encode()))
505
+ elif isinstance(data, bytes):
506
+ self.writer.write(self.encode(0x2, data))
507
+ else:
508
+ self.writer.write(self.encode(0x1, json.dumps(data).encode()))
509
+ await self.writer.drain()
510
+ except (ConnectionError, OSError, RuntimeError) as e:
511
+ raise self._disconnected(e) from e
417
512
 
418
513
  async def _instance_close(self, code: int = 1000, message: str = ''):
419
- self.writer.write(self.encode(0x8, struct.pack('!H', code) + message.encode('utf-8')))
420
- await self.writer.drain()
514
+ if self._is_disconnected or self.writer is None:
515
+ return
516
+ try:
517
+ self.writer.write(self.encode(0x8, struct.pack('!H', code) + message.encode('utf-8')))
518
+ await self.writer.drain()
519
+ except (ConnectionError, OSError, RuntimeError) as e:
520
+ raise self._disconnected(e) from e
421
521
 
422
522
  async def ping(self):
423
- self.writer.write(self.encode(0x9, b''))
424
- await self.writer.drain()
523
+ self._check_running()
524
+ try:
525
+ self.writer.write(self.encode(0x9, b''))
526
+ await self.writer.drain()
527
+ except (ConnectionError, OSError, RuntimeError) as e:
528
+ raise self._disconnected(e) from e
425
529
 
426
530
  async def pong(self):
427
- self.writer.write(self.encode(0xA, b''))
428
- await self.writer.drain()
531
+ self._check_running()
532
+ try:
533
+ self.writer.write(self.encode(0xA, b''))
534
+ await self.writer.drain()
535
+ except (ConnectionError, OSError, RuntimeError) as e:
536
+ raise self._disconnected(e) from e
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: CheeseAPI
3
- Version: 2.0.8b7
3
+ Version: 2.0.8-beta.8
4
4
  Summary: 一款web协程框架
5
5
  Project-URL: Source, https://github.com/CheeseUnknown/CheeseAPI
6
6
  Author-email: Cheese Unknown <cheese@cheese.ren>
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "CheeseAPI"
7
- version = "2.0.8-beta.7"
7
+ version = "2.0.8-beta.8"
8
8
  description = "一款web协程框架"
9
9
  readme = "README.md"
10
10
  license-files = { paths = [ "LICENSE" ] }
File without changes
File without changes
File without changes