[PM-14863] Force unlock when keys are cleared / on first unlock and fix account switching behavior

apps/desktop/desktop_native/core/Cargo.toml

@@ -10,22 +10,22 @@ default = ["sys"]
 manual_test = []
 
 sys = [
-  "dep:widestring",
-  "dep:windows",
-  "dep:core-foundation",
-  "dep:security-framework",
-  "dep:security-framework-sys",
-  "dep:gio",
-  "dep:libsecret",
-  "dep:zbus",
-  "dep:zbus_polkit",
+    "dep:widestring",
+    "dep:windows",
+    "dep:core-foundation",
+    "dep:security-framework",
+    "dep:security-framework-sys",
+    "dep:gio",
+    "dep:libsecret",
+    "dep:zbus",
+    "dep:zbus_polkit",
 ]
 
 [dependencies]
 aes = "=0.8.4"
 anyhow = "=1.0.93"
 arboard = { version = "=3.4.1", default-features = false, features = [
-  "wayland-data-control",
+    "wayland-data-control",
 ] }
 async-stream = "=0.3.5"
 base64 = "=0.22.1"
@@ -50,7 +50,7 @@ ssh-key = { version = "=0.6.6", default-features = false, features = [
     "rsa",
     "getrandom",
 ] }
-bitwarden-russh = { git = "https://github.com/bitwarden/bitwarden-russh.git", branch = "km/PM-14863/allow-unlock-on-list" }
+bitwarden-russh = { git = "https://github.com/bitwarden/bitwarden-russh.git", rev = "b4e7f2f" }
 tokio = { version = "=1.40.0", features = ["io-util", "sync", "macros", "net"] }
 tokio-stream = { version = "=0.1.15", features = ["net"] }
 tokio-util = { version = "0.7.11", features = ["codec"] }
@@ -64,15 +64,15 @@ ed25519 = { version = "=2.2.3", features = ["pkcs8"] }
 [target.'cfg(windows)'.dependencies]
 widestring = { version = "=1.1.0", optional = true }
 windows = { version = "=0.57.0", features = [
-  "Foundation",
-  "Security_Credentials_UI",
-  "Security_Cryptography",
-  "Storage_Streams",
-  "Win32_Foundation",
-  "Win32_Security_Credentials",
-  "Win32_System_WinRT",
-  "Win32_UI_Input_KeyboardAndMouse",
-  "Win32_UI_WindowsAndMessaging",
+    "Foundation",
+    "Security_Credentials_UI",
+    "Security_Cryptography",
+    "Storage_Streams",
+    "Win32_Foundation",
+    "Win32_Security_Credentials",
+    "Win32_System_WinRT",
+    "Win32_UI_Input_KeyboardAndMouse",
+    "Win32_UI_WindowsAndMessaging",
 ], optional = true }
 
 [target.'cfg(windows)'.dev-dependencies]

apps/desktop/desktop_native/core/src/ssh_agent/mod.rs

@@ -17,19 +17,12 @@ pub mod importer;
 pub struct BitwardenDesktopAgent {
     keystore: ssh_agent::KeyStore,
     cancellation_token: CancellationToken,
-    show_ui_request_tx: tokio::sync::mpsc::Sender<(u32, String)>,
+    show_ui_request_tx: tokio::sync::mpsc::Sender<(u32, (String, bool))>,
     get_ui_response_rx: Arc<Mutex<tokio::sync::broadcast::Receiver<(u32, bool)>>>,
     request_id: Arc<Mutex<u32>>,
-    // before first unlock, or after account switching, listing keys should require an unlock to get a list of public keys
+    /// before first unlock, or after account switching, listing keys should require an unlock to get a list of public keys
     needs_unlock: Arc<Mutex<bool>>,
-}
-
-impl BitwardenDesktopAgent {
-    async fn get_request_id(&self) -> u32 {
-        let mut request_id = self.request_id.lock().await;
-        *request_id += 1;
-        *request_id
-    }
+    is_running: Arc<tokio::sync::Mutex<bool>>,
 }
 
 impl ssh_agent::Agent for BitwardenDesktopAgent {
@@ -42,8 +35,9 @@ impl ssh_agent::Agent for BitwardenDesktopAgent {
         let request_id = self.get_request_id().await;
 
         let mut rx_channel = self.get_ui_response_rx.lock().await.resubscribe();
+        let message = (request_id, (ssh_key.cipher_uuid.clone(), false));
         self.show_ui_request_tx
-            .send((request_id, ssh_key.cipher_uuid.clone()))
+            .send(message)
             .await
             .expect("Should send request to ui");
         while let Ok((id, response)) = rx_channel.recv().await {
@@ -55,17 +49,16 @@ impl ssh_agent::Agent for BitwardenDesktopAgent {
     }
 
     async fn can_list(&self) -> bool {
-        let needs_unlock = self.needs_unlock.lock().await;
-        if !*needs_unlock {
+        if !*self.needs_unlock.lock().await{
             return true;
         }
-        drop(needs_unlock);
 
         let request_id = self.get_request_id().await;
 
         let mut rx_channel = self.get_ui_response_rx.lock().await.resubscribe();
+        let message = (request_id, ("".to_string(), true));
         self.show_ui_request_tx
-            .send((request_id, "".to_string()))
+            .send(message)
             .await
             .expect("Should send request to ui");
         while let Ok((id, response)) = rx_channel.recv().await {
@@ -106,9 +99,7 @@ impl BitwardenDesktopAgent {
         let keystore = &mut self.keystore;
         keystore.0.write().expect("RwLock is not poisoned").clear();
 
-        let mut needs_unlock = self.needs_unlock.blocking_lock();
-        *needs_unlock = false;
-        drop(needs_unlock);
+        *self.needs_unlock.blocking_lock() = false;
 
         for (key, name, cipher_id) in new_keys.iter() {
             match parse_key_safe(&key) {
@@ -157,12 +148,25 @@ impl BitwardenDesktopAgent {
     pub fn clear_keys(&mut self) -> Result<(), anyhow::Error> {
         let keystore = &mut self.keystore;
         keystore.0.write().expect("RwLock is not poisoned").clear();
-        let mut needs_unlock = self.needs_unlock.blocking_lock();
-        *needs_unlock = true;
-        drop(needs_unlock);
+        *self.needs_unlock.blocking_lock() = true;
 
         Ok(())
     }
+
+    async fn get_request_id(&self) -> u32 {
+        if !*self.is_running.lock().await {
+            println!("[BitwardenDesktopAgent] Agent is not running, but tried to get request id");
+            return 0;
+        }
+
+        let mut request_id = self.request_id.lock().await;
+        *request_id += 1;
+        *request_id
+    }
+
+    pub fn is_running(self) -> bool {
+        return self.is_running.blocking_lock().clone();
+    }
 }
 
 fn parse_key_safe(pem: &str) -> Result<ssh_key::private::PrivateKey, anyhow::Error> {

apps/desktop/desktop_native/core/src/ssh_agent/unix.rs

@@ -14,7 +14,7 @@ use super::BitwardenDesktopAgent;
 
 impl BitwardenDesktopAgent {
     pub async fn start_server(
-        auth_request_tx: tokio::sync::mpsc::Sender<(u32, String)>,
+        auth_request_tx: tokio::sync::mpsc::Sender<(u32, (String, bool))>,
         auth_response_rx: Arc<Mutex<tokio::sync::broadcast::Receiver<(u32, bool)>>>,
     ) -> Result<Self, anyhow::Error> {
         let agent = BitwardenDesktopAgent {
@@ -24,6 +24,7 @@ impl BitwardenDesktopAgent {
             get_ui_response_rx: auth_response_rx,
             request_id: Arc::new(tokio::sync::Mutex::new(0)),
             needs_unlock: Arc::new(tokio::sync::Mutex::new(true)),
+            is_running: Arc::new(tokio::sync::Mutex::new(false)),
         };
         let cloned_agent_state = agent.clone();
         tokio::spawn(async move {

apps/desktop/desktop_native/core/src/ssh_agent/windows.rs

@@ -12,7 +12,7 @@ use super::BitwardenDesktopAgent;
 
 impl BitwardenDesktopAgent {
     pub async fn start_server(
-        auth_request_tx: tokio::sync::mpsc::Sender<(u32, String)>,
+        auth_request_tx: tokio::sync::mpsc::Sender<(u32, (String, bool))>,
         auth_response_rx: Arc<Mutex<tokio::sync::broadcast::Receiver<(u32, bool)>>>,
     ) -> Result<Self, anyhow::Error> {
         let agent_state = BitwardenDesktopAgent {
@@ -22,6 +22,7 @@ impl BitwardenDesktopAgent {
             cancellation_token: CancellationToken::new(),
             request_id: Arc::new(tokio::sync::Mutex::new(0)),
             needs_unlock: Arc::new(tokio::sync::Mutex::new(true)),
+            is_running: Arc::new(tokio::sync::Mutex::new(true)),
         };
         let stream = named_pipe_listener_stream::NamedPipeServerStream::new(
             agent_state.cancellation_token.clone(),

apps/desktop/desktop_native/napi/index.d.ts

@@ -69,7 +69,7 @@ export declare namespace sshagent {
     status: SshKeyImportStatus
     sshKey?: SshKey
   }
-  export function serve(callback: (err: Error | null, arg: string) => any): Promise<SshAgentState>
+  export function serve(callback: (err: Error | null, arg0: string, arg1: boolean) => any): Promise<SshAgentState>
   export function stop(agentState: SshAgentState): void
   export function isRunning(agentState: SshAgentState): boolean
   export function setKeys(agentState: SshAgentState, newKeys: Array<PrivateKey>): void

apps/desktop/desktop_native/napi/src/lib.rs

@@ -247,15 +247,15 @@ pub mod sshagent {
 
     #[napi]
     pub async fn serve(
-        callback: ThreadsafeFunction<String, CalleeHandled>,
+        callback: ThreadsafeFunction<(String, bool), CalleeHandled>,
     ) -> napi::Result<SshAgentState> {
-        let (auth_request_tx, mut auth_request_rx) = tokio::sync::mpsc::channel::<(u32, String)>(32);
+        let (auth_request_tx, mut auth_request_rx) = tokio::sync::mpsc::channel::<(u32, (String, bool))>(32);
         let (auth_response_tx, auth_response_rx) = tokio::sync::broadcast::channel::<(u32, bool)>(32);
         let auth_response_tx_arc = Arc::new(Mutex::new(auth_response_tx));
         tokio::spawn(async move {
             let _ = auth_response_rx;
 
-            while let Some((request_id, cipher_uuid)) = auth_request_rx.recv().await {
+            while let Some((request_id, (cipher_uuid, is_list_request))) = auth_request_rx.recv().await {
                 let cloned_request_id = request_id.clone();
                 let cloned_cipher_uuid = cipher_uuid.clone();
                 let cloned_response_tx_arc = auth_response_tx_arc.clone();
@@ -266,7 +266,7 @@ pub mod sshagent {
                     let auth_response_tx_arc = cloned_response_tx_arc;
                     let callback = cloned_callback;
                     let promise_result: Result<Promise<bool>, napi::Error> =
-                        callback.call_async(Ok(cipher_uuid)).await;
+                        callback.call_async(Ok((cipher_uuid, is_list_request))).await;
                     match promise_result {
                         Ok(promise_result) => match promise_result.await {
                             Ok(result) => {

apps/desktop/src/platform/main/main-ssh-agent.service.ts

@@ -27,7 +27,7 @@
   init() {
     // handle sign request passing to UI
     sshagent
-      .serve(async (err: Error, cipherId: string) => {
+      .serve(async (err: Error, cipherId: string, isListRequest: boolean) => {
         // clear all old (> SIGN_TIMEOUT) requests
         this.requestResponses = this.requestResponses.filter(
           (response) => response.timestamp > new Date(Date.now() - this.SIGN_TIMEOUT),
@@ -37,6 +37,7 @@
         const id_for_this_request = this.request_id;
         this.messagingService.send("sshagent.signrequest", {
           cipherId,
+          isListRequest,
           requestId: id_for_this_request,
         });
 
@@ -112,9 +113,9 @@
       }
     });
 
     ipcMain.handle("sshagent.clearkeys", async (event: any) => {
       if (this.agentState != null) {
         sshagent.clearKeys(this.agentState);
       }
     });
   }

apps/desktop/src/platform/services/ssh-agent.service.ts

@@ -18,7 +18,7 @@
   withLatestFrom,
 } from "rxjs";
 
 import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
 import { AuthService } from "@bitwarden/common/auth/abstractions/auth.service";
 import { AuthenticationStatus } from "@bitwarden/common/auth/enums/authentication-status";
 import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
@@ -54,7 +54,7 @@
     private i18nService: I18nService,
     private desktopSettingsService: DesktopSettingsService,
     private configService: ConfigService,
     private accountService: AccountService,
   ) {}
 
   async init() {
@@ -115,39 +115,36 @@
             ),
           ),
           // This concatMap handles showing the dialog to approve the request.
-          concatMap(([message, ciphers]) => {
+          switchMap(([message, ciphers]) => {
             const cipherId = message.cipherId as string;
+            const isListRequest = message.isListRequest as boolean;
             const requestId = message.requestId as number;
 
-            // cipherid is empty has the special meaning of "unlock in order to list public keys"
-            if (cipherId == "") {
-              return of(true).pipe(
-                switchMap(async (result) => {
-                  const sshCiphers = ciphers.filter(
-                    (cipher) => cipher.type === CipherType.SshKey && !cipher.isDeleted,
-                  );
-                  const keys = sshCiphers.map((cipher) => {
-                    return {
-                      name: cipher.name,
-                      privateKey: cipher.sshKey.privateKey,
-                      cipherId: cipher.id,
-                    };
-                  });
-                  await ipc.platform.sshAgent.setKeys(keys);
-                  await ipc.platform.sshAgent.signRequestResponse(requestId, Boolean(result));
-                }),
-              );
+            if (isListRequest) {
+              (async () => {
+                const sshCiphers = ciphers.filter(
+                  (cipher) => cipher.type === CipherType.SshKey && !cipher.isDeleted,
+                );
+                const keys = sshCiphers.map((cipher) => {
+                  return {
+                    name: cipher.name,
+                    privateKey: cipher.sshKey.privateKey,
+                    cipherId: cipher.id,
+                  };
+                });
+                await ipc.platform.sshAgent.setKeys(keys);
+                await ipc.platform.sshAgent.signRequestResponse(requestId, true);
+              })().catch((e) => this.logService.error("Failed to respond to SSH request", e));
+              return;
             }
 
             if (ciphers === undefined) {
-              return of(false).pipe(
-                switchMap((result) =>
-                  ipc.platform.sshAgent.signRequestResponse(requestId, Boolean(result)),
-                ),
-              );
+              ipc.platform.sshAgent
+                .signRequestResponse(requestId, false)
+                .catch((e) => this.logService.error("Failed to respond to SSH request", e));
             }
 
             const cipher = ciphers.find((cipher) => cipher.id == cipherId);
 
             ipc.platform.focusWindow();
             const dialogRef = ApproveSshRequestComponent.open(
@@ -158,22 +155,34 @@
 
             return dialogRef.closed.pipe(
               switchMap((result) => {
-                return ipc.platform.sshAgent.signRequestResponse(requestId, Boolean(result));
+                return ipc.platform.sshAgent.signRequestResponse(requestId, result);
               }),
             );
           }),
           takeUntil(this.destroy$),
         )
         .subscribe();
 
-      this.accountService.activeAccount$
-        .pipe(skip(1), takeUntil(this.destroy$))
-        .subscribe((account) => {
+      this.accountService.activeAccount$.pipe(skip(1), takeUntil(this.destroy$)).subscribe({
+        next: (account) => {
           this.logService.info("Active account changed, clearing SSH keys");
           ipc.platform.sshAgent
             .clearKeys()
             .catch((e) => this.logService.error("Failed to clear SSH keys", e));
-        });
+        },
+        error: (e: unknown) => {
+          this.logService.error("Error in active account observable", e);
+          ipc.platform.sshAgent
+            .clearKeys()
+            .catch((e) => this.logService.error("Failed to clear SSH keys", e));
+        },
+        complete: () => {
+          this.logService.info("Active account observable completed, clearing SSH keys");
+          ipc.platform.sshAgent
+            .clearKeys()
+            .catch((e) => this.logService.error("Failed to clear SSH keys", e));
+        },
+      });
 
       combineLatest([
         timer(0, this.SSH_REFRESH_INTERVAL),
@@ -182,7 +191,7 @@
         .pipe(
           concatMap(async ([, enabled]) => {
             if (!enabled) {
               await ipc.platform.sshAgent.clearKeys();
               return;
             }