Archived
feat(routes): 迁移下载/批量打包/播客下载(账本事务)+ 修密钥绑定缺陷
迁移端点:
- POST /api/v1/download(单曲:file 后台任务 / json 同步交付+结算)
- POST /api/v1/download/batch_zip(按去重歌曲数计费)
- POST /api/v1/podcast/download、/podcast/download/batch_zip
- 新增 web/billing.py(预扣+开账同事务 / 只退 open 行 / rename 承接)
与 web/tasks_registry.py(任务登记/容量闸/账本撞名重试,与 main 同口径)
- worker 仍复用 main 既有实现(无 flask 依赖),账本与交付语义不变
修复我引入的真实缺陷:
- download 参数校验分支(ID/音质/格式非法)漏退款 → 预扣泄漏。改为统一 _fail() 退款
- web/routers/{auth,register}.py 在导入期绑定 deps.secret_key,导致 patch web.deps.secret_key
无效、且测试顺序相关。改为调用期 deps.secret_key() 解析
测试:tests/test_download_billing.py 7 项沙箱用例如实断言账本守恒
(预扣必退、无 open 行残留、次数不足不建任务);
全套 pytest 69/69(顺序无关);ruff All checks passed;pyright 0 errors;奇偶 36/36。
This commit is contained in:
1 parent
fb40b2132c
commit
9df0028195
9 files changed
+808
-7
No files matched your search
@@ -40,7 +40,6 @@ def app_env(tmp_path, monkeypatch):
|
|||||||
conn.close()
|
conn.close()
|
||||||
|
|
||||||
monkeypatch.setattr('web.deps.secret_key', lambda: SECRET)
|
monkeypatch.setattr('web.deps.secret_key', lambda: SECRET)
|
||||||
monkeypatch.setattr('web.routers.auth.secret_key', lambda: SECRET)
|
|
||||||
|
|
||||||
# 清空登录失败锁定/2FA 的模块级内存态,避免测试间串扰
|
# 清空登录失败锁定/2FA 的模块级内存态,避免测试间串扰
|
||||||
import auth_routes
|
import auth_routes
|
||||||
|
|||||||
@@ -0,0 +1,141 @@
|
|||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""下载/计费端点沙箱测试(不触碰生产 data/site.db)。
|
||||||
|
|
||||||
|
重点验证账本守恒:
|
||||||
|
- 未登录 401
|
||||||
|
- 无效歌曲 ID:400 且**预扣已退**(账本不留 open 行、余额不变)
|
||||||
|
- 次数不足:401/expired 且不建任务
|
||||||
|
- 批量超 500:400 且退款
|
||||||
|
- 参数校验失败各分支均按「只退 open 行」退费
|
||||||
|
"""
|
||||||
|
import pytest
|
||||||
|
from fastapi.testclient import TestClient
|
||||||
|
|
||||||
|
import auth_system as auth
|
||||||
|
import core.db as db
|
||||||
|
|
||||||
|
SECRET = 'download-test-secret'
|
||||||
|
PW = 'Passw0rd12'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture()
|
||||||
|
def env(tmp_path, monkeypatch):
|
||||||
|
db_path = tmp_path / 'site.db'
|
||||||
|
monkeypatch.setattr(auth, 'DATA_DIR', tmp_path)
|
||||||
|
monkeypatch.setattr(auth, 'DB_PATH', db_path)
|
||||||
|
monkeypatch.setattr(auth, 'BACKUP_DIR', tmp_path / 'backups')
|
||||||
|
monkeypatch.setattr(auth, 'SECRET_KEY_PATH', tmp_path / 'secret.key')
|
||||||
|
monkeypatch.setattr(db, 'DATA_DIR', tmp_path)
|
||||||
|
monkeypatch.setattr(db, 'DB_PATH', db_path)
|
||||||
|
auth.init_db()
|
||||||
|
# 单一密钥源:web 层统一经 web.deps.secret_key 取密钥(调用期解析),
|
||||||
|
# patch 它即让登录签发与依赖校验同源
|
||||||
|
monkeypatch.setattr('web.deps.secret_key', lambda: SECRET)
|
||||||
|
|
||||||
|
# 造 3 次余额用户 + 0 次余额用户
|
||||||
|
auth.create_user('payer', PW, email='p@example.com', credits=3)
|
||||||
|
auth.create_user('broke', PW, email='b@example.com', credits=0)
|
||||||
|
auth.create_user('rich', PW, email='r@example.com', credits=600)
|
||||||
|
|
||||||
|
from web.app import create_app
|
||||||
|
return TestClient(create_app())
|
||||||
|
|
||||||
|
|
||||||
|
def _login(env, username):
|
||||||
|
"""登录并返回 (cookie_value, body)。
|
||||||
|
|
||||||
|
直接取 Set-Cookie 原始值并用 Cookie 头携带——httpx 的 cookie jar 会对含
|
||||||
|
base64 padding('=') 的值加双引号,导致服务端签名解析失败(真实浏览器不会),
|
||||||
|
故测试避开 jar。
|
||||||
|
"""
|
||||||
|
r = env.post('/api/v1/auth/login', json={'username': username, 'password': PW})
|
||||||
|
assert r.status_code == 200, r.text
|
||||||
|
raw = r.headers.get('set-cookie', '')
|
||||||
|
cookie = ''
|
||||||
|
for part in raw.split(';'):
|
||||||
|
if part.strip().startswith('session='):
|
||||||
|
cookie = part.strip()[len('session='):]
|
||||||
|
break
|
||||||
|
assert cookie, f'no session cookie in: {raw!r}'
|
||||||
|
return cookie, r.json()
|
||||||
|
|
||||||
|
|
||||||
|
def _balance(username):
|
||||||
|
u = auth.get_user_by_name(username)
|
||||||
|
return int(u['duration_seconds'] or 0)
|
||||||
|
|
||||||
|
|
||||||
|
def _open_ledger_rows():
|
||||||
|
conn = auth.get_conn()
|
||||||
|
rows = conn.execute("SELECT task_id, reserved, status FROM billing_ledger WHERE status='open'").fetchall()
|
||||||
|
conn.close()
|
||||||
|
return [dict(r) for r in rows]
|
||||||
|
|
||||||
|
|
||||||
|
def test_download_requires_login(env):
|
||||||
|
r = env.post('/api/v1/download', json={'id': '347230'})
|
||||||
|
assert r.status_code == 401
|
||||||
|
|
||||||
|
|
||||||
|
def test_download_bad_id_refunds(env):
|
||||||
|
ck, _ = _login(env, 'payer')
|
||||||
|
H = {'Cookie': f'session={ck}'}
|
||||||
|
before = _balance('payer')
|
||||||
|
r = env.post('/api/v1/download', json={'id': 'not-a-number'}, headers=H)
|
||||||
|
assert r.status_code == 400
|
||||||
|
assert '无效' in r.json()['message']
|
||||||
|
# 预扣已退回,且无 open 账本行残留
|
||||||
|
assert _balance('payer') == before
|
||||||
|
assert _open_ledger_rows() == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_download_bad_quality_refunds(env):
|
||||||
|
ck, _ = _login(env, 'payer')
|
||||||
|
H = {'Cookie': f'session={ck}'}
|
||||||
|
before = _balance('payer')
|
||||||
|
r = env.post('/api/v1/download', json={'id': '347230', 'quality': 'nope'}, headers=H)
|
||||||
|
assert r.status_code == 400
|
||||||
|
assert _balance('payer') == before
|
||||||
|
assert _open_ledger_rows() == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_download_bad_format_refunds(env):
|
||||||
|
ck, _ = _login(env, 'payer')
|
||||||
|
H = {'Cookie': f'session={ck}'}
|
||||||
|
before = _balance('payer')
|
||||||
|
r = env.post('/api/v1/download', json={'id': '347230', 'format': 'xml'}, headers=H)
|
||||||
|
assert r.status_code == 400
|
||||||
|
assert _balance('payer') == before
|
||||||
|
|
||||||
|
|
||||||
|
def test_download_insufficient_credits(env):
|
||||||
|
ck, _ = _login(env, 'broke')
|
||||||
|
H = {'Cookie': f'session={ck}'}
|
||||||
|
r = env.post('/api/v1/download', json={'id': '347230'}, headers=H)
|
||||||
|
assert r.status_code == 401
|
||||||
|
assert r.json().get('error_code') == 'expired'
|
||||||
|
# 未开账
|
||||||
|
assert _open_ledger_rows() == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_batch_zip_too_many_refunds(env):
|
||||||
|
# 需足够余额才能为「封顶 500」预扣成功,再验证 400 时把 500 全额退回
|
||||||
|
ck, _ = _login(env, 'rich')
|
||||||
|
H = {'Cookie': f'session={ck}'}
|
||||||
|
before = _balance('rich')
|
||||||
|
ids = ','.join(str(1000 + i) for i in range(501))
|
||||||
|
r = env.post('/api/v1/download/batch_zip', json={'ids': ids}, headers=H)
|
||||||
|
assert r.status_code == 400
|
||||||
|
assert '500' in r.json()['message']
|
||||||
|
# 预扣(封顶 500)已全额退回,余额回到 600
|
||||||
|
assert _balance('rich') == before
|
||||||
|
assert _open_ledger_rows() == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_batch_zip_missing_ids_refunds(env):
|
||||||
|
ck, _ = _login(env, 'payer')
|
||||||
|
H = {'Cookie': f'session={ck}'}
|
||||||
|
before = _balance('payer')
|
||||||
|
r = env.post('/api/v1/download/batch_zip', json={}, headers=H)
|
||||||
|
assert r.status_code == 400
|
||||||
|
assert _balance('payer') == before
|
||||||
@@ -30,8 +30,6 @@ def env(tmp_path, monkeypatch):
|
|||||||
monkeypatch.setattr(db, 'DB_PATH', db_path)
|
monkeypatch.setattr(db, 'DB_PATH', db_path)
|
||||||
auth.init_db()
|
auth.init_db()
|
||||||
monkeypatch.setattr('web.deps.secret_key', lambda: SECRET)
|
monkeypatch.setattr('web.deps.secret_key', lambda: SECRET)
|
||||||
monkeypatch.setattr('web.routers.auth.secret_key', lambda: SECRET)
|
|
||||||
monkeypatch.setattr('web.routers.register.secret_key', lambda: SECRET)
|
|
||||||
import auth_routes
|
import auth_routes
|
||||||
auth_routes._SEND_CODE_IP.clear()
|
auth_routes._SEND_CODE_IP.clear()
|
||||||
from web.app import create_app
|
from web.app import create_app
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ from web.routers import admin_site as admin_site_router
|
|||||||
from web.routers import auth as auth_router
|
from web.routers import auth as auth_router
|
||||||
from web.routers import auth_status as auth_status_router
|
from web.routers import auth_status as auth_status_router
|
||||||
from web.routers import delivery as delivery_router
|
from web.routers import delivery as delivery_router
|
||||||
|
from web.routers import download as download_router
|
||||||
from web.routers import index as index_router
|
from web.routers import index as index_router
|
||||||
from web.routers import music as music_router
|
from web.routers import music as music_router
|
||||||
from web.routers import podcast as podcast_router
|
from web.routers import podcast as podcast_router
|
||||||
@@ -99,6 +100,7 @@ def create_app(serve_flask_fallback: bool = False) -> FastAPI:
|
|||||||
app.include_router(admin_router.router, prefix='/api/v1', tags=['admin'])
|
app.include_router(admin_router.router, prefix='/api/v1', tags=['admin'])
|
||||||
app.include_router(register_router.router, prefix='/api/v1', tags=['auth'])
|
app.include_router(register_router.router, prefix='/api/v1', tags=['auth'])
|
||||||
app.include_router(admin_site_router.router, prefix='/api/v1', tags=['admin'])
|
app.include_router(admin_site_router.router, prefix='/api/v1', tags=['admin'])
|
||||||
|
app.include_router(download_router.router, prefix='/api/v1', tags=['download'])
|
||||||
# 首页与静态资源(迁移完成形态;fallback 模式下同路径也能由 Flask 提供)
|
# 首页与静态资源(迁移完成形态;fallback 模式下同路径也能由 Flask 提供)
|
||||||
app.include_router(index_router.router, tags=['index'])
|
app.include_router(index_router.router, tags=['index'])
|
||||||
|
|
||||||
|
|||||||
+107
@@ -0,0 +1,107 @@
|
|||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""FastAPI 计费依赖(账本预扣/退款),与 Flask login_required 同口径。
|
||||||
|
|
||||||
|
Flask 版 login_required 用 flask.g 暂存请求级账本键;FastAPI 无 g,改用
|
||||||
|
「依赖返回值携带 ledger_key + 路由内显式退款」的方式,语义完全一致:
|
||||||
|
- 预扣与开账同事务:auth_system.reserve_and_open(A2/B8)
|
||||||
|
- 失败退款:只退「仍 open 的账本行」(与装饰器 _refund_open 同口径,杜绝双退)
|
||||||
|
- 任务接手:auth_system.ledger_rename 把请求级行改挂 task_id(worker 结算唯一入口)
|
||||||
|
|
||||||
|
金额相关逻辑一律走 auth_system 既有函数,本模块不新造账务语义。
|
||||||
|
"""
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
import secrets
|
||||||
|
from collections.abc import Callable
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
import auth_system
|
||||||
|
from web.deps import SiteUser
|
||||||
|
from web.errors import AppError
|
||||||
|
|
||||||
|
logger = logging.getLogger('web.billing')
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class Billing:
|
||||||
|
"""一次请求的计费上下文。"""
|
||||||
|
user: SiteUser
|
||||||
|
ledger_key: str | None # 请求级账本键(unlimited 用户为 None)
|
||||||
|
need: int # 本次预扣次数(free 路径为 0)
|
||||||
|
free: bool
|
||||||
|
|
||||||
|
async def refund_open(self, reason: str = 'rejected') -> None:
|
||||||
|
"""仅对仍 open 的账本行全额退款(被 rename/结算过则跳过)。"""
|
||||||
|
key = self.ledger_key
|
||||||
|
if not key:
|
||||||
|
return
|
||||||
|
|
||||||
|
def _do() -> None:
|
||||||
|
try:
|
||||||
|
row = auth_system.ledger_get(key)
|
||||||
|
if row and row['status'] == 'open':
|
||||||
|
auth_system.ledger_settle(key, 0, reason)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
await asyncio.to_thread(_do)
|
||||||
|
|
||||||
|
async def settle(self, delivered: int, status: str = 'delivered') -> None:
|
||||||
|
"""交付结算(json 模式等无 task_id 的场景)。"""
|
||||||
|
key = self.ledger_key
|
||||||
|
if not key:
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
await asyncio.to_thread(auth_system.ledger_settle, key, delivered, status)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error('账本结算失败(重启对账兜底): %s %s', key, e)
|
||||||
|
|
||||||
|
async def rename_to(self, task_id: str) -> bool:
|
||||||
|
"""把请求级账本行改挂 task_id(任务接手账目)。无请求级行时开 reserved=0 的账。"""
|
||||||
|
if not self.ledger_key:
|
||||||
|
await asyncio.to_thread(auth_system.ledger_open, task_id, self.user.id, 0)
|
||||||
|
return True
|
||||||
|
return await asyncio.to_thread(auth_system.ledger_rename, self.ledger_key, task_id)
|
||||||
|
|
||||||
|
|
||||||
|
async def require_billing(
|
||||||
|
user: SiteUser | None,
|
||||||
|
*,
|
||||||
|
action: str = '',
|
||||||
|
free: bool = False,
|
||||||
|
cost: Callable[[], int] | None = None,
|
||||||
|
ip: str = '-',
|
||||||
|
) -> Billing:
|
||||||
|
"""登录 + 预扣闸(对应 Flask login_required)。
|
||||||
|
|
||||||
|
free=True:仅要求登录,不扣次数(解析/试听类)。
|
||||||
|
cost: 预扣次数计算(下载类:单曲 1 / 批量按去重后歌曲数)。
|
||||||
|
"""
|
||||||
|
if user is None:
|
||||||
|
raise AppError('请先登录后再使用解析功能', 401, error_code='unauthorized')
|
||||||
|
|
||||||
|
fresh = await asyncio.to_thread(auth_system.get_user_by_id, user.id)
|
||||||
|
if fresh is None:
|
||||||
|
raise AppError('登录已失效,请重新登录', 401, error_code='unauthorized')
|
||||||
|
|
||||||
|
credits = int(fresh['duration_seconds'] or 0)
|
||||||
|
unlimited = fresh['role'] == 'admin' or credits == -1
|
||||||
|
|
||||||
|
if free or unlimited:
|
||||||
|
return Billing(user=user, ledger_key=None, need=0, free=free)
|
||||||
|
|
||||||
|
need = 1
|
||||||
|
if cost is not None:
|
||||||
|
try:
|
||||||
|
need = max(1, int(cost()))
|
||||||
|
except Exception:
|
||||||
|
need = 1
|
||||||
|
|
||||||
|
key = 'req-' + secrets.token_hex(8)
|
||||||
|
remaining = await asyncio.to_thread(auth_system.reserve_and_open, fresh['id'], need, key)
|
||||||
|
if remaining is None:
|
||||||
|
await asyncio.to_thread(auth_system.add_log, 'fail', fresh['username'], ip,
|
||||||
|
action or '下载', f'剩余次数不足(本次需 {need} 次)')
|
||||||
|
raise AppError(f'下载次数已用尽(本次需 {need} 次,当前 {credits} 次),'
|
||||||
|
'请充值卡密后继续下载', 401, error_code='expired')
|
||||||
|
return Billing(user=user, ledger_key=key, need=need, free=False)
|
||||||
+3
-2
@@ -23,10 +23,11 @@ from fastapi import APIRouter, Depends, Request
|
|||||||
from fastapi.responses import JSONResponse
|
from fastapi.responses import JSONResponse
|
||||||
|
|
||||||
import auth_system as auth
|
import auth_system as auth
|
||||||
|
import web.deps as deps
|
||||||
from core import db, repo
|
from core import db, repo
|
||||||
from core import session as cs
|
from core import session as cs
|
||||||
from web import ratelimit
|
from web import ratelimit
|
||||||
from web.deps import SiteUser, current_user, secret_key
|
from web.deps import SiteUser, current_user
|
||||||
from web.response import error, success
|
from web.response import error, success
|
||||||
|
|
||||||
router = APIRouter()
|
router = APIRouter()
|
||||||
@@ -85,7 +86,7 @@ def _cookie_kwargs() -> dict:
|
|||||||
def _login_response(payload: dict, message: str, uid: int, seq: int) -> JSONResponse:
|
def _login_response(payload: dict, message: str, uid: int, seq: int) -> JSONResponse:
|
||||||
"""构造带会话 Cookie 的成功响应(Cookie 格式与 Flask 一致)。"""
|
"""构造带会话 Cookie 的成功响应(Cookie 格式与 Flask 一致)。"""
|
||||||
resp = JSONResponse(content={'status': 200, 'success': True, 'message': message, 'data': payload})
|
resp = JSONResponse(content={'status': 200, 'success': True, 'message': message, 'data': payload})
|
||||||
cookie = cs.dumps(secret_key(), {
|
cookie = cs.dumps(deps.secret_key(), {
|
||||||
'site_uid': uid, 'site_uid_seq': seq, 'site_uid_ts': int(time.time()),
|
'site_uid': uid, 'site_uid_seq': seq, 'site_uid_ts': int(time.time()),
|
||||||
})
|
})
|
||||||
resp.set_cookie(SESSION_COOKIE, cookie, **_cookie_kwargs())
|
resp.set_cookie(SESSION_COOKIE, cookie, **_cookie_kwargs())
|
||||||
|
|||||||
@@ -0,0 +1,473 @@
|
|||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""下载提交端点(FastAPI 栈)——与 Flask 侧逐字段对齐。
|
||||||
|
|
||||||
|
已迁移:
|
||||||
|
- POST /api/v1/download ↔ Flask /download(计费 1 次;file/json 两模式)
|
||||||
|
- POST /api/v1/download/batch_zip ↔ Flask /download/batch_zip(按去重歌曲数计费)
|
||||||
|
- POST /api/v1/podcast/download ↔ Flask /podcast/download(计费 1 次)
|
||||||
|
- POST /api/v1/podcast/download/batch_zip ↔ Flask /podcast/download/batch_zip
|
||||||
|
|
||||||
|
计费:经 web.billing.require_billing(预扣+开账同事务 / 失败只退 open 行),
|
||||||
|
任务登记与启动复用 web.tasks_registry(与 main._register_download_task 同口径);
|
||||||
|
worker 仍复用 main 的既有实现(无 flask 依赖),故账本与交付语义不变。
|
||||||
|
"""
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
from typing import Annotated, Any
|
||||||
|
|
||||||
|
from fastapi import APIRouter, Depends, Request
|
||||||
|
|
||||||
|
from core import media
|
||||||
|
from core.config import VALID_QUALITY_LEVELS
|
||||||
|
from core.state import REQ_LOG
|
||||||
|
from music_api import APIException
|
||||||
|
from web import ratelimit
|
||||||
|
from web.billing import require_billing
|
||||||
|
from web.deps import SiteUser, current_user, require_login
|
||||||
|
from web.response import error, success
|
||||||
|
from web.tasks_registry import register_download_task, start_task
|
||||||
|
|
||||||
|
router = APIRouter()
|
||||||
|
|
||||||
|
LoginUser = Annotated[SiteUser, Depends(require_login)]
|
||||||
|
|
||||||
|
logger = logging.getLogger('web.download')
|
||||||
|
|
||||||
|
|
||||||
|
async def _request_data(request: Request) -> dict[str, Any]:
|
||||||
|
data: dict[str, Any] = {}
|
||||||
|
if request.method == 'GET':
|
||||||
|
data.update(dict(request.query_params))
|
||||||
|
return data
|
||||||
|
try:
|
||||||
|
body = await request.json()
|
||||||
|
if isinstance(body, dict):
|
||||||
|
data.update(body)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
if not data:
|
||||||
|
try:
|
||||||
|
form = await request.form()
|
||||||
|
data.update(dict(form.items()))
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return data
|
||||||
|
|
||||||
|
|
||||||
|
def _extract_id(raw: Any) -> str:
|
||||||
|
from web.routers.music import _extract_id as _ex
|
||||||
|
return _ex(raw)
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_meta_fields(raw: Any):
|
||||||
|
from core.utils import parse_meta_fields
|
||||||
|
return parse_meta_fields(raw)
|
||||||
|
|
||||||
|
|
||||||
|
def _record_log(action: str, ip: str, song: str) -> None:
|
||||||
|
REQ_LOG.append({'action': action, 'ip': ip, 'song': str(song or '')[:100],
|
||||||
|
'time': time.strftime('%Y-%m-%d %H:%M:%S')})
|
||||||
|
if len(REQ_LOG) > 2000:
|
||||||
|
del REQ_LOG[:len(REQ_LOG) - 2000]
|
||||||
|
|
||||||
|
|
||||||
|
def _dedup_ids(raw_ids: Any) -> list[str]:
|
||||||
|
if isinstance(raw_ids, str):
|
||||||
|
raw_list = raw_ids.split(',')
|
||||||
|
else:
|
||||||
|
raw_list = [str(x) for x in (raw_ids or [])]
|
||||||
|
out, seen = [], set()
|
||||||
|
for x in raw_list:
|
||||||
|
mid = _extract_id(x)
|
||||||
|
s = str(mid).strip()
|
||||||
|
if s and s not in seen:
|
||||||
|
seen.add(s)
|
||||||
|
out.append(s)
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def _cookies(data: dict[str, Any]) -> dict[str, str]:
|
||||||
|
import music_api
|
||||||
|
return music_api.load_cookies()
|
||||||
|
|
||||||
|
|
||||||
|
async def _resolve_personal(data: dict[str, Any], user: SiteUser | None, request: Request):
|
||||||
|
from web.routers.music import _resolve_user_cookies
|
||||||
|
return await _resolve_user_cookies(data, user, request)
|
||||||
|
|
||||||
|
|
||||||
|
# ============ 单曲下载 ============
|
||||||
|
|
||||||
|
@router.post('/download')
|
||||||
|
async def download(user: Annotated[SiteUser | None, Depends(current_user)], request: Request):
|
||||||
|
"""下载单曲(file=后台任务 / json=同步返回)。"""
|
||||||
|
await ratelimit.guard(request, bucket='dl', limit=20, window=60)
|
||||||
|
data = await _request_data(request)
|
||||||
|
music_id = data.get('id')
|
||||||
|
quality = data.get('quality', 'lossless')
|
||||||
|
return_format = data.get('format', 'file')
|
||||||
|
ip = ratelimit.client_ip(request)
|
||||||
|
|
||||||
|
billing = await require_billing(user, action='音乐下载', ip=ip)
|
||||||
|
|
||||||
|
async def _fail(msg: str, code: int = 400):
|
||||||
|
"""参数/校验失败:退回本次预扣后返回错误(账本守恒)。"""
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error(msg, code)
|
||||||
|
|
||||||
|
try:
|
||||||
|
if not music_id:
|
||||||
|
return await _fail("参数 'music_id' 不能为空", 400)
|
||||||
|
music_id = str(music_id).strip()
|
||||||
|
if not music_id.isdigit():
|
||||||
|
return await _fail('无效的歌曲 ID')
|
||||||
|
if quality not in VALID_QUALITY_LEVELS:
|
||||||
|
return await _fail(f"无效的音质参数,支持: {', '.join(VALID_QUALITY_LEVELS)}")
|
||||||
|
if return_format not in ('file', 'json'):
|
||||||
|
return await _fail("返回格式只支持 'file' 或 'json'")
|
||||||
|
music_id = _extract_id(music_id)
|
||||||
|
|
||||||
|
account = data.get('account', 'admin')
|
||||||
|
if account == 'user':
|
||||||
|
cookies = await _resolve_personal(data, user, request)
|
||||||
|
if not cookies:
|
||||||
|
return await _fail('请先扫码登录个人账号后再使用个人解析', 401)
|
||||||
|
else:
|
||||||
|
cookies = _cookies(data)
|
||||||
|
|
||||||
|
filename_format = data.get('filename_format', 'artist_song')
|
||||||
|
write_meta = str(data.get('write_meta', 'true')).lower() in ('1', 'true', 'yes')
|
||||||
|
meta_fields = _parse_meta_fields(data.get('meta_fields'))
|
||||||
|
lyric_format = data.get('lyric_format', 'lrc')
|
||||||
|
if lyric_format not in ('lrc', 'srt'):
|
||||||
|
lyric_format = 'lrc'
|
||||||
|
need_zip = str(data.get('zip', 'false')).lower() in ('1', 'true', 'yes')
|
||||||
|
|
||||||
|
import main as legacy
|
||||||
|
downloader = legacy.api_service.downloader
|
||||||
|
|
||||||
|
if return_format == 'json':
|
||||||
|
try:
|
||||||
|
info = await asyncio.to_thread(
|
||||||
|
downloader.get_music_info, int(music_id), quality, cookies)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error('获取音乐信息失败: %s', e)
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('获取失败:歌曲不存在或暂无版权', 404)
|
||||||
|
base = await asyncio.to_thread(
|
||||||
|
legacy._build_download_base, filename_format,
|
||||||
|
{'name': info.name, 'artist_string': info.artists})
|
||||||
|
filename = f'{base}.{info.file_type}'
|
||||||
|
try:
|
||||||
|
dl = await asyncio.to_thread(
|
||||||
|
downloader.download_music_file, int(music_id), quality,
|
||||||
|
write_meta=write_meta, meta_fields=meta_fields,
|
||||||
|
cookies=cookies, music_info=info)
|
||||||
|
if not dl.success:
|
||||||
|
logger.error('下载失败: %s', dl.error_message)
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('下载失败,请稍后重试或更换音质', 500)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error('下载异常: %s', e)
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('服务暂时不可用,请稍后再试', 500)
|
||||||
|
try:
|
||||||
|
_record_log('下载', ip, info.name)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
response_data = {
|
||||||
|
'music_id': music_id,
|
||||||
|
'name': info.name,
|
||||||
|
'artist': info.artists.replace('/', ', '),
|
||||||
|
'album': info.album,
|
||||||
|
'quality': quality,
|
||||||
|
'quality_name': media.quality_display_name(quality),
|
||||||
|
'file_type': info.file_type,
|
||||||
|
'file_size': info.file_size,
|
||||||
|
'file_size_formatted': media.format_file_size(info.file_size),
|
||||||
|
'filename': filename,
|
||||||
|
'duration': info.duration * 1000,
|
||||||
|
}
|
||||||
|
await billing.settle(1, 'delivered')
|
||||||
|
return success(response_data, '下载完成')
|
||||||
|
|
||||||
|
# file 模式:注册后台任务
|
||||||
|
device_id = str(data.get('device_id') or '').strip()[:64]
|
||||||
|
await asyncio.to_thread(legacy._touch_device, device_id)
|
||||||
|
task = _single_task(data, music_id, quality, filename_format, write_meta,
|
||||||
|
meta_fields, lyric_format, need_zip, cookies, ip, device_id)
|
||||||
|
reg_err = await register_download_task(billing, task, str(data.get('task_id') or '').strip())
|
||||||
|
if reg_err:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error(reg_err, 429)
|
||||||
|
task_id = task['task_id']
|
||||||
|
start_err = await start_task(task_id, legacy._single_download_worker, (task_id,))
|
||||||
|
if start_err:
|
||||||
|
return error(start_err, 500)
|
||||||
|
return success({'task_id': task_id, 'status': 'queued'}, '下载任务已提交')
|
||||||
|
except APIException:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('服务暂时不可用,请稍后再试', 500)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error('下载音乐异常: %s', e)
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('服务暂时不可用,请稍后再试', 500)
|
||||||
|
|
||||||
|
|
||||||
|
def _single_task(data, music_id, quality, filename_format, write_meta, meta_fields,
|
||||||
|
lyric_format, need_zip, cookies, ip, device_id) -> dict[str, Any]:
|
||||||
|
return {
|
||||||
|
'type': 'single',
|
||||||
|
'total': 1, 'done': 0, 'failed': 0,
|
||||||
|
'status': 'queued', 'start': time.time(), 'msg': '排队中...',
|
||||||
|
'cancel': threading.Event(), 'pause': threading.Event(),
|
||||||
|
'device_id': device_id, 'result': None,
|
||||||
|
'params': {
|
||||||
|
'music_id': music_id, 'quality': quality,
|
||||||
|
'filename_format': filename_format, 'write_meta': write_meta,
|
||||||
|
'meta_fields': meta_fields, 'lyric_format': lyric_format,
|
||||||
|
'need_zip': need_zip, 'cookies': cookies, 'remote_ip': ip,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
# ============ 批量 ZIP ============
|
||||||
|
|
||||||
|
@router.post('/download/batch_zip')
|
||||||
|
async def download_batch_zip(user: Annotated[SiteUser | None, Depends(current_user)], request: Request):
|
||||||
|
"""批量打包 ZIP(按去重歌曲数计费)。"""
|
||||||
|
await ratelimit.guard(request, bucket='dl', limit=20, window=60)
|
||||||
|
data = await _request_data(request)
|
||||||
|
ip = ratelimit.client_ip(request)
|
||||||
|
|
||||||
|
def _cost() -> int:
|
||||||
|
return max(1, min(len(_dedup_ids(data.get('ids') or data.get('id') or [])), 500))
|
||||||
|
|
||||||
|
billing = await require_billing(user, action='批量ZIP下载', cost=_cost, ip=ip)
|
||||||
|
try:
|
||||||
|
raw_ids = data.get('ids') or data.get('id')
|
||||||
|
ids = _dedup_ids(raw_ids)
|
||||||
|
if not ids:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('缺少 ids 参数')
|
||||||
|
non_digit = [x for x in ids if not x.isdigit()]
|
||||||
|
if non_digit:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('ids 中含无效歌曲 ID')
|
||||||
|
if len(ids) > 500:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('单个批量任务最多 500 首,请分批提交')
|
||||||
|
|
||||||
|
quality = data.get('quality', 'lossless')
|
||||||
|
if quality not in VALID_QUALITY_LEVELS:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error(f"无效的音质参数,支持: {', '.join(VALID_QUALITY_LEVELS)}")
|
||||||
|
filename_format = data.get('filename_format', 'artist_song')
|
||||||
|
write_meta = str(data.get('write_meta', 'true')).lower() in ('1', 'true', 'yes')
|
||||||
|
meta_fields = _parse_meta_fields(data.get('meta_fields'))
|
||||||
|
lyric_format = data.get('lyric_format', 'lrc')
|
||||||
|
if lyric_format not in ('lrc', 'srt'):
|
||||||
|
lyric_format = 'lrc'
|
||||||
|
try:
|
||||||
|
concurrency = int(data.get('concurrency', 8) or 8)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
concurrency = 8
|
||||||
|
concurrency = max(1, min(10, concurrency))
|
||||||
|
|
||||||
|
cookies = _cookies(data)
|
||||||
|
device_id = str(data.get('device_id') or '').strip()[:64]
|
||||||
|
import main as legacy
|
||||||
|
await asyncio.to_thread(legacy._touch_device, device_id)
|
||||||
|
task = {
|
||||||
|
'type': 'batch',
|
||||||
|
'total': len(ids), 'done': 0, 'failed': 0,
|
||||||
|
'status': 'queued', 'start': time.time(), 'msg': '排队中...',
|
||||||
|
'cancel': threading.Event(), 'pause': threading.Event(),
|
||||||
|
'device_id': device_id, 'result': None,
|
||||||
|
'params': {
|
||||||
|
'ids': ids, 'quality': quality, 'filename_format': filename_format,
|
||||||
|
'write_meta': write_meta, 'meta_fields': meta_fields,
|
||||||
|
'lyric_format': lyric_format, 'concurrency': concurrency,
|
||||||
|
'cookies': cookies, 'remote_ip': ip,
|
||||||
|
'want_zip': True,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
reg_err = await register_download_task(billing, task, str(data.get('task_id') or '').strip())
|
||||||
|
if reg_err:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error(reg_err, 429)
|
||||||
|
task_id = task['task_id']
|
||||||
|
start_err = await start_task(task_id, legacy._batch_zip_worker, (task_id,))
|
||||||
|
if start_err:
|
||||||
|
return error(start_err, 500)
|
||||||
|
return success({'task_id': task_id, 'status': 'queued', 'total': len(ids)},
|
||||||
|
'批量打包任务已提交')
|
||||||
|
except Exception as e:
|
||||||
|
logger.error('批量打包提交异常: %s', e)
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('服务暂时不可用,请稍后再试', 500)
|
||||||
|
|
||||||
|
|
||||||
|
# ============ 播客单集下载 ============
|
||||||
|
|
||||||
|
def _podcast_billing_count(data: dict[str, Any]) -> int:
|
||||||
|
"""播客打包计费(与 Flask _podcast_billing_count 同口径)。"""
|
||||||
|
if data.get('radio_id'):
|
||||||
|
return 1
|
||||||
|
limit = _safe_int(data.get('limit', 200) or 200, 200)
|
||||||
|
limit = max(1, min(limit, 500))
|
||||||
|
ids = _dedup_ids(data.get('ids'))
|
||||||
|
return max(1, min(len(ids), limit))
|
||||||
|
|
||||||
|
|
||||||
|
def _safe_int(value: Any, default: int) -> int:
|
||||||
|
try:
|
||||||
|
return int(value)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return default
|
||||||
|
|
||||||
|
|
||||||
|
def _podcast_quality(data: dict[str, Any]) -> str:
|
||||||
|
q = str(data.get('quality') or 'auto')
|
||||||
|
return q if q in ('auto', 'standard', 'exhigh', 'lossless', 'hires') else 'auto'
|
||||||
|
|
||||||
|
|
||||||
|
@router.post('/podcast/download')
|
||||||
|
async def podcast_download(user: Annotated[SiteUser | None, Depends(current_user)], request: Request):
|
||||||
|
"""播客单集下载(计费 1 次)。"""
|
||||||
|
await ratelimit.guard(request, bucket='dl', limit=20, window=60)
|
||||||
|
data = await _request_data(request)
|
||||||
|
raw = data.get('id')
|
||||||
|
ip = ratelimit.client_ip(request)
|
||||||
|
|
||||||
|
billing = await require_billing(user, action='播客单集下载', ip=ip)
|
||||||
|
try:
|
||||||
|
if not raw:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error("参数 'id' 不能为空", 400)
|
||||||
|
import podcast_api
|
||||||
|
quality = _podcast_quality(data)
|
||||||
|
need_zip = str(data.get('zip', 'false')).lower() in ('1', 'true', 'yes')
|
||||||
|
ref = podcast_api.extract_podcast_ref(raw)
|
||||||
|
if ref['kind'] == 'radio':
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('这是电台链接,请先解析电台后选择单集下载,或使用「打包整个播客」')
|
||||||
|
|
||||||
|
account = data.get('account', 'admin')
|
||||||
|
if account == 'user':
|
||||||
|
cookies = await _resolve_personal(data, user, request)
|
||||||
|
if not cookies:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('请先扫码登录个人账号后再使用个人解析', 401)
|
||||||
|
else:
|
||||||
|
cookies = _cookies(data)
|
||||||
|
|
||||||
|
device_id = str(data.get('device_id') or '').strip()[:64]
|
||||||
|
import main as legacy
|
||||||
|
await asyncio.to_thread(legacy._touch_device, device_id)
|
||||||
|
task = {
|
||||||
|
'type': 'single',
|
||||||
|
'total': 1, 'done': 0, 'failed': 0,
|
||||||
|
'status': 'queued', 'start': time.time(), 'msg': '排队中...',
|
||||||
|
'cancel': threading.Event(), 'pause': threading.Event(),
|
||||||
|
'device_id': device_id, 'result': None,
|
||||||
|
'params': {'program_id': ref['id'], 'quality': quality,
|
||||||
|
'need_zip': need_zip, 'cookies': cookies, 'remote_ip': ip},
|
||||||
|
}
|
||||||
|
reg_err = await register_download_task(billing, task, str(data.get('task_id') or '').strip())
|
||||||
|
if reg_err:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error(reg_err, 429)
|
||||||
|
task_id = task['task_id']
|
||||||
|
start_err = await start_task(task_id, legacy._podcast_download_worker, (task_id,))
|
||||||
|
if start_err:
|
||||||
|
return error(start_err, 500)
|
||||||
|
return success({'task_id': task_id, 'status': 'queued'}, '播客单集下载任务已提交')
|
||||||
|
except APIException as e:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error(str(e), 404)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error('播客单集下载提交异常: %s', e)
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('服务暂时不可用,请稍后再试', 500)
|
||||||
|
|
||||||
|
|
||||||
|
@router.post('/podcast/download/batch_zip')
|
||||||
|
async def podcast_batch_zip(user: Annotated[SiteUser | None, Depends(current_user)], request: Request):
|
||||||
|
"""整台播客打包(先扣 1(电台模式)或按去重单集数扣,worker 按实际集数补扣差额)。"""
|
||||||
|
await ratelimit.guard(request, bucket='dl', limit=20, window=60)
|
||||||
|
data = await _request_data(request)
|
||||||
|
ip = ratelimit.client_ip(request)
|
||||||
|
|
||||||
|
def _cost() -> int:
|
||||||
|
return _podcast_billing_count(data)
|
||||||
|
|
||||||
|
billing = await require_billing(user, action='播客整台打包', cost=_cost, ip=ip)
|
||||||
|
try:
|
||||||
|
raw_radio = data.get('radio_id')
|
||||||
|
raw_ids = data.get('ids')
|
||||||
|
if not raw_radio and not raw_ids:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('缺少 radio_id 或 ids 参数')
|
||||||
|
limit = max(1, min(_safe_int(data.get('limit', 200) or 200, 200), 500))
|
||||||
|
quality = _podcast_quality(data)
|
||||||
|
cookies = _cookies(data)
|
||||||
|
|
||||||
|
import podcast_api
|
||||||
|
if raw_radio:
|
||||||
|
ref = podcast_api.extract_podcast_ref(raw_radio)
|
||||||
|
if ref['kind'] == 'program':
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('这是单集链接,请使用「下载单集」或直接输入电台链接/ID')
|
||||||
|
radio_id = ref['id']
|
||||||
|
radio_name = ''
|
||||||
|
try:
|
||||||
|
radio_name = podcast_api.get_radio_detail(radio_id, cookies).get('name', '')
|
||||||
|
except APIException:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
id_list = _dedup_ids(raw_ids)[:limit]
|
||||||
|
if not id_list:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('缺少 ids 参数')
|
||||||
|
radio_id = None
|
||||||
|
radio_name = str(data.get('name') or '播客单集')
|
||||||
|
|
||||||
|
billed = _podcast_billing_count(data)
|
||||||
|
device_id = str(data.get('device_id') or '').strip()[:64]
|
||||||
|
import main as legacy
|
||||||
|
await asyncio.to_thread(legacy._touch_device, device_id)
|
||||||
|
task = {
|
||||||
|
'type': 'batch',
|
||||||
|
'total': 0, 'done': 0, 'failed': 0,
|
||||||
|
'status': 'queued', 'start': time.time(), 'msg': '排队中...',
|
||||||
|
'cancel': threading.Event(), 'pause': threading.Event(),
|
||||||
|
'device_id': device_id, 'result': None,
|
||||||
|
'params': {
|
||||||
|
'radio_id': radio_id,
|
||||||
|
'ids': None if radio_id else id_list,
|
||||||
|
'limit': limit, 'quality': quality, 'radio_name': radio_name,
|
||||||
|
'cookies': cookies, 'remote_ip': ip,
|
||||||
|
# 整台模式:worker 解析完实际集数后按差额补扣(一首 = 1 次)
|
||||||
|
'billing_uid': user.id if user else None,
|
||||||
|
'billed': billed,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
reg_err = await register_download_task(billing, task, str(data.get('task_id') or '').strip())
|
||||||
|
if reg_err:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error(reg_err, 429)
|
||||||
|
task_id = task['task_id']
|
||||||
|
start_err = await start_task(task_id, legacy._podcast_zip_worker, (task_id,))
|
||||||
|
if start_err:
|
||||||
|
return error(start_err, 500)
|
||||||
|
return success({'task_id': task_id, 'status': 'queued'}, '播客打包任务已提交')
|
||||||
|
except APIException as e:
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error(str(e), 404)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error('播客打包提交异常: %s', e)
|
||||||
|
await billing.refund_open('rejected')
|
||||||
|
return error('服务暂时不可用,请稍后再试', 500)
|
||||||
@@ -24,10 +24,11 @@ from fastapi import APIRouter, Depends, Request
|
|||||||
from fastapi.responses import JSONResponse
|
from fastapi.responses import JSONResponse
|
||||||
|
|
||||||
import auth_system as auth
|
import auth_system as auth
|
||||||
|
import web.deps as deps
|
||||||
from core import db, repo
|
from core import db, repo
|
||||||
from core import session as cs
|
from core import session as cs
|
||||||
from web import ratelimit
|
from web import ratelimit
|
||||||
from web.deps import SiteUser, current_user, secret_key
|
from web.deps import SiteUser, current_user
|
||||||
from web.response import error, success
|
from web.response import error, success
|
||||||
|
|
||||||
router = APIRouter()
|
router = APIRouter()
|
||||||
@@ -72,7 +73,7 @@ def _user_public(row: dict) -> dict:
|
|||||||
def _login_cookie_response(payload: dict, message: str, uid: int, seq: int,
|
def _login_cookie_response(payload: dict, message: str, uid: int, seq: int,
|
||||||
secure: bool) -> JSONResponse:
|
secure: bool) -> JSONResponse:
|
||||||
resp = JSONResponse(content={'status': 200, 'success': True, 'message': message, 'data': payload})
|
resp = JSONResponse(content={'status': 200, 'success': True, 'message': message, 'data': payload})
|
||||||
cookie = cs.dumps(secret_key(), {'site_uid': uid, 'site_uid_seq': seq,
|
cookie = cs.dumps(deps.secret_key(), {'site_uid': uid, 'site_uid_seq': seq,
|
||||||
'site_uid_ts': int(time.time())})
|
'site_uid_ts': int(time.time())})
|
||||||
resp.set_cookie(SESSION_COOKIE, cookie, max_age=_SESSION_MAX_AGE, httponly=True,
|
resp.set_cookie(SESSION_COOKIE, cookie, max_age=_SESSION_MAX_AGE, httponly=True,
|
||||||
samesite='lax', path='/', secure=secure)
|
samesite='lax', path='/', secure=secure)
|
||||||
|
|||||||
@@ -0,0 +1,79 @@
|
|||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""下载/播客任务登记(FastAPI 栈)——镜像 main._register_download_task。
|
||||||
|
|
||||||
|
与 Flask 版的唯一差异:不依赖 flask.g,请求级账本键由 Billing 显式携带。
|
||||||
|
其余(归属、容量闸、task_id 白名单与碰撞换新、账本承接与撞名重试)逐条对齐。
|
||||||
|
"""
|
||||||
|
import logging
|
||||||
|
import re
|
||||||
|
import uuid
|
||||||
|
from collections.abc import Callable
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from core.config import MAX_ACTIVE_TASKS_PER_USER, MAX_TERMINAL_TASKS, MAX_TOTAL_TASKS
|
||||||
|
from core.state import DOWNLOAD_TASKS, DOWNLOAD_TASKS_LOCK
|
||||||
|
from web.billing import Billing
|
||||||
|
|
||||||
|
logger = logging.getLogger('web.tasks')
|
||||||
|
|
||||||
|
_TASK_ID_RE = re.compile(r'^[A-Za-z0-9_-]{1,32}$')
|
||||||
|
|
||||||
|
|
||||||
|
async def register_download_task(billing: Billing, task: dict[str, Any],
|
||||||
|
task_id_hint: str = '') -> str | None:
|
||||||
|
"""任务登记:失败返回错误描述,成功返回 None。task['task_id'] 为最终 id。"""
|
||||||
|
user_id = billing.user.id
|
||||||
|
task['user_id'] = user_id
|
||||||
|
task['billing_reserved'] = billing.need
|
||||||
|
|
||||||
|
with DOWNLOAD_TASKS_LOCK:
|
||||||
|
tid = (task_id_hint or '').strip() or uuid.uuid4().hex[:12]
|
||||||
|
if not _TASK_ID_RE.match(tid):
|
||||||
|
tid = uuid.uuid4().hex[:12]
|
||||||
|
while tid in DOWNLOAD_TASKS:
|
||||||
|
tid = uuid.uuid4().hex[:12]
|
||||||
|
active = sum(1 for t in DOWNLOAD_TASKS.values()
|
||||||
|
if t.get('user_id') == user_id
|
||||||
|
and t.get('status') in ('queued', 'running', 'paused'))
|
||||||
|
if active >= MAX_ACTIVE_TASKS_PER_USER:
|
||||||
|
return f'您有 {active} 个任务正在进行,请等待完成或取消后再提交'
|
||||||
|
active_total = sum(1 for t in DOWNLOAD_TASKS.values()
|
||||||
|
if t.get('status') in ('queued', 'running', 'paused'))
|
||||||
|
if active_total >= MAX_TOTAL_TASKS:
|
||||||
|
return '当前任务队列已满,请稍后再试'
|
||||||
|
if len(DOWNLOAD_TASKS) - active_total >= MAX_TERMINAL_TASKS:
|
||||||
|
return '当前任务队列已满,请稍后再试'
|
||||||
|
task['task_id'] = tid
|
||||||
|
DOWNLOAD_TASKS[tid] = task
|
||||||
|
|
||||||
|
# 账本承接(撞名重试,与 Flask 同口径)
|
||||||
|
try:
|
||||||
|
renamed = await billing.rename_to(tid)
|
||||||
|
retries = 0
|
||||||
|
while not renamed and retries < 5:
|
||||||
|
retries += 1
|
||||||
|
with DOWNLOAD_TASKS_LOCK:
|
||||||
|
DOWNLOAD_TASKS.pop(tid, None)
|
||||||
|
tid = uuid.uuid4().hex[:12]
|
||||||
|
while tid in DOWNLOAD_TASKS:
|
||||||
|
tid = uuid.uuid4().hex[:12]
|
||||||
|
task['task_id'] = tid
|
||||||
|
DOWNLOAD_TASKS[tid] = task
|
||||||
|
renamed = await billing.rename_to(tid)
|
||||||
|
if not renamed:
|
||||||
|
raise RuntimeError('任务账本键撞名,请重新提交')
|
||||||
|
billing.ledger_key = tid
|
||||||
|
except Exception:
|
||||||
|
with DOWNLOAD_TASKS_LOCK:
|
||||||
|
if DOWNLOAD_TASKS.get(tid) is task:
|
||||||
|
DOWNLOAD_TASKS.pop(tid, None)
|
||||||
|
raise
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
async def start_task(task_id: str, worker: Callable, args: tuple = ()) -> str | None:
|
||||||
|
"""启动 worker 线程(复用 main._start_task_thread 的实现,失败收敛终态并退费)。"""
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
import main as legacy
|
||||||
|
return await asyncio.to_thread(legacy._start_task_thread, task_id, worker, args)
|
||||||
Reference in new issue
Block a user