yggdrasil/api/database/backup/
restore.rs

1//! COPY 感知的 SQL 预处理、原子恢复、临时文件清理与缓存失效。
2
3use super::commands::psql_restore_command;
4use super::files::backup_path;
5use crate::api::database::tasks::{self, TaskStatus};
6use std::path::{Path, PathBuf};
7
8fn sql_line_without_eol(mut line: &[u8]) -> &[u8] {
9    if let Some(stripped) = line.strip_suffix(b"\n") {
10        line = stripped;
11    }
12    if let Some(stripped) = line.strip_suffix(b"\r") {
13        line = stripped;
14    }
15    line
16}
17
18fn is_legacy_pg_dump_owner_statement(line: &[u8]) -> bool {
19    const PREFIXES: [&[u8]; 3] = [b"ALTER FUNCTION ", b"ALTER SEQUENCE ", b"ALTER TABLE "];
20    let line = sql_line_without_eol(line);
21    PREFIXES.iter().any(|prefix| line.starts_with(prefix))
22        && line
23            .windows(b" OWNER TO ".len())
24            .any(|window| window == b" OWNER TO ")
25        && line.ends_with(b";")
26}
27
28/// 去除恢复脚本行首空白，确保带缩进的 psql 元命令也进入安全检查。
29fn trim_restore_line_start(line: &[u8]) -> &[u8] {
30    let line = sql_line_without_eol(line);
31    let start = line
32        .iter()
33        .position(|byte| !byte.is_ascii_whitespace())
34        .unwrap_or(line.len());
35    &line[start..]
36}
37
38fn is_allowed_restore_meta_command(line: &[u8]) -> bool {
39    let line = trim_restore_line_start(line);
40    let Some(rest) = line.strip_prefix(b"\\") else {
41        return false;
42    };
43    let mut fields = rest
44        .split(|byte| byte.is_ascii_whitespace())
45        .filter(|field| !field.is_empty());
46    let Some(command) = fields.next() else {
47        return false;
48    };
49    if command != b"restrict" && command != b"unrestrict" {
50        return false;
51    }
52    let Some(token) = fields.next() else {
53        return false;
54    };
55    !token.is_empty()
56        && token
57            .iter()
58            .all(|byte| byte.is_ascii_alphanumeric() || *byte == b'_')
59        && fields.next().is_none()
60}
61
62fn write_owner_neutral_restore_sql(
63    mut input: impl std::io::BufRead,
64    mut output: impl std::io::Write,
65) -> std::io::Result<usize> {
66    let mut line = Vec::new();
67    let mut in_copy_data = false;
68    let mut removed = 0;
69    loop {
70        line.clear();
71        if input.read_until(b'\n', &mut line)? == 0 {
72            break;
73        }
74        let sql_line = sql_line_without_eol(&line);
75        if in_copy_data {
76            output.write_all(&line)?;
77            if sql_line == b"\\." {
78                in_copy_data = false;
79            }
80            continue;
81        }
82        if sql_line.starts_with(b"COPY ") && sql_line.ends_with(b" FROM stdin;") {
83            in_copy_data = true;
84            output.write_all(&line)?;
85        } else if is_legacy_pg_dump_owner_statement(&line) {
86            removed += 1;
87        } else if trim_restore_line_start(sql_line).first() == Some(&b'\\')
88            && !is_allowed_restore_meta_command(&line)
89        {
90            return Err(std::io::Error::new(
91                std::io::ErrorKind::InvalidData,
92                "unsupported psql meta-command in restore file",
93            ));
94        } else {
95            output.write_all(&line)?;
96        }
97    }
98
99    output.flush()?;
100    Ok(removed)
101}
102
103struct PreparedRestoreSql {
104    path: PathBuf,
105    removed_owner_statements: usize,
106}
107
108impl PreparedRestoreSql {
109    fn prepare(source: &Path) -> std::io::Result<Self> {
110        let input = std::io::BufReader::new(std::fs::File::open(source)?);
111        let path =
112            std::env::temp_dir().join(format!("yggdrasil-restore-{}.sql", uuid::Uuid::new_v4()));
113        let mut options = std::fs::OpenOptions::new();
114        options.write(true).create_new(true);
115        #[cfg(unix)]
116        {
117            use std::os::unix::fs::OpenOptionsExt;
118            options.mode(0o600);
119        }
120        let output = std::io::BufWriter::new(options.open(&path)?);
121        match write_owner_neutral_restore_sql(input, output) {
122            Ok(removed_owner_statements) => Ok(Self {
123                path,
124                removed_owner_statements,
125            }),
126            Err(error) => {
127                let _ = std::fs::remove_file(&path);
128                Err(error)
129            }
130        }
131    }
132}
133
134impl Drop for PreparedRestoreSql {
135    fn drop(&mut self) {
136        if let Err(error) = std::fs::remove_file(&self.path) {
137            tracing::warn!(
138                path = %self.path.display(),
139                "restore: failed to delete prepared SQL: {error}"
140            );
141        }
142    }
143}
144/// 后台执行恢复：探测 psql，可用则原子恢复，不可用则报告。
145///
146/// **跨环境恢复**：旧备份可能含 `ALTER ... OWNER TO ...`，而目标集群未必存在
147/// 源库角色。恢复前流式生成临时副本，只移除 pg_dump 生成的属主语句；COPY 数据
148/// 原样保留。新备份已由 `--no-owner --no-privileges` 从源头去掉角色元数据。
149///
150/// **原子恢复**：备份由 `pg_dump --clean --if-exists` 生成，脚本自带
151/// `DROP ... IF EXISTS`；psql 使用 `--single-transaction`，任一 SQL 错误都会
152/// 回滚先前的 DROP/CREATE/COPY，不留下半恢复数据库。
153///
154/// **错误中止**：`-v ON_ERROR_STOP=1` 让 psql 在第一条 SQL 错误时立即退出
155/// （退出码 3）。否则 psql 即使满屏 ERROR 也返回 0，导致 `status.success()`
156/// 误判为成功——这是「恢复完成却无任何数据变更」假成功的根因。
157///
158/// **恢复成功后失效全量缓存**：恢复会用备份时刻的数据重建 posts 等表，
159/// 现有 moka 缓存（列表/标签/单篇/统计/搜索）与 SSR 世代号必须一并冲刷，
160/// 否则前端仍读旧数据。
161pub(super) async fn run_restore(task_id: &str, filename: &str) {
162    let path = backup_path(filename);
163    let db_url = match std::env::var("DATABASE_URL") {
164        Ok(u) if !u.is_empty() => u,
165        _ => {
166            tasks::update(
167                task_id,
168                "DATABASE_URL 未配置",
169                100,
170                TaskStatus::Failed,
171                None,
172                Some("恢复需要 DATABASE_URL".to_string()),
173                None,
174            );
175            return;
176        }
177    };
178    let psql_ok = tokio::task::spawn_blocking(|| {
179        std::process::Command::new("psql")
180            .arg("--version")
181            .output()
182            .is_ok()
183    })
184    .await
185    .unwrap_or(false);
186    if !psql_ok {
187        tasks::update(
188            task_id,
189            "psql 不可用",
190            100,
191            TaskStatus::Failed,
192            None,
193            Some("恢复需要 psql，但当前环境未安装 psql".to_string()),
194            None,
195        );
196        return;
197    }
198    tasks::update(
199        task_id,
200        "正在用 psql 恢复",
201        50,
202        TaskStatus::Running,
203        None,
204        None,
205        None,
206    );
207    // 预处理备份与 psql 恢复均为阻塞文件/子进程操作，整体移出 tokio worker。
208    // 临时 SQL 由 RAII 守卫在成功、失败和启动错误路径统一删除。
209    let restore_result = tokio::task::spawn_blocking(move || {
210        let prepared = PreparedRestoreSql::prepare(&path)?;
211        if prepared.removed_owner_statements > 0 {
212            tracing::info!(
213                removed = prepared.removed_owner_statements,
214                "restore: removed legacy pg_dump owner statements"
215            );
216        }
217
218        let mut command = psql_restore_command(&db_url, &prepared.path)?;
219        command.output()
220    })
221    .await
222    .unwrap_or_else(|join_e| Err(std::io::Error::other(join_e.to_string())));
223    match restore_result {
224        Ok(o) if o.status.success() => {
225            // 恢复会覆盖文章、评论、用户、设置与素材元数据；所有内存/磁盘缓存
226            // 必须一起冲刷。generation 仅用于观测，不能替代物理 SSR 删除。
227            crate::cache::invalidate_all_post_caches();
228            crate::cache::invalidate_search_results();
229            crate::cache::invalidate_all_comments();
230            crate::cache::invalidate_friend_links();
231            crate::cache::invalidate_site_settings();
232            crate::cache::invalidate_security_settings();
233            crate::cache::invalidate_image_cache_settings();
234            crate::cache::invalidate_log_targets();
235            crate::cache::SESSION_CACHE.invalidate_all();
236            crate::api::image::invalidate_all_caches().await;
237            crate::ssr_cache::invalidate_ssr_all_public();
238            crate::ssr_cache::bump_global_generation();
239            tasks::update(task_id, "恢复完成", 100, TaskStatus::Done, None, None, None);
240        }
241        Ok(o) => {
242            let stderr = String::from_utf8_lossy(&o.stderr).to_string();
243            tasks::update(
244                task_id,
245                "恢复失败",
246                100,
247                TaskStatus::Failed,
248                None,
249                Some(stderr),
250                None,
251            );
252        }
253        Err(e) => {
254            tasks::update(
255                task_id,
256                "psql 启动失败",
257                100,
258                TaskStatus::Failed,
259                None,
260                Some(e.to_string()),
261                None,
262            );
263        }
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270
271    #[test]
272    fn legacy_owner_statements_are_removed_without_touching_copy_data() {
273        let sql = concat!(
274            "-- YGGDRASIL BACKUP v1\n",
275            "ALTER FUNCTION public.f() OWNER TO yggdrasil;\n",
276            "COPY public.lines (value) FROM stdin;\n",
277            "ALTER TABLE public.literal OWNER TO text;\n",
278            "\\.\n",
279            "ALTER TABLE public.users OWNER TO yggdrasil;\n",
280            "ALTER SEQUENCE public.users_id_seq OWNER TO yggdrasil;\n",
281            "ALTER TABLE public.users ENABLE ROW LEVEL SECURITY;\n",
282        );
283        let mut output = Vec::new();
284        let removed = write_owner_neutral_restore_sql(std::io::Cursor::new(sql), &mut output)
285            .expect("过滤测试 SQL 应成功");
286        let output = String::from_utf8(output).expect("测试 SQL 是 UTF-8");
287
288        assert_eq!(removed, 3);
289        assert!(
290            output.contains("ALTER TABLE public.literal OWNER TO text;"),
291            "COPY 数据不得被当成 SQL 删除: {output}"
292        );
293        assert!(!output.contains("OWNER TO yggdrasil"), "{output}");
294        assert!(
295            output.contains("ALTER TABLE public.users ENABLE ROW LEVEL SECURITY;"),
296            "{output}"
297        );
298    }
299
300    #[test]
301    fn restore_rejects_shell_meta_commands() {
302        let sql = concat!(
303            "-- YGGDRASIL BACKUP v1\n",
304            "\\! touch /tmp/yggdrasil-pwned\n",
305        );
306        let mut output = Vec::new();
307        let error =
308            write_owner_neutral_restore_sql(std::io::Cursor::new(sql), &mut output).unwrap_err();
309        assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
310    }
311
312    #[test]
313    fn restore_keeps_generated_restrict_meta_commands() {
314        let sql = concat!(
315            "\\restrict abc123\n",
316            "SELECT 1;\n",
317            "\\unrestrict abc123\n",
318        );
319        let mut output = Vec::new();
320        write_owner_neutral_restore_sql(std::io::Cursor::new(sql), &mut output)
321            .expect("generated restrict commands are allowed");
322        let output = String::from_utf8(output).expect("test SQL is UTF-8");
323        assert!(output.contains("\\restrict abc123"));
324        assert!(output.contains("\\unrestrict abc123"));
325    }
326    /// Real pg_dump/psql regression. Requires a disposable ygg_backup_test database;
327    /// uses its own URL so it can run independently of the application's global pool.
328    #[tokio::test]
329    #[ignore = "requires pg_dump, psql, and disposable YGG_BACKUP_TEST_DATABASE_URL"]
330    async fn backup_round_trip_and_failed_restore_rolls_back_database() {
331        use super::super::commands::pg_dump_command;
332        use super::super::files::{has_valid_signature, read_first_line, BACKUP_SIGNATURE};
333
334        let db_url = std::env::var("YGG_BACKUP_TEST_DATABASE_URL")
335            .expect("set YGG_BACKUP_TEST_DATABASE_URL to a disposable ygg_backup_test database");
336        let config: tokio_postgres::Config = db_url.parse().expect("valid test database URL");
337        assert_eq!(
338            config.get_dbname(),
339            Some("ygg_backup_test"),
340            "backup regression only runs on the disposable ygg_backup_test database"
341        );
342        let (client, connection) = config
343            .connect(tokio_postgres::NoTls)
344            .await
345            .expect("connect to disposable backup database");
346        tokio::spawn(async move {
347            connection.await.expect("backup regression connection");
348        });
349        client
350            .batch_execute(
351                "DROP TABLE IF EXISTS public.ygg_backup_regression; \
352                 CREATE TABLE public.ygg_backup_regression (id integer PRIMARY KEY, value text NOT NULL)",
353            )
354            .await
355            .expect("create backup fixture");
356        let values = [
357            "ALTER TABLE public.literal OWNER TO absent_role;",
358            "\\! COPY data stays literal\n正文保持换行",
359        ];
360        for (index, value) in values.iter().enumerate() {
361            client
362                .execute(
363                    "INSERT INTO public.ygg_backup_regression VALUES ($1, $2)",
364                    &[&(index as i32), value],
365                )
366                .await
367                .expect("write backup fixture");
368        }
369
370        struct TestDirectory(PathBuf);
371        impl Drop for TestDirectory {
372            fn drop(&mut self) {
373                let _ = std::fs::remove_dir_all(&self.0);
374            }
375        }
376        let dir = TestDirectory(std::env::temp_dir().join(format!(
377            "yggdrasil-backup-regression-{}",
378            uuid::Uuid::new_v4()
379        )));
380        std::fs::create_dir(&dir.0).expect("create regression directory");
381        let source = dir.0.join("backup.sql");
382        let dump = pg_dump_command(&db_url)
383            .expect("build pg_dump command")
384            .output()
385            .expect("run pg_dump");
386        assert!(dump.status.success(), "pg_dump must succeed");
387        let mut sql = format!("{BACKUP_SIGNATURE}\n-- mode: pg_dump\n").into_bytes();
388        sql.extend_from_slice(&dump.stdout);
389        // Exercise legacy cross-role compatibility even on a new owner-neutral dump.
390        sql.extend_from_slice(
391            b"\nALTER TABLE public.ygg_backup_regression OWNER TO ygg_nonexistent_restore_role;\n",
392        );
393        std::fs::write(&source, &sql).expect("write signed backup");
394        assert!(has_valid_signature(
395            &read_first_line(&source).expect("read backup signature")
396        ));
397        client
398            .execute(
399                "UPDATE public.ygg_backup_regression SET value = 'changed'",
400                &[],
401            )
402            .await
403            .expect("change fixture before restoring");
404        let prepared = PreparedRestoreSql::prepare(&source).expect("prepare real dump");
405        assert_eq!(prepared.removed_owner_statements, 1);
406        let prepared_path = prepared.path.clone();
407        #[cfg(unix)]
408        {
409            use std::os::unix::fs::PermissionsExt;
410            assert_eq!(
411                std::fs::metadata(&prepared_path)
412                    .expect("prepared file metadata")
413                    .permissions()
414                    .mode()
415                    & 0o777,
416                0o600
417            );
418        }
419        let restored = psql_restore_command(&db_url, &prepared_path)
420            .expect("build restore command")
421            .output()
422            .expect("run psql restore");
423        assert!(restored.status.success(), "valid backup must restore");
424        drop(prepared);
425        assert!(
426            !prepared_path.exists(),
427            "successful restore cleans temporary SQL"
428        );
429        let actual: Vec<String> = client
430            .query(
431                "SELECT value FROM public.ygg_backup_regression ORDER BY id",
432                &[],
433            )
434            .await
435            .expect("read restored fixture")
436            .iter()
437            .map(|row| row.get(0))
438            .collect();
439        assert_eq!(actual, values);
440
441        // A SQL failure after DROP/CREATE/COPY must preserve the current live state.
442        client
443            .execute(
444                "UPDATE public.ygg_backup_regression SET value = 'live after backup'",
445                &[],
446            )
447            .await
448            .expect("change live fixture");
449        sql.extend_from_slice(b"SELECT * FROM public.ygg_backup_missing_table;\n");
450        std::fs::write(&source, &sql).expect("write failing restore fixture");
451        let prepared = PreparedRestoreSql::prepare(&source).expect("prepare failing SQL");
452        let prepared_path = prepared.path.clone();
453        let failed = psql_restore_command(&db_url, &prepared_path)
454            .expect("build failing restore command")
455            .output()
456            .expect("run failing restore");
457        assert!(!failed.status.success(), "SQL error must stop restore");
458        drop(prepared);
459        assert!(
460            !prepared_path.exists(),
461            "failed restore cleans temporary SQL"
462        );
463        let actual: Vec<String> = client
464            .query(
465                "SELECT value FROM public.ygg_backup_regression ORDER BY id",
466                &[],
467            )
468            .await
469            .expect("read fixture after failed restore")
470            .iter()
471            .map(|row| row.get(0))
472            .collect();
473        assert_eq!(actual, vec!["live after backup"; values.len()]);
474        client
475            .batch_execute("DROP TABLE public.ygg_backup_regression")
476            .await
477            .expect("clean backup fixture");
478    }
479}
yggdrasil/api/database/backup/restore.rs

yggdrasil/api/database/backup/
restore.rs
