From 420d844e1f38b64b1b0aeb76373803b1fee5677b Mon Sep 17 00:00:00 2001 From: wcqqq1214 Date: Tue, 27 Oct 2026 01:44:53 +0800 Subject: [PATCH] fix: restore backup attachments after changing data directories --- astrbot/core/backup/importer.py | 39 +++++++-- tests/test_backup.py | 143 ++++++++++++++++++++++++++++++++ 2 files changed, 174 insertions(+), 8 deletions(-) diff --git a/astrbot/core/backup/importer.py b/astrbot/core/backup/importer.py index 70de7b46d2..6451a56e40 100644 --- a/astrbot/core/backup/importer.py +++ b/astrbot/core/backup/importer.py @@ -16,11 +16,12 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -from sqlalchemy import delete +from sqlalchemy import delete, update from astrbot.core import logger from astrbot.core.config.default import VERSION from astrbot.core.db import BaseDatabase +from astrbot.core.db.po import Attachment from astrbot.core.utils.astrbot_path import ( get_astrbot_data_path, get_astrbot_knowledge_base_path, @@ -496,7 +497,7 @@ async def import_all( await progress_callback("attachments", 0, 100, "正在导入附件...") attachment_count = await self._import_attachments( - zf, main_data.get("attachments", []) + zf, main_data.get("attachments", []), result ) result.imported_files["attachments"] = attachment_count @@ -842,18 +843,28 @@ async def _import_attachments( self, zf: zipfile.ZipFile, attachments: list[dict], + result: ImportResult, ) -> int: - """导入附件文件""" + """Restore attachments and rebase paths when the data directory changes. + + Args: + zf: Open backup archive. + attachments: Attachment records containing their original paths. + result: Import result used to report attachment failures. + + Returns: + Number of files restored with their database paths updated. + """ count = 0 - attachments_dir = Path(self.config_path).parent / "attachments" + attachments_dir = (Path(self.config_path).parent / "attachments").resolve() attachments_dir.mkdir(parents=True, exist_ok=True) attachment_prefix = "files/attachments/" for name in zf.namelist(): if name.startswith(attachment_prefix) and name != attachment_prefix: try: - # 从附件记录中找到原始路径 + # Preserve paths within this instance; rebase migrated files. attachment_id = os.path.splitext(os.path.basename(name))[0] original_path = None for att in attachments: @@ -861,22 +872,34 @@ async def _import_attachments( original_path = att.get("path") break - if original_path: + if original_path and _validate_path_within( + Path(original_path), attachments_dir + ): target_path = Path(original_path) else: target_path = attachments_dir / os.path.basename(name) # Validate path is within attachments directory (CWE-22) if not _validate_path_within(target_path, attachments_dir): - logger.warning(f"附件路径越界，已跳过: {target_path}") + result.add_error( + f"Attachment path is outside the attachments directory: {target_path}" + ) continue target_path.parent.mkdir(parents=True, exist_ok=True) with zf.open(name) as src, open(target_path, "wb") as dst: dst.write(src.read()) + if original_path != str(target_path): + async with self.main_db.get_db() as session: + async with session.begin(): + await session.execute( + update(Attachment) + .where(Attachment.attachment_id == attachment_id) + .values(path=str(target_path)) + ) count += 1 except Exception as e: - logger.warning(f"导入附件 {name} 失败: {e}") + result.add_error(f"Failed to restore attachment {name}: {e}") return count diff --git a/tests/test_backup.py b/tests/test_backup.py index 7fbd93a50e..c771db9832 100644 --- a/tests/test_backup.py +++ b/tests/test_backup.py @@ -628,6 +628,91 @@ async def test_import_file_not_exists(self, mock_main_db, tmp_path): assert result.success is False assert any("不存在" in err for err in result.errors) + @pytest.mark.asyncio + async def test_attachment_write_failure_marks_restore_failed( + self, mock_main_db, tmp_path + ): + """Report attachment write failures through the overall import result. + + Args: + mock_main_db: Database double for the unrelated table import. + tmp_path: Temporary directory for the archive and restore target. + """ + attachment_dir = tmp_path / "attachments" + (attachment_dir / "test-id.txt").mkdir(parents=True) + archive = tmp_path / "backup.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("manifest.json", json.dumps({"astrbot_version": VERSION})) + zf.writestr( + "databases/main_db.json", + json.dumps( + { + "attachments": [ + { + "attachment_id": "test-id", + "path": "/old-install/data/attachments/uploaded.txt", + } + ] + } + ), + ) + zf.writestr("files/attachments/test-id.txt", "payload") + importer = AstrBotImporter( + mock_main_db, config_path=str(tmp_path / "cmd_config.json") + ) + importer._clear_main_db = AsyncMock() + importer._import_main_database = AsyncMock(return_value={"attachments": 1}) + + result = await importer.import_all(str(archive)) + + assert not result.success + assert result.imported_files["attachments"] == 0 + assert any("Failed to restore attachment" in error for error in result.errors) + + @pytest.mark.asyncio + async def test_attachment_relocation_rejects_symlink_escape( + self, mock_main_db, tmp_path, require_symlink + ): + """Keep relocated attachment writes inside the target attachment directory. + + Args: + mock_main_db: Database double that must not be used for a rejected file. + tmp_path: Temporary directory for the archive and restore target. + require_symlink: Fixture that skips platforms without symlink support. + """ + attachment_dir = tmp_path / "attachments" + attachment_dir.mkdir() + outside = tmp_path / "outside.txt" + outside.write_bytes(b"keep existing bytes") + (attachment_dir / "test-id.txt").symlink_to(outside) + archive = tmp_path / "backup.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("files/attachments/test-id.txt", "replacement") + importer = AstrBotImporter( + mock_main_db, config_path=str(tmp_path / "cmd_config.json") + ) + result = ImportResult() + + with zipfile.ZipFile(archive) as zf: + count = await importer._import_attachments( + zf, + [ + { + "attachment_id": "test-id", + "path": "/old-install/data/attachments/uploaded.txt", + } + ], + result, + ) + + assert count == 0 + assert not result.success + assert any( + "outside the attachments directory" in error for error in result.errors + ) + assert outside.read_bytes() == b"keep existing bytes" + mock_main_db.get_db.assert_not_called() + @pytest.mark.asyncio async def test_import_invalid_zip(self, mock_main_db, tmp_path): """测试导入无效的 ZIP 文件""" @@ -1073,6 +1158,64 @@ def test_kb_metadata_models_contain_expected_tables(self): class TestBackupIntegration: """备份集成测试""" + @pytest.mark.asyncio + @pytest.mark.parametrize("relocate", [False, True], ids=["same-root", "new-root"]) + async def test_attachment_restore_paths(self, tmp_path, monkeypatch, relocate): + """Restore attachment bytes and usable database paths after migration. + + Args: + tmp_path: Temporary directory for isolated instances and backups. + monkeypatch: Fixture used to select each instance's data root. + relocate: Whether to restore into a different installation directory. + """ + source_root = tmp_path / "source" + target_root = tmp_path / "target" if relocate else source_root + monkeypatch.setenv("ASTRBOT_ROOT", str(source_root)) + source_data = source_root / "data" + target_data = target_root / "data" + original = source_data / "attachments" / "uploaded.txt" + original.parent.mkdir(parents=True) + target_data.mkdir(parents=True, exist_ok=True) + original.write_bytes(b"attachment payload") + source = SQLiteDatabase(str(source_data / "source.db")) + target = SQLiteDatabase(str(target_data / "target.db")) + try: + attachment = await source.insert_attachment( + str(original), "file", "text/plain" + ) + archive = await AstrBotExporter( + source, config_path=str(source_data / "cmd_config.json") + ).export_all(output_dir=str(tmp_path / "backups")) + with zipfile.ZipFile(archive) as zf: + assert ( + zf.read(f"files/attachments/{attachment.attachment_id}.txt") + == b"attachment payload" + ) + original.unlink() + monkeypatch.setenv("ASTRBOT_ROOT", str(target_root)) + + result = await AstrBotImporter( + target, config_path=str(target_data / "cmd_config.json") + ).import_all(archive) + + assert result.success, result.errors + assert result.warnings == [] + assert result.imported_tables["attachments"] == 1 + assert result.imported_files["attachments"] == 1 + restored = await target.get_attachment_by_id(attachment.attachment_id) + assert restored is not None + restored_path = Path(restored.path) + assert restored_path.is_absolute() + assert restored_path.is_relative_to(target_data / "attachments") + assert restored_path.read_bytes() == b"attachment payload" + if relocate: + assert not original.exists() + else: + assert restored_path == original + finally: + await source.engine.dispose() + await target.engine.dispose() + @pytest.mark.asyncio @pytest.mark.parametrize( ("model", "data"),