obsidian-livesync/src/modules/essential/ModuleMigration.ts

import {
    LOG_LEVEL_INFO,
    LOG_LEVEL_NOTICE,
    LOG_LEVEL_VERBOSE,
    Logger,
} from "@vrtmrz/livesync-commonlib/compat/common/logger";
import { EVENT_REQUEST_RUN_DOCTOR, EVENT_REQUEST_RUN_FIX_INCOMPLETE, eventHub } from "@/common/events.ts";
import { AbstractModule } from "@/modules/AbstractModule.ts";
import { $msg } from "@/common/translation";
import { performDoctorConsultation, RebuildOptions } from "@vrtmrz/livesync-commonlib/compat/common/configForDoc";
import { isValidPath } from "@/common/utils.ts";
import { isMetaEntry } from "@vrtmrz/livesync-commonlib/compat/common/types";
import {
    isDeletedEntry,
    isDocContentSame,
    isLoadedEntry,
    readAsBlob,
} from "@vrtmrz/livesync-commonlib/compat/common/utils";
import { countCompromisedChunks } from "@vrtmrz/livesync-commonlib/compat/pouchdb/negotiation";
import type { LiveSyncCore } from "@/main.ts";
import { SetupManager } from "@/modules/features/SetupManager.ts";
import { showOnboardingInvitation } from "@/serviceFeatures/setupObsidian/setupManagerHandlers.ts";
import {
    runConfiguredStartupLifecycle,
    runStartupEntryLifecycle,
} from "@/serviceFeatures/configuredStartupLifecycle.ts";
import { disableLegacyBulkChunkPreSend } from "@/common/compatibilitySettings.ts";

type ErrorInfo = {
    path: string;
    recordedSize: number;
    actualSize: number;
    storageSize: number;
    contentMatched: boolean;
    isConflicted?: boolean;
};

const INCOMPLETE_DOCUMENT_NOTICE_GROUP = "startup-integrity-check";

interface CompromisedChunkCounter {
    countCompromisedChunks(): Promise<number | boolean>;
}

function hasCompromisedChunkCounter(value: object | undefined): value is CompromisedChunkCounter {
    return (
        value !== undefined && "countCompromisedChunks" in value && typeof value.countCompromisedChunks === "function"
    );
}

export class ModuleMigration extends AbstractModule<LiveSyncCore> {
    constructor(
        core: LiveSyncCore,
        private readonly waitForCompatibilityReview: () => Promise<void> = () => Promise.resolve()
    ) {
        super(core);
    }

    async migrateUsingDoctor(skipRebuild: boolean = false, activateReason = "updated", forceRescan = false) {
        const { shouldRebuild, shouldRebuildLocal, isModified, settings } = await performDoctorConsultation(
            {
                confirm: this.core.confirm,
                translate: this.services.context.translate,
            },
            this.settings,
            {
                localRebuild: skipRebuild ? RebuildOptions.SkipEvenIfRequired : RebuildOptions.AutomaticAcceptable,
                remoteRebuild: skipRebuild ? RebuildOptions.SkipEvenIfRequired : RebuildOptions.AutomaticAcceptable,
                activateReason,
                forceRescan,
            }
        );
        if (isModified) {
            this.settings = settings;
            await this.saveSettings();
        }
        if (!skipRebuild) {
            if (shouldRebuild) {
                await this.core.rebuilder.scheduleRebuild();
                this.services.appLifecycle.performRestart();
                return false;
            } else if (shouldRebuildLocal) {
                await this.core.rebuilder.scheduleFetch();
                this.services.appLifecycle.performRestart();
                return false;
            }
        }
        return true;
    }

    async migrateDisableBulkSend() {
        if (disableLegacyBulkChunkPreSend(this.settings)) {
            this._log($msg("moduleMigration.logBulkSendCorrupted"), LOG_LEVEL_NOTICE);
            await this.saveSettings();
        }
    }

    initialMessage() {
        const manager = this.core.getModule(SetupManager);
        showOnboardingInvitation(this.core, manager);
    }

    async hasIncompleteDocs(force: boolean = false): Promise<boolean> {
        const incompleteDocsChecked = (await this.core.kvDB.get<boolean>("checkIncompleteDocs")) || false;
        if (incompleteDocsChecked && !force) {
            this._log("Incomplete docs check already done, skipping.", LOG_LEVEL_VERBOSE);
            return Promise.resolve(true);
        }

        const noticeGroups = this.core.services.context.noticeGroups;
        noticeGroups.setItem(INCOMPLETE_DOCUMENT_NOTICE_GROUP, "checking", {
            message: "Checking for incomplete documents...",
        });
        this._log("Checking for incomplete documents...", LOG_LEVEL_VERBOSE);

        try {
            const errorFiles = [] as ErrorInfo[];
            for await (const metaDoc of this.localDatabase.findAllNormalDocs({ conflicts: true })) {
                const path = this.getPath(metaDoc);

                if (!isValidPath(path)) {
                    continue;
                }
                if (!(await this.services.vault.isTargetFile(path))) {
                    continue;
                }
                if (!isMetaEntry(metaDoc)) {
                    continue;
                }

                const doc = await this.localDatabase.getDBEntryFromMeta(metaDoc);
                if (!doc || !isLoadedEntry(doc)) {
                    continue;
                }
                if (isDeletedEntry(doc)) {
                    continue;
                }
                const isConflicted = metaDoc?._conflicts && metaDoc._conflicts.length > 0;

                let storageFileContent;
                try {
                    storageFileContent = await this.core.storageAccess.readHiddenFileBinary(path);
                } catch (e) {
                    Logger(`Failed to read file ${path}: Possibly unprocessed or missing`);
                    Logger(e, LOG_LEVEL_VERBOSE);
                    continue;
                }
                // const storageFileBlob = createBlob(storageFileContent);
                const sizeOnStorage = storageFileContent.byteLength;
                const recordedSize = doc.size;
                const docBlob = readAsBlob(doc);
                const actualSize = docBlob.size;
                if (
                    recordedSize !== actualSize ||
                    sizeOnStorage !== actualSize ||
                    sizeOnStorage !== recordedSize ||
                    isConflicted
                ) {
                    const contentMatched = await isDocContentSame(doc.data, storageFileContent);
                    errorFiles.push({
                        path,
                        recordedSize,
                        actualSize,
                        storageSize: sizeOnStorage,
                        contentMatched,
                        isConflicted,
                    });
                    Logger(
                        `Size mismatch for ${path}: ${recordedSize} (DB Recorded) , ${actualSize} (DB Stored) , ${sizeOnStorage} (Storage Stored), ${contentMatched ? "Content Matched" : "Content Mismatched"} ${isConflicted ? "Conflicted" : "Not Conflicted"}`
                    );
                }
            }
            if (errorFiles.length == 0) {
                Logger("No size mismatches found", LOG_LEVEL_INFO);
                noticeGroups.setItem(INCOMPLETE_DOCUMENT_NOTICE_GROUP, "result", {
                    message: "No size mismatches found",
                });
                await this.core.kvDB.set("checkIncompleteDocs", true);
                return Promise.resolve(true);
            }
            Logger(`Found ${errorFiles.length} size mismatches`, LOG_LEVEL_INFO);
            noticeGroups.setItem(INCOMPLETE_DOCUMENT_NOTICE_GROUP, "result", {
                message: `Found ${errorFiles.length} size mismatches`,
            });
            // We have to repair them following rules and situations:
            // A. DB Recorded != DB Stored
            //   A.1. DB Recorded == Storage Stored
            //        Possibly recoverable from storage. Just overwrite the DB content with storage content.
            //   A.2. Neither
            //        Probably it cannot be resolved on this device. Even if the storage content is larger than DB Recorded, it possibly corrupted.
            //        We do not fix it automatically. Leave it as is. Possibly other device can do this.
            // B. DB Recorded == DB Stored ,  < Storage Stored
            //   Very fragile, if DB Recorded size is less than Storage Stored size, we possibly repair the content (The issue was `unexpectedly shortened file`).
            //   We do not fix it automatically, but it will be automatically overwritten in other process.
            // C. DB Recorded == DB Stored ,  > Storage Stored
            //   Probably restored by the user by resolving A or B on other device, We should overwrite the storage
            //   Also do not fix it automatically. It should be overwritten by replication.
            const recoverable = errorFiles.filter((e) => {
                return e.recordedSize === e.storageSize && !e.isConflicted;
            });
            const unrecoverable = errorFiles.filter((e) => {
                return e.recordedSize !== e.storageSize || e.isConflicted;
            });
            const fileInfo = (e: (typeof errorFiles)[0]) => {
                return `${e.path} (M: ${e.recordedSize}, A: ${e.actualSize}, S: ${e.storageSize}) ${e.isConflicted ? "(Conflicted)" : ""}`;
            };
            const messageUnrecoverable =
                unrecoverable.length > 0
                    ? $msg("moduleMigration.fix0256.messageUnrecoverable", {
                          filesNotRecoverable: unrecoverable.map((e) => `- ${fileInfo(e)}`).join("\n"),
                      })
                    : "";

            const message = $msg("moduleMigration.fix0256.message", {
                files: recoverable.map((e) => `- ${fileInfo(e)}`).join("\n"),
                messageUnrecoverable,
            });
            const CHECK_IT_LATER = $msg("moduleMigration.fix0256.buttons.checkItLater");
            const FIX = $msg("moduleMigration.fix0256.buttons.fix");
            const DISMISS = $msg("moduleMigration.fix0256.buttons.DismissForever");
            const ret = await this.core.confirm.askSelectStringDialogue(message, [CHECK_IT_LATER, FIX, DISMISS], {
                title: $msg("moduleMigration.fix0256.title"),
                defaultAction: CHECK_IT_LATER,
            });
            if (ret == FIX) {
                for (const file of recoverable) {
                    // Overwrite the database with the files on the storage
                    const stubFile = await this.core.storageAccess.getFileStub(file.path);
                    if (stubFile == null) {
                        Logger(`Could not find stub file for ${file.path}`, LOG_LEVEL_NOTICE);
                        continue;
                    }

                    stubFile.stat.mtime = Date.now();
                    const result = await this.core.fileHandler.storeFileToDB(stubFile, true, false);
                    if (result) {
                        Logger(`Successfully restored ${file.path} from storage`);
                    } else {
                        Logger(`Failed to restore ${file.path} from storage`, LOG_LEVEL_NOTICE);
                    }
                }
            } else if (ret === DISMISS) {
                // User chose to dismiss the issue
                await this.core.kvDB.set("checkIncompleteDocs", true);
            }

            return Promise.resolve(true);
        } catch (error) {
            noticeGroups.setItem(INCOMPLETE_DOCUMENT_NOTICE_GROUP, "result", {
                message: "The incomplete document check could not be completed.",
            });
            throw error;
        } finally {
            noticeGroups.finish(INCOMPLETE_DOCUMENT_NOTICE_GROUP);
        }
    }

    async hasCompromisedChunks(): Promise<boolean> {
        Logger(`Checking for compromised chunks...`, LOG_LEVEL_VERBOSE);
        if (!this.settings.encrypt) {
            // If not encrypted, we do not need to check for compromised chunks.
            return true;
        }
        // Check local database for compromised chunks
        const localCompromised = await countCompromisedChunks(this.localDatabase.localDatabase);
        const remote = this.services.replicator.getActiveReplicator();
        const remoteCompromised =
            this.services.API.isOnline && hasCompromisedChunkCounter(remote)
                ? await remote.countCompromisedChunks()
                : 0;
        if (localCompromised === false) {
            Logger(`Failed to count compromised chunks in local database`, LOG_LEVEL_NOTICE);
            return false;
        }
        if (remoteCompromised === false) {
            Logger(`Failed to count compromised chunks in remote database`, LOG_LEVEL_NOTICE);
            return false;
        }
        if (remoteCompromised === 0 && localCompromised === 0) {
            return true;
        }
        Logger(
            `Found compromised chunks : ${localCompromised} in local, ${remoteCompromised} in remote`,
            LOG_LEVEL_NOTICE
        );
        const title = $msg("moduleMigration.insecureChunkExist.title");
        const msg = $msg("moduleMigration.insecureChunkExist.message");
        const REBUILD = $msg("moduleMigration.insecureChunkExist.buttons.rebuild");
        const FETCH = $msg("moduleMigration.insecureChunkExist.buttons.fetch");
        const DISMISS = $msg("moduleMigration.insecureChunkExist.buttons.later");
        const buttons = [REBUILD, FETCH, DISMISS];
        if (remoteCompromised != 0) {
            buttons.splice(buttons.indexOf(FETCH), 1);
        }
        const result = await this.core.confirm.askSelectStringDialogue(msg, buttons, {
            title,
            defaultAction: DISMISS,
            timeout: 0,
        });
        if (result === REBUILD) {
            // Rebuild the database
            await this.core.rebuilder.scheduleRebuild();
            this.services.appLifecycle.performRestart();
            return false;
        } else if (result === FETCH) {
            // Fetch the latest data from remote
            await this.core.rebuilder.scheduleFetch();
            this.services.appLifecycle.performRestart();
            return false;
        } else {
            // User chose to dismiss the issue
            this._log($msg("moduleMigration.insecureChunkExist.laterMessage"), LOG_LEVEL_NOTICE);
        }
        return true;
    }

    async _everyOnFirstInitialize(): Promise<boolean> {
        return await runConfiguredStartupLifecycle({
            databaseReady: this.localDatabase.isReady,
            reportDatabaseNotReady: () => this._log($msg("moduleMigration.logLocalDatabaseNotReady"), LOG_LEVEL_NOTICE),
            hasCompromisedChunks: () => this.hasCompromisedChunks(),
            hasIncompleteDocuments: () => this.hasIncompleteDocs(),
            waitForCompatibilityReview: () => this.waitForCompatibilityReview(),
            runDoctor: () => this.migrateUsingDoctor(false),
            migrateBulkSend: () => this.migrateDisableBulkSend(),
        });
    }
    _everyOnLayoutReady(): Promise<boolean> {
        const shouldInitialiseDatabase = runStartupEntryLifecycle({
            configured: this.settings.isConfigured === true,
            inviteToOnboarding: () => this.initialMessage(),
        });
        if (!shouldInitialiseDatabase) return Promise.resolve(false);
        eventHub.onEvent(EVENT_REQUEST_RUN_DOCTOR, async (reason) => {
            await this.migrateUsingDoctor(false, reason, true);
        });
        eventHub.onEvent(EVENT_REQUEST_RUN_FIX_INCOMPLETE, async () => {
            await this.hasIncompleteDocs(true);
        });
        return Promise.resolve(true);
    }
    override onBindFunction(core: LiveSyncCore, services: typeof core.services): void {
        super.onBindFunction(core, services);
        services.appLifecycle.onLayoutReady.addHandler(this._everyOnLayoutReady.bind(this));
        services.appLifecycle.onFirstInitialise.addHandler(this._everyOnFirstInitialize.bind(this));
    }
}