feat: strengthen backup and mobile workflows

2026-09-16 21:12:52 +08:00
parent 6f38190c92
commit 6c234d7d82
64 changed files with 5059 additions and 635 deletions
@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import hashlib
+import secrets
+import shutil
+from datetime import timedelta
+from pathlib import Path
+
+from fastapi import APIRouter, Depends, File, Query, UploadFile
+from fastapi.responses import FileResponse
+from sqlalchemy import func, select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+from starlette.background import BackgroundTask
+
+from backend.auth import current_user, hash_token
+from backend.config import get_settings
+from backend.db import get_db
+from backend.models import BackupPreflight, Countdown, User, new_id, utcnow
+
+from .archive import MAX_ARCHIVE_BYTES as ARCHIVE_DEFAULT_LIMIT
+from .archive import backup_error, parse_archive_path
+from .schemas import RestoreRequest
+from .service import BackupRepairPending, export_v2, restore_v2, validate_archive
+
+router = APIRouter(prefix="/api/v1/backup", tags=["backup"])
+MAX_ARCHIVE_BYTES = ARCHIVE_DEFAULT_LIMIT
+MAX_PENDING_PREFLIGHTS_PER_USER = 3
+_READ_CHUNK = 1024 * 1024
+
+
+def _staging_root() -> Path:
+    root = Path(get_settings().backup_staging_dir).resolve()
+    root.mkdir(parents=True, exist_ok=True)
+    return root
+
+
+def _safe_staging_path(value: str) -> Path:
+    root = _staging_root()
+    path = Path(value).resolve()
+    if path.parent != root:
+        raise backup_error("backup_preflight_invalid", "预检暂存无效", 409)
+    return path
+
+
+def _path_size(path: str | None) -> int:
+    if not path:
+        return 0
+    root = Path(path)
+    if not root.exists():
+        return 0
+    if root.is_file():
+        return root.stat().st_size
+    return sum(item.stat().st_size for item in root.rglob("*") if item.is_file())
+
+
+async def _prune(db: AsyncSession) -> None:
+    now = utcnow()
+    rows = list((await db.scalars(select(BackupPreflight).where(
+        BackupPreflight.expires_at <= now,
+        BackupPreflight.status.in_((
+            "reserved", "pending", "failed", "consuming", "cleanup_pending",
+            "repair_pending",
+        )),
+    ).with_for_update())).all())
+    from .storage import remove_quarantine, restore_quarantine_dir
+
+    attachment_root = Path(get_settings().attachment_dir).resolve()
+    for row in rows:
+        try:
+            if row.status == "repair_pending":
+                if row.cleanup_path:
+                    restore_quarantine_dir(attachment_root, Path(row.cleanup_path))
+                    row.cleanup_path = None
+                _safe_staging_path(row.staging_path).unlink(missing_ok=True)
+                row.status = "failed"
+                continue
+            if row.cleanup_path:
+                remove_quarantine(Path(row.cleanup_path))
+                row.cleanup_path = None
+            _safe_staging_path(row.staging_path).unlink(missing_ok=True)
+        except (OSError, ValueError):
+            continue
+        if row.status == "cleanup_pending":
+            row.status = "consumed"
+            row.consumed_at = row.consumed_at or now
+        else:
+            await db.delete(row)
+    if rows:
+        await db.commit()
+
+
+@router.get("/export.zip")
+async def export_zip(user: User = Depends(current_user), db: AsyncSession = Depends(get_db)):
+    path = await export_v2(db, user)
+    await db.commit()
+    return FileResponse(
+        path,
+        media_type="application/zip",
+        filename="dodo-backup-v2.zip",
+        background=BackgroundTask(path.unlink, missing_ok=True),
+    )
+
+
+@router.post("/preflight")
+async def preflight(
+    mode: str = Query(pattern="^(merge|replace)$"),
+    file: UploadFile = File(...),
+    user: User = Depends(current_user),
+    db: AsyncSession = Depends(get_db),
+):
+    await _prune(db)
+    settings = get_settings()
+    configured_limit = settings.backup_max_archive_mb * 1024 * 1024
+    limit = min(MAX_ARCHIVE_BYTES, configured_limit)
+    max_pending = min(MAX_PENDING_PREFLIGHTS_PER_USER, settings.backup_max_pending_per_user)
+    max_staged = settings.backup_max_staged_mb_per_user * 1024 * 1024
+    token = secrets.token_urlsafe(32)
+    staging = _staging_root() / f"{new_id()}.zip"
+    reservation = BackupPreflight(
+        token_hash=hash_token(token), user_id=user.id, backup_id=new_id(),
+        archive_sha256="0" * 64, archive_size=limit, staging_path=str(staging), mode=mode,
+        status="reserved", expires_at=utcnow() + timedelta(seconds=settings.backup_preflight_ttl_seconds),
+    )
+    # A user-row write serializes quota decisions across workers on PostgreSQL;
+    # SQLite serializes writers at the database level.
+    await db.execute(update(User).where(User.id == user.id).values(username=User.username))
+    pending_count, pending_bytes = (await db.execute(select(
+        func.count(BackupPreflight.id), func.coalesce(func.sum(BackupPreflight.archive_size), 0)
+    ).where(
+        BackupPreflight.user_id == user.id,
+        BackupPreflight.status.in_((
+            "reserved", "pending", "failed", "consuming", "cleanup_pending",
+            "repair_pending",
+        )),
+    ))).one()
+    pending_rows = list((await db.scalars(select(BackupPreflight).where(
+        BackupPreflight.user_id == user.id,
+        BackupPreflight.status.in_((
+            "reserved", "pending", "failed", "consuming", "cleanup_pending",
+            "repair_pending",
+        )),
+    ))).all())
+    pending_bytes += sum(_path_size(item.cleanup_path) for item in pending_rows)
+    if pending_count >= max_pending or pending_bytes + limit > max_staged:
+        await db.rollback()
+        raise backup_error("backup_preflight_quota", "待处理预检配额已达上限", 429)
+    db.add(reservation)
+    await db.flush()
+    reservation_id = reservation.id
+    await db.commit()
+
+    size = 0
+    digest = hashlib.sha256()
+    archive = None
+    try:
+        with staging.open("xb") as output:
+            while chunk := await file.read(_READ_CHUNK):
+                size += len(chunk)
+                if size > limit:
+                    raise backup_error("backup_size_invalid", "备份文件大小无效")
+                digest.update(chunk)
+                output.write(chunk)
+        archive = parse_archive_path(staging, max_archive_bytes=limit)
+        validate_archive(archive)
+        if mode == "merge":
+            incoming_pin_ids = {
+                str(item["id"])
+                for item in archive.entities["countdowns"]
+                if item.get("pinned") is True and item.get("archived_at") is None
+            }
+            existing_pin_ids = set((await db.scalars(select(Countdown.id).where(
+                Countdown.user_id == user.id,
+                Countdown.pinned.is_(True),
+                Countdown.archived_at.is_(None),
+            ))).all())
+            if (
+                incoming_pin_ids
+                and existing_pin_ids
+                and incoming_pin_ids != {str(item) for item in existing_pin_ids}
+            ):
+                raise backup_error("backup_constraint_invalid", "合并恢复会产生多个置顶倒数日")
+        shutil.rmtree(archive.staging_dir, ignore_errors=True)
+        if archive.archive_sha256 != digest.hexdigest():
+            raise backup_error("backup_checksum_mismatch", "备份校验和不匹配")
+        reservation.backup_id = archive.backup_id
+        reservation.archive_sha256 = archive.archive_sha256
+        reservation.archive_size = size
+        reservation.status = "pending"
+        await db.commit()
+    except Exception:
+        await db.rollback()
+        if archive is not None:
+            shutil.rmtree(archive.staging_dir, ignore_errors=True)
+        failed = await db.get(BackupPreflight, reservation_id)
+        if failed is not None:
+            await db.delete(failed)
+            await db.commit()
+        staging.unlink(missing_ok=True)
+        raise
+    return {"valid": True, "preflight_token": token, "backup_id": archive.backup_id,
+            "archive_sha256": archive.archive_sha256,
+            "entities": {name: len(rows) for name, rows in archive.entities.items()}}
+
+
+@router.post("/restore")
+async def restore(
+    payload: RestoreRequest,
+    user: User = Depends(current_user),
+    db: AsyncSession = Depends(get_db),
+):
+    await _prune(db)
+    token_hash = hash_token(payload.preflight_token)
+    row = await db.scalar(select(BackupPreflight).where(
+        BackupPreflight.token_hash == token_hash, BackupPreflight.user_id == user.id
+    ).with_for_update())
+    if row is None or row.mode != payload.mode or row.expires_at <= utcnow():
+        raise backup_error("backup_preflight_invalid", "预检令牌无效或已过期", 409)
+    if row.status == "repair_pending" and row.cleanup_path:
+        from .storage import restore_quarantine_dir
+        root = Path(get_settings().attachment_dir).resolve()
+        try:
+            restore_quarantine_dir(root, Path(row.cleanup_path))
+        except OSError as exc:
+            raise backup_error("backup_repair_pending", "原附件复原尚未完成，请稍后重试", 503) from exc
+        row.status = "failed"
+        row.cleanup_path = None
+        await db.commit()
+        raise backup_error("backup_restore_retry", "附件已复原，请重新预检后重试", 409)
+    if row.status == "cleanup_pending":
+        from .storage import remove_quarantine
+        try:
+            if row.cleanup_path:
+                remove_quarantine(Path(row.cleanup_path))
+            _safe_staging_path(row.staging_path).unlink(missing_ok=True)
+        except OSError as exc:
+            raise backup_error("backup_cleanup_pending", "清理尚未完成，请稍后重试", 503) from exc
+        row.status = "consumed"
+        row.cleanup_path = None
+        await db.commit()
+        return {"restored": 0, "mode": payload.mode, "already_imported": False, "cleanup_retried": True}
+    claimed = await db.execute(update(BackupPreflight).where(
+        BackupPreflight.id == row.id, BackupPreflight.status.in_(("pending", "failed"))
+    ).values(status="consuming", consumed_at=utcnow()))
+    if claimed.rowcount != 1:
+        await db.rollback()
+        raise backup_error("backup_preflight_invalid", "预检令牌已使用", 409)
+    await db.commit()
+    staging = _safe_staging_path(row.staging_path)
+    archive = parse_archive_path(staging)
+    try:
+        if archive.archive_sha256 != row.archive_sha256 or archive.backup_id != row.backup_id:
+            row.status = "failed"
+            await db.commit()
+            raise backup_error("backup_preflight_invalid", "预检暂存已改变", 409)
+        result = await restore_v2(db, user, archive, payload.mode, operation=row)
+    except BackupRepairPending:
+        raise
+    except Exception:
+        await db.refresh(row)
+        if row.status not in {"cleanup_pending", "repair_pending"}:
+            row.status = "failed"
+            await db.commit()
+        raise
+    finally:
+        shutil.rmtree(archive.staging_dir, ignore_errors=True)
+    row.status = "cleanup_pending"
+    await db.commit()
+    try:
+        staging.unlink(missing_ok=True)
+    except OSError as exc:
+        raise backup_error("backup_cleanup_pending", "数据已恢复，但暂存清理未完成；请使用同一令牌重试", 500) from exc
+    row.status = "consumed"
+    await db.commit()
+    return result