From 8ac282cc37f6cc60128ed65f07a8a27fed94113a Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 28 Jul 2026 12:19:48 +0800 Subject: [PATCH] fix(web): support SSH host key prompts --- .../components/ssh/SshHostKeyPromptDialog.vue | 104 +++++++---- .../__tests__/SshHostKeyPromptDialog.spec.ts | 163 ++++++++++++++++++ .../backend/__tests__/sshPromptApi.spec.ts | 25 +++ apps/desktop/src/lib/backend/api.ts | 1 + apps/desktop/src/lib/backend/http.ts | 5 + apps/desktop/src/lib/backend/tauri.ts | 11 ++ crates/dbx-web/src/main.rs | 6 + crates/dbx-web/src/routes/connection.rs | 1 + crates/dbx-web/src/routes/mod.rs | 1 + crates/dbx-web/src/routes/mongo.rs | 1 + crates/dbx-web/src/routes/mq.rs | 1 + crates/dbx-web/src/routes/ssh_prompt.rs | 59 +++++++ crates/dbx-web/src/ssh_prompt.rs | 162 +++++++++++++++++ crates/dbx-web/src/state.rs | 1 + 14 files changed, 509 insertions(+), 32 deletions(-) create mode 100644 apps/desktop/src/components/ssh/__tests__/SshHostKeyPromptDialog.spec.ts create mode 100644 apps/desktop/src/lib/backend/__tests__/sshPromptApi.spec.ts create mode 100644 crates/dbx-web/src/routes/ssh_prompt.rs create mode 100644 crates/dbx-web/src/ssh_prompt.rs diff --git a/apps/desktop/src/components/ssh/SshHostKeyPromptDialog.vue b/apps/desktop/src/components/ssh/SshHostKeyPromptDialog.vue index 3dd66034f9..94317e3853 100644 --- a/apps/desktop/src/components/ssh/SshHostKeyPromptDialog.vue +++ b/apps/desktop/src/components/ssh/SshHostKeyPromptDialog.vue @@ -5,6 +5,9 @@ import { invoke } from "@tauri-apps/api/core"; import { Dialog, DialogContent, DialogDescription, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { Button } from "@/components/ui/button"; import { useToast } from "@/composables/useToast"; +import { resolveSshPrompt } from "@/lib/backend/api"; +import { isTauriRuntime } from "@/lib/backend/tauriRuntime"; +import { apiUrl } from "@/lib/common/webPath"; interface SshPromptRequest { id: string; @@ -38,15 +41,51 @@ const remember = ref(true); const secretCode = ref(""); const resolving = ref(false); const unlisteners: Array<() => void> = []; +let webEventSource: EventSource | null = null; let mounted = true; -onMounted(async () => { +function enqueuePrompt(request: SshPromptRequest) { + if (queue.value.some((prompt) => prompt.id === request.id)) return; + queue.value.push(request); + visible.value = true; +} + +function handleNotice(n: SshHostKeyNotice) { + let message: string; + if (n.kind === "Changed") { + message = t("connection.sshHostKeyNoticeChanged", { host: n.host, port: n.port }); + } else if (n.kind === "Rejected") { + message = t("connection.sshHostKeyNoticeRejected", { host: n.host, port: n.port }); + } else if (n.kind === "LearnFailed") { + message = t("connection.sshHostKeyNoticeLearnFailed", { host: n.host, port: n.port }); + } else { + message = n.message || t("connection.sshHostKeyNoticeGeneric"); + } + toast(message, 6000); +} + +function handleWebEvent(data: string) { + try { + const event = JSON.parse(data) as { type: "sync" | "prompt" | "notice" | "dismiss"; pendingIds?: string[]; request?: SshPromptRequest; notice?: SshHostKeyNotice; id?: string }; + if (event.type === "sync" && event.pendingIds) { + const currentId = queue.value[0]?.id; + queue.value = queue.value.filter((prompt) => event.pendingIds?.includes(prompt.id)); + if (currentId && !queue.value.some((prompt) => prompt.id === currentId)) resetPromptState(); + visible.value = queue.value.length > 0; + } else if (event.type === "prompt" && event.request) enqueuePrompt(event.request); + else if (event.type === "notice" && event.notice) handleNotice(event.notice); + else if (event.type === "dismiss" && event.id) dismissPrompt(event.id); + } catch (error) { + console.error("[DBX] invalid SSH prompt event:", error); + } +} + +async function setupTauriPromptBridge() { try { const { listen } = await import("@tauri-apps/api/event"); if (!mounted) return; const promptHandle = await listen("ssh-prompt", (event) => { - queue.value.push(event.payload); - visible.value = true; + enqueuePrompt(event.payload); }); if (!mounted) { promptHandle(); @@ -55,18 +94,7 @@ onMounted(async () => { unlisteners.push(promptHandle); const noticeHandle = await listen("ssh-host-key-notice", (event) => { - const n = event.payload; - let message: string; - if (n.kind === "Changed") { - message = t("connection.sshHostKeyNoticeChanged", { host: n.host, port: n.port }); - } else if (n.kind === "Rejected") { - message = t("connection.sshHostKeyNoticeRejected", { host: n.host, port: n.port }); - } else if (n.kind === "LearnFailed") { - message = t("connection.sshHostKeyNoticeLearnFailed", { host: n.host, port: n.port }); - } else { - message = n.message || t("connection.sshHostKeyNoticeGeneric"); - } - toast(message, 6000); + handleNotice(event.payload); }); if (!mounted) { noticeHandle(); @@ -93,14 +121,29 @@ onMounted(async () => { void invoke("ssh_prompt_not_ready").catch(() => undefined); } } catch { - // Not running inside Tauri (e.g. web preview) — no prompts to handle. + // The desktop bridge is unavailable, so the connection remains fail-closed. } +} + +function setupWebPromptBridge() { + webEventSource = new EventSource(apiUrl("/api/ssh/prompts")); + webEventSource.onmessage = (event) => handleWebEvent(event.data); + webEventSource.onerror = () => { + // EventSource automatically retries. Do not close it while the app is mounted. + }; +} + +onMounted(() => { + if (isTauriRuntime()) void setupTauriPromptBridge(); + else setupWebPromptBridge(); }); onBeforeUnmount(() => { mounted = false; unlisteners.forEach((u) => u()); - void invoke("ssh_prompt_not_ready").catch(() => undefined); + webEventSource?.close(); + webEventSource = null; + if (isTauriRuntime()) void import("@tauri-apps/api/core").then(({ invoke }) => invoke("ssh_prompt_not_ready")).catch(() => undefined); }); async function resolve(action: "accept" | "reject" | "secret") { @@ -108,19 +151,14 @@ async function resolve(action: "accept" | "reject" | "secret") { if (!prompt || resolving.value) return; resolving.value = true; try { - await invoke("resolve_ssh_prompt", { - resolution: { - id: prompt.id, - action, - remember: action === "accept" ? remember.value : undefined, - secret: action === "secret" ? secretCode.value : undefined, - }, + await resolveSshPrompt({ + id: prompt.id, + action, + remember: action === "accept" ? remember.value : undefined, + secret: action === "secret" ? secretCode.value : undefined, }); // Advance to the next queued prompt (if any) and reset per-prompt state. - queue.value.shift(); - remember.value = true; - secretCode.value = ""; - if (queue.value.length === 0) visible.value = false; + dismissPrompt(prompt.id); } catch (e) { console.error("[DBX] resolve_ssh_prompt failed:", e); // "No pending" / "already cancelled" mean the backend already dropped this @@ -138,13 +176,15 @@ function dismissPrompt(id: string) { const idx = queue.value.findIndex((p) => p.id === id); if (idx === -1) return; queue.value.splice(idx, 1); - if (idx === 0) { - remember.value = true; - secretCode.value = ""; - } + if (idx === 0) resetPromptState(); if (queue.value.length === 0) visible.value = false; } +function resetPromptState() { + remember.value = true; + secretCode.value = ""; +} + function accept() { void resolve("accept"); } diff --git a/apps/desktop/src/components/ssh/__tests__/SshHostKeyPromptDialog.spec.ts b/apps/desktop/src/components/ssh/__tests__/SshHostKeyPromptDialog.spec.ts new file mode 100644 index 0000000000..661903f44e --- /dev/null +++ b/apps/desktop/src/components/ssh/__tests__/SshHostKeyPromptDialog.spec.ts @@ -0,0 +1,163 @@ +// @vitest-environment happy-dom + +import { createApp, defineComponent, h, nextTick, type App } from "vue"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import i18n from "@/i18n"; + +const { resolveSshPromptMock } = vi.hoisted(() => ({ + resolveSshPromptMock: vi.fn(), +})); + +vi.mock("@/lib/backend/api", () => ({ resolveSshPrompt: resolveSshPromptMock })); +vi.mock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); +vi.mock("@/composables/useToast", () => ({ useToast: () => ({ toast: vi.fn() }) })); + +vi.mock("@/components/ui/dialog", async () => { + const { defineComponent, h } = await import("vue"); + const passthrough = defineComponent({ + setup(_props, { slots }) { + return () => h("div", slots.default?.()); + }, + }); + const dialog = defineComponent({ + props: { open: Boolean }, + setup(props, { slots }) { + return () => (props.open ? h("div", slots.default?.()) : null); + }, + }); + return { + Dialog: dialog, + DialogContent: passthrough, + DialogDescription: passthrough, + DialogFooter: passthrough, + DialogHeader: passthrough, + DialogTitle: passthrough, + }; +}); + +vi.mock("@/components/ui/button", async () => { + const { defineComponent, h } = await import("vue"); + return { + Button: defineComponent({ + inheritAttrs: false, + setup(_props, { attrs, slots }) { + return () => h("button", attrs, slots.default?.()); + }, + }), + }; +}); + +import SshHostKeyPromptDialog from "@/components/ssh/SshHostKeyPromptDialog.vue"; + +class MockEventSource { + static instances: MockEventSource[] = []; + + readonly url: string; + onmessage: ((event: MessageEvent) => void) | null = null; + onerror: (() => void) | null = null; + closed = false; + + constructor(url: string | URL) { + this.url = String(url); + MockEventSource.instances.push(this); + } + + emit(data: unknown) { + this.onmessage?.(new MessageEvent("message", { data: JSON.stringify(data) })); + } + + close() { + this.closed = true; + } +} + +const mountedApps: App[] = []; + +beforeEach(() => { + resolveSshPromptMock.mockReset().mockResolvedValue(undefined); + MockEventSource.instances = []; + vi.stubGlobal("EventSource", MockEventSource as unknown as typeof EventSource); + i18n.global.locale.value = "en"; +}); + +afterEach(() => { + for (const app of mountedApps.splice(0)) app.unmount(); + document.body.innerHTML = ""; + vi.unstubAllGlobals(); +}); + +async function mountDialog() { + const container = document.createElement("div"); + document.body.append(container); + const app = createApp( + defineComponent({ + setup() { + return () => h(SshHostKeyPromptDialog); + }, + }), + ); + mountedApps.push(app); + app.use(i18n); + app.mount(container); + await nextTick(); +} + +describe("SshHostKeyPromptDialog web bridge", () => { + it("shows an SSE host-key prompt and posts the user's acceptance", async () => { + await mountDialog(); + + const eventSource = MockEventSource.instances[0]; + expect(eventSource?.url).toBe("/api/ssh/prompts"); + + eventSource?.emit({ + type: "prompt", + request: { + id: "prompt-1", + kind: "HostKeyVerify", + host: "ssh.example.test", + port: 22, + key_type: "ssh-ed25519", + fingerprint: "SHA256:test-fingerprint", + }, + }); + await nextTick(); + + expect(document.body.textContent).toContain("ssh.example.test:22"); + expect(document.body.textContent).toContain("SHA256:test-fingerprint"); + + const buttons = document.body.querySelectorAll("button"); + buttons.item(buttons.length - 1).click(); + + await vi.waitFor(() => { + expect(resolveSshPromptMock).toHaveBeenCalledWith({ + id: "prompt-1", + action: "accept", + remember: true, + secret: undefined, + }); + }); + }); + + it("clears prompts that are no longer pending after an SSE reconnect snapshot", async () => { + await mountDialog(); + + const eventSource = MockEventSource.instances[0]; + eventSource?.emit({ + type: "prompt", + request: { + id: "prompt-2", + kind: "HostKeyVerify", + host: "stale.example.test", + port: 22, + key_type: "ssh-ed25519", + fingerprint: "SHA256:stale", + }, + }); + await nextTick(); + expect(document.body.textContent).toContain("stale.example.test:22"); + + eventSource?.emit({ type: "sync", pendingIds: [] }); + await nextTick(); + expect(document.body.textContent).not.toContain("stale.example.test:22"); + }); +}); diff --git a/apps/desktop/src/lib/backend/__tests__/sshPromptApi.spec.ts b/apps/desktop/src/lib/backend/__tests__/sshPromptApi.spec.ts new file mode 100644 index 0000000000..b542d4291f --- /dev/null +++ b/apps/desktop/src/lib/backend/__tests__/sshPromptApi.spec.ts @@ -0,0 +1,25 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +afterEach(() => { + vi.unstubAllGlobals(); + vi.resetModules(); +}); + +describe("web SSH prompt API", () => { + it("posts the host-key resolution to the authenticated web endpoint", async () => { + const fetchMock = vi.fn().mockResolvedValue({ + ok: true, + json: vi.fn().mockResolvedValue(null), + }); + vi.stubGlobal("fetch", fetchMock); + + const { resolveSshPrompt } = await import("@/lib/backend/http"); + await resolveSshPrompt({ id: "prompt-1", action: "accept", remember: true }); + + expect(fetchMock).toHaveBeenCalledWith("/api/ssh/prompts/resolve", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ id: "prompt-1", action: "accept", remember: true }), + }); + }); +}); diff --git a/apps/desktop/src/lib/backend/api.ts b/apps/desktop/src/lib/backend/api.ts index 3c825c6c15..22dd670bac 100644 --- a/apps/desktop/src/lib/backend/api.ts +++ b/apps/desktop/src/lib/backend/api.ts @@ -67,6 +67,7 @@ export const loadConnections = forward("loadConnections"); export const loadTunnelProfiles = forward("loadTunnelProfiles"); export const saveTunnelProfiles = forward("saveTunnelProfiles"); export const testTunnelProfile = forward("testTunnelProfile"); +export const resolveSshPrompt = forward("resolveSshPrompt"); export const readKeychainPassword = forward("readKeychainPassword"); export const readKeychainPasswords = forward("readKeychainPasswords"); export const decryptConfig = forward("decryptConfig"); diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index 2125a4a903..7f2fe63009 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -121,6 +121,7 @@ import type { MongoGridFsFileInfo, AppSupportInfo, PromptTemplate, + SshPromptResolution, } from "@/lib/backend/tauri"; import type { QueryEditability } from "@/lib/sql/sqlAnalysis"; import { isTerminalTransferProgress } from "@/lib/backend/transferProgress"; @@ -320,6 +321,10 @@ export async function testTunnelProfile(profile: TunnelProfile): Promise return post("/api/tunnel-profiles/test", profile); } +export async function resolveSshPrompt(resolution: SshPromptResolution): Promise { + await post("/api/ssh/prompts/resolve", resolution); +} + export async function readKeychainPassword(_service: string): Promise { return ""; // Not available in web backend } diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index d232e9a481..fcfdca860a 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -74,6 +74,13 @@ import type { CreateDatabaseSqlOptions } from "@/lib/database/createDatabaseSql" import type { DatabaseNameSqlOptions, DatabasePropertyEditSqlOptions, DropTableChildObjectSqlOptions, DropObjectSqlOptions, DuplicateTableStructureSqlOptions, CopyTableDataSqlOptions, SchemaNameSqlOptions, TableAdminSqlOptions } from "@/lib/database/dbAdminSql"; import type { BuildDatabaseSqlExportOptions, BuildExportInsertStatementsOptions } from "@/lib/export/databaseExport"; +export interface SshPromptResolution { + id: string; + action: "accept" | "reject" | "secret"; + remember?: boolean; + secret?: string; +} + export interface AgentDriverInfo { db_type: string; label: string; @@ -1280,6 +1287,10 @@ export async function testTunnelProfile(profile: TunnelProfile): Promise return invoke("test_tunnel_profile", { profile }); } +export async function resolveSshPrompt(resolution: SshPromptResolution): Promise { + await invoke("resolve_ssh_prompt", { resolution }); +} + export async function readKeychainPassword(service: string): Promise { return invoke("read_keychain_password", { service, account: null }); } diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index cf2359870f..ce858caf3c 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -2,6 +2,7 @@ mod auth; mod error; mod routes; mod sse; +mod ssh_prompt; mod state; use std::collections::{HashMap, HashSet}; @@ -239,8 +240,11 @@ async fn main() { nacos_imports: RwLock::new(HashMap::new()), login_rate_limit: tokio::sync::Mutex::new(state::LoginRateLimit { fail_count: 0, locked_until: None }), export_files: RwLock::new(HashMap::new()), + ssh_prompts: Arc::new(ssh_prompt::SshPromptHub::new()), }); + ssh_prompt::install_web_ssh_prompt_bridge(web_state.ssh_prompts.clone()); + // API routes let api = Router::new() // Auth @@ -283,6 +287,8 @@ async fn main() { // System .route("/system/fonts", get(routes::jdbc::list_system_fonts)) .route("/ssh/config-hosts", get(routes::ssh_config::list_ssh_config_hosts)) + .route("/ssh/prompts", get(routes::ssh_prompt::stream_ssh_prompts)) + .route("/ssh/prompts/resolve", post(routes::ssh_prompt::resolve_ssh_prompt)) // Tunnel profiles .route("/tunnel-profiles/list", get(routes::tunnel_profiles::load_tunnel_profiles)) .route("/tunnel-profiles/save", post(routes::tunnel_profiles::save_tunnel_profiles)) diff --git a/crates/dbx-web/src/routes/connection.rs b/crates/dbx-web/src/routes/connection.rs index 3ddd28aade..8ff8616a17 100644 --- a/crates/dbx-web/src/routes/connection.rs +++ b/crates/dbx-web/src/routes/connection.rs @@ -500,6 +500,7 @@ mod tests { nacos_imports: RwLock::new(HashMap::new()), login_rate_limit: Mutex::new(LoginRateLimit { fail_count: 0, locked_until: None }), export_files: RwLock::new(HashMap::new()), + ssh_prompts: Arc::new(crate::ssh_prompt::SshPromptHub::new()), }); (state, dir) } diff --git a/crates/dbx-web/src/routes/mod.rs b/crates/dbx-web/src/routes/mod.rs index f5b4d0da74..e7b32fdc54 100644 --- a/crates/dbx-web/src/routes/mod.rs +++ b/crates/dbx-web/src/routes/mod.rs @@ -28,6 +28,7 @@ pub mod schema_cache; pub mod schema_diff; pub mod sql_file; pub mod ssh_config; +pub mod ssh_prompt; pub mod tab_runtime_cache; pub mod table_export; pub mod table_import; diff --git a/crates/dbx-web/src/routes/mongo.rs b/crates/dbx-web/src/routes/mongo.rs index 25b044dca5..b885637505 100644 --- a/crates/dbx-web/src/routes/mongo.rs +++ b/crates/dbx-web/src/routes/mongo.rs @@ -840,6 +840,7 @@ mod tests { nacos_imports: RwLock::new(HashMap::new()), login_rate_limit: Mutex::new(LoginRateLimit { fail_count: 0, locked_until: None }), export_files: RwLock::new(HashMap::new()), + ssh_prompts: Arc::new(crate::ssh_prompt::SshPromptHub::new()), }); (state, dir) } diff --git a/crates/dbx-web/src/routes/mq.rs b/crates/dbx-web/src/routes/mq.rs index 6a65c85a89..f78d3c4132 100644 --- a/crates/dbx-web/src/routes/mq.rs +++ b/crates/dbx-web/src/routes/mq.rs @@ -1546,6 +1546,7 @@ mod tests { nacos_imports: RwLock::new(HashMap::new()), login_rate_limit: Mutex::new(LoginRateLimit { fail_count: 0, locked_until: None }), export_files: RwLock::new(HashMap::new()), + ssh_prompts: Arc::new(crate::ssh_prompt::SshPromptHub::new()), }); (state, dir) } diff --git a/crates/dbx-web/src/routes/ssh_prompt.rs b/crates/dbx-web/src/routes/ssh_prompt.rs new file mode 100644 index 0000000000..f847489235 --- /dev/null +++ b/crates/dbx-web/src/routes/ssh_prompt.rs @@ -0,0 +1,59 @@ +use axum::extract::State; +use axum::response::sse::{Event, KeepAlive, Sse}; +use axum::Json; +use dbx_core::db::ssh_prompt::SshPromptAnswer; +use futures::Stream; +use serde::Deserialize; +use std::convert::Infallible; +use std::sync::Arc; + +use crate::error::AppError; +use crate::ssh_prompt::SshPromptEvent; +use crate::state::WebState; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ResolveSshPromptRequest { + pub id: String, + pub action: String, + pub remember: Option, + pub secret: Option, +} + +pub async fn stream_ssh_prompts( + State(state): State>, +) -> Sse>> { + let (replay, mut receiver) = state.ssh_prompts.subscribe(); + let stream = async_stream::stream! { + for event in replay { + yield Ok(ssh_prompt_event(event)); + } + loop { + match receiver.recv().await { + Ok(event) => yield Ok(ssh_prompt_event(event)), + Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => break, + Err(tokio::sync::broadcast::error::RecvError::Closed) => break, + } + } + }; + Sse::new(stream).keep_alive(KeepAlive::default()) +} + +pub async fn resolve_ssh_prompt( + State(state): State>, + Json(request): Json, +) -> Result, AppError> { + let answer = match request.action.as_str() { + "accept" => SshPromptAnswer::Accept { remember: request.remember.unwrap_or(false) }, + "reject" => SshPromptAnswer::Reject, + "secret" => SshPromptAnswer::Secret(request.secret.unwrap_or_default()), + other => return Err(AppError::bad_request(format!("Unknown SSH prompt action: {other}"))), + }; + + state.ssh_prompts.resolve(&request.id, answer).map_err(AppError::bad_request)?; + Ok(Json(())) +} + +fn ssh_prompt_event(event: SshPromptEvent) -> Event { + Event::default().data(serde_json::to_string(&event).expect("SSH prompt events must serialize")) +} diff --git a/crates/dbx-web/src/ssh_prompt.rs b/crates/dbx-web/src/ssh_prompt.rs new file mode 100644 index 0000000000..8f443791cd --- /dev/null +++ b/crates/dbx-web/src/ssh_prompt.rs @@ -0,0 +1,162 @@ +use dbx_core::db::ssh_prompt::{SshHostKeyNotice, SshPromptAnswer, SshPromptEnvelope, SshPromptRequest}; +use serde::Serialize; +use std::sync::{Arc, Mutex}; +use tokio::sync::{broadcast, mpsc}; + +const EVENT_BUFFER_SIZE: usize = 32; +const PROMPT_CHANNEL_SIZE: usize = 16; + +#[derive(Clone, Debug, Serialize)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum SshPromptEvent { + Sync { pending_ids: Vec }, + Prompt { request: SshPromptRequest }, + Notice { notice: SshHostKeyNotice }, + Dismiss { id: String }, +} + +struct PendingPrompt { + responder: tokio::sync::oneshot::Sender, + request: SshPromptRequest, +} + +pub struct SshPromptHub { + pending: Mutex>, + events: broadcast::Sender, +} + +impl SshPromptHub { + pub fn new() -> Self { + let (events, _) = broadcast::channel(EVENT_BUFFER_SIZE); + Self { pending: Mutex::new(Vec::new()), events } + } + + pub fn register(&self, envelope: SshPromptEnvelope) { + let request = envelope.request; + self.pending.lock().unwrap().push(PendingPrompt { responder: envelope.responder, request: request.clone() }); + let _ = self.events.send(SshPromptEvent::Prompt { request }); + } + + pub fn subscribe(&self) -> (Vec, broadcast::Receiver) { + let receiver = self.events.subscribe(); + let pending = self.pending.lock().unwrap(); + let mut replay = vec![SshPromptEvent::Sync { + pending_ids: pending.iter().map(|prompt| prompt.request.id.clone()).collect(), + }]; + replay.extend(pending.iter().map(|prompt| SshPromptEvent::Prompt { request: prompt.request.clone() })); + (replay, receiver) + } + + pub fn resolve(&self, id: &str, answer: SshPromptAnswer) -> Result<(), String> { + let pending = { + let mut pending = self.pending.lock().unwrap(); + let index = pending + .iter() + .position(|prompt| prompt.request.id == id) + .ok_or_else(|| format!("No pending SSH prompt with id {id}"))?; + pending.remove(index) + }; + let result = pending.responder.send(answer); + let _ = self.events.send(SshPromptEvent::Dismiss { id: id.to_string() }); + result.map_err(|_| format!("SSH prompt {id} was already cancelled")) + } + + pub fn notify(&self, notice: SshHostKeyNotice) { + let _ = self.events.send(SshPromptEvent::Notice { notice }); + } + + pub fn sweep_cancelled(&self) { + let cancelled = { + let mut pending = self.pending.lock().unwrap(); + let cancelled = pending + .iter() + .filter(|prompt| prompt.responder.is_closed()) + .map(|prompt| prompt.request.id.clone()) + .collect::>(); + pending.retain(|prompt| !prompt.responder.is_closed()); + cancelled + }; + + for id in cancelled { + let _ = self.events.send(SshPromptEvent::Dismiss { id }); + } + } +} + +pub fn install_web_ssh_prompt_bridge(hub: Arc) { + let (prompt_tx, mut prompt_rx) = mpsc::channel(PROMPT_CHANNEL_SIZE); + dbx_core::db::ssh_prompt::install_ssh_prompt_gateway(prompt_tx); + let prompt_hub = hub.clone(); + tokio::spawn(async move { + while let Some(envelope) = prompt_rx.recv().await { + prompt_hub.register(envelope); + } + }); + + let (notice_tx, mut notice_rx) = mpsc::channel(PROMPT_CHANNEL_SIZE); + dbx_core::db::ssh_prompt::install_ssh_notice_gateway(notice_tx); + let notice_hub = hub.clone(); + tokio::spawn(async move { + while let Some(notice) = notice_rx.recv().await { + notice_hub.notify(notice); + } + }); + + tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(10)); + loop { + interval.tick().await; + hub.sweep_cancelled(); + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use dbx_core::db::ssh_prompt::SshPromptKind; + + fn request(id: &str) -> SshPromptRequest { + SshPromptRequest { + id: id.to_string(), + kind: SshPromptKind::HostKeyVerify, + host: "example.test".to_string(), + port: 22, + key_type: Some("ssh-ed25519".to_string()), + fingerprint: Some("SHA256:test".to_string()), + prompt: None, + } + } + + #[tokio::test] + async fn late_subscriber_replays_pending_prompt_and_resolution() { + let hub = SshPromptHub::new(); + let (responder, receiver) = tokio::sync::oneshot::channel(); + hub.register(SshPromptEnvelope { request: request("prompt-1"), responder }); + + let (replay, mut events) = hub.subscribe(); + assert!( + matches!(replay.as_slice(), [SshPromptEvent::Sync { pending_ids }, SshPromptEvent::Prompt { request }] if pending_ids == &["prompt-1"] && request.id == "prompt-1") + ); + + hub.resolve("prompt-1", SshPromptAnswer::Accept { remember: true }).unwrap(); + assert!(matches!(receiver.await.unwrap(), SshPromptAnswer::Accept { remember: true })); + assert!(matches!(events.recv().await.unwrap(), SshPromptEvent::Dismiss { id } if id == "prompt-1")); + } + + #[tokio::test] + async fn invalid_resolution_keeps_prompt_pending() { + let hub = SshPromptHub::new(); + let (responder, receiver) = tokio::sync::oneshot::channel(); + hub.register(SshPromptEnvelope { request: request("prompt-2"), responder }); + + assert!(hub.resolve("missing", SshPromptAnswer::Reject).is_err()); + let (replay, _) = hub.subscribe(); + assert!( + matches!(replay.as_slice(), [SshPromptEvent::Sync { pending_ids }, SshPromptEvent::Prompt { request }] if pending_ids == &["prompt-2"] && request.id == "prompt-2") + ); + + hub.resolve("prompt-2", SshPromptAnswer::Reject).unwrap(); + assert!(matches!(receiver.await.unwrap(), SshPromptAnswer::Reject)); + } +} diff --git a/crates/dbx-web/src/state.rs b/crates/dbx-web/src/state.rs index a8d8ed956d..c5784b5c83 100644 --- a/crates/dbx-web/src/state.rs +++ b/crates/dbx-web/src/state.rs @@ -35,6 +35,7 @@ pub struct WebState { pub login_rate_limit: Mutex, /// Table export temp files: export_id -> (file_path, format) pub export_files: RwLock>, + pub ssh_prompts: Arc, } impl WebState {