Skip to main content

yggdrasil/api/database/
backup.rs

1#![allow(clippy::unused_unit, deprecated)]
2
3//! 备份与恢复(读写,最高风险)。
4//!
5//! 备份:探测 pg_dump 可用性——可用则子进程生成完整、角色无关的 .sql
6//! (`--clean --if-exists --no-owner --no-privileges`),不可用则回退纯 SQL(仅数据)。
7//! 备份文件含签名头。每次备份按设置附带 uploads 素材打包(`_uploads.tar.gz`,
8//! 排除可重建的 `.cache/`),与 .sql 成对展示/下载/删除。
9//!
10//! 来源标记:手动备份文件名前缀 `backup_`,定时任务自动备份前缀 `auto_`;
11//! 自动备份按 `backup_retention_count` 轮转(只删自动,手动永不自动删)。
12//!
13//! 恢复:仅接受本系统生成的备份(签名校验)+ 二次确认 + 路径穿越防护;
14//! 旧备份的 pg_dump 属主语句会在 COPY 感知的预处理阶段移除,再由
15//! `psql --single-transaction -v ON_ERROR_STOP=1` 原子重放,失败时完整回滚;
16//! 仅恢复数据库——uploads 需从配对 tar.gz 手动还原;成功后全量失效缓存。
17//! 长耗时操作走后台任务 + 进度轮询(见 [`crate::api::database::tasks`])。
18//!
19//! 导入:`POST /api/database/backups/import`(纯 Axum 路由,multipart 流式落盘)
20//! 把本机备份回灌进 backups/。两步式——导入只入库,恢复仍走上述管线。
21//! 导入时即强制签名校验(外来 SQL 拒收不留盘);保留原始文件名、同名冲突拒绝;
22//! `.tmp` 落盘 + 原子 rename,半截文件永不出现在列表。上限 `BACKUP_IMPORT_MAX_MB`
23//! (默认 512MB),路由层 DefaultBodyLimit + 流式计数双保险。
24//!
25//! **兼容性提示**:`--clean --if-exists` 是后加的备份参数。本修复之前生成的备份
26//! 文件不含 DROP,对其执行恢复会在第一条「relation already exists」处中止并报
27//! 失败(行为正确,但无法恢复数据)——需重新创建备份才能恢复。
28
29// Component/PathBuf/chrono::Utc 仅 server 构建的备份逻辑用到。
30#[cfg(feature = "server")]
31use std::path::{Component, Path, PathBuf};
32
33#[cfg(feature = "server")]
34use chrono::Utc;
35use dioxus::prelude::*;
36use serde::{Deserialize, Serialize};
37
38// admin 鉴权 + AppError + tasks 进度表仅在 server 构建里被 server function 体引用。
39#[cfg(feature = "server")]
40use crate::api::auth::get_current_admin_user;
41#[cfg(feature = "server")]
42use crate::api::database::tasks::{self, TaskKind, TaskStatus};
43#[cfg(feature = "server")]
44use crate::api::error::AppError;
45
46// 以下常量仅被 server 构建的备份/恢复逻辑引用(WASM 构建里相关函数体被 cfg 剥掉,
47// 故常量也需 gate,否则非 server 构建会报 dead_code)。
48
49/// 备份目录(项目根,与 uploads/ 平级,gitignored)。
50#[cfg(feature = "server")]
51const BACKUP_DIR: &str = "backups";
52/// 文件名白名单正则:仅字母数字下划线点连字符(防路径穿越)。
53#[cfg(feature = "server")]
54const FILENAME_RE: &str = r"^[a-zA-Z0-9_.\-]+$";
55/// 备份文件签名头(恢复时校验,拒绝非本系统文件)。
56#[cfg(feature = "server")]
57const BACKUP_SIGNATURE: &str = "-- YGGDRASIL BACKUP v1";
58/// 导入单文件上限默认值(MB),env `BACKUP_IMPORT_MAX_MB` 覆盖。
59#[cfg(feature = "server")]
60const DEFAULT_IMPORT_MAX_MB: u64 = 512;
61/// multipart 框架开销宽限(boundary/头部),加在路由 body limit 上,
62/// 避免恰好等于上限的文件被框架字节误杀(流式计数只计载荷字节,才是权威)。
63#[cfg(feature = "server")]
64pub(crate) const MULTIPART_FRAME_SLACK: u64 = 1024 * 1024;
65
66/// 备份文件元信息(列表展示用)。
67#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
68pub struct BackupInfo {
69    pub filename: String,
70    pub size_bytes: u64,
71    /// 备份模式:pg_dump / sql-fallback(从签名头解析)。
72    pub mode: String,
73    pub created_at: Option<String>,
74    /// 来源:manual / auto(从文件名前缀解析)。
75    pub origin: String,
76    /// 配对的 uploads 打包文件名与大小(无则 None——老备份或禁用了素材打包)。
77    pub uploads_filename: Option<String>,
78    pub uploads_size_bytes: Option<u64>,
79}
80
81/// 备份来源:决定文件名前缀与是否参与自动轮转。
82#[cfg(feature = "server")]
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
84pub(crate) enum BackupOrigin {
85    /// admin 面板手动触发(`backup_` 前缀,永不自动删除)。
86    Manual,
87    /// 定时任务触发(`auto_` 前缀,按 retention_count 轮转)。
88    Auto,
89}
90
91#[cfg(feature = "server")]
92impl BackupOrigin {
93    fn file_prefix(self) -> &'static str {
94        match self {
95            Self::Manual => "backup",
96            Self::Auto => "auto",
97        }
98    }
99}
100
101/// 一次备份运行的结果(调度任务据此落库 last_run)。
102#[cfg(feature = "server")]
103#[derive(Debug)]
104pub(crate) struct BackupRunOutcome {
105    /// 数据库 SQL 文件名。
106    pub sql_filename: String,
107    /// uploads 打包文件名(禁用了素材打包时为 None)。
108    pub uploads_filename: Option<String>,
109    /// 非致命告警(如 DB 成功但 uploads 打包失败)。
110    pub warning: Option<String>,
111}
112
113/// 发起备份,立即返回 task_id,后台任务执行。
114#[server(CreateBackup, "/api")]
115pub async fn create_backup() -> Result<String, ServerFnError> {
116    let _user = get_current_admin_user().await?;
117
118    #[cfg(feature = "server")]
119    {
120        let task_id = uuid::Uuid::new_v4().to_string();
121        tasks::insert(task_id.clone(), TaskKind::Backup);
122        let tid = task_id.clone();
123        tokio::spawn(async move {
124            // 手动触发:进度经任务表轮询展示,返回值无人消费。
125            let _ = run_backup(&tid, BackupOrigin::Manual).await;
126        });
127        Ok(task_id)
128    }
129    #[cfg(not(feature = "server"))]
130    {
131        Ok(String::new())
132    }
133}
134
135/// 定时任务入口:注册任务后执行自动备份,返回结果供调度器落库 last_run。
136#[cfg(feature = "server")]
137pub(crate) async fn run_auto_backup() -> Result<BackupRunOutcome, String> {
138    let task_id = uuid::Uuid::new_v4().to_string();
139    tasks::insert(task_id.clone(), TaskKind::Backup);
140    run_backup(&task_id, BackupOrigin::Auto).await
141}
142
143/// 后台执行备份:pg_dump 优先,不可用回退纯 SQL;成功后按设置打包 uploads,
144/// 自动备份再做轮转。Ok/Err 均已同步进任务进度表(供前端轮询),返回值供
145/// 调度器落库 last_run。
146#[cfg(feature = "server")]
147async fn run_backup(task_id: &str, origin: BackupOrigin) -> Result<BackupRunOutcome, String> {
148    let _ = std::fs::create_dir_all(BACKUP_DIR);
149    let timestamp = Utc::now().format("%Y%m%d_%H%M%S").to_string();
150    let prefix = origin.file_prefix();
151
152    // 读取备份设置(include_uploads / retention_count);读取失败用默认值——
153    // DB 不可达时 pg_dump 自身也会失败并按原路径上报。
154    let settings = match crate::db::pool::get_conn().await {
155        Ok(conn) => crate::api::settings::load_backup_settings(&conn)
156            .await
157            .unwrap_or_default(),
158        Err(_) => crate::models::settings::BackupSettings::default(),
159    };
160
161    // 探测 pg_dump(fork+exec+wait 仍是阻塞系统调用,移出 tokio worker 线程)
162    let pg_dump_ok = tokio::task::spawn_blocking(|| {
163        std::process::Command::new("pg_dump")
164            .arg("--version")
165            .output()
166            .is_ok()
167    })
168    .await
169    .unwrap_or(false);
170
171    let sql_result = if pg_dump_ok {
172        run_pg_dump_backup(task_id, prefix, &timestamp).await
173    } else {
174        run_sql_fallback_backup(task_id, prefix, &timestamp).await
175    };
176    let sql_filename = match sql_result {
177        Ok(f) => f,
178        // 子流程已把具体错误写进任务进度表,这里仅透传给调度器。
179        Err(e) => return Err(e),
180    };
181
182    // uploads 素材打包:失败不拖垮已成功的 DB 备份,降级为告警。
183    let mut uploads_filename = None;
184    let mut warning = None;
185    if settings.include_uploads {
186        tasks::update(
187            task_id,
188            "正在打包 uploads 素材",
189            92,
190            TaskStatus::Running,
191            None,
192            None,
193            None,
194        );
195        let tar_name = uploads_archive_name(&sql_filename);
196        let tar_path = backup_path(&tar_name);
197        match tokio::task::spawn_blocking(move || {
198            create_uploads_tarball(Path::new("uploads"), &tar_path)
199        })
200        .await
201        {
202            Ok(Ok(())) => uploads_filename = Some(tar_name),
203            Ok(Err(e)) => warning = Some(format!("uploads 打包失败: {e}")),
204            Err(e) => warning = Some(format!("uploads 打包任务panic: {e}")),
205        }
206        if warning.is_some() {
207            tracing::warn!("backup uploads tarball failed: {:?}", warning);
208        }
209    }
210
211    // 自动备份轮转:只删 auto_ 前缀的最旧者,手动备份永不触碰。
212    if origin == BackupOrigin::Auto {
213        rotate_auto_backups(settings.retention_count);
214    }
215
216    tasks::update(
217        task_id,
218        "完成",
219        100,
220        TaskStatus::Done,
221        warning.clone(),
222        None,
223        Some(sql_filename.clone()),
224    );
225    Ok(BackupRunOutcome {
226        sql_filename,
227        uploads_filename,
228        warning,
229    })
230}
231
232/// SQL 文件名 → 配对 uploads 打包文件名(`X.sql` → `X_uploads.tar.gz`)。
233#[cfg(feature = "server")]
234fn uploads_archive_name(sql_filename: &str) -> String {
235    format!("{}_uploads.tar.gz", sql_filename.trim_end_matches(".sql"))
236}
237
238/// 把 uploads 目录打成 tar.gz(排除可重建的 .cache/ 与 VCS 占位 .gitkeep)。
239/// 阻塞 IO + CPU(gzip),调用方须放 spawn_blocking。
240/// uploads_dir 显式传入(而非写死相对路径)便于测试注入临时目录。
241#[cfg(feature = "server")]
242fn create_uploads_tarball(uploads_dir: &Path, out_path: &Path) -> std::io::Result<()> {
243    let file = std::fs::File::create(out_path)?;
244    let gz = flate2::write::GzEncoder::new(file, flate2::Compression::default());
245    let mut builder = tar::Builder::new(gz);
246    if uploads_dir.is_dir() {
247        for entry in std::fs::read_dir(uploads_dir)? {
248            let entry = entry?;
249            let name = entry.file_name();
250            if name == ".cache" || name == ".gitkeep" {
251                continue;
252            }
253            let path = entry.path();
254            if path.is_dir() {
255                builder.append_dir_all(&name, &path)?;
256            } else if path.is_file() {
257                builder.append_path_with_name(&path, &name)?;
258            }
259        }
260    }
261    let gz = builder.into_inner()?;
262    gz.finish()?;
263    Ok(())
264}
265
266/// 从自动备份 SQL 文件名列表挑出超出保留份数、应删除的最旧者。
267/// 文件名内嵌定宽时间戳,字典序即时间序。纯函数便于单测。
268#[cfg(feature = "server")]
269fn select_expired_auto_backups(names: &[String], keep: usize) -> Vec<String> {
270    let mut autos: Vec<&String> = names
271        .iter()
272        .filter(|n| n.starts_with("auto_") && n.ends_with(".sql"))
273        .collect();
274    autos.sort();
275    let excess = autos.len().saturating_sub(keep);
276    autos.into_iter().take(excess).cloned().collect()
277}
278
279/// 执行自动备份轮转:删除超龄自动备份的 .sql 与配对 tar.gz。
280/// 单个删除失败只记日志,不中断其余。
281#[cfg(feature = "server")]
282fn rotate_auto_backups(keep: i32) {
283    let names: Vec<String> = match std::fs::read_dir(BACKUP_DIR) {
284        Ok(entries) => entries
285            .flatten()
286            .map(|e| e.file_name().to_string_lossy().to_string())
287            .collect(),
288        Err(e) => {
289            tracing::warn!("backup rotation: cannot read {BACKUP_DIR}: {e}");
290            return;
291        }
292    };
293    for sql_name in select_expired_auto_backups(&names, keep.max(0) as usize) {
294        for path in [
295            backup_path(&sql_name),
296            backup_path(&uploads_archive_name(&sql_name)),
297        ] {
298            match std::fs::remove_file(&path) {
299                Ok(()) => tracing::info!("backup rotation: deleted {}", path.display()),
300                Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
301                Err(e) => {
302                    tracing::warn!("backup rotation: failed to delete {}: {e}", path.display())
303                }
304            }
305        }
306    }
307}
308
309/// pg_dump 模式:子进程生成完整备份(含 schema),前置签名头。
310///
311/// `--clean --if-exists` 生成幂等的删后重建脚本;`--no-owner --no-privileges`
312/// 去掉集群角色元数据,使备份可恢复到使用不同数据库角色的环境。详见 `run_restore`。
313///
314/// 成功返回 SQL 文件名(最终「完成」进度由调用方统一上报)。
315#[cfg(feature = "server")]
316fn pg_command_with_db_config(
317    program: &str,
318    db_url: &str,
319) -> std::io::Result<std::process::Command> {
320    use std::str::FromStr;
321    use tokio_postgres::config::{Host, SslMode};
322
323    let config = tokio_postgres::Config::from_str(db_url).map_err(|error| {
324        std::io::Error::new(
325            std::io::ErrorKind::InvalidInput,
326            format!("DATABASE_URL 解析失败: {error}"),
327        )
328    })?;
329    let mut command = std::process::Command::new(program);
330    // pg_dump/psql 不读取 DATABASE_URL;移除继承的敏感环境变量,只保留
331    // 拆出的 PG* 连接参数,避免子进程环境再次暴露完整 URI。
332    command.env_remove("DATABASE_URL");
333
334    let hosts: Vec<String> = config
335        .get_hosts()
336        .iter()
337        .map(|host| match host {
338            Host::Tcp(host) => host.clone(),
339            Host::Unix(path) => path.to_string_lossy().into_owned(),
340        })
341        .collect();
342    if !hosts.is_empty() {
343        command.env("PGHOST", hosts.join(","));
344    }
345    if !config.get_hostaddrs().is_empty() {
346        command.env(
347            "PGHOSTADDR",
348            config
349                .get_hostaddrs()
350                .iter()
351                .map(ToString::to_string)
352                .collect::<Vec<_>>()
353                .join(","),
354        );
355    }
356    if !config.get_ports().is_empty() {
357        command.env(
358            "PGPORT",
359            config
360                .get_ports()
361                .iter()
362                .map(ToString::to_string)
363                .collect::<Vec<_>>()
364                .join(","),
365        );
366    }
367    if let Some(user) = config.get_user() {
368        command.env("PGUSER", user);
369    }
370    if let Some(password) = config.get_password() {
371        command.env("PGPASSWORD", String::from_utf8_lossy(password).into_owned());
372    }
373    if let Some(dbname) = config.get_dbname() {
374        command.env("PGDATABASE", dbname);
375    }
376    if let Some(options) = config.get_options() {
377        command.env("PGOPTIONS", options);
378    }
379    if let Some(application_name) = config.get_application_name() {
380        command.env("PGAPPNAME", application_name);
381    }
382    let ssl_mode = match config.get_ssl_mode() {
383        SslMode::Disable => "disable",
384        SslMode::Prefer => "prefer",
385        SslMode::Require => "require",
386        _ => "prefer",
387    };
388    command.env("PGSSLMODE", ssl_mode);
389    if let Some(timeout) = config.get_connect_timeout() {
390        command.env("PGCONNECT_TIMEOUT", timeout.as_secs().max(1).to_string());
391    }
392    Ok(command)
393}
394
395#[cfg(feature = "server")]
396fn pg_dump_command(db_url: &str) -> std::io::Result<std::process::Command> {
397    let mut command = pg_command_with_db_config("pg_dump", db_url)?;
398    command.args(["--clean", "--if-exists", "--no-owner", "--no-privileges"]);
399    Ok(command)
400}
401
402#[cfg(feature = "server")]
403async fn run_pg_dump_backup(
404    task_id: &str,
405    prefix: &str,
406    timestamp: &str,
407) -> Result<String, String> {
408    tasks::update(
409        task_id,
410        "正在用 pg_dump 导出",
411        10,
412        TaskStatus::Running,
413        None,
414        None,
415        None,
416    );
417    let filename = format!("{prefix}_{timestamp}.sql");
418    let path = backup_path(&filename);
419    let db_url = match std::env::var("DATABASE_URL") {
420        Ok(u) if !u.is_empty() => u,
421        _ => {
422            let msg = "pg_dump 备份需要 DATABASE_URL".to_string();
423            tasks::update(
424                task_id,
425                "DATABASE_URL 未配置",
426                100,
427                TaskStatus::Failed,
428                None,
429                Some(msg.clone()),
430                None,
431            );
432            return Err(msg);
433        }
434    };
435
436    let mut header = String::new();
437    header.push_str(&format!("{}\n", BACKUP_SIGNATURE));
438    header.push_str(&format!("-- created_at: {}\n", Utc::now()));
439    header.push_str("-- mode: pg_dump\n");
440
441    // 先写签名头,再追加 pg_dump 输出。
442    if let Err(e) = std::fs::write(&path, &header) {
443        let msg = format!("无法写入备份目录: {e}");
444        tasks::update(
445            task_id,
446            "写入备份文件失败",
447            100,
448            TaskStatus::Failed,
449            None,
450            Some(msg.clone()),
451            None,
452        );
453        return Err(msg);
454    }
455
456    let stdout_file = match std::fs::OpenOptions::new().append(true).open(&path) {
457        Ok(f) => f,
458        Err(e) => {
459            let msg = e.to_string();
460            tasks::update(
461                task_id,
462                "pg_dump 启动失败",
463                100,
464                TaskStatus::Failed,
465                None,
466                Some(msg.clone()),
467                None,
468            );
469            return Err(msg);
470        }
471    };
472    // pg_dump 导出可持续数十秒到数分钟,整个子进程生命周期(spawn + wait_with_output)
473    // 移入 spawn_blocking,避免阻塞 tokio worker 线程。注意 stdout 重定向到备份文件,
474    // 故不能用 .output()(它会用 piped 覆盖 stdout 配置);保留 spawn() + wait_with_output()
475    // 两段式。闭包返回 Result<Output, (bool, io::Error)>:true=启动(spawn)阶段失败,
476    // false=等待(wait)阶段失败,分别对应原有「启动失败」「执行失败」两条上报路径;
477    // 闭包 panic(JoinError)按「执行失败」处理。
478    let dump_result = tokio::task::spawn_blocking(
479        move || -> Result<std::process::Output, (bool, std::io::Error)> {
480            let mut command = pg_dump_command(&db_url).map_err(|error| (true, error))?;
481            command
482                .stdout(std::process::Stdio::from(stdout_file))
483                .stderr(std::process::Stdio::piped())
484                .spawn()
485                .map_err(|e| (true, e))?
486                .wait_with_output()
487                .map_err(|e| (false, e))
488        },
489    )
490    .await
491    .unwrap_or_else(|join_e| Err((false, std::io::Error::other(join_e.to_string()))));
492    match dump_result {
493        Ok(o) if o.status.success() => Ok(filename),
494        Ok(o) => {
495            let msg = String::from_utf8_lossy(&o.stderr).to_string();
496            tasks::update(
497                task_id,
498                "pg_dump 失败",
499                100,
500                TaskStatus::Failed,
501                None,
502                Some(msg.clone()),
503                None,
504            );
505            Err(msg)
506        }
507        Err((true, e)) => {
508            let msg = e.to_string();
509            tasks::update(
510                task_id,
511                "pg_dump 启动失败",
512                100,
513                TaskStatus::Failed,
514                None,
515                Some(msg.clone()),
516                None,
517            );
518            Err(msg)
519        }
520        Err((false, e)) => {
521            let msg = e.to_string();
522            tasks::update(
523                task_id,
524                "pg_dump 执行失败",
525                100,
526                TaskStatus::Failed,
527                None,
528                Some(msg.clone()),
529                None,
530            );
531            Err(msg)
532        }
533    }
534}
535
536/// 纯 SQL 回退:仅备份数据(不含 schema),按表计数精确进度。
537///
538/// 注意:该产物不是合法 SQL 脚本(COPY TO STDOUT 的 CSV 原文拼接),
539/// 不能经 psql 恢复——仅供极端环境下抢救数据。生产镜像自带 pg_dump,
540/// 正常不会走到这里。
541///
542/// 成功返回 SQL 文件名(最终「完成」进度由调用方统一上报)。
543#[cfg(feature = "server")]
544async fn run_sql_fallback_backup(
545    task_id: &str,
546    prefix: &str,
547    timestamp: &str,
548) -> Result<String, String> {
549    tasks::update(
550        task_id,
551        "pg_dump 不可用,使用纯 SQL 回退(仅数据)",
552        10,
553        TaskStatus::Running,
554        Some("仅备份数据,不含 schema/索引/触发器,且不可经 psql 恢复".to_string()),
555        None,
556        None,
557    );
558    let filename = format!("{prefix}_{timestamp}_sqlfallback.sql");
559    let path = backup_path(&filename);
560
561    let client = match crate::db::pool::get_conn().await {
562        Ok(c) => c,
563        Err(e) => {
564            let msg = e.to_string();
565            tasks::update(
566                task_id,
567                "数据库连接失败",
568                100,
569                TaskStatus::Failed,
570                None,
571                Some(msg.clone()),
572                None,
573            );
574            return Err(msg);
575        }
576    };
577
578    // 取 public schema 下所有表名
579    let tables: Vec<String> = match client
580        .query(
581            "SELECT tablename FROM pg_tables WHERE schemaname = 'public' ORDER BY tablename",
582            &[],
583        )
584        .await
585    {
586        Ok(rows) => rows.into_iter().map(|r| r.get(0)).collect(),
587        Err(e) => {
588            let msg = e.to_string();
589            tasks::update(
590                task_id,
591                "读取表清单失败",
592                100,
593                TaskStatus::Failed,
594                None,
595                Some(msg.clone()),
596                None,
597            );
598            return Err(msg);
599        }
600    };
601    let total = tables.len().max(1);
602
603    let mut out = String::new();
604    out.push_str(&format!("{}\n", BACKUP_SIGNATURE));
605    out.push_str(&format!("-- created_at: {}\n", Utc::now()));
606    out.push_str("-- mode: sql-fallback\n\n");
607
608    for (i, table) in tables.iter().enumerate() {
609        out.push_str(&format!("\n-- table: {}\n", table));
610        let copy_stmt = format!("COPY \"{}\" TO STDOUT WITH CSV", table);
611        match client.copy_out(&copy_stmt).await {
612            Ok(stream) => {
613                use futures::StreamExt;
614                // CopyOutStream 是 !Unpin,必须 pin 才能调 next。
615                tokio::pin!(stream);
616                while let Some(chunk) = stream.next().await {
617                    if let Ok(bytes) = chunk {
618                        out.push_str(&String::from_utf8_lossy(&bytes));
619                    }
620                }
621            }
622            Err(e) => {
623                out.push_str(&format!("-- 导出失败: {}\n", e));
624            }
625        }
626        // 按表更新进度(用 u32 避免大 schema 下的截断/溢出)
627        tasks::update(
628            task_id,
629            &format!("导出表 {}/{}", i + 1, total),
630            (10 + (i + 1) as u32 * 90 / total as u32).min(99) as u8,
631            TaskStatus::Running,
632            None,
633            None,
634            None,
635        );
636    }
637
638    if let Err(e) = std::fs::write(&path, out) {
639        let msg = format!("无法写入备份目录: {e}");
640        tasks::update(
641            task_id,
642            "写入备份文件失败",
643            100,
644            TaskStatus::Failed,
645            None,
646            Some(msg.clone()),
647            None,
648        );
649        return Err(msg);
650    }
651    Ok(filename)
652}
653
654/// 发起恢复:校验签名 + 路径穿越防护 + 二次确认,立即返回 task_id。
655#[server(RestoreBackup, "/api")]
656pub async fn restore_backup(filename: String, confirm: bool) -> Result<String, ServerFnError> {
657    let _user = get_current_admin_user().await?;
658
659    // 全部校验都在 server cfg 块内:confirm/regex/backup_path/std::fs 都是 server-only。
660    // WASM 侧的 server-function 客户端桩只返回 Ok(String::new())。
661    #[cfg(feature = "server")]
662    {
663        if !confirm {
664            return Err(AppError::BadRequest("需确认恢复(会覆盖现有数据)".to_string()).into());
665        }
666        // 路径穿越防护
667        if !is_valid_backup_filename(&filename) {
668            return Err(AppError::BadRequest("无效的文件名".to_string()).into());
669        }
670        let path = backup_path(&filename);
671        if !path.exists() {
672            return Err(AppError::NotFound("备份文件不存在").into());
673        }
674
675        // 签名校验:仅读取首行(备份文件可达数十 MB,无需整文件读入内存)。
676        let first_line = read_first_line(&path).unwrap_or_default();
677        if !has_valid_signature(&first_line) {
678            return Err(
679                AppError::BadRequest("非本系统生成的备份文件,拒绝恢复".to_string()).into(),
680            );
681        }
682
683        let task_id = uuid::Uuid::new_v4().to_string();
684        tasks::insert(task_id.clone(), TaskKind::Restore);
685        let tid = task_id.clone();
686        let f = filename;
687        tokio::spawn(async move {
688            run_restore(&tid, &f).await;
689        });
690        Ok(task_id)
691    }
692    #[cfg(not(feature = "server"))]
693    {
694        // WASM 客户端桩:忽略参数,返回空 task_id。
695        let _ = (filename, confirm);
696        Ok(String::new())
697    }
698}
699
700/// 后台执行恢复:探测 psql,可用则原子恢复,不可用则报告。
701///
702/// **跨环境恢复**:旧备份可能含 `ALTER ... OWNER TO ...`,而目标集群未必存在
703/// 源库角色。恢复前流式生成临时副本,只移除 pg_dump 生成的属主语句;COPY 数据
704/// 原样保留。新备份已由 `--no-owner --no-privileges` 从源头去掉角色元数据。
705///
706/// **原子恢复**:备份由 `pg_dump --clean --if-exists` 生成,脚本自带
707/// `DROP ... IF EXISTS`;psql 使用 `--single-transaction`,任一 SQL 错误都会
708/// 回滚先前的 DROP/CREATE/COPY,不留下半恢复数据库。
709///
710/// **错误中止**:`-v ON_ERROR_STOP=1` 让 psql 在第一条 SQL 错误时立即退出
711/// (退出码 3)。否则 psql 即使满屏 ERROR 也返回 0,导致 `status.success()`
712/// 误判为成功——这是「恢复完成却无任何数据变更」假成功的根因。
713///
714/// **恢复成功后失效全量缓存**:恢复会用备份时刻的数据重建 posts 等表,
715/// 现有 moka 缓存(列表/标签/单篇/统计/搜索)与 SSR 世代号必须一并冲刷,
716/// 否则前端仍读旧数据。
717#[cfg(feature = "server")]
718fn sql_line_without_eol(mut line: &[u8]) -> &[u8] {
719    if let Some(stripped) = line.strip_suffix(b"\n") {
720        line = stripped;
721    }
722    if let Some(stripped) = line.strip_suffix(b"\r") {
723        line = stripped;
724    }
725    line
726}
727
728#[cfg(feature = "server")]
729fn is_legacy_pg_dump_owner_statement(line: &[u8]) -> bool {
730    const PREFIXES: [&[u8]; 3] = [b"ALTER FUNCTION ", b"ALTER SEQUENCE ", b"ALTER TABLE "];
731    let line = sql_line_without_eol(line);
732    PREFIXES.iter().any(|prefix| line.starts_with(prefix))
733        && line
734            .windows(b" OWNER TO ".len())
735            .any(|window| window == b" OWNER TO ")
736        && line.ends_with(b";")
737}
738
739/// 去除恢复脚本行首空白,确保带缩进的 psql 元命令也进入安全检查。
740#[cfg(feature = "server")]
741fn trim_restore_line_start(line: &[u8]) -> &[u8] {
742    let line = sql_line_without_eol(line);
743    let start = line
744        .iter()
745        .position(|byte| !byte.is_ascii_whitespace())
746        .unwrap_or(line.len());
747    &line[start..]
748}
749
750#[cfg(feature = "server")]
751fn is_allowed_restore_meta_command(line: &[u8]) -> bool {
752    let line = trim_restore_line_start(line);
753    let Some(rest) = line.strip_prefix(b"\\") else {
754        return false;
755    };
756    let mut fields = rest
757        .split(|byte| byte.is_ascii_whitespace())
758        .filter(|field| !field.is_empty());
759    let Some(command) = fields.next() else {
760        return false;
761    };
762    if command != b"restrict" && command != b"unrestrict" {
763        return false;
764    }
765    let Some(token) = fields.next() else {
766        return false;
767    };
768    !token.is_empty()
769        && token
770            .iter()
771            .all(|byte| byte.is_ascii_alphanumeric() || *byte == b'_')
772        && fields.next().is_none()
773}
774
775#[cfg(feature = "server")]
776fn write_owner_neutral_restore_sql(
777    mut input: impl std::io::BufRead,
778    mut output: impl std::io::Write,
779) -> std::io::Result<usize> {
780    let mut line = Vec::new();
781    let mut in_copy_data = false;
782    let mut removed = 0;
783    loop {
784        line.clear();
785        if input.read_until(b'\n', &mut line)? == 0 {
786            break;
787        }
788        let sql_line = sql_line_without_eol(&line);
789        if in_copy_data {
790            output.write_all(&line)?;
791            if sql_line == b"\\." {
792                in_copy_data = false;
793            }
794            continue;
795        }
796        if sql_line.starts_with(b"COPY ") && sql_line.ends_with(b" FROM stdin;") {
797            in_copy_data = true;
798            output.write_all(&line)?;
799        } else if is_legacy_pg_dump_owner_statement(&line) {
800            removed += 1;
801        } else if trim_restore_line_start(sql_line).first() == Some(&b'\\')
802            && !is_allowed_restore_meta_command(&line)
803        {
804            return Err(std::io::Error::new(
805                std::io::ErrorKind::InvalidData,
806                "unsupported psql meta-command in restore file",
807            ));
808        } else {
809            output.write_all(&line)?;
810        }
811    }
812
813    output.flush()?;
814    Ok(removed)
815}
816
817#[cfg(feature = "server")]
818struct PreparedRestoreSql {
819    path: PathBuf,
820    removed_owner_statements: usize,
821}
822
823#[cfg(feature = "server")]
824impl PreparedRestoreSql {
825    fn prepare(source: &Path) -> std::io::Result<Self> {
826        let input = std::io::BufReader::new(std::fs::File::open(source)?);
827        let path =
828            std::env::temp_dir().join(format!("yggdrasil-restore-{}.sql", uuid::Uuid::new_v4()));
829        let mut options = std::fs::OpenOptions::new();
830        options.write(true).create_new(true);
831        #[cfg(unix)]
832        {
833            use std::os::unix::fs::OpenOptionsExt;
834            options.mode(0o600);
835        }
836        let output = std::io::BufWriter::new(options.open(&path)?);
837        match write_owner_neutral_restore_sql(input, output) {
838            Ok(removed_owner_statements) => Ok(Self {
839                path,
840                removed_owner_statements,
841            }),
842            Err(error) => {
843                let _ = std::fs::remove_file(&path);
844                Err(error)
845            }
846        }
847    }
848}
849
850#[cfg(feature = "server")]
851impl Drop for PreparedRestoreSql {
852    fn drop(&mut self) {
853        if let Err(error) = std::fs::remove_file(&self.path) {
854            tracing::warn!(
855                path = %self.path.display(),
856                "restore: failed to delete prepared SQL: {error}"
857            );
858        }
859    }
860}
861
862#[cfg(feature = "server")]
863fn psql_restore_command(db_url: &str, path: &Path) -> std::io::Result<std::process::Command> {
864    let mut command = pg_command_with_db_config("psql", db_url)?;
865    command
866        .args(["--single-transaction", "-v", "ON_ERROR_STOP=1", "-f"])
867        .arg(path)
868        .stdout(std::process::Stdio::null())
869        .stderr(std::process::Stdio::piped());
870    Ok(command)
871}
872
873#[cfg(feature = "server")]
874async fn run_restore(task_id: &str, filename: &str) {
875    let path = backup_path(filename);
876    let db_url = match std::env::var("DATABASE_URL") {
877        Ok(u) if !u.is_empty() => u,
878        _ => {
879            tasks::update(
880                task_id,
881                "DATABASE_URL 未配置",
882                100,
883                TaskStatus::Failed,
884                None,
885                Some("恢复需要 DATABASE_URL".to_string()),
886                None,
887            );
888            return;
889        }
890    };
891    let psql_ok = tokio::task::spawn_blocking(|| {
892        std::process::Command::new("psql")
893            .arg("--version")
894            .output()
895            .is_ok()
896    })
897    .await
898    .unwrap_or(false);
899    if !psql_ok {
900        tasks::update(
901            task_id,
902            "psql 不可用",
903            100,
904            TaskStatus::Failed,
905            None,
906            Some("恢复需要 psql,但当前环境未安装 psql".to_string()),
907            None,
908        );
909        return;
910    }
911    tasks::update(
912        task_id,
913        "正在用 psql 恢复",
914        50,
915        TaskStatus::Running,
916        None,
917        None,
918        None,
919    );
920    // 预处理备份与 psql 恢复均为阻塞文件/子进程操作,整体移出 tokio worker。
921    // 临时 SQL 由 RAII 守卫在成功、失败和启动错误路径统一删除。
922    let restore_result = tokio::task::spawn_blocking(move || {
923        let prepared = PreparedRestoreSql::prepare(&path)?;
924        if prepared.removed_owner_statements > 0 {
925            tracing::info!(
926                removed = prepared.removed_owner_statements,
927                "restore: removed legacy pg_dump owner statements"
928            );
929        }
930
931        let mut command = psql_restore_command(&db_url, &prepared.path)?;
932        command.output()
933    })
934    .await
935    .unwrap_or_else(|join_e| Err(std::io::Error::other(join_e.to_string())));
936    match restore_result {
937        Ok(o) if o.status.success() => {
938            // 恢复会覆盖文章、评论、用户、设置与素材元数据;所有内存/磁盘缓存
939            // 必须一起冲刷。generation 仅用于观测,不能替代物理 SSR 删除。
940            crate::cache::invalidate_all_post_caches();
941            crate::cache::invalidate_search_results();
942            crate::cache::invalidate_all_comments();
943            crate::cache::invalidate_friend_links();
944            crate::cache::invalidate_site_settings();
945            crate::cache::invalidate_security_settings();
946            crate::cache::invalidate_image_cache_settings();
947            crate::cache::invalidate_log_targets();
948            crate::cache::SESSION_CACHE.invalidate_all();
949            crate::api::image::invalidate_all_caches().await;
950            crate::ssr_cache::invalidate_ssr_all_public();
951            crate::ssr_cache::bump_global_generation();
952            tasks::update(task_id, "恢复完成", 100, TaskStatus::Done, None, None, None);
953        }
954        Ok(o) => {
955            let stderr = String::from_utf8_lossy(&o.stderr).to_string();
956            tasks::update(
957                task_id,
958                "恢复失败",
959                100,
960                TaskStatus::Failed,
961                None,
962                Some(stderr),
963                None,
964            );
965        }
966        Err(e) => {
967            tasks::update(
968                task_id,
969                "psql 启动失败",
970                100,
971                TaskStatus::Failed,
972                None,
973                Some(e.to_string()),
974                None,
975            );
976        }
977    }
978}
979
980/// 列出 backups/ 目录下的备份文件元信息。
981///
982/// uploads 打包(`*_uploads.tar.gz`)不单独成行,按文件名配对挂到对应 .sql 上。
983#[server(ListBackups, "/api")]
984pub async fn list_backups() -> Result<Vec<BackupInfo>, ServerFnError> {
985    let _user = get_current_admin_user().await?;
986    #[cfg(feature = "server")]
987    {
988        let mut infos: Vec<BackupInfo> = Vec::new();
989        // 先收集 uploads 包:文件名 → 大小,供配对查询。
990        let mut tarballs: std::collections::HashMap<String, u64> = std::collections::HashMap::new();
991        if let Ok(entries) = std::fs::read_dir(BACKUP_DIR) {
992            for entry in entries.flatten() {
993                let name = entry.file_name().to_string_lossy().to_string();
994                if name.ends_with("_uploads.tar.gz") {
995                    if let Ok(meta) = entry.metadata() {
996                        tarballs.insert(name, meta.len());
997                    }
998                }
999            }
1000        }
1001        if let Ok(entries) = std::fs::read_dir(BACKUP_DIR) {
1002            for entry in entries.flatten() {
1003                let name = entry.file_name().to_string_lossy().to_string();
1004                if !name.ends_with(".sql") {
1005                    continue;
1006                }
1007                let meta = match entry.metadata() {
1008                    Ok(m) => m,
1009                    Err(_) => continue,
1010                };
1011                // 仅读取前 3 行(签名/created_at/mode),避免把每个(可能数十 MB
1012                // 的)备份文件整文件读入内存只为取 `-- mode:` 行。
1013                let mode = read_first_lines(entry.path(), 3)
1014                    .map(|lines| parse_backup_mode(&lines.join("\n")))
1015                    .unwrap_or_else(|_| "unknown".to_string());
1016                let created_at = meta
1017                    .modified()
1018                    .ok()
1019                    .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
1020                    .map(|d| {
1021                        chrono::DateTime::<Utc>::from_timestamp(d.as_secs() as i64, 0)
1022                            .map(|dt| dt.to_rfc3339())
1023                            .unwrap_or_default()
1024                    });
1025                let tar_name = uploads_archive_name(&name);
1026                let (uploads_filename, uploads_size_bytes) = tarballs
1027                    .get_key_value(&tar_name)
1028                    .map(|(k, v)| (Some(k.clone()), Some(*v)))
1029                    .unwrap_or((None, None));
1030                infos.push(BackupInfo {
1031                    origin: if name.starts_with("auto_") {
1032                        "auto"
1033                    } else {
1034                        "manual"
1035                    }
1036                    .to_string(),
1037                    filename: name,
1038                    size_bytes: meta.len(),
1039                    mode,
1040                    created_at,
1041                    uploads_filename,
1042                    uploads_size_bytes,
1043                });
1044            }
1045        }
1046        // 按创建时间降序(新的在前)
1047        infos.sort_by(|a, b| b.created_at.cmp(&a.created_at));
1048        Ok(infos)
1049    }
1050    #[cfg(not(feature = "server"))]
1051    {
1052        Ok(vec![])
1053    }
1054}
1055
1056/// 删除备份文件(连同配对的 uploads 打包;配对缺失不视为错误)。
1057#[server(DeleteBackup, "/api")]
1058pub async fn delete_backup(filename: String) -> Result<(), ServerFnError> {
1059    let _user = get_current_admin_user().await?;
1060    #[cfg(feature = "server")]
1061    {
1062        if !is_valid_backup_filename(&filename) {
1063            return Err(AppError::BadRequest("无效的文件名".to_string()).into());
1064        }
1065        let path = backup_path(&filename);
1066        if !path.exists() {
1067            return Err(AppError::NotFound("备份文件不存在").into());
1068        }
1069        std::fs::remove_file(&path).map_err(|_| AppError::Internal("删除失败"))?;
1070        // 成对清理:传 .sql 删配对 tar.gz,传 tar.gz 删配对 .sql。
1071        let pair = if filename.ends_with(".sql") {
1072            Some(uploads_archive_name(&filename))
1073        } else {
1074            filename
1075                .strip_suffix("_uploads.tar.gz")
1076                .map(|stem| format!("{stem}.sql"))
1077        };
1078        if let Some(pair_name) = pair {
1079            match std::fs::remove_file(backup_path(&pair_name)) {
1080                Ok(()) | Err(_) => {} // 配对缺失/删除失败不影响主文件删除结果
1081            }
1082        }
1083        Ok(())
1084    }
1085    #[cfg(not(feature = "server"))]
1086    {
1087        Ok(())
1088    }
1089}
1090
1091/// 构造 backups/ 下的安全路径(额外防御:校验规范化后仍在 BACKUP_DIR 内)。
1092///
1093/// 纵深防御:即便第一道白名单 `is_valid_backup_filename` 被绕过,这里也要
1094/// 保证结果不逃出 BACKUP_DIR。直接对 filename 做 components 检查——
1095/// 含 `..`(ParentDir)、绝对路径前缀(RootDir/Prefix,如 `/etc` 或 `C:\`)
1096/// 的 filename 一律降级为 BACKUP_DIR 本身。
1097///
1098/// 注意:不能用 `[BACKUP_DIR, filename].collect::<PathBuf>()` 后再检——
1099/// 当 filename 是绝对路径时,PathBuf 语义会丢弃 BACKUP_DIR 前缀(如
1100/// `["backups", "/etc/passwd"]` → `/etc/passwd`),导致 components 检查
1101/// 在错位的路径上运行而漏判。必须先检 filename 本身。
1102#[cfg(feature = "server")]
1103fn backup_path(filename: &str) -> PathBuf {
1104    // 直接检查 filename 的 components:只允许 Normal 段。
1105    let filename_is_safe = std::path::Path::new(filename)
1106        .components()
1107        .all(|c| matches!(c, Component::Normal(_)));
1108    if filename_is_safe {
1109        let mut p = PathBuf::from(BACKUP_DIR);
1110        p.push(filename);
1111        p
1112    } else {
1113        // 命中 ParentDir/RootDir/Prefix/CurDir → 降级为 BACKUP_DIR
1114        PathBuf::from(BACKUP_DIR)
1115    }
1116}
1117
1118/// 校验备份文件名是否符合白名单(仅字母数字下划线点连字符)。
1119/// 返回 true 表示安全可用。提取为纯函数便于单测覆盖路径穿越边界。
1120#[cfg(feature = "server")]
1121fn is_valid_backup_filename(filename: &str) -> bool {
1122    // regex::Regex::new 在 FILENAME_RE 是常量正则,编译期可验证不会 panic。
1123    regex::Regex::new(FILENAME_RE)
1124        .map(|re| re.is_match(filename))
1125        .unwrap_or(false)
1126}
1127
1128/// 导入单文件上限(字节):env `BACKUP_IMPORT_MAX_MB`(MB 为单位),默认 512MB。
1129/// main.rs 路由 DefaultBodyLimit 与 handler 流式计数共用此值。
1130#[cfg(feature = "server")]
1131pub(crate) fn import_max_bytes() -> u64 {
1132    std::env::var("BACKUP_IMPORT_MAX_MB")
1133        .ok()
1134        .and_then(|s| s.parse::<u64>().ok())
1135        .filter(|mb| *mb > 0)
1136        .unwrap_or(DEFAULT_IMPORT_MAX_MB)
1137        .saturating_mul(1024 * 1024)
1138}
1139
1140/// 校验并清洗导入文件名:剥离路径分量(旧浏览器可能带 `C:\fakepath\` 前缀),
1141/// 要求 `.sql` 后缀 + 非隐藏文件 + 备份文件名白名单。纯函数便于单测。
1142#[cfg(feature = "server")]
1143fn sanitize_import_filename(raw: &str) -> Option<String> {
1144    let name = raw.rsplit(['/', '\\']).next().unwrap_or(raw);
1145    if name.len() > 255
1146        || !name.ends_with(".sql")
1147        || name.starts_with('.')
1148        || !is_valid_backup_filename(name)
1149    {
1150        return None;
1151    }
1152    Some(name.to_string())
1153}
1154
1155/// 查询 backups/ 所在文件系统的可用空间(取最长挂载点前缀的磁盘)。
1156/// 预检用途:探测失败返回 None,降级为跳过预检(写入失败仍会清理并报错)。
1157#[cfg(feature = "server")]
1158fn backup_partition_free_space() -> Option<u64> {
1159    let dir = std::fs::canonicalize(BACKUP_DIR).ok()?;
1160    let disks = sysinfo::Disks::new_with_refreshed_list();
1161    disks
1162        .iter()
1163        .filter(|d| dir.starts_with(d.mount_point()))
1164        .max_by_key(|d| d.mount_point().as_os_str().len())
1165        .map(|d| d.available_space())
1166}
1167
1168/// 从备份文件全文提取 `-- mode: <value>` 行的值(如 "pg_dump"/"sql-fallback")。
1169/// 提取为纯函数:把文件内容作为参数传入,便于单测。
1170/// 缺失或格式不符返回 "unknown"。
1171#[cfg(feature = "server")]
1172fn parse_backup_mode(content: &str) -> String {
1173    content
1174        .lines()
1175        .find(|l| l.starts_with("-- mode:"))
1176        .map(|l| l.trim_start_matches("-- mode:").trim().to_string())
1177        .filter(|s| !s.is_empty())
1178        .unwrap_or_else(|| "unknown".to_string())
1179}
1180
1181/// 校验备份文件首行是否含本系统签名头。
1182/// 提取为纯函数:把首行(或全文)作为参数传入,便于单测。
1183#[cfg(feature = "server")]
1184fn has_valid_signature(content: &str) -> bool {
1185    content
1186        .lines()
1187        .next()
1188        .map(|l| l.trim().contains(BACKUP_SIGNATURE))
1189        .unwrap_or(false)
1190}
1191
1192/// 仅读取文件首行(用于签名校验,避免把整个备份文件读入内存)。
1193#[cfg(feature = "server")]
1194fn read_first_line(path: impl AsRef<Path>) -> std::io::Result<String> {
1195    use std::io::BufRead;
1196    let mut reader = std::io::BufReader::new(std::fs::File::open(path)?);
1197    let mut line = String::new();
1198    reader.read_line(&mut line)?;
1199    Ok(line)
1200}
1201
1202/// 读取文件前 `n` 行(用于解析 `-- mode:` 等头部元信息,避免整文件读入内存)。
1203/// 不足 `n` 行时返回实际读到的行;任一行读取失败则整体返回该错误。
1204#[cfg(feature = "server")]
1205fn read_first_lines(path: impl AsRef<Path>, n: usize) -> std::io::Result<Vec<String>> {
1206    use std::io::BufRead;
1207    let reader = std::io::BufReader::new(std::fs::File::open(path)?);
1208    reader.lines().take(n).collect()
1209}
1210
1211/// Axum 处理器:导入备份(multipart 流式落盘,admin 鉴权 + 签名校验 + 同名拒绝)。
1212///
1213/// 设计要点(见 `docs/adr/0001-backup-import.md`):
1214/// - 两步式:导入只入库 backups/ 并出现在列表,恢复仍走 [`restore_backup`] 管线;
1215/// - 导入时即校验签名(拒绝外来 SQL 入库),与恢复时的签名校验是两层独立防线;
1216/// - 流式写 `.tmp` + 原子 rename:半截文件永不入库,任一失败路径清理 tmp;
1217/// - 体积双保险:路由 `DefaultBodyLimit`(上限 + [`MULTIPART_FRAME_SLACK`])
1218///   兜整个请求体,此处流式计数只计载荷字节,超限 413 早断;
1219/// - 仅 server 构建:纯 Axum 路由(在 main.rs 注册),无 WASM 消费者。
1220#[cfg(feature = "server")]
1221pub async fn import_backup(
1222    connect_info: Option<
1223        axum::extract::Extension<axum::extract::ConnectInfo<std::net::SocketAddr>>,
1224    >,
1225    headers: axum::http::HeaderMap,
1226    mut multipart: axum::extract::Multipart,
1227) -> Result<axum::Json<serde_json::Value>, (axum::http::StatusCode, axum::Json<serde_json::Value>)>
1228{
1229    use crate::api::upload::upload_error;
1230    use axum::http::StatusCode;
1231    use tokio::io::AsyncWriteExt;
1232
1233    // 0. 限流:复用 UPLOAD 桶(同为 admin 文件上传,不为低频导入单开配置面)。
1234    let peer = connect_info.map(|axum::extract::Extension(axum::extract::ConnectInfo(addr))| addr);
1235    let ip = crate::api::rate_limit::get_client_ip_with_peer(&headers, peer).await;
1236    if let Err(msg) = crate::api::rate_limit::check_upload_limit(&ip) {
1237        return Err(upload_error(StatusCode::TOO_MANY_REQUESTS, msg));
1238    }
1239
1240    // 1. cookie session → admin(与 upload_image 同一校验链)。
1241    let cookie_header = headers
1242        .get("cookie")
1243        .and_then(|h| h.to_str().ok())
1244        .unwrap_or("");
1245    let token = match crate::auth::session::parse_session_token(cookie_header) {
1246        Some(t) => t,
1247        None => return Err(upload_error(StatusCode::UNAUTHORIZED, "未登录")),
1248    };
1249    let user = match crate::api::auth::get_user_by_token(token).await {
1250        Ok(Some(u)) => u,
1251        _ => return Err(upload_error(StatusCode::UNAUTHORIZED, "会话已过期")),
1252    };
1253    if user.role != crate::models::user::UserRole::Admin {
1254        return Err(upload_error(StatusCode::FORBIDDEN, "权限不足"));
1255    }
1256
1257    let max_bytes = import_max_bytes();
1258
1259    // 2. Content-Length 早拒:明显超限不读 body 直接 413。
1260    let content_length = headers
1261        .get(axum::http::header::CONTENT_LENGTH)
1262        .and_then(|v| v.to_str().ok())
1263        .and_then(|s| s.parse::<u64>().ok());
1264    if let Some(cl) = content_length {
1265        if cl > max_bytes + MULTIPART_FRAME_SLACK {
1266            return Err(upload_error(
1267                StatusCode::PAYLOAD_TOO_LARGE,
1268                "文件超过导入上限",
1269            ));
1270        }
1271    }
1272
1273    // 3. 取 multipart 字段并清洗文件名。
1274    let mut field = match multipart.next_field().await {
1275        Ok(Some(f)) => f,
1276        Ok(None) => return Err(upload_error(StatusCode::BAD_REQUEST, "未找到文件")),
1277        Err(e) => {
1278            tracing::error!("backup import multipart error: {e:?}");
1279            return Err(upload_error(StatusCode::BAD_REQUEST, "文件读取失败"));
1280        }
1281    };
1282    let filename = match sanitize_import_filename(field.file_name().unwrap_or_default()) {
1283        Some(n) => n,
1284        None => {
1285            return Err(upload_error(
1286                StatusCode::BAD_REQUEST,
1287                "文件名不合法:仅接受以 .sql 结尾的备份文件名(字母/数字/下划线/点/连字符)",
1288            ))
1289        }
1290    };
1291
1292    // 4. 同名冲突拒绝(不覆盖——覆盖是破坏性操作;用户可先删除旧文件再导入)。
1293    if let Err(e) = std::fs::create_dir_all(BACKUP_DIR) {
1294        tracing::error!("backup import: create dir failed: {e}");
1295        return Err(upload_error(
1296            StatusCode::INTERNAL_SERVER_ERROR,
1297            "无法创建备份目录",
1298        ));
1299    }
1300    let final_path = backup_path(&filename);
1301    if final_path.exists() {
1302        return Err(upload_error(StatusCode::CONFLICT, "已存在同名备份文件"));
1303    }
1304
1305    // 5. 磁盘空间预检(Content-Length 已知且探测成功时;否则降级跳过)。
1306    if let (Some(cl), Some(free)) = (content_length, backup_partition_free_space()) {
1307        if cl > free {
1308            return Err(upload_error(
1309                StatusCode::INSUFFICIENT_STORAGE,
1310                "磁盘空间不足",
1311            ));
1312        }
1313    }
1314
1315    // 6. 流式落盘 tmp:按载荷字节精确计数,超限/写失败/读失败都清理半截文件。
1316    let tmp_name = format!(
1317        ".import-{}-{}.tmp",
1318        std::process::id(),
1319        std::time::SystemTime::now()
1320            .duration_since(std::time::UNIX_EPOCH)
1321            .map(|d| d.as_nanos())
1322            .unwrap_or(0)
1323    );
1324    let tmp_path = backup_path(&tmp_name);
1325    let mut out = match tokio::fs::File::create(&tmp_path).await {
1326        Ok(f) => f,
1327        Err(e) => {
1328            tracing::error!("backup import: create tmp failed: {e}");
1329            return Err(upload_error(
1330                StatusCode::INTERNAL_SERVER_ERROR,
1331                "无法写入备份目录",
1332            ));
1333        }
1334    };
1335    let mut written: u64 = 0;
1336    let stream_result: Result<(), (StatusCode, &'static str)> = loop {
1337        match field.chunk().await {
1338            Ok(Some(chunk)) => {
1339                written += chunk.len() as u64;
1340                if written > max_bytes {
1341                    break Err((StatusCode::PAYLOAD_TOO_LARGE, "文件超过导入上限"));
1342                }
1343                if let Err(e) = out.write_all(&chunk).await {
1344                    tracing::error!("backup import: write failed: {e}");
1345                    break Err((
1346                        StatusCode::INTERNAL_SERVER_ERROR,
1347                        "写入失败(磁盘可能已满)",
1348                    ));
1349                }
1350            }
1351            Ok(None) => break Ok(()),
1352            Err(e) => {
1353                tracing::error!("backup import: chunk error: {e:?}");
1354                break Err((StatusCode::BAD_REQUEST, "文件读取失败"));
1355            }
1356        }
1357    };
1358    drop(out);
1359    if let Err((status, msg)) = stream_result {
1360        let _ = std::fs::remove_file(&tmp_path);
1361        return Err(upload_error(status, msg));
1362    }
1363
1364    // 7. 导入时即签名校验:非本系统文件拒收且不留盘。
1365    let first_line = read_first_line(&tmp_path).unwrap_or_default();
1366    if !has_valid_signature(&first_line) {
1367        let _ = std::fs::remove_file(&tmp_path);
1368        return Err(upload_error(
1369            StatusCode::BAD_REQUEST,
1370            "非本系统生成的备份文件,拒绝导入",
1371        ));
1372    }
1373
1374    // 8. 原子 rename 入库;rename 前二次冲突检查,收窄两个并发同名导入的竞态窗。
1375    if final_path.exists() {
1376        let _ = std::fs::remove_file(&tmp_path);
1377        return Err(upload_error(StatusCode::CONFLICT, "已存在同名备份文件"));
1378    }
1379    if let Err(e) = std::fs::rename(&tmp_path, &final_path) {
1380        let _ = std::fs::remove_file(&tmp_path);
1381        tracing::error!("backup import: rename failed: {e}");
1382        return Err(upload_error(StatusCode::INTERNAL_SERVER_ERROR, "入库失败"));
1383    }
1384
1385    tracing::info!(
1386        operator = %user.username,
1387        filename = %filename,
1388        size_bytes = written,
1389        "备份导入成功"
1390    );
1391    Ok(axum::Json(
1392        serde_json::json!({ "success": true, "filename": filename }),
1393    ))
1394}
1395
1396/// Axum 处理器:下载备份文件(admin 鉴权 + 路径白名单)。
1397/// 仅 server 构建:纯 Axum 路由(在 main.rs 注册),无 WASM 消费者。
1398#[cfg(feature = "server")]
1399pub async fn download_backup(
1400    axum::extract::Path(filename): axum::extract::Path<String>,
1401    headers: axum::http::HeaderMap,
1402) -> Result<impl axum::response::IntoResponse, (axum::http::StatusCode, String)> {
1403    use axum::http::{header, StatusCode};
1404
1405    // 鉴权
1406    let cookie_header = headers
1407        .get("cookie")
1408        .and_then(|h| h.to_str().ok())
1409        .unwrap_or("");
1410    let token = crate::auth::session::parse_session_token(cookie_header).map(str::to_string);
1411    let token = match token {
1412        Some(t) => t,
1413        None => return Err((StatusCode::UNAUTHORIZED, "未登录".to_string())),
1414    };
1415    let user = match crate::api::auth::get_user_by_token(&token).await {
1416        Ok(Some(u)) => u,
1417        _ => return Err((StatusCode::UNAUTHORIZED, "会话已过期".to_string())),
1418    };
1419    if user.role != crate::models::user::UserRole::Admin {
1420        return Err((StatusCode::FORBIDDEN, "权限不足".to_string()));
1421    }
1422
1423    // 路径白名单
1424    if !is_valid_backup_filename(&filename) {
1425        return Err((StatusCode::BAD_REQUEST, "无效的文件名".to_string()));
1426    }
1427    let path = backup_path(&filename);
1428    let bytes = tokio::fs::read(&path)
1429        .await
1430        .map_err(|_| (StatusCode::NOT_FOUND, "文件不存在".to_string()))?;
1431    let disposition = format!("attachment; filename=\"{}\"", filename);
1432    let content_type = if filename.ends_with(".tar.gz") {
1433        "application/gzip"
1434    } else {
1435        "application/sql; charset=utf-8"
1436    };
1437    Ok((
1438        StatusCode::OK,
1439        [
1440            (
1441                header::CONTENT_TYPE,
1442                axum::http::HeaderValue::from_static(content_type),
1443            ),
1444            (
1445                header::CONTENT_DISPOSITION,
1446                axum::http::HeaderValue::from_str(&disposition)
1447                    .unwrap_or_else(|_| axum::http::HeaderValue::from_static("attachment")),
1448            ),
1449        ],
1450        axum::body::Body::from(bytes),
1451    ))
1452}
1453
1454#[cfg(all(test, feature = "server"))]
1455mod tests {
1456    use super::*;
1457
1458    // ── is_valid_backup_filename:文件名白名单(路径穿越第一道防线) ──
1459
1460    #[test]
1461    fn filename_accepts_normal_names() {
1462        for name in [
1463            "backup_20260702_120000.sql",
1464            "backup_20260702_120000_sqlfallback.sql",
1465            "a.sql",
1466            "A-B_C.123",
1467        ] {
1468            assert!(is_valid_backup_filename(name), "正常文件名应通过: {name}");
1469        }
1470    }
1471
1472    #[test]
1473    fn filename_rejects_path_traversal() {
1474        // 路径穿越:白名单只允许字母数字下划线点连字符,/ 和 .. 都应被拒。
1475        for evil in [
1476            "../etc/passwd",
1477            "..\\windows\\win.ini",
1478            "/etc/passwd",
1479            "a/../../b",
1480            "backup.sql/../../etc",
1481        ] {
1482            assert!(!is_valid_backup_filename(evil), "路径穿越应被拒: {evil}");
1483        }
1484    }
1485
1486    #[test]
1487    fn filename_rejects_spaces_and_special_chars() {
1488        // 空格、中文、shell 元字符等都不在白名单。
1489        for evil in [
1490            "backup with space.sql",
1491            "备份.sql",
1492            "a;rm -rf.sql",
1493            r"a\$b.sql",
1494            "a`b`.sql",
1495            "",
1496        ] {
1497            assert!(!is_valid_backup_filename(evil), "特殊字符应被拒: {evil:?}");
1498        }
1499    }
1500
1501    // ── backup_path:路径穿越纵深防御(白名单之外的二次防御) ────────
1502
1503    #[test]
1504    fn backup_path_stays_in_backup_dir_for_normal_name() {
1505        let p = backup_path("backup_20260702.sql");
1506        assert!(p.starts_with(BACKUP_DIR), "应在 {BACKUP_DIR}/ 下");
1507        assert_eq!(
1508            p.file_name().and_then(|n| n.to_str()),
1509            Some("backup_20260702.sql")
1510        );
1511    }
1512
1513    #[test]
1514    fn backup_path_collapses_traversal_to_backup_dir() {
1515        // 即便绕过白名单调用 backup_path(纵深防御),../ 也应被规约回 BACKUP_DIR,
1516        // 而非指向 backups/ 之外。Component::ParentDir / RootDir 命中即降级。
1517        for evil in ["../etc/passwd", "../../etc/shadow"] {
1518            let p = backup_path(evil);
1519            // 不应逃出 BACKUP_DIR(应为 BACKUP_DIR 本身,不含文件名)
1520            assert_eq!(
1521                p,
1522                PathBuf::from(BACKUP_DIR),
1523                "穿越应被规约回 {BACKUP_DIR}: {evil}"
1524            );
1525        }
1526    }
1527
1528    #[test]
1529    fn backup_path_rejects_absolute_path() {
1530        // Component::RootDir 命中也应降级。
1531        let p = backup_path("/etc/passwd");
1532        assert_eq!(p, PathBuf::from(BACKUP_DIR));
1533    }
1534
1535    // ── has_valid_signature:备份签名校验(拒绝非本系统文件) ───────
1536
1537    #[test]
1538    fn signature_matches_exact_header() {
1539        let content = "-- YGGDRASIL BACKUP v1\n-- mode: pg_dump\nSELECT 1;\n";
1540        assert!(has_valid_signature(content));
1541    }
1542
1543    #[test]
1544    fn signature_matches_with_leading_whitespace() {
1545        // 首行允许前导空白(trim 后匹配),容忍编辑器缩进。
1546        let content = "  -- YGGDRASIL BACKUP v1\nrest\n";
1547        assert!(has_valid_signature(content));
1548    }
1549
1550    #[test]
1551    fn signature_rejects_non_system_file() {
1552        // 普通 SQL 文件首行不含签名 → 拒绝恢复(防任意文件读取/执行)。
1553        let content = "SELECT * FROM users;\n-- YGGDRASIL BACKUP v1\n";
1554        // 注意:签名必须在首行。第二行有签名不算。
1555        assert!(!has_valid_signature(content));
1556    }
1557
1558    #[test]
1559    fn signature_rejects_empty_and_garbage() {
1560        assert!(!has_valid_signature(""));
1561        assert!(!has_valid_signature("garbage\n"));
1562        assert!(!has_valid_signature("\n\n-- YGGDRASIL BACKUP v1"));
1563    }
1564
1565    // ── parse_backup_mode:模式解析(列表展示用) ───────────────────
1566
1567    #[test]
1568    fn parse_mode_pg_dump() {
1569        let content = "-- YGGDRASIL BACKUP v1\n-- mode: pg_dump\n...\n";
1570        assert_eq!(parse_backup_mode(content), "pg_dump");
1571    }
1572
1573    #[test]
1574    fn parse_mode_sql_fallback() {
1575        let content = "-- YGGDRASIL BACKUP v1\n-- mode: sql-fallback\n\n-- table: posts\n";
1576        assert_eq!(parse_backup_mode(content), "sql-fallback");
1577    }
1578
1579    #[test]
1580    fn parse_mode_unknown_when_absent() {
1581        let content = "-- YGGDRASIL BACKUP v1\nSELECT 1;\n";
1582        assert_eq!(parse_backup_mode(content), "unknown");
1583    }
1584
1585    #[test]
1586    fn parse_mode_unknown_when_empty_value() {
1587        // "-- mode:" 后无值 → unknown(防空字符串显示)
1588        let content = "-- mode:\nrest\n";
1589        assert_eq!(parse_backup_mode(content), "unknown");
1590    }
1591
1592    #[test]
1593    fn parse_mode_only_matches_first_occurrence() {
1594        // 多个 -- mode: 行取第一个。
1595        let content = "-- mode: pg_dump\n-- mode: sql-fallback\n";
1596        assert_eq!(parse_backup_mode(content), "pg_dump");
1597    }
1598
1599    // ── uploads_archive_name:SQL ↔ uploads 包配对规则 ─────────────
1600
1601    #[test]
1602    fn uploads_archive_name_pairs_with_sql() {
1603        assert_eq!(
1604            uploads_archive_name("backup_20260810_040000.sql"),
1605            "backup_20260810_040000_uploads.tar.gz"
1606        );
1607        assert_eq!(
1608            uploads_archive_name("auto_20260810_040000_sqlfallback.sql"),
1609            "auto_20260810_040000_sqlfallback_uploads.tar.gz"
1610        );
1611    }
1612
1613    // ── select_expired_auto_backups:轮转选择(只动 auto,保最新 N 份) ──
1614
1615    #[test]
1616    fn rotation_keeps_newest_and_deletes_oldest() {
1617        let names: Vec<String> = [
1618            "auto_20260808_040000.sql",
1619            "auto_20260810_040000.sql",
1620            "auto_20260809_040000.sql",
1621        ]
1622        .iter()
1623        .map(|s| s.to_string())
1624        .collect();
1625        // 文件名时间戳定宽,乱序输入也应按时间排序。
1626        assert_eq!(
1627            select_expired_auto_backups(&names, 2),
1628            vec!["auto_20260808_040000.sql".to_string()]
1629        );
1630    }
1631
1632    #[test]
1633    fn rotation_ignores_manual_backups_and_tarballs() {
1634        let names: Vec<String> = [
1635            "backup_20260101_000000.sql",          // 手动:永不轮转
1636            "auto_20260808_040000_uploads.tar.gz", // tar.gz:随配对 sql 删除,不单独计数
1637            "auto_20260808_040000.sql",
1638            "auto_20260809_040000.sql",
1639        ]
1640        .iter()
1641        .map(|s| s.to_string())
1642        .collect();
1643        assert!(select_expired_auto_backups(&names, 2).is_empty());
1644        assert_eq!(
1645            select_expired_auto_backups(&names, 1),
1646            vec!["auto_20260808_040000.sql".to_string()]
1647        );
1648    }
1649
1650    #[test]
1651    fn rotation_empty_and_exact_keep() {
1652        assert!(select_expired_auto_backups(&[], 5).is_empty());
1653        let names: Vec<String> = ["auto_20260808_040000.sql"]
1654            .iter()
1655            .map(|s| s.to_string())
1656            .collect();
1657        assert!(select_expired_auto_backups(&names, 1).is_empty());
1658        // keep=0 全删
1659        assert_eq!(select_expired_auto_backups(&names, 0).len(), 1);
1660    }
1661
1662    // ── create_uploads_tarball:打包排除 .cache/.gitkeep ───────────
1663
1664    #[test]
1665    fn tarball_excludes_cache_and_gitkeep() {
1666        let nanos = std::time::SystemTime::now()
1667            .duration_since(std::time::UNIX_EPOCH)
1668            .expect("系统时间必然晚于 UNIX_EPOCH")
1669            .as_nanos();
1670        let dir = std::env::temp_dir().join(format!(
1671            "yggdrasil_backup_tar_test_{}_{}",
1672            nanos,
1673            std::process::id()
1674        ));
1675        let uploads = dir.join("uploads");
1676        std::fs::create_dir_all(uploads.join("2026")).expect("创建测试目录");
1677        std::fs::create_dir_all(uploads.join(".cache")).expect("创建测试目录");
1678        std::fs::write(uploads.join("2026/pic.webp"), b"img").expect("写测试文件");
1679        std::fs::write(uploads.join(".cache/x.webp"), b"cache").expect("写测试文件");
1680        std::fs::write(uploads.join(".gitkeep"), b"").expect("写测试文件");
1681
1682        let out = dir.join("out.tar.gz");
1683        create_uploads_tarball(&uploads, &out).expect("打包应成功");
1684
1685        let file = std::fs::File::open(&out).expect("打开打包产物");
1686        let gz = flate2::read::GzDecoder::new(file);
1687        let mut archive = tar::Archive::new(gz);
1688        let entries: Vec<String> = archive
1689            .entries()
1690            .expect("读取 tar 条目")
1691            .map(|e| {
1692                e.expect("tar 条目有效")
1693                    .path()
1694                    .expect("路径有效")
1695                    .to_string_lossy()
1696                    .to_string()
1697            })
1698            .collect();
1699        assert!(
1700            entries.iter().any(|p| p.contains("2026/pic.webp")),
1701            "应包含素材文件: {entries:?}"
1702        );
1703        assert!(
1704            !entries
1705                .iter()
1706                .any(|p| p.contains(".cache") || p.contains(".gitkeep")),
1707            "应排除 .cache 与 .gitkeep: {entries:?}"
1708        );
1709        std::fs::remove_dir_all(&dir).expect("清理测试目录");
1710    }
1711
1712    // ── sanitize_import_filename:导入文件名校验 ─────────────────
1713
1714    #[test]
1715    fn import_filename_accepts_plain_backup_name() {
1716        assert_eq!(
1717            sanitize_import_filename("backup_20260816_200000.sql"),
1718            Some("backup_20260816_200000.sql".to_string())
1719        );
1720        // 255 字节边界接受
1721        let max_ok = format!("{}.sql", "a".repeat(251));
1722        assert!(sanitize_import_filename(&max_ok).is_some());
1723    }
1724
1725    #[test]
1726    fn import_filename_strips_path_components() {
1727        // 旧浏览器可能送 `C:\fakepath\` 前缀;POSIX 路径同样剥掉,只留 basename。
1728        assert_eq!(
1729            sanitize_import_filename("C:\\fakepath\\auto_20260816_200000.sql"),
1730            Some("auto_20260816_200000.sql".to_string())
1731        );
1732        assert_eq!(
1733            sanitize_import_filename("/tmp/x/backup_1.sql"),
1734            Some("backup_1.sql".to_string())
1735        );
1736    }
1737
1738    #[test]
1739    fn import_filename_rejects_non_sql_and_hidden() {
1740        // 配对 tar.gz 不可经导入入口塞进 backups/
1741        assert_eq!(sanitize_import_filename("x_uploads.tar.gz"), None);
1742        assert_eq!(sanitize_import_filename("noext"), None);
1743        // 隐藏文件与裸 ".sql" 拒绝(避开 tmp 文件命名空间)
1744        assert_eq!(sanitize_import_filename(".hidden.sql"), None);
1745        assert_eq!(sanitize_import_filename(".sql"), None);
1746        assert_eq!(sanitize_import_filename(""), None);
1747    }
1748
1749    #[test]
1750    fn import_filename_rejects_illegal_chars_and_overlong() {
1751        assert_eq!(sanitize_import_filename("带中文.sql"), None);
1752        assert_eq!(sanitize_import_filename("has space.sql"), None);
1753        let long = format!("{}.sql", "a".repeat(252));
1754        assert_eq!(sanitize_import_filename(&long), None);
1755    }
1756    #[test]
1757    fn pg_dump_is_owner_acl_neutral_without_db_url_in_args() {
1758        let db_url = "postgres://user:secret@example.com:5432/dbname";
1759        let command = pg_dump_command(db_url).expect("valid database URL");
1760        let args: Vec<_> = command
1761            .get_args()
1762            .map(|arg| arg.to_string_lossy().into_owned())
1763            .collect();
1764        assert!(args.iter().any(|arg| arg == "--no-owner"), "{args:?}");
1765        assert!(args.iter().any(|arg| arg == "--no-privileges"), "{args:?}");
1766        assert!(!args.iter().any(|arg| arg.contains("secret")), "{args:?}");
1767        assert!(command
1768            .get_envs()
1769            .any(|(key, value)| key == "PGPASSWORD" && value.is_some()));
1770    }
1771
1772    #[test]
1773    fn psql_restore_is_atomic_without_db_url_in_args() {
1774        let db_url = "postgres://user:secret@example.com:5432/dbname";
1775        let command =
1776            psql_restore_command(db_url, Path::new("backup.sql")).expect("valid database URL");
1777        let args: Vec<_> = command
1778            .get_args()
1779            .map(|arg| arg.to_string_lossy().into_owned())
1780            .collect();
1781        assert!(
1782            args.iter().any(|arg| arg == "--single-transaction"),
1783            "{args:?}"
1784        );
1785        assert!(!args.iter().any(|arg| arg.contains("secret")), "{args:?}");
1786    }
1787
1788    #[test]
1789    fn legacy_owner_statements_are_removed_without_touching_copy_data() {
1790        let sql = concat!(
1791            "-- YGGDRASIL BACKUP v1\n",
1792            "ALTER FUNCTION public.f() OWNER TO yggdrasil;\n",
1793            "COPY public.lines (value) FROM stdin;\n",
1794            "ALTER TABLE public.literal OWNER TO text;\n",
1795            "\\.\n",
1796            "ALTER TABLE public.users OWNER TO yggdrasil;\n",
1797            "ALTER SEQUENCE public.users_id_seq OWNER TO yggdrasil;\n",
1798            "ALTER TABLE public.users ENABLE ROW LEVEL SECURITY;\n",
1799        );
1800        let mut output = Vec::new();
1801        let removed = write_owner_neutral_restore_sql(std::io::Cursor::new(sql), &mut output)
1802            .expect("过滤测试 SQL 应成功");
1803        let output = String::from_utf8(output).expect("测试 SQL 是 UTF-8");
1804
1805        assert_eq!(removed, 3);
1806        assert!(
1807            output.contains("ALTER TABLE public.literal OWNER TO text;"),
1808            "COPY 数据不得被当成 SQL 删除: {output}"
1809        );
1810        assert!(!output.contains("OWNER TO yggdrasil"), "{output}");
1811        assert!(
1812            output.contains("ALTER TABLE public.users ENABLE ROW LEVEL SECURITY;"),
1813            "{output}"
1814        );
1815    }
1816    #[test]
1817    fn restore_rejects_shell_meta_commands() {
1818        let sql = concat!(
1819            "-- YGGDRASIL BACKUP v1\n",
1820            "\\! touch /tmp/yggdrasil-pwned\n",
1821        );
1822        let mut output = Vec::new();
1823        let error =
1824            write_owner_neutral_restore_sql(std::io::Cursor::new(sql), &mut output).unwrap_err();
1825        assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
1826    }
1827
1828    #[test]
1829    fn restore_keeps_generated_restrict_meta_commands() {
1830        let sql = concat!(
1831            "\\restrict abc123\n",
1832            "SELECT 1;\n",
1833            "\\unrestrict abc123\n",
1834        );
1835        let mut output = Vec::new();
1836        write_owner_neutral_restore_sql(std::io::Cursor::new(sql), &mut output)
1837            .expect("generated restrict commands are allowed");
1838        let output = String::from_utf8(output).expect("test SQL is UTF-8");
1839        assert!(output.contains("\\restrict abc123"));
1840        assert!(output.contains("\\unrestrict abc123"));
1841    }
1842}