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

apps/desktop/desktop_native/core/Cargo.toml

@@ -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-10098/clean-russh-implementation" }
+bitwarden-russh = { git = "https://github.com/bitwarden/bitwarden-russh.git", branch = "km/PM-14863/allow-unlock-on-list" }
 tokio = { version = "=1.40.0", features = ["io-util", "sync", "macros", "net"] }
 tokio-stream = { version = "=0.1.15", features = ["net"] }
 tokio-util = "=0.7.12"

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

@@ -20,6 +20,8 @@ pub struct BitwardenDesktopAgent {
     show_ui_request_tx: tokio::sync::mpsc::Sender<(u32, String)>,
     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
+    needs_unlock: Arc<Mutex<bool>>,
 }
 
 impl BitwardenDesktopAgent {
@@ -46,6 +48,28 @@ impl ssh_agent::Agent for BitwardenDesktopAgent {
         }
         false
     }
+
+    async fn can_list(&self) -> bool {
+        let needs_unlock = self.needs_unlock.lock().await;
+        if !*needs_unlock {
+            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();
+        self.show_ui_request_tx
+            .send((request_id, "".to_string()))
+            .await
+            .expect("Should send request to ui");
+        while let Ok((id, response)) = rx_channel.recv().await {
+            if id == request_id {
+                return response;
+            }
+        }
+        false
+    }
 }
 
 impl BitwardenDesktopAgent {
@@ -65,6 +89,10 @@ 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);
+
         for (key, name, cipher_id) in new_keys.iter() {
             match parse_key_safe(&key) {
                 Ok(private_key) => {
@@ -102,6 +130,16 @@ impl BitwardenDesktopAgent {
             });
         Ok(())
     }
+
+    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);
+
+        Ok(())
+    }
 }
 
 fn parse_key_safe(pem: &str) -> Result<ssh_key::private::PrivateKey, anyhow::Error> {

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

@@ -23,6 +23,7 @@ impl BitwardenDesktopAgent {
             show_ui_request_tx: auth_request_tx,
             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)),
         };
         let cloned_agent_state = agent.clone();
         tokio::spawn(async move {

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

@@ -21,6 +21,7 @@ impl BitwardenDesktopAgent {
             get_ui_response_rx: auth_response_rx,
             cancellation_token: CancellationToken::new(),
             request_id: Arc::new(tokio::sync::Mutex::new(0)),
+            needs_unlock: 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

@@ -74,6 +74,7 @@ export declare namespace sshagent {
   export function setKeys(agentState: SshAgentState, newKeys: Array<PrivateKey>): void
   export function lock(agentState: SshAgentState): void
   export function importKey(encodedKey: string, password: string): SshKeyImportResult
+  export function clearKeys(agentState: SshAgentState): void
   export function generateKeypair(keyAlgorithm: string): Promise<SshKey>
   export class SshAgentState {   }
 }

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

@@ -339,6 +339,12 @@ pub mod sshagent {
         Ok(result.into())
     }
 
+    #[napi]
+    pub fn clear_keys(agent_state: &mut SshAgentState) -> napi::Result<()> {
+        let bitwarden_agent_state = &mut agent_state.state;
+        bitwarden_agent_state.clear_keys().map_err(|e| napi::Error::from_reason(e.to_string()))
+    }
+
     #[napi]
     pub async fn generate_keypair(key_algorithm: String) -> napi::Result<SshKey> {
         desktop_core::ssh_agent::generator::generate_keypair(key_algorithm)

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

@@ -111,5 +111,11 @@
         sshagent.lock(this.agentState);
       }
     });
+
+    ipcMain.handle("sshagent.clearkeys", async (event: any) => {
+      if (this.agentState != null) {
+        sshagent.clearKeys(this.agentState);
+      }
+    });
   }
 }

apps/desktop/src/platform/preload.ts

@@ -56,6 +56,9 @@
   lock: async () => {
     return await ipcRenderer.invoke("sshagent.lock");
   },
+  clearKeys: async () => {
+    return await ipcRenderer.invoke("sshagent.clearkeys");
+  },
   importKey: async (key: string, password: string): Promise<ssh.SshKeyImportResult> => {
     const res = await ipcRenderer.invoke("sshagent.importkey", {
       privateKey: key,

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

@@ -8,6 +8,7 @@
   from,
   map,
   of,
+  skip,
   Subject,
   switchMap,
   takeUntil,
@@ -17,6 +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";
@@ -52,6 +54,7 @@
     private i18nService: I18nService,
     private desktopSettingsService: DesktopSettingsService,
     private configService: ConfigService,
+    private accountService: AccountService,
   ) {}
 
   async init() {
@@ -110,19 +113,39 @@
             ),
           ),
           // This concatMap handles showing the dialog to approve the request.
-          concatMap(([message, decryptedCiphers]) => {
+          concatMap(([message, ciphers]) => {
             const cipherId = message.cipherId as string;
             const requestId = message.requestId as number;
 
-            if (decryptedCiphers === undefined) {
+            // 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 (ciphers === undefined) {
               return of(false).pipe(
                 switchMap((result) =>
                   ipc.platform.sshAgent.signRequestResponse(requestId, Boolean(result)),
                 ),
               );
             }
 
-            const cipher = decryptedCiphers.find((cipher) => cipher.id == cipherId);
+            const cipher = ciphers.find((cipher) => cipher.id == cipherId);
 
             ipc.platform.focusWindow();
             const dialogRef = ApproveSshRequestComponent.open(
@@ -141,14 +164,23 @@
         )
         .subscribe();
 
+      this.accountService.activeAccount$
+        .pipe(skip(1), takeUntil(this.destroy$))
+        .subscribe((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));
+        });
+
       combineLatest([
         timer(0, this.SSH_REFRESH_INTERVAL),
         this.desktopSettingsService.sshAgentEnabled$,
       ])
         .pipe(
           concatMap(async ([, enabled]) => {
             if (!enabled) {
-              await ipc.platform.sshAgent.setKeys([]);
+              await ipc.platform.sshAgent.clearKeys();
               return;
             }