Enforce TOTP verification during desktop sync setup

apps/client/src/setup.ts

@@ -17,6 +17,9 @@ class SetupController {
     private syncServerHostInput: HTMLInputElement;
     private syncProxyInput: HTMLInputElement;
     private passwordInput: HTMLInputElement;
+    private totpTokenInput: HTMLInputElement;
+    private totpSection: HTMLElement;
+    private totpEnabled = false;
     private sections: Record<SetupStep, HTMLElement>;
 
     constructor(rootNode: HTMLElement, syncInProgress: boolean) {
@@ -29,6 +32,8 @@ class SetupController {
         this.syncServerHostInput = mustGetElement("sync-server-host", HTMLInputElement);
         this.syncProxyInput = mustGetElement("sync-proxy", HTMLInputElement);
         this.passwordInput = mustGetElement("password", HTMLInputElement);
+        this.totpTokenInput = mustGetElement("totp-token", HTMLInputElement);
+        this.totpSection = mustGetElement("totp-section", HTMLElement);
         this.sections = {
             "setup-type": mustGetElement("setup-type-section", HTMLElement),
             "new-document-in-progress": mustGetElement("new-document-in-progress-section", HTMLElement),
@@ -56,6 +61,10 @@ class SetupController {
             });
         }
 
+        this.syncServerHostInput.addEventListener("blur", () => {
+            void this.checkTotpStatus();
+        });
+
         for (const backButton of document.querySelectorAll<HTMLElement>("[data-action='back']")) {
             backButton.addEventListener("click", () => {
                 this.back();
@@ -87,9 +96,40 @@ class SetupController {
         }
     }
 
+    private async checkTotpStatus() {
+        const syncServerHost = this.syncServerHostInput.value.trim();
+
+        if (!syncServerHost) {
+            this.setTotpEnabled(false);
+            return;
+        }
+
+        try {
+            const resp = await $.post("api/setup/check-server-totp", {
+                syncServerHost
+            });
+
+            this.setTotpEnabled(!!resp.totpEnabled);
+        } catch {
+            // If we can't reach the server, don't show the TOTP field yet.
+            this.setTotpEnabled(false);
+        }
+    }
+
+    private setTotpEnabled(enabled: boolean) {
+        this.totpEnabled = enabled;
+
+        if (!enabled) {
+            this.totpTokenInput.value = "";
+        }
+
+        this.render();
+    }
+
     private back() {
         this.setStep("setup-type");
         this.setupType = "";
+        this.setTotpEnabled(false);
 
         for (const input of this.setupTypeInputs) {
             input.checked = false;
@@ -113,11 +153,21 @@ class SetupController {
             return;
         }
 
+        await this.checkTotpStatus();
+
+        const totpToken = this.totpTokenInput.value.trim();
+
+        if (this.totpEnabled && !totpToken) {
+            showAlert("TOTP token can't be empty when two-factor authentication is enabled");
+            return;
+        }
+
         // not using server.js because it loads too many dependencies
         const resp = await $.post("api/setup/sync-from-server", {
             syncServerHost,
             syncProxy,
-            password
+            password,
+            totpToken
         });
 
         if (resp.result === "success") {
@@ -139,13 +189,10 @@ class SetupController {
             section.style.display = step === this.step ? "" : "none";
         }
 
+        this.totpSection.style.display = this.totpEnabled ? "" : "none";
         this.setupTypeNextButton.disabled = !this.setupType;
     }
 
-    private getSelectedSetupType(): SetupType {
-        return (this.setupTypeInputs.find((input) => input.checked)?.value ?? "") as SetupType;
-    }
-
     private startSyncPolling() {
         if (this.syncPollIntervalId !== null) {
             return;
@@ -196,7 +243,7 @@ function mustGetElement<T extends typeof HTMLElement>(id: string, ctor: T): Inst
     return element as InstanceType<T>;
 }
 
-addEventListener("DOMContentLoaded", (event) => {
+addEventListener("DOMContentLoaded", () => {
     const rootNode = document.getElementById("setup-dialog");
     if (!rootNode || !(rootNode instanceof HTMLElement)) return;
 

apps/server/src/assets/translations/cn/server.json

@@ -155,6 +155,8 @@
     "proxy-instruction": "如果您将代理设置留空，将使用系统代理（仅适用于桌面应用）",
     "password": "密码",
     "password-placeholder": "密码",
+    "totp-token": "TOTP 验证码",
+    "totp-token-placeholder": "请输入 TOTP 验证码",
     "back": "返回",
     "finish-setup": "完成设置"
   },

apps/server/src/assets/translations/en/server.json

@@ -252,6 +252,8 @@
     "proxy-instruction": "If you leave proxy setting blank, system proxy will be used (applies to the desktop application only)",
     "password": "Password",
     "password-placeholder": "Password",
+    "totp-token": "TOTP Token",
+    "totp-token-placeholder": "Enter your TOTP code",
     "back": "Back",
     "finish-setup": "Finish setup"
   },

apps/server/src/assets/translations/tw/server.json

@@ -155,6 +155,8 @@
     "proxy-instruction": "如果您將代理設定留空，將使用系統代理（僅適用於桌面版）",
     "password": "密碼",
     "password-placeholder": "密碼",
+    "totp-token": "TOTP 驗證碼",
+    "totp-token-placeholder": "請輸入 TOTP 驗證碼",
     "back": "返回",
     "finish-setup": "完成設定"
   },

apps/server/src/assets/views/setup.ejs

@@ -141,6 +141,10 @@
                 <label for="password"><%= t("setup_sync-from-server.password") %></label>
                 <input type="password" id="password" class="form-control" placeholder="<%= t("setup_sync-from-server.password-placeholder") %>">
             </div>
+            <div id="totp-section" class="form-group" style="margin-bottom: 8px; display: none;">
+                <label for="totp-token"><%= t("setup_sync-from-server.totp-token") %></label>
+                <input type="text" id="totp-token" class="form-control" placeholder="<%= t("setup_sync-from-server.totp-token-placeholder") %>" autocomplete="one-time-code">
+            </div>
 
             <button type="button" data-action="back" class="btn btn-secondary"><%= t("setup_sync-from-server.back") %></button>
 

apps/server/src/express.d.ts

@@ -8,6 +8,7 @@ export declare module "express-serve-static-core" {
 
             authorization?: string;
             "trilium-cred"?: string;
+            "trilium-totp"?: string;
             "x-csrf-token"?: string;
 
             "trilium-component-id"?: string;

apps/server/src/routes/api/setup.ts

@@ -1,16 +1,19 @@
-"use strict";
 
-import sqlInit from "../../services/sql_init.js";
-import setupService from "../../services/setup.js";
-import log from "../../services/log.js";
-import appInfo from "../../services/app_info.js";
+
 import type { Request } from "express";
 
+import appInfo from "../../services/app_info.js";
+import log from "../../services/log.js";
+import setupService from "../../services/setup.js";
+import sqlInit from "../../services/sql_init.js";
+import totp from "../../services/totp.js";
+
 function getStatus() {
     return {
         isInitialized: sqlInit.isDbInitialized(),
         schemaExists: sqlInit.schemaExists(),
-        syncVersion: appInfo.syncVersion
+        syncVersion: appInfo.syncVersion,
+        totpEnabled: totp.isTotpEnabled()
     };
 }
 
@@ -19,9 +22,9 @@ async function setupNewDocument() {
 }
 
 function setupSyncFromServer(req: Request) {
-    const { syncServerHost, syncProxy, password } = req.body;
+    const { syncServerHost, syncProxy, password, totpToken } = req.body;
 
-    return setupService.setupSyncFromSyncServer(syncServerHost, syncProxy, password);
+    return setupService.setupSyncFromSyncServer(syncServerHost, syncProxy, password, totpToken);
 }
 
 function saveSyncSeed(req: Request) {
@@ -82,10 +85,26 @@ function getSyncSeed() {
     };
 }
 
+async function checkServerTotpStatus(req: Request) {
+    const { syncServerHost } = req.body;
+
+    if (!syncServerHost) {
+        return { totpEnabled: false };
+    }
+
+    try {
+        const resp = await setupService.checkRemoteTotpStatus(syncServerHost);
+        return { totpEnabled: !!resp.totpEnabled };
+    } catch {
+        return { totpEnabled: false };
+    }
+}
+
 export default {
     getStatus,
     setupNewDocument,
     setupSyncFromServer,
     getSyncSeed,
-    saveSyncSeed
+    saveSyncSeed,
+    checkServerTotpStatus
 };

apps/server/src/routes/routes.ts

@@ -244,6 +244,7 @@ function register(app: express.Application) {
     asyncRoute(PST, "/api/setup/sync-from-server", [auth.checkAppNotInitialized], setupApiRoute.setupSyncFromServer, apiResultHandler);
     route(GET, "/api/setup/sync-seed", [loginRateLimiter, auth.checkCredentials], setupApiRoute.getSyncSeed, apiResultHandler);
     asyncRoute(PST, "/api/setup/sync-seed", [auth.checkAppNotInitialized], setupApiRoute.saveSyncSeed, apiResultHandler);
+    asyncRoute(PST, "/api/setup/check-server-totp", [auth.checkAppNotInitialized], setupApiRoute.checkServerTotpStatus, apiResultHandler);
 
     apiRoute(GET, "/api/autocomplete", autocompleteApiRoute.getAutocomplete);
     apiRoute(GET, "/api/autocomplete/notesCount", autocompleteApiRoute.getNotesCount);

apps/server/src/services/api-interface.ts

@@ -6,6 +6,7 @@ import type { OptionRow } from "@triliumnext/commons";
 export interface SetupStatusResponse {
     syncVersion: number;
     schemaExists: boolean;
+    totpEnabled: boolean;
 }
 
 /**

apps/server/src/services/auth.spec.ts

@@ -9,6 +9,10 @@ import options from "./options";
 
 let app: Application;
 
+function encodeCred(password: string): string {
+    return Buffer.from(`dummy:${password}`).toString("base64");
+}
+
 describe("Auth", () => {
     beforeAll(async () => {
         const buildApp = (await (import("../../src/app.js"))).default;
@@ -72,4 +76,49 @@ describe("Auth", () => {
                 .expect(200);
         });
     });
+
+    describe("Setup status endpoint", () => {
+        it("returns totpEnabled: true when TOTP is enabled", async () => {
+            cls.init(() => {
+                options.setOption("mfaEnabled", "true");
+                options.setOption("mfaMethod", "totp");
+                options.setOption("totpVerificationHash", "hi");
+            });
+            const response = await supertest(app)
+                .get("/api/setup/status")
+                .expect(200);
+            expect(response.body.totpEnabled).toBe(true);
+        });
+
+        it("returns totpEnabled: false when TOTP is disabled", async () => {
+            cls.init(() => {
+                options.setOption("mfaEnabled", "false");
+            });
+            const response = await supertest(app)
+                .get("/api/setup/status")
+                .expect(200);
+            expect(response.body.totpEnabled).toBe(false);
+        });
+    });
+
+    describe("checkCredentials TOTP enforcement", () => {
+        beforeAll(() => {
+            config.General.noAuthentication = false;
+            refreshAuth();
+        });
+
+        it("does not require TOTP token when TOTP is disabled", async () => {
+            cls.init(() => {
+                options.setOption("mfaEnabled", "false");
+            });
+            // Will still fail with 401 due to wrong password, but NOT because of missing TOTP
+            const response = await supertest(app)
+                .get("/api/setup/sync-seed")
+                .set("trilium-cred", encodeCred("wrongpassword"))
+                .expect(401);
+            // The error should be about password, not TOTP
+            expect(response.text).toContain("Incorrect password");
+        });
+    });
 }, 60_000);
+

apps/server/src/services/auth.ts

@@ -1,15 +1,17 @@
-import etapiTokenService from "./etapi_tokens.js";
-import log from "./log.js";
-import sqlInit from "./sql_init.js";
-import { isElectron } from "./utils.js";
-import passwordEncryptionService from "./encryption/password_encryption.js";
+import type { NextFunction, Request, Response } from "express";
+
+import attributes from "./attributes.js";
 import config from "./config.js";
 import passwordService from "./encryption/password.js";
-import totp from "./totp.js";
+import passwordEncryptionService from "./encryption/password_encryption.js";
+import recoveryCodeService from "./encryption/recovery_codes.js";
+import etapiTokenService from "./etapi_tokens.js";
+import log from "./log.js";
 import openID from "./open_id.js";
 import options from "./options.js";
-import attributes from "./attributes.js";
-import type { NextFunction, Request, Response } from "express";
+import sqlInit from "./sql_init.js";
+import totp from "./totp.js";
+import { isElectron } from "./utils.js";
 
 let noAuthentication = false;
 refreshAuth();
@@ -161,9 +163,28 @@ function checkCredentials(req: Request, res: Response, next: NextFunction) {
     if (!passwordEncryptionService.verifyPassword(password)) {
         res.setHeader("Content-Type", "text/plain").status(401).send("Incorrect password");
         log.info(`WARNING: Wrong password from ${req.ip}, rejecting.`);
-    } else {
-        next();
+        return;
+    }
+
+    // Verify TOTP if enabled
+    if (totp.isTotpEnabled()) {
+        const totpHeader = req.headers["trilium-totp"];
+        const totpToken = Array.isArray(totpHeader) ? totpHeader[0] : totpHeader;
+        if (typeof totpToken !== "string" || !totpToken) {
+            res.setHeader("Content-Type", "text/plain").status(401).send("TOTP token is required");
+            log.info(`WARNING: Missing or invalid TOTP token from ${req.ip}, rejecting.`);
+            return;
+        }
+
+        // Accept TOTP code or recovery code
+        if (!totp.validateTOTP(totpToken) && !recoveryCodeService.verifyRecoveryCode(totpToken)) {
+            res.setHeader("Content-Type", "text/plain").status(401).send("Incorrect TOTP token");
+            log.info(`WARNING: Wrong TOTP token from ${req.ip}, rejecting.`);
+            return;
+        }
     }
+
+    next();
 }
 
 export default {

apps/server/src/services/encryption/totp_encryption.ts

@@ -1,8 +1,9 @@
+import type { OptionNames } from "@triliumnext/commons";
+
 import optionService from "../options.js";
-import myScryptService from "./my_scrypt.js";
-import { randomSecureToken, toBase64, constantTimeCompare } from "../utils.js";
+import { constantTimeCompare,randomSecureToken, toBase64 } from "../utils.js";
 import dataEncryptionService from "./data_encryption.js";
-import type { OptionNames } from "@triliumnext/commons";
+import myScryptService from "./my_scrypt.js";
 
 const TOTP_OPTIONS: Record<string, OptionNames> = {
     SALT: "totpEncryptionSalt",

apps/server/src/services/request.ts

@@ -62,6 +62,9 @@ async function exec<T>(opts: ExecOpts): Promise<T> {
 
             if (opts.auth) {
                 headers["trilium-cred"] = Buffer.from(`dummy:${opts.auth.password}`).toString("base64");
+                if (opts.auth.totpToken) {
+                    headers["trilium-totp"] = opts.auth.totpToken;
+                }
             }
 
             const request = (await client).request({

apps/server/src/services/request_interface.ts

@@ -14,6 +14,7 @@ export interface ExecOpts {
     cookieJar?: CookieJar;
     auth?: {
         password?: string;
+        totpToken?: string;
     };
     timeout: number;
     body?: string | {};