fix(native): cleanup deleted docs and blobs (#14689)

packages/frontend/apps/electron/src/helper/dialog/dialog.ts

@@ -1,4 +1,4 @@
-import { parse } from 'node:path';
+import { parse, resolve } from 'node:path';
 
 import { DocStorage, ValidationResult } from '@affine/native';
 import { parseUniversalId } from '@affine/nbstore';
@@ -71,10 +71,34 @@ function getDefaultDBFileName(name: string, id: string) {
   return fileName.replace(/[/\\?%*:|"<>]/g, '-');
 }
 
+async function resolveExistingPath(path: string) {
+  if (!(await fs.pathExists(path))) {
+    return null;
+  }
+  try {
+    return await fs.realpath(path);
+  } catch {
+    return resolve(path);
+  }
+}
+
+async function isSameFilePath(sourcePath: string, targetPath: string) {
+  if (resolve(sourcePath) === resolve(targetPath)) {
+    return true;
+  }
+
+  const [sourceRealPath, targetRealPath] = await Promise.all([
+    resolveExistingPath(sourcePath),
+    resolveExistingPath(targetPath),
+  ]);
+
+  return !!sourceRealPath && sourceRealPath === targetRealPath;
+}
+
 /**
  * This function is called when the user clicks the "Save" button in the "Save Workspace" dialog.
  *
- * It will just copy the file to the given path
+ * It will export a compacted database file to the given path
  */
 export async function saveDBFileAs(
   universalId: string,
@@ -115,12 +139,26 @@ export async function saveDBFileAs(
 
     const filePath = ret.filePath;
     if (ret.canceled || !filePath) {
-      return {
-        canceled: true,
-      };
+      return { canceled: true };
     }
 
-    await fs.copyFile(dbPath, filePath);
+    if (await isSameFilePath(dbPath, filePath)) {
+      return { error: 'DB_FILE_PATH_INVALID' };
+    }
+
+    const tempFilePath = `${filePath}.${nanoid(6)}.tmp`;
+    if (await fs.pathExists(tempFilePath)) {
+      await fs.remove(tempFilePath);
+    }
+
+    try {
+      await pool.vacuumInto(universalId, tempFilePath);
+      await fs.move(tempFilePath, filePath, { overwrite: true });
+    } finally {
+      if (await fs.pathExists(tempFilePath)) {
+        await fs.remove(tempFilePath);
+      }
+    }
     logger.log('saved', filePath);
     if (!fakedResult) {
       mainRPC.showItemInFolder(filePath).catch(err => {
@@ -183,11 +221,7 @@ export async function loadDBFile(
     const provided =
       getFakedResult() ??
       (dbFilePath
-        ? {
-            filePath: dbFilePath,
-            filePaths: [dbFilePath],
-            canceled: false,
-          }
+        ? { filePath: dbFilePath, filePaths: [dbFilePath], canceled: false }
         : undefined);
     const ret =
       provided ??
@@ -224,6 +258,10 @@ export async function loadDBFile(
       return await cpV1DBFile(originalPath, workspaceId);
     }
 
+    if (!(await storage.validateImportSchema())) {
+      return { error: 'DB_FILE_INVALID' };
+    }
+
     // v2 import logic
     const internalFilePath = await getSpaceDBPath(
       'local',
@@ -231,8 +269,8 @@ export async function loadDBFile(
       workspaceId
     );
     await fs.ensureDir(parse(internalFilePath).dir);
-    await fs.copy(originalPath, internalFilePath);
-    logger.info(`loadDBFile, copy: ${originalPath} -> ${internalFilePath}`);
+    await storage.vacuumInto(internalFilePath);
+    logger.info(`loadDBFile, vacuum: ${originalPath} -> ${internalFilePath}`);
 
     storage = new DocStorage(internalFilePath);
     await storage.setSpaceId(workspaceId);
@@ -260,17 +298,16 @@ async function cpV1DBFile(
     return { error: 'DB_FILE_INVALID' }; // invalid db file
   }
 
-  // checkout to make sure wal is flushed
   const connection = new SqliteConnection(originalPath);
-  await connection.connect();
-  await connection.checkpoint();
-  await connection.close();
+  if (!(await connection.validateImportSchema())) {
+    return { error: 'DB_FILE_INVALID' };
+  }
 
   const internalFilePath = await getWorkspaceDBPath('workspace', workspaceId);
 
-  await fs.ensureDir(await getWorkspacesBasePath());
-  await fs.copy(originalPath, internalFilePath);
-  logger.info(`loadDBFile, copy: ${originalPath} -> ${internalFilePath}`);
+  await fs.ensureDir(parse(internalFilePath).dir);
+  await connection.vacuumInto(internalFilePath);
+  logger.info(`loadDBFile, vacuum: ${originalPath} -> ${internalFilePath}`);
 
   await storeWorkspaceMeta(workspaceId, {
     id: workspaceId,

packages/frontend/apps/electron/test/dialog/dialog.spec.ts

@@ -0,0 +1,268 @@
+import { afterEach, describe, expect, test, vi } from 'vitest';
+
+const connect = vi.fn();
+const checkpoint = vi.fn();
+const poolVacuumInto = vi.fn();
+const pathExists = vi.fn();
+const remove = vi.fn();
+const move = vi.fn();
+const realpath = vi.fn();
+const copyFile = vi.fn();
+const ensureDir = vi.fn();
+const copy = vi.fn();
+const storeWorkspaceMeta = vi.fn();
+const getSpaceDBPath = vi.fn();
+const getWorkspaceDBPath = vi.fn();
+const getWorkspacesBasePath = vi.fn();
+const docValidate = vi.fn();
+const docValidateImportSchema = vi.fn();
+const docVacuumInto = vi.fn();
+const docSetSpaceId = vi.fn();
+const sqliteValidate = vi.fn();
+const sqliteValidateImportSchema = vi.fn();
+const sqliteVacuumInto = vi.fn();
+
+vi.doMock('nanoid', () => ({
+  nanoid: () => 'workspace-1',
+}));
+
+vi.doMock('@affine/native', () => {
+  const ValidationResult = {
+    MissingTables: 'MissingTables',
+    MissingDocIdColumn: 'MissingDocIdColumn',
+    MissingVersionColumn: 'MissingVersionColumn',
+    GeneralError: 'GeneralError',
+    Valid: 'Valid',
+  };
+
+  return {
+    ValidationResult,
+    DocStorage: class {
+      constructor(private readonly path: string) {}
+
+      validate() {
+        return docValidate(this.path);
+      }
+
+      validateImportSchema() {
+        return docValidateImportSchema(this.path);
+      }
+
+      vacuumInto(path: string) {
+        return docVacuumInto(this.path, path);
+      }
+
+      setSpaceId(spaceId: string) {
+        return docSetSpaceId(this.path, spaceId);
+      }
+    },
+    SqliteConnection: class {
+      static validate(path: string) {
+        return sqliteValidate(path);
+      }
+
+      constructor(private readonly path: string) {}
+
+      validateImportSchema() {
+        return sqliteValidateImportSchema(this.path);
+      }
+
+      vacuumInto(path: string) {
+        return sqliteVacuumInto(this.path, path);
+      }
+    },
+  };
+});
+
+vi.doMock('@affine/electron/helper/nbstore', () => ({
+  getDocStoragePool: () => ({
+    connect,
+    checkpoint,
+    vacuumInto: poolVacuumInto,
+  }),
+}));
+
+vi.doMock('@affine/electron/helper/main-rpc', () => ({
+  mainRPC: {
+    showItemInFolder: vi.fn(),
+  },
+}));
+
+vi.doMock('@affine/electron/helper/workspace/meta', () => ({
+  getSpaceDBPath,
+  getWorkspaceDBPath,
+  getWorkspacesBasePath,
+}));
+
+vi.doMock('@affine/electron/helper/workspace', () => ({
+  storeWorkspaceMeta,
+}));
+
+vi.doMock('fs-extra', () => ({
+  default: {
+    pathExists,
+    remove,
+    move,
+    realpath,
+    copyFile,
+    ensureDir,
+    copy,
+  },
+}));
+
+afterEach(() => {
+  vi.clearAllMocks();
+  vi.resetModules();
+});
+
+describe('dialog export', () => {
+  test('saveDBFileAs exports a vacuumed backup instead of copying the live db', async () => {
+    const dbPath = '/tmp/workspace/storage.db';
+    const exportPath = '/tmp/export.affine';
+    const tempExportPath = '/tmp/export.affine.workspace-1.tmp';
+    const id = '@peer(local);@type(workspace);@id(workspace-1);';
+
+    pathExists.mockImplementation(async path => path === dbPath);
+    realpath.mockImplementation(async path => path);
+    getSpaceDBPath.mockResolvedValue(dbPath);
+    move.mockResolvedValue(undefined);
+
+    const { saveDBFileAs, setFakeDialogResult } =
+      await import('@affine/electron/helper/dialog/dialog');
+
+    setFakeDialogResult({ filePath: exportPath });
+
+    const result = await saveDBFileAs(id, 'My Space');
+
+    expect(result).toEqual({ filePath: exportPath });
+    expect(connect).toHaveBeenCalledWith(id, dbPath);
+    expect(checkpoint).toHaveBeenCalledWith(id);
+    expect(poolVacuumInto).toHaveBeenCalledWith(id, tempExportPath);
+    expect(move).toHaveBeenCalledWith(tempExportPath, exportPath, {
+      overwrite: true,
+    });
+    expect(remove).not.toHaveBeenCalledWith(exportPath);
+    expect(copyFile).not.toHaveBeenCalled();
+  });
+
+  test('saveDBFileAs rejects exporting over the live database path', async () => {
+    const dbPath = '/tmp/workspace/storage.db';
+    const id = '@peer(local);@type(workspace);@id(workspace-1);';
+
+    pathExists.mockResolvedValue(false);
+    getSpaceDBPath.mockResolvedValue(dbPath);
+
+    const { saveDBFileAs, setFakeDialogResult } =
+      await import('@affine/electron/helper/dialog/dialog');
+
+    setFakeDialogResult({ filePath: dbPath });
+
+    const result = await saveDBFileAs(id, 'My Space');
+
+    expect(result).toEqual({ error: 'DB_FILE_PATH_INVALID' });
+    expect(poolVacuumInto).not.toHaveBeenCalled();
+    expect(copyFile).not.toHaveBeenCalled();
+  });
+
+  test('saveDBFileAs rejects exporting to a symlink alias of the live database', async () => {
+    const dbPath = '/tmp/workspace/storage.db';
+    const exportPath = '/tmp/alias.affine';
+    const id = '@peer(local);@type(workspace);@id(workspace-1);';
+
+    pathExists.mockResolvedValue(true);
+    realpath.mockImplementation(async path =>
+      path === exportPath ? dbPath : path
+    );
+    getSpaceDBPath.mockResolvedValue(dbPath);
+
+    const { saveDBFileAs, setFakeDialogResult } =
+      await import('@affine/electron/helper/dialog/dialog');
+
+    setFakeDialogResult({ filePath: exportPath });
+
+    const result = await saveDBFileAs(id, 'My Space');
+
+    expect(result).toEqual({ error: 'DB_FILE_PATH_INVALID' });
+    expect(poolVacuumInto).not.toHaveBeenCalled();
+    expect(move).not.toHaveBeenCalled();
+  });
+});
+
+describe('dialog import', () => {
+  test('loadDBFile validates schema and vacuums v2 imports into internal storage', async () => {
+    const originalPath = '/tmp/import.affine';
+    const internalPath = '/app/workspaces/local/workspace-1/storage.db';
+
+    getWorkspacesBasePath.mockResolvedValue('/app/workspaces');
+    getSpaceDBPath.mockResolvedValue(internalPath);
+    docValidate.mockResolvedValue(true);
+    docValidateImportSchema.mockResolvedValue(true);
+    docVacuumInto.mockResolvedValue(undefined);
+    docSetSpaceId.mockResolvedValue(undefined);
+    ensureDir.mockResolvedValue(undefined);
+
+    const { loadDBFile, setFakeDialogResult } =
+      await import('@affine/electron/helper/dialog/dialog');
+
+    setFakeDialogResult({ filePath: originalPath });
+
+    const result = await loadDBFile();
+
+    expect(result).toEqual({ workspaceId: 'workspace-1' });
+    expect(docValidate).toHaveBeenCalledWith(originalPath);
+    expect(docValidateImportSchema).toHaveBeenCalledWith(originalPath);
+    expect(docVacuumInto).toHaveBeenCalledWith(originalPath, internalPath);
+    expect(docSetSpaceId).toHaveBeenCalledWith(internalPath, 'workspace-1');
+    expect(copy).not.toHaveBeenCalled();
+  });
+
+  test('loadDBFile rejects v2 imports with unexpected schema objects', async () => {
+    const originalPath = '/tmp/import.affine';
+
+    getWorkspacesBasePath.mockResolvedValue('/app/workspaces');
+    docValidate.mockResolvedValue(true);
+    docValidateImportSchema.mockResolvedValue(false);
+
+    const { loadDBFile, setFakeDialogResult } =
+      await import('@affine/electron/helper/dialog/dialog');
+
+    setFakeDialogResult({ filePath: originalPath });
+
+    const result = await loadDBFile();
+
+    expect(result).toEqual({ error: 'DB_FILE_INVALID' });
+    expect(docVacuumInto).not.toHaveBeenCalled();
+    expect(copy).not.toHaveBeenCalled();
+  });
+
+  test('loadDBFile validates schema and vacuums v1 imports into internal storage', async () => {
+    const originalPath = '/tmp/import-v1.affine';
+    const internalPath = '/app/workspaces/workspace-1/storage.db';
+
+    getWorkspacesBasePath.mockResolvedValue('/app/workspaces');
+    getWorkspaceDBPath.mockResolvedValue(internalPath);
+    docValidate.mockResolvedValue(false);
+    sqliteValidate.mockResolvedValue('Valid');
+    sqliteValidateImportSchema.mockResolvedValue(true);
+    sqliteVacuumInto.mockResolvedValue(undefined);
+    ensureDir.mockResolvedValue(undefined);
+
+    const { loadDBFile, setFakeDialogResult } =
+      await import('@affine/electron/helper/dialog/dialog');
+
+    setFakeDialogResult({ filePath: originalPath });
+
+    const result = await loadDBFile();
+
+    expect(result).toEqual({ workspaceId: 'workspace-1' });
+    expect(sqliteValidate).toHaveBeenCalledWith(originalPath);
+    expect(sqliteValidateImportSchema).toHaveBeenCalledWith(originalPath);
+    expect(ensureDir).toHaveBeenCalledWith('/app/workspaces/workspace-1');
+    expect(sqliteVacuumInto).toHaveBeenCalledWith(originalPath, internalPath);
+    expect(storeWorkspaceMeta).toHaveBeenCalledWith('workspace-1', {
+      id: 'workspace-1',
+      mainDBPath: internalPath,
+    });
+    expect(copy).not.toHaveBeenCalled();
+  });
+});