From 3def78fc5a01002515415b57bb4f917d69e1aa4c Mon Sep 17 00:00:00 2001 From: Guoyu Su Date: Thu, 15 Oct 2026 10:35:15 +0800 Subject: [PATCH] fix(connection): forget transient password on user close --- .../connectionStore.savePassword.spec.ts | 145 ++++++++++++++++++ apps/desktop/src/stores/connectionStore.ts | 67 +++++++- 2 files changed, 204 insertions(+), 8 deletions(-) diff --git a/apps/desktop/src/stores/__tests__/connectionStore.savePassword.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.savePassword.spec.ts index b20067f04a..b21932e981 100644 --- a/apps/desktop/src/stores/__tests__/connectionStore.savePassword.spec.ts +++ b/apps/desktop/src/stores/__tests__/connectionStore.savePassword.spec.ts @@ -76,9 +76,154 @@ describe("connectionStore save_password opt-out", () => { }); afterEach(() => { + vi.useRealTimers(); vi.restoreAllMocks(); }); + it.each(["postgres", "mysql", "redis"] as const)("ordinary close retires the %s transient password before disconnect", async (db_type) => { + const calls: string[] = []; + installApiMocks({ + sessionCredentialStatus: vi.fn().mockResolvedValue(true), + forgetSessionCredential: vi.fn(async () => { + calls.push("forget"); + }), + disconnectDb: vi.fn(async () => { + calls.push("disconnect"); + }), + }); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [postgresConnection({ db_type, save_password: false })]; + await store.disconnect("pg-1"); + expect(calls).toEqual(["forget", "disconnect"]); + }); + + it.each([true, undefined])("preserves ordinary-close behavior for save_password=%s", async (save_password) => { + installApiMocks(); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [postgresConnection({ save_password })]; + await store.disconnect("pg-1"); + const api = await import("@/lib/backend/api"); + expect(api.sessionCredentialStatus).not.toHaveBeenCalled(); + expect(api.forgetSessionCredential).not.toHaveBeenCalled(); + expect(api.disconnectDb).toHaveBeenCalledTimes(1); + }); + + it("system disconnect does not retire a transient credential", async () => { + installApiMocks(); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [postgresConnection({ save_password: false })]; + await store.disconnect("pg-1", { markClosedByUser: false }); + const api = await import("@/lib/backend/api"); + expect(api.sessionCredentialStatus).not.toHaveBeenCalled(); + expect(api.forgetSessionCredential).not.toHaveBeenCalled(); + }); + + it("ordinary and explicit repeated closes are safe when there is no credential", async () => { + installApiMocks(); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [postgresConnection({ save_password: false })]; + await store.disconnect("pg-1"); + await store.disconnectAndForgetConnectionPassword("pg-1"); + const api = await import("@/lib/backend/api"); + expect(api.forgetSessionCredential).not.toHaveBeenCalled(); + expect(api.disconnectDb).toHaveBeenCalledTimes(2); + }); + + it("coalesces concurrent closes so an old forget cannot erase a newly entered password", async () => { + const forget = deferred(); + installApiMocks({ sessionCredentialStatus: vi.fn().mockResolvedValue(true), forgetSessionCredential: vi.fn(() => forget.promise) }); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [postgresConnection({ save_password: false })]; + const closes = [store.disconnect("pg-1"), store.disconnectAndForgetConnectionPassword("pg-1")]; + await vi.waitFor(async () => { + const api = await import("@/lib/backend/api"); + expect(api.forgetSessionCredential).toHaveBeenCalledTimes(1); + }); + forget.resolve(); + await Promise.all(closes); + const api = await import("@/lib/backend/api"); + expect(api.forgetSessionCredential).toHaveBeenCalledTimes(1); + expect(api.disconnectDb).toHaveBeenCalledTimes(1); + }); + + it.each([false, true])("accepts exact already-gone error only if absence is confirmed: remaining=%s", async (remaining) => { + const error = new Error("Connection has no transient session credential to forget: pg-1"); + installApiMocks({ sessionCredentialStatus: vi.fn().mockResolvedValueOnce(true).mockResolvedValue(remaining), forgetSessionCredential: vi.fn().mockRejectedValue(error) }); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [postgresConnection({ save_password: false })]; + const closing = store.disconnect("pg-1"); + if (remaining) await expect(closing).rejects.toThrow(error.message); + else await closing; + const api = await import("@/lib/backend/api"); + expect(api.disconnectDb).toHaveBeenCalledTimes(1); + expect(api.sessionCredentialStatus).toHaveBeenCalledTimes(2); + }); + + it.each(["status", "forget"])("propagates %s failure, still disconnects and forces fresh input on retry", async (failure) => { + const status = vi.fn().mockResolvedValue(true); + const forget = vi.fn().mockResolvedValue(undefined); + (failure === "status" ? status : forget).mockRejectedValueOnce(new Error("transport failed")); + installApiMocks({ sessionCredentialStatus: status, forgetSessionCredential: forget }); + installPasswordPromptMock(); + requestPassword.mockResolvedValue({ password: "fresh", rememberPassword: false }); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + const config = postgresConnection({ save_password: false }); + store.connections = [config]; + await expect(store.disconnect(config.id)).rejects.toThrow("transport failed"); + const api = await import("@/lib/backend/api"); + expect(api.disconnectDb).toHaveBeenCalledTimes(1); + await store.connect(config); + expect(requestPassword).toHaveBeenCalledTimes(1); + }); + + it("keeps the pool disconnect timeout after retiring the credential", async () => { + vi.useFakeTimers(); + installApiMocks({ sessionCredentialStatus: vi.fn().mockResolvedValue(true), disconnectDb: vi.fn(() => new Promise(() => {})) }); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [postgresConnection({ save_password: false })]; + const closing = store.disconnect("pg-1"); + await vi.advanceTimersByTimeAsync(6_000); + await closing; + const api = await import("@/lib/backend/api"); + expect(api.forgetSessionCredential).toHaveBeenCalledTimes(1); + expect(api.disconnectDb).toHaveBeenCalledTimes(1); + }); + + it.each(["connect", "ensureConnected"] as const)("immediate %s waits beyond the pool timeout for retirement and prompts even with stale status", async (reopen) => { + vi.useFakeTimers(); + const forget = deferred(); + const disconnectDb = vi.fn().mockResolvedValue(undefined); + installApiMocks({ sessionCredentialStatus: vi.fn().mockResolvedValue(true), forgetSessionCredential: vi.fn(() => forget.promise), disconnectDb }); + installPasswordPromptMock(); + requestPassword.mockResolvedValue({ password: "new-input", rememberPassword: false }); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + const config = postgresConnection({ save_password: false }); + store.connections = [config]; + await store.connect(config); + const api = await import("@/lib/backend/api"); + vi.mocked(api.connectDb).mockClear(); + const closing = store.disconnect(config.id); + const reopening = reopen === "connect" ? store.connect(config) : store.ensureConnected(config.id); + await vi.advanceTimersByTimeAsync(6_000); + expect(disconnectDb).not.toHaveBeenCalled(); + expect(api.connectDb).not.toHaveBeenCalled(); + expect(requestPassword).not.toHaveBeenCalled(); + forget.resolve(); + await closing; + await reopening; + expect(requestPassword).toHaveBeenCalledTimes(1); + expect(api.connectDb).toHaveBeenCalledWith(expect.objectContaining({ password: "new-input" }), expect.any(Number)); + }); + it("addConnection blanks the in-memory password when save_password is false", async () => { installApiMocks(); const { useConnectionStore } = await import("@/stores/connectionStore"); diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index e36d504da3..0c0f8d531f 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -709,6 +709,7 @@ export const useConnectionStore = defineStore("connection", () => { }; const connectInFlight = new Map>(); const disconnectInFlight = new Map>(); + const credentialCloseInFlight = new Map>(); const disconnectInFlightScoped = new Map(); const passwordChangeReconnectRequired = new Map(); const cancelDisconnectInFlight = new Map>(); @@ -1117,12 +1118,14 @@ export const useConnectionStore = defineStore("connection", () => { } function getBlockingDisconnectInFlight(connectionId: string): Promise | undefined { + const credentialClose = credentialCloseInFlight.get(connectionId); + if (credentialClose) return credentialClose; return disconnectInFlightScoped.get(connectionId) ? undefined : disconnectInFlight.get(connectionId); } /** True while any disconnect for the connection is in flight (full or scoped). */ function hasDisconnectInFlight(connectionId: string): boolean { - return disconnectInFlight.has(connectionId); + return credentialCloseInFlight.has(connectionId) || disconnectInFlight.has(connectionId); } /** 用户显式关闭过、且此后没有重新连上的连接(见 `userClosedConnectionIds`)。 */ @@ -1170,13 +1173,43 @@ export const useConnectionStore = defineStore("connection", () => { return bounded; } - function startDisconnectRequest(connectionId: string, metadataDrain?: Promise, expectedConfigFingerprint?: string): Promise { + function startDisconnectRequest(connectionId: string, metadataDrain?: Promise, expectedConfigFingerprint?: string, beforeDisconnect?: () => Promise): Promise { + const pendingCredentialClose = credentialCloseInFlight.get(connectionId); + if (pendingCredentialClose) return pendingCredentialClose; const clientAttempt = activeLocalConnectionAttempts.get(connectionId) ?? successfulLocalConnectionAttempts.get(connectionId); const disconnectIfCurrent = () => { const config = getConfig(connectionId); if (expectedConfigFingerprint && (!config || connectionConfigFingerprint(config) !== expectedConfigFingerprint)) return; return api.disconnectDb(connectionId, clientAttempt); }; + if (beforeDisconnect) { + // Register before any asynchronous retirement work. A scoped pool teardown + // may normally overlap reconnect, but a pending forget must never do so. + const request = Promise.resolve().then(async () => { + let retirementFailed = false; + let retirementError: unknown; + try { + await beforeDisconnect(); + } catch (error) { + retirementFailed = true; + retirementError = error; + } + try { + const poolRequest = metadataDrain ? metadataDrain.then(disconnectIfCurrent) : Promise.resolve(disconnectIfCurrent()); + await trackDisconnectRequest(connectionId, poolRequest, clientAttempt != null); + } catch (error) { + if (!retirementFailed) throw error; + } + if (retirementFailed) throw retirementError; + }); + const tracked = request.finally(() => { + if (credentialCloseInFlight.get(connectionId) === tracked) credentialCloseInFlight.delete(connectionId); + }); + // Disconnect also awaits UI persistence; observe an early rejection now. + void tracked.catch(() => undefined); + credentialCloseInFlight.set(connectionId, tracked); + return tracked; + } let request: Promise; try { request = metadataDrain ? metadataDrain.then(disconnectIfCurrent) : Promise.resolve(disconnectIfCurrent()); @@ -4600,6 +4633,17 @@ export const useConnectionStore = defineStore("connection", () => { } } + async function retireSessionCredentialIfPresent(connectionId: string): Promise { + if (!(await api.sessionCredentialStatus(connectionId))) return; + try { + await api.forgetSessionCredential(connectionId); + } catch (error) { + const alreadyGone = `Connection has no transient session credential to forget: ${connectionId}`; + if (connectionErrorMessage(error) === alreadyGone && !(await api.sessionCredentialStatus(connectionId))) return; + throw error; + } + } + /** * Prompt for a transient password when saving is disabled, or when a server * confirms that a metadata-only synced connection sent no password. The @@ -4688,7 +4732,7 @@ export const useConnectionStore = defineStore("connection", () => { const prompted = await ensureConnectionPassword({ ...config, password: "", save_password: false }, true); config = prompted.config; rememberPassword = prompted.rememberPassword; - } else if (connectionNeedsPasswordPrompt(config) && (await pluginConnectionPasswordPromptNeeded(config)) && !(await hasSessionCredential(config.id))) { + } else if (connectionNeedsPasswordPrompt(config) && (await pluginConnectionPasswordPromptNeeded(config)) && (isConnectionClosedByUser(config.id) || !(await hasSessionCredential(config.id)))) { const prompted = await ensureConnectionPassword(config); config = prompted.config; rememberPassword = prompted.rememberPassword; @@ -4882,11 +4926,18 @@ export const useConnectionStore = defineStore("connection", () => { // 重新挂载,也不该在这一刻把连接连回来。 if (options.markClosedByUser !== false) markConnectionClosedByUser(connectionId); const shouldRemoveOneTimeConnection = getConfig(connectionId)?.one_time === true; - if (hasSqlServerActivityTraceForConnection(connectionId)) { + const shouldRetireCredential = options.markClosedByUser !== false && getConfig(connectionId)?.save_password === false; + if (!shouldRetireCredential && hasSqlServerActivityTraceForConnection(connectionId)) { await disposeSqlServerActivityTracesForConnection(connectionId); if (!isExpectedConfigCurrent()) return; } - const disconnectRequest = startDisconnectRequest(connectionId, metadataDrain, options.expectedConfigFingerprint); + const beforeDisconnect = shouldRetireCredential + ? async () => { + await retireSessionCredentialIfPresent(connectionId); + if (hasSqlServerActivityTraceForConnection(connectionId)) await disposeSqlServerActivityTracesForConnection(connectionId); + } + : undefined; + const disconnectRequest = startDisconnectRequest(connectionId, metadataDrain, options.expectedConfigFingerprint, beforeDisconnect); cancelLocalConnectionAttempt(connectionId); connectedIds.value.delete(connectionId); @@ -4956,11 +5007,11 @@ export const useConnectionStore = defineStore("connection", () => { /** * "断开并忘记本次密码":关闭连接池(保留其它已保存连接的凭据),并清除该 * `save_password === false` 连接本次运行期的会话密码,使下一次连接必须重新输入。 - * 若该连接本无会话凭据,后端会返回错误并在此抛出(不静默吞掉),避免误报成功。 + * 普通关闭已清除凭据时可安全重复;真实 API 错误仍向调用方传播。 */ async function disconnectAndForgetConnectionPassword(connectionId: string) { await disconnect(connectionId); - await api.forgetSessionCredential(connectionId); + if (getConfig(connectionId)?.save_password !== false) await retireSessionCredentialIfPresent(connectionId); } async function retirePasswordAfterChange(connectionId: string, expectedConfigFingerprint?: string): Promise<{ status: "retired"; fingerprint: string } | { status: "skipped" }> { @@ -5105,7 +5156,7 @@ export const useConnectionStore = defineStore("connection", () => { // Fast-path the common case (password saved or no password needed) so the // in-flight dedup above keeps its exact microtask cadence; only await the // interactive prompt when the connection actually needs a typed password. - if (connectionNeedsPasswordPrompt(config) && (await pluginConnectionPasswordPromptNeeded(config)) && !(await hasSessionCredential(connectionId))) { + if (connectionNeedsPasswordPrompt(config) && (await pluginConnectionPasswordPromptNeeded(config)) && (isConnectionClosedByUser(connectionId) || !(await hasSessionCredential(connectionId)))) { if (options.allowPasswordPrompt === false) throw new Error(CONNECTION_PASSWORD_REQUIRED_MESSAGE); const prompted = await ensureConnectionPassword(config); config = prompted.config;