From 6af77ec0d92163685292085a143284dfffd75fd3 Mon Sep 17 00:00:00 2001 From: "zed-zippy[bot]" <234243425+zed-zippy[bot]@users.noreply.github.com> Date: Mon, 29 Dec 2025 20:33:18 +0000 Subject: [PATCH] Fix worktree trust handling of multiple projects on the same remote host (#45834) (cherry-pick to preview) (#45839) Cherry-pick of #45834 to preview ---- Closes https://github.com/zed-industries/zed/issues/45630 Remote host location alone is not enough to distinguish between remote worktrees: different remote projects open in different windows will have the same remote host location and _will_ have the same `WorktreeId`. Thus, require an associated `WorktreeStore` with all `WorktreeId`-related trust questions, and store those IDs based on the store key. Release Notes: - Fixed worktree trust handling of multiple projects on the same remote host Co-authored-by: Kirill Bulatov --- .../remote_editing_collaboration_tests.rs | 44 +- crates/editor/src/editor_tests.rs | 12 +- crates/git_ui/src/worktree_picker.rs | 14 +- crates/project/src/lsp_store.rs | 2 +- crates/project/src/project.rs | 21 +- crates/project/src/project_settings.rs | 2 +- crates/project/src/trusted_worktrees.rs | 666 ++++++++++-------- crates/remote_server/src/headless_project.rs | 14 +- crates/remote_server/src/unix.rs | 4 +- crates/rpc/src/proto_client.rs | 11 +- crates/workspace/src/persistence.rs | 31 +- crates/workspace/src/security_modal.rs | 8 +- crates/workspace/src/workspace.rs | 51 +- crates/zed/src/main.rs | 4 +- 14 files changed, 486 insertions(+), 398 deletions(-) diff --git a/crates/collab/src/tests/remote_editing_collaboration_tests.rs b/crates/collab/src/tests/remote_editing_collaboration_tests.rs index 5342b0bbd4b11afb24ccbaa6d4bf17df036cec76..e022c0a8e77b513bb165ca3595104ed2e5af47b7 100644 --- a/crates/collab/src/tests/remote_editing_collaboration_tests.rs +++ b/crates/collab/src/tests/remote_editing_collaboration_tests.rs @@ -855,8 +855,6 @@ async fn test_slow_adapter_startup_retries( #[gpui::test] async fn test_ssh_remote_worktree_trust(cx_a: &mut TestAppContext, server_cx: &mut TestAppContext) { - use project::trusted_worktrees::RemoteHostLocation; - cx_a.update(|cx| { release_channel::init(semver::Version::new(0, 0, 0), cx); project::trusted_worktrees::init(HashMap::default(), None, None, cx); @@ -991,23 +989,19 @@ async fn test_ssh_remote_worktree_trust(cx_a: &mut TestAppContext, server_cx: &m }); assert_eq!(worktree_ids.len(), 2); - let remote_host = project_a.read_with(cx_a, |project, cx| { - project - .remote_connection_options(cx) - .map(RemoteHostLocation::from) - }); - let trusted_worktrees = cx_a.update(|cx| TrustedWorktrees::try_get_global(cx).expect("trust global should exist")); + let worktree_store = project_a.read_with(cx_a, |project, _| project.worktree_store()); - let can_trust_a = - trusted_worktrees.update(cx_a, |store, cx| store.can_trust(worktree_ids[0], cx)); - let can_trust_b = - trusted_worktrees.update(cx_a, |store, cx| store.can_trust(worktree_ids[1], cx)); + let can_trust_a = trusted_worktrees.update(cx_a, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[0], cx) + }); + let can_trust_b = trusted_worktrees.update(cx_a, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[1], cx) + }); assert!(!can_trust_a, "project_a should be restricted initially"); assert!(!can_trust_b, "project_b should be restricted initially"); - let worktree_store = project_a.read_with(cx_a, |project, _| project.worktree_store()); let has_restricted = trusted_worktrees.read_with(cx_a, |store, cx| { store.has_restricted_worktrees(&worktree_store, cx) }); @@ -1054,8 +1048,8 @@ async fn test_ssh_remote_worktree_trust(cx_a: &mut TestAppContext, server_cx: &m trusted_worktrees.update(cx_a, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::Worktree(worktree_ids[0])]), - remote_host.clone(), cx, ); }); @@ -1080,25 +1074,29 @@ async fn test_ssh_remote_worktree_trust(cx_a: &mut TestAppContext, server_cx: &m "inlay hints should be queried after trust approval" ); - let can_trust_a = - trusted_worktrees.update(cx_a, |store, cx| store.can_trust(worktree_ids[0], cx)); - let can_trust_b = - trusted_worktrees.update(cx_a, |store, cx| store.can_trust(worktree_ids[1], cx)); + let can_trust_a = trusted_worktrees.update(cx_a, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[0], cx) + }); + let can_trust_b = trusted_worktrees.update(cx_a, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[1], cx) + }); assert!(can_trust_a, "project_a should be trusted after trust()"); assert!(!can_trust_b, "project_b should still be restricted"); trusted_worktrees.update(cx_a, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::Worktree(worktree_ids[1])]), - remote_host.clone(), cx, ); }); - let can_trust_a = - trusted_worktrees.update(cx_a, |store, cx| store.can_trust(worktree_ids[0], cx)); - let can_trust_b = - trusted_worktrees.update(cx_a, |store, cx| store.can_trust(worktree_ids[1], cx)); + let can_trust_a = trusted_worktrees.update(cx_a, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[0], cx) + }); + let can_trust_b = trusted_worktrees.update(cx_a, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[1], cx) + }); assert!(can_trust_a, "project_a should remain trusted"); assert!(can_trust_b, "project_b should now be trusted"); diff --git a/crates/editor/src/editor_tests.rs b/crates/editor/src/editor_tests.rs index 8548b28d971ba48a17d7a4ac7d5fc522392b75a2..97bf380a47187c7e5edc978bfd34d3e801b7eb5c 100644 --- a/crates/editor/src/editor_tests.rs +++ b/crates/editor/src/editor_tests.rs @@ -29419,11 +29419,14 @@ async fn test_local_worktree_trust(cx: &mut TestAppContext) { .map(|wt| wt.read(cx).id()) .expect("should have a worktree") }); + let worktree_store = project.read_with(cx, |project, _| project.worktree_store()); let trusted_worktrees = cx.update(|cx| TrustedWorktrees::try_get_global(cx).expect("trust global should exist")); - let can_trust = trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!(!can_trust, "worktree should be restricted initially"); let buffer_before_approval = project @@ -29469,8 +29472,8 @@ async fn test_local_worktree_trust(cx: &mut TestAppContext) { trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, std::collections::HashSet::from_iter([PathTrust::Worktree(worktree_id)]), - None, cx, ); }); @@ -29497,7 +29500,8 @@ async fn test_local_worktree_trust(cx: &mut TestAppContext) { "inlay hints should be queried after trust approval" ); - let can_trust_after = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust_after = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!(can_trust_after, "worktree should be trusted after trust()"); } diff --git a/crates/git_ui/src/worktree_picker.rs b/crates/git_ui/src/worktree_picker.rs index fef5e16c80ddd26ae6dd0b2a5c0ad1d8e5b21b2c..40ec67b3b9b92268a56acc892f8cbfb46e6209ec 100644 --- a/crates/git_ui/src/worktree_picker.rs +++ b/crates/git_ui/src/worktree_picker.rs @@ -13,7 +13,7 @@ use picker::{Picker, PickerDelegate, PickerEditorPosition}; use project::{ DirectoryLister, git_store::Repository, - trusted_worktrees::{PathTrust, RemoteHostLocation, TrustedWorktrees}, + trusted_worktrees::{PathTrust, TrustedWorktrees}, }; use recent_projects::{RemoteConnectionModal, connect}; use remote::{RemoteConnectionOptions, remote_client::ConnectionIdentifier}; @@ -271,16 +271,18 @@ impl WorktreeListDelegate { if let Some((parent_worktree, _)) = project.read(cx).find_worktree(repo_path, cx) { + let worktree_store = project.read(cx).worktree_store(); trusted_worktrees.update(cx, |trusted_worktrees, cx| { - if trusted_worktrees.can_trust(parent_worktree.read(cx).id(), cx) { + if trusted_worktrees.can_trust( + &worktree_store, + parent_worktree.read(cx).id(), + cx, + ) { trusted_worktrees.trust( + &worktree_store, HashSet::from_iter([PathTrust::AbsPath( new_worktree_path.clone(), )]), - project - .read(cx) - .remote_connection_options(cx) - .map(RemoteHostLocation::from), cx, ); } diff --git a/crates/project/src/lsp_store.rs b/crates/project/src/lsp_store.rs index 2ea3dbf70fcb4359f3f5985cc6cd3bb4db7df009..92f1e0eb35691d9fd4695f0496d18f80647490cd 100644 --- a/crates/project/src/lsp_store.rs +++ b/crates/project/src/lsp_store.rs @@ -386,7 +386,7 @@ impl LocalLspStore { let untrusted_worktree_task = TrustedWorktrees::try_get_global(cx).and_then(|trusted_worktrees| { let can_trust = trusted_worktrees.update(cx, |trusted_worktrees, cx| { - trusted_worktrees.can_trust(worktree_id, cx) + trusted_worktrees.can_trust(&self.worktree_store, worktree_id, cx) }); if can_trust { self.restricted_worktrees_tasks.remove(&worktree_id); diff --git a/crates/project/src/project.rs b/crates/project/src/project.rs index 12c56f35f594a46fb35613c2432bf6519759389a..615146e896df377d51d0fa5e1837f1f3880eec02 100644 --- a/crates/project/src/project.rs +++ b/crates/project/src/project.rs @@ -52,7 +52,9 @@ pub use project_search::Search; use anyhow::{Context as _, Result, anyhow}; use buffer_store::{BufferStore, BufferStoreEvent}; -use client::{Client, Collaborator, PendingEntitySubscription, TypedEnvelope, UserStore, proto}; +use client::{ + Client, Collaborator, PendingEntitySubscription, ProjectId, TypedEnvelope, UserStore, proto, +}; use clock::ReplicaId; use dap::client::DebugAdapterClient; @@ -1295,7 +1297,7 @@ impl Project { worktree_store.clone(), Some(RemoteHostLocation::from(connection_options)), None, - Some((remote_proto.clone(), REMOTE_SERVER_PROJECT_ID)), + Some((remote_proto.clone(), ProjectId(REMOTE_SERVER_PROJECT_ID))), cx, ); } @@ -4814,7 +4816,7 @@ impl Project { let worktree_id = WorktreeId::from_proto(envelope.payload.worktree_id); if let Some(trusted_worktrees) = TrustedWorktrees::try_get_global(cx) { trusted_worktrees.update(cx, |trusted_worktrees, cx| { - trusted_worktrees.can_trust(worktree_id, cx) + trusted_worktrees.can_trust(&project.worktree_store, worktree_id, cx) }); } if let Some(worktree) = project.worktree_for_id(worktree_id, cx) { @@ -4853,18 +4855,14 @@ impl Project { .update(|cx| TrustedWorktrees::try_get_global(cx))? .context("missing trusted worktrees")?; trusted_worktrees.update(&mut cx, |trusted_worktrees, cx| { - let remote_host = this - .read(cx) - .remote_connection_options(cx) - .map(RemoteHostLocation::from); trusted_worktrees.trust( + &this.read(cx).worktree_store(), envelope .payload .trusted_paths .into_iter() .filter_map(|proto_path| PathTrust::from_proto(proto_path)) .collect(), - remote_host, cx, ); })?; @@ -4880,6 +4878,7 @@ impl Project { .update(|cx| TrustedWorktrees::try_get_global(cx))? .context("missing trusted worktrees")?; trusted_worktrees.update(&mut cx, |trusted_worktrees, cx| { + let worktree_store = this.read(cx).worktree_store().downgrade(); let restricted_paths = envelope .payload .worktree_ids @@ -4887,11 +4886,7 @@ impl Project { .map(WorktreeId::from_proto) .map(PathTrust::Worktree) .collect::>(); - let remote_host = this - .read(cx) - .remote_connection_options(cx) - .map(RemoteHostLocation::from); - trusted_worktrees.restrict(restricted_paths, remote_host, cx); + trusted_worktrees.restrict(worktree_store, restricted_paths, cx); })?; Ok(proto::Ack {}) } diff --git a/crates/project/src/project_settings.rs b/crates/project/src/project_settings.rs index 6d95411681d5d350271e7071b752f27d0807f60d..ea028d8ae7b683d453ab54afeac51e6e37db0dba 100644 --- a/crates/project/src/project_settings.rs +++ b/crates/project/src/project_settings.rs @@ -1046,7 +1046,7 @@ impl SettingsObserver { if *can_trust_worktree.get_or_init(|| { if let Some(trusted_worktrees) = TrustedWorktrees::try_get_global(cx) { trusted_worktrees.update(cx, |trusted_worktrees, cx| { - trusted_worktrees.can_trust(worktree_id, cx) + trusted_worktrees.can_trust(&self.worktree_store, worktree_id, cx) }) } else { true diff --git a/crates/project/src/trusted_worktrees.rs b/crates/project/src/trusted_worktrees.rs index 5eb64498a426ed435629d3507189df9b33d5d015..45d060a75070381629119a6363fd66be01482f9a 100644 --- a/crates/project/src/trusted_worktrees.rs +++ b/crates/project/src/trusted_worktrees.rs @@ -39,8 +39,11 @@ //! To ease trusting multiple directory worktrees at once, it's possible to trust a parent directory of a certain directory worktree opened in Zed. //! Trusting a directory means trusting all its subdirectories as well, including all current and potential directory worktrees. +use client::ProjectId; use collections::{HashMap, HashSet}; -use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Global, SharedString, WeakEntity}; +use gpui::{ + App, AppContext as _, Context, Entity, EventEmitter, Global, SharedString, Task, WeakEntity, +}; use remote::RemoteConnectionOptions; use rpc::{AnyProtoClient, proto}; use settings::{Settings as _, WorktreeId}; @@ -53,20 +56,14 @@ use util::debug_panic; use crate::{project_settings::ProjectSettings, worktree_store::WorktreeStore}; pub fn init( - db_trusted_paths: TrustedPaths, - downstream_client: Option<(AnyProtoClient, u64)>, - upstream_client: Option<(AnyProtoClient, u64)>, + db_trusted_paths: DbTrustedPaths, + downstream_client: Option<(AnyProtoClient, ProjectId)>, + upstream_client: Option<(AnyProtoClient, ProjectId)>, cx: &mut App, ) { if TrustedWorktrees::try_get_global(cx).is_none() { let trusted_worktrees = cx.new(|_| { - TrustedWorktreesStore::new( - db_trusted_paths, - None, - None, - downstream_client, - upstream_client, - ) + TrustedWorktreesStore::new(db_trusted_paths, downstream_client, upstream_client) }); cx.set_global(TrustedWorktrees(trusted_worktrees)) } @@ -76,38 +73,36 @@ pub fn init( pub fn track_worktree_trust( worktree_store: Entity, remote_host: Option, - downstream_client: Option<(AnyProtoClient, u64)>, - upstream_client: Option<(AnyProtoClient, u64)>, + downstream_client: Option<(AnyProtoClient, ProjectId)>, + upstream_client: Option<(AnyProtoClient, ProjectId)>, cx: &mut App, ) { match TrustedWorktrees::try_get_global(cx) { Some(trusted_worktrees) => { trusted_worktrees.update(cx, |trusted_worktrees, cx| { - let sync_upstream = trusted_worktrees.upstream_client.as_ref().map(|(_, id)| id) - != upstream_client.as_ref().map(|(_, id)| id); - trusted_worktrees.downstream_client = downstream_client; - trusted_worktrees.upstream_client = upstream_client; + if let Some(downstream_client) = downstream_client { + trusted_worktrees.downstream_clients.push(downstream_client); + } + if let Some(upstream_client) = upstream_client.clone() { + trusted_worktrees.upstream_clients.push(upstream_client); + } trusted_worktrees.add_worktree_store(worktree_store, remote_host, cx); - if sync_upstream { - if let Some((upstream_client, upstream_project_id)) = - &trusted_worktrees.upstream_client - { - let trusted_paths = trusted_worktrees - .trusted_paths - .iter() - .flat_map(|(_, paths)| { - paths.iter().map(|trusted_path| trusted_path.to_proto()) + if let Some((upstream_client, upstream_project_id)) = upstream_client { + let trusted_paths = trusted_worktrees + .trusted_paths + .iter() + .flat_map(|(_, paths)| { + paths.iter().map(|trusted_path| trusted_path.to_proto()) + }) + .collect::>(); + if !trusted_paths.is_empty() { + upstream_client + .send(proto::TrustWorktrees { + project_id: upstream_project_id.0, + trusted_paths, }) - .collect::>(); - if !trusted_paths.is_empty() { - upstream_client - .send(proto::TrustWorktrees { - project_id: *upstream_project_id, - trusted_paths, - }) - .ok(); - } + .ok(); } } }); @@ -132,12 +127,15 @@ impl TrustedWorktrees { /// /// Emits an event each time the worktree was checked and found not trusted, /// or a certain worktree had been trusted. +#[derive(Debug)] pub struct TrustedWorktreesStore { - downstream_client: Option<(AnyProtoClient, u64)>, - upstream_client: Option<(AnyProtoClient, u64)>, + downstream_clients: Vec<(AnyProtoClient, ProjectId)>, + upstream_clients: Vec<(AnyProtoClient, ProjectId)>, worktree_stores: HashMap, Option>, + db_trusted_paths: DbTrustedPaths, trusted_paths: TrustedPaths, - restricted: HashSet, + restricted: HashMap, HashSet>, + worktree_trust_serialization: Task<()>, } /// An identifier of a host to split the trust questions by. @@ -165,7 +163,7 @@ impl From for RemoteHostLocation { SharedString::new(docker_connection_options.container_id), ), }; - RemoteHostLocation { + Self { user_name, host_identifier: host_name, } @@ -214,48 +212,50 @@ impl PathTrust { /// A change of trust on a certain host. #[derive(Debug)] pub enum TrustedWorktreesEvent { - Trusted(Option, HashSet), - Restricted(Option, HashSet), + Trusted(WeakEntity, HashSet), + Restricted(WeakEntity, HashSet), } impl EventEmitter for TrustedWorktreesStore {} -pub type TrustedPaths = HashMap, HashSet>; +type TrustedPaths = HashMap, HashSet>; +pub type DbTrustedPaths = HashMap, HashSet>; impl TrustedWorktreesStore { fn new( - trusted_paths: TrustedPaths, - worktree_store: Option>, - remote_host: Option, - downstream_client: Option<(AnyProtoClient, u64)>, - upstream_client: Option<(AnyProtoClient, u64)>, + db_trusted_paths: DbTrustedPaths, + downstream_client: Option<(AnyProtoClient, ProjectId)>, + upstream_client: Option<(AnyProtoClient, ProjectId)>, ) -> Self { if let Some((upstream_client, upstream_project_id)) = &upstream_client { - let trusted_paths = trusted_paths + let trusted_paths = db_trusted_paths .iter() - .flat_map(|(_, paths)| paths.iter().map(|trusted_path| trusted_path.to_proto())) + .flat_map(|(_, paths)| { + paths + .iter() + .cloned() + .map(PathTrust::AbsPath) + .map(|trusted_path| trusted_path.to_proto()) + }) .collect::>(); if !trusted_paths.is_empty() { upstream_client .send(proto::TrustWorktrees { - project_id: *upstream_project_id, + project_id: upstream_project_id.0, trusted_paths, }) .ok(); } } - let worktree_stores = match worktree_store { - Some(worktree_store) => HashMap::from_iter([(worktree_store.downgrade(), remote_host)]), - None => HashMap::default(), - }; - Self { - trusted_paths, - downstream_client, - upstream_client, - restricted: HashSet::default(), - worktree_stores, + db_trusted_paths, + downstream_clients: downstream_client.into_iter().collect(), + upstream_clients: upstream_client.into_iter().collect(), + trusted_paths: HashMap::default(), + worktree_stores: HashMap::default(), + restricted: HashMap::default(), + worktree_trust_serialization: Task::ready(()), } } @@ -265,13 +265,15 @@ impl TrustedWorktreesStore { worktree_store: &Entity, cx: &App, ) -> bool { - self.worktree_stores - .contains_key(&worktree_store.downgrade()) - && self.restricted.iter().any(|restricted_worktree| { - worktree_store - .read(cx) - .worktree_for_id(*restricted_worktree, cx) - .is_some() + self.restricted + .get(&worktree_store.downgrade()) + .is_some_and(|restricted_worktrees| { + restricted_worktrees.iter().any(|restricted_worktree| { + worktree_store + .read(cx) + .worktree_for_id(*restricted_worktree, cx) + .is_some() + }) }) } @@ -280,40 +282,55 @@ impl TrustedWorktreesStore { /// and the ones that got auto trusted based on trust hierarchy (see module-level docs). pub fn trust( &mut self, + worktree_store: &Entity, mut trusted_paths: HashSet, - remote_host: Option, cx: &mut Context, ) { + let weak_worktree_store = worktree_store.downgrade(); let mut new_trusted_single_file_worktrees = HashSet::default(); let mut new_trusted_other_worktrees = HashSet::default(); let mut new_trusted_abs_paths = HashSet::default(); for trusted_path in trusted_paths.iter().chain( self.trusted_paths - .remove(&remote_host) + .remove(&weak_worktree_store) .iter() .flat_map(|current_trusted| current_trusted.iter()), ) { match trusted_path { PathTrust::Worktree(worktree_id) => { - self.restricted.remove(worktree_id); - if let Some((abs_path, is_file, host)) = - self.find_worktree_data(*worktree_id, cx) + if let Some(restricted_worktrees) = + self.restricted.get_mut(&weak_worktree_store) + { + restricted_worktrees.remove(worktree_id); + }; + + if let Some(worktree) = + worktree_store.read(cx).worktree_for_id(*worktree_id, cx) { - if host == remote_host { - if is_file { - new_trusted_single_file_worktrees.insert(*worktree_id); - } else { - new_trusted_other_worktrees.insert((abs_path, *worktree_id)); - } + if worktree.read(cx).is_single_file() { + new_trusted_single_file_worktrees.insert(*worktree_id); + } else { + new_trusted_other_worktrees + .insert((worktree.read(cx).abs_path(), *worktree_id)); } } } - PathTrust::AbsPath(path) => { + PathTrust::AbsPath(abs_path) => { debug_assert!( - path.is_absolute(), - "Cannot trust non-absolute path {path:?}" + abs_path.is_absolute(), + "Cannot trust non-absolute path {abs_path:?}" ); - new_trusted_abs_paths.insert(path.clone()); + if let Some((worktree_id, is_file)) = + find_worktree_in_store(worktree_store.read(cx), abs_path, cx) + { + if is_file { + new_trusted_single_file_worktrees.insert(worktree_id); + } else { + new_trusted_other_worktrees + .insert((Arc::from(abs_path.as_path()), worktree_id)); + } + } + new_trusted_abs_paths.insert(abs_path.clone()); } } } @@ -326,37 +343,45 @@ impl TrustedWorktreesStore { if !new_trusted_other_worktrees.is_empty() { new_trusted_single_file_worktrees.clear(); } - self.restricted = std::mem::take(&mut self.restricted) - .into_iter() - .filter(|restricted_worktree| { - let Some((restricted_worktree_path, is_file, restricted_host)) = - self.find_worktree_data(*restricted_worktree, cx) - else { - return false; - }; - if restricted_host != remote_host { - return true; - } - // When trusting an abs path on the host, we transitively trust all single file worktrees on this host too. - if is_file && !new_trusted_abs_paths.is_empty() { - trusted_paths.insert(PathTrust::Worktree(*restricted_worktree)); - return false; - } + if let Some(restricted_worktrees) = self.restricted.remove(&weak_worktree_store) { + let new_restricted_worktrees = restricted_worktrees + .into_iter() + .filter(|restricted_worktree| { + let Some(worktree) = worktree_store + .read(cx) + .worktree_for_id(*restricted_worktree, cx) + else { + return false; + }; + let is_file = worktree.read(cx).is_single_file(); + + // When trusting an abs path on the host, we transitively trust all single file worktrees on this host too. + if is_file && !new_trusted_abs_paths.is_empty() { + trusted_paths.insert(PathTrust::Worktree(*restricted_worktree)); + return false; + } - let retain = (!is_file || new_trusted_other_worktrees.is_empty()) - && new_trusted_abs_paths.iter().all(|new_trusted_path| { - !restricted_worktree_path.starts_with(new_trusted_path) - }); - if !retain { - trusted_paths.insert(PathTrust::Worktree(*restricted_worktree)); - } - retain - }) - .collect(); + let restricted_worktree_path = worktree.read(cx).abs_path(); + let retain = (!is_file || new_trusted_other_worktrees.is_empty()) + && new_trusted_abs_paths.iter().all(|new_trusted_path| { + !restricted_worktree_path.starts_with(new_trusted_path) + }); + if !retain { + trusted_paths.insert(PathTrust::Worktree(*restricted_worktree)); + } + retain + }) + .collect(); + self.restricted + .insert(weak_worktree_store.clone(), new_restricted_worktrees); + } { - let trusted_paths = self.trusted_paths.entry(remote_host.clone()).or_default(); + let trusted_paths = self + .trusted_paths + .entry(weak_worktree_store.clone()) + .or_default(); trusted_paths.extend(new_trusted_abs_paths.into_iter().map(PathTrust::AbsPath)); trusted_paths.extend( new_trusted_other_worktrees @@ -371,11 +396,11 @@ impl TrustedWorktreesStore { } cx.emit(TrustedWorktreesEvent::Trusted( - remote_host, + weak_worktree_store, trusted_paths.clone(), )); - if let Some((upstream_client, upstream_project_id)) = &self.upstream_client { + for (upstream_client, upstream_project_id) in &self.upstream_clients { let trusted_paths = trusted_paths .iter() .map(|trusted_path| trusted_path.to_proto()) @@ -383,7 +408,7 @@ impl TrustedWorktreesStore { if !trusted_paths.is_empty() { upstream_client .send(proto::TrustWorktrees { - project_id: *upstream_project_id, + project_id: upstream_project_id.0, trusted_paths, }) .ok(); @@ -395,63 +420,81 @@ impl TrustedWorktreesStore { /// This will emit [`TrustedWorktreesEvent::Restricted`] event for all passed entries. pub fn restrict( &mut self, + worktree_store: WeakEntity, restricted_paths: HashSet, - remote_host: Option, cx: &mut Context, ) { + let mut restricted = HashSet::default(); for restricted_path in restricted_paths { match restricted_path { PathTrust::Worktree(worktree_id) => { - self.restricted.insert(worktree_id); - cx.emit(TrustedWorktreesEvent::Restricted( - remote_host.clone(), - HashSet::from_iter([PathTrust::Worktree(worktree_id)]), - )); + self.restricted + .entry(worktree_store.clone()) + .or_default() + .insert(worktree_id); + restricted.insert(PathTrust::Worktree(worktree_id)); } PathTrust::AbsPath(..) => debug_panic!("Unexpected: cannot restrict an abs path"), } } + + cx.emit(TrustedWorktreesEvent::Restricted( + worktree_store, + restricted, + )); } /// Erases all trust information. /// Requires Zed's restart to take proper effect. pub fn clear_trusted_paths(&mut self) { self.trusted_paths.clear(); + self.db_trusted_paths.clear(); } /// Checks whether a certain worktree is trusted (or on a larger trust level). /// If not, emits [`TrustedWorktreesEvent::Restricted`] event if for the first time and not trusted, or no corresponding worktree store was found. /// /// No events or data adjustment happens when `trust_all_worktrees` auto trust is enabled. - pub fn can_trust(&mut self, worktree_id: WorktreeId, cx: &mut Context) -> bool { + pub fn can_trust( + &mut self, + worktree_store: &Entity, + worktree_id: WorktreeId, + cx: &mut Context, + ) -> bool { if ProjectSettings::get_global(cx).session.trust_all_worktrees { return true; } - if self.restricted.contains(&worktree_id) { - return false; - } - let Some((worktree_path, is_file, remote_host)) = self.find_worktree_data(worktree_id, cx) - else { + let weak_worktree_store = worktree_store.downgrade(); + let Some(worktree) = worktree_store.read(cx).worktree_for_id(worktree_id, cx) else { return false; }; + let worktree_path = worktree.read(cx).abs_path(); + let is_file = worktree.read(cx).is_single_file(); + if self + .restricted + .get(&weak_worktree_store) + .is_some_and(|restricted_worktrees| restricted_worktrees.contains(&worktree_id)) + { + return false; + } if self .trusted_paths - .get(&remote_host) + .get(&weak_worktree_store) .is_some_and(|trusted_paths| trusted_paths.contains(&PathTrust::Worktree(worktree_id))) { return true; } // See module documentation for details on trust level. - if is_file && self.trusted_paths.contains_key(&remote_host) { + if is_file && self.trusted_paths.contains_key(&weak_worktree_store) { return true; } let parent_path_trusted = self.trusted_paths - .get(&remote_host) + .get(&weak_worktree_store) .is_some_and(|trusted_paths| { trusted_paths.iter().any(|trusted_path| { let PathTrust::AbsPath(trusted_path) = trusted_path else { @@ -464,23 +507,26 @@ impl TrustedWorktreesStore { return true; } - self.restricted.insert(worktree_id); + self.restricted + .entry(weak_worktree_store.clone()) + .or_default() + .insert(worktree_id); cx.emit(TrustedWorktreesEvent::Restricted( - remote_host, + weak_worktree_store, HashSet::from_iter([PathTrust::Worktree(worktree_id)]), )); - if let Some((downstream_client, downstream_project_id)) = &self.downstream_client { + for (downstream_client, downstream_project_id) in &self.downstream_clients { downstream_client .send(proto::RestrictWorktrees { - project_id: *downstream_project_id, + project_id: downstream_project_id.0, worktree_ids: vec![worktree_id.to_proto()], }) .ok(); } - if let Some((upstream_client, upstream_project_id)) = &self.upstream_client { + for (upstream_client, upstream_project_id) in &self.upstream_clients { upstream_client .send(proto::RestrictWorktrees { - project_id: *upstream_project_id, + project_id: upstream_project_id.0, worktree_ids: vec![worktree_id.to_proto()], }) .ok(); @@ -491,15 +537,20 @@ impl TrustedWorktreesStore { /// Lists all explicitly restricted worktrees (via [`TrustedWorktreesStore::can_trust`] method calls) for a particular worktree store on a particular host. pub fn restricted_worktrees( &self, - worktree_store: &WorktreeStore, + worktree_store: &Entity, cx: &App, ) -> HashSet<(WorktreeId, Arc)> { let mut single_file_paths = HashSet::default(); + let other_paths = self .restricted - .iter() + .get(&worktree_store.downgrade()) + .into_iter() + .flatten() .filter_map(|&restricted_worktree_id| { - let worktree = worktree_store.worktree_for_id(restricted_worktree_id, cx)?; + let worktree = worktree_store + .read(cx) + .worktree_for_id(restricted_worktree_id, cx)?; let worktree = worktree.read(cx); let abs_path = worktree.abs_path(); if worktree.is_single_file() { @@ -521,74 +572,58 @@ impl TrustedWorktreesStore { /// Switches the "trust nothing" mode to "automatically trust everything". /// This does not influence already persisted data, but stops adding new worktrees there. pub fn auto_trust_all(&mut self, cx: &mut Context) { - for (remote_host, worktrees) in std::mem::take(&mut self.restricted) - .into_iter() - .flat_map(|restricted_worktree| { - let (_, _, host) = self.find_worktree_data(restricted_worktree, cx)?; - Some((restricted_worktree, host)) - }) - .fold(HashMap::default(), |mut acc, (worktree_id, remote_host)| { + for (worktree_store, worktrees) in std::mem::take(&mut self.restricted).into_iter().fold( + HashMap::default(), + |mut acc, (remote_host, worktrees)| { acc.entry(remote_host) .or_insert_with(HashSet::default) - .insert(PathTrust::Worktree(worktree_id)); + .extend(worktrees.into_iter().map(PathTrust::Worktree)); acc - }) - { - self.trust(worktrees, remote_host, cx); + }, + ) { + if let Some(worktree_store) = worktree_store.upgrade() { + self.trust(&worktree_store, worktrees, cx); + } } } - /// Returns a normalized representation of the trusted paths to store in the DB. - pub fn trusted_paths_for_serialization( + pub fn schedule_serialization(&mut self, cx: &mut Context, serialize: S) + where + S: FnOnce(HashMap, HashSet>, &App) -> Task<()> + + 'static, + { + self.worktree_trust_serialization = serialize(self.trusted_paths_for_serialization(cx), cx); + } + + fn trusted_paths_for_serialization( &mut self, cx: &mut Context, ) -> HashMap, HashSet> { - let new_trusted_worktrees = self - .trusted_paths - .clone() - .into_iter() - .map(|(host, paths)| { + self.trusted_paths + .iter() + .filter_map(|(worktree_store, paths)| { + let host = self.worktree_stores.get(&worktree_store)?.clone(); let abs_paths = paths - .into_iter() + .iter() .flat_map(|path| match path { - PathTrust::Worktree(worktree_id) => self - .find_worktree_data(worktree_id, cx) - .map(|(abs_path, ..)| abs_path.to_path_buf()), - PathTrust::AbsPath(abs_path) => Some(abs_path), + PathTrust::Worktree(worktree_id) => worktree_store + .upgrade() + .and_then(|worktree_store| { + worktree_store.read(cx).worktree_for_id(*worktree_id, cx) + }) + .map(|worktree| worktree.read(cx).abs_path().to_path_buf()), + PathTrust::AbsPath(abs_path) => Some(abs_path.clone()), }) - .collect(); - (host, abs_paths) + .collect::>(); + Some((host, abs_paths)) + }) + .chain(self.db_trusted_paths.clone()) + .fold(HashMap::default(), |mut acc, (host, paths)| { + acc.entry(host) + .or_insert_with(HashSet::default) + .extend(paths); + acc }) - .collect(); - new_trusted_worktrees - } - - fn find_worktree_data( - &mut self, - worktree_id: WorktreeId, - cx: &mut Context, - ) -> Option<(Arc, bool, Option)> { - let mut worktree_data = None; - self.worktree_stores.retain( - |worktree_store, remote_host| match worktree_store.upgrade() { - Some(worktree_store) => { - if worktree_data.is_none() { - if let Some(worktree) = - worktree_store.read(cx).worktree_for_id(worktree_id, cx) - { - worktree_data = Some(( - worktree.read(cx).abs_path(), - worktree.read(cx).is_single_file(), - remote_host.clone(), - )); - } - } - true - } - None => false, - }, - ); - worktree_data } fn add_worktree_store( @@ -597,18 +632,26 @@ impl TrustedWorktreesStore { remote_host: Option, cx: &mut Context, ) { + let weak_worktree_store = worktree_store.downgrade(); self.worktree_stores - .insert(worktree_store.downgrade(), remote_host.clone()); + .insert(weak_worktree_store.clone(), remote_host.clone()); - if let Some(trusted_paths) = self.trusted_paths.remove(&remote_host) { + let mut new_trusted_paths = HashSet::default(); + if let Some(db_trusted_paths) = self.db_trusted_paths.get(&remote_host) { + new_trusted_paths.extend(db_trusted_paths.clone().into_iter().map(PathTrust::AbsPath)); + } + if let Some(trusted_paths) = self.trusted_paths.remove(&weak_worktree_store) { + new_trusted_paths.extend(trusted_paths); + } + if !new_trusted_paths.is_empty() { self.trusted_paths.insert( - remote_host.clone(), - trusted_paths + weak_worktree_store, + new_trusted_paths .into_iter() .map(|path_trust| match path_trust { PathTrust::AbsPath(abs_path) => { find_worktree_in_store(worktree_store.read(cx), &abs_path, cx) - .map(PathTrust::Worktree) + .map(|(worktree_id, _)| PathTrust::Worktree(worktree_id)) .unwrap_or_else(|| PathTrust::AbsPath(abs_path)) } other => other, @@ -619,14 +662,14 @@ impl TrustedWorktreesStore { } } -pub fn find_worktree_in_store( +fn find_worktree_in_store( worktree_store: &WorktreeStore, abs_path: &Path, cx: &App, -) -> Option { +) -> Option<(WorktreeId, bool)> { let (worktree, path_in_worktree) = worktree_store.find_worktree(&abs_path, cx)?; if path_in_worktree.is_empty() { - Some(worktree.read(cx).id()) + Some((worktree.read(cx).id(), worktree.read(cx).is_single_file())) } else { None } @@ -703,15 +746,17 @@ mod tests { }) .detach(); - let can_trust = trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!(!can_trust, "worktree should be restricted by default"); { let events = events.borrow(); assert_eq!(events.len(), 1); match &events[0] { - TrustedWorktreesEvent::Restricted(host, paths) => { - assert!(host.is_none()); + TrustedWorktreesEvent::Restricted(event_worktree_store, paths) => { + assert_eq!(event_worktree_store, &worktree_store.downgrade()); assert!(paths.contains(&PathTrust::Worktree(worktree_id))); } _ => panic!("expected Restricted event"), @@ -723,15 +768,16 @@ mod tests { }); assert!(has_restricted, "should have restricted worktrees"); - let restricted = worktree_store.read_with(cx, |ws, cx| { - trusted_worktrees.read(cx).restricted_worktrees(ws, cx) + let restricted = trusted_worktrees.read_with(cx, |trusted_worktrees, cx| { + trusted_worktrees.restricted_worktrees(&worktree_store, cx) }); assert!(restricted.iter().any(|(id, _)| *id == worktree_id)); events.borrow_mut().clear(); - let can_trust_again = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust_again = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!(!can_trust_again, "worktree should still be restricted"); assert!( events.borrow().is_empty(), @@ -740,8 +786,8 @@ mod tests { trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::Worktree(worktree_id)]), - None, cx, ); }); @@ -750,16 +796,17 @@ mod tests { let events = events.borrow(); assert_eq!(events.len(), 1); match &events[0] { - TrustedWorktreesEvent::Trusted(host, paths) => { - assert!(host.is_none()); + TrustedWorktreesEvent::Trusted(event_worktree_store, paths) => { + assert_eq!(event_worktree_store, &worktree_store.downgrade()); assert!(paths.contains(&PathTrust::Worktree(worktree_id))); } _ => panic!("expected Trusted event"), } } - let can_trust_after = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust_after = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!(can_trust_after, "worktree should be trusted after trust()"); let has_restricted_after = trusted_worktrees.read_with(cx, |store, cx| { @@ -770,8 +817,8 @@ mod tests { "should have no restricted worktrees after trust" ); - let restricted_after = worktree_store.read_with(cx, |ws, cx| { - trusted_worktrees.read(cx).restricted_worktrees(ws, cx) + let restricted_after = trusted_worktrees.read_with(cx, |trusted_worktrees, cx| { + trusted_worktrees.restricted_worktrees(&worktree_store, cx) }); assert!( restricted_after.is_empty(), @@ -796,7 +843,7 @@ mod tests { worktree.id() }); - let trusted_worktrees = init_trust_global(worktree_store, cx); + let trusted_worktrees = init_trust_global(worktree_store.clone(), cx); let events: Rc>> = Rc::default(); cx.update({ @@ -816,7 +863,9 @@ mod tests { }) .detach(); - let can_trust = trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!( !can_trust, "single-file worktree should be restricted by default" @@ -826,8 +875,8 @@ mod tests { let events = events.borrow(); assert_eq!(events.len(), 1); match &events[0] { - TrustedWorktreesEvent::Restricted(host, paths) => { - assert!(host.is_none()); + TrustedWorktreesEvent::Restricted(event_worktree_store, paths) => { + assert_eq!(event_worktree_store, &worktree_store.downgrade()); assert!(paths.contains(&PathTrust::Worktree(worktree_id))); } _ => panic!("expected Restricted event"), @@ -838,8 +887,8 @@ mod tests { trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::Worktree(worktree_id)]), - None, cx, ); }); @@ -848,16 +897,17 @@ mod tests { let events = events.borrow(); assert_eq!(events.len(), 1); match &events[0] { - TrustedWorktreesEvent::Trusted(host, paths) => { - assert!(host.is_none()); + TrustedWorktreesEvent::Trusted(event_worktree_store, paths) => { + assert_eq!(event_worktree_store, &worktree_store.downgrade()); assert!(paths.contains(&PathTrust::Worktree(worktree_id))); } _ => panic!("expected Trusted event"), } } - let can_trust_after = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust_after = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!( can_trust_after, "single-file worktree should be trusted after trust()" @@ -902,11 +952,12 @@ mod tests { }); assert_eq!(worktree_ids.len(), 3); - let trusted_worktrees = init_trust_global(worktree_store, cx); + let trusted_worktrees = init_trust_global(worktree_store.clone(), cx); for &worktree_id in &worktree_ids { - let can_trust = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!( !can_trust, "worktree {worktree_id:?} should be restricted initially" @@ -915,18 +966,21 @@ mod tests { trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::Worktree(worktree_ids[1])]), - None, cx, ); }); - let can_trust_0 = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_ids[0], cx)); - let can_trust_1 = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_ids[1], cx)); - let can_trust_2 = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_ids[2], cx)); + let can_trust_0 = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[0], cx) + }); + let can_trust_1 = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[1], cx) + }); + let can_trust_2 = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[2], cx) + }); assert!(!can_trust_0, "worktree 0 should still be restricted"); assert!(can_trust_1, "worktree 1 should be trusted"); @@ -969,42 +1023,48 @@ mod tests { }); assert_eq!(worktree_ids.len(), 2); - let trusted_worktrees = init_trust_global(worktree_store, cx); + let trusted_worktrees = init_trust_global(worktree_store.clone(), cx); - let can_trust_a = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_ids[0], cx)); - let can_trust_b = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_ids[1], cx)); + let can_trust_a = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[0], cx) + }); + let can_trust_b = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[1], cx) + }); assert!(!can_trust_a, "project_a should be restricted initially"); assert!(!can_trust_b, "project_b should be restricted initially"); trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::Worktree(worktree_ids[0])]), - None, cx, ); }); - let can_trust_a = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_ids[0], cx)); - let can_trust_b = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_ids[1], cx)); + let can_trust_a = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[0], cx) + }); + let can_trust_b = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[1], cx) + }); assert!(can_trust_a, "project_a should be trusted after trust()"); assert!(!can_trust_b, "project_b should still be restricted"); trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::Worktree(worktree_ids[1])]), - None, cx, ); }); - let can_trust_a = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_ids[0], cx)); - let can_trust_b = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_ids[1], cx)); + let can_trust_a = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[0], cx) + }); + let can_trust_b = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_ids[1], cx) + }); assert!(can_trust_a, "project_a should remain trusted"); assert!(can_trust_b, "project_b should now be trusted"); } @@ -1043,17 +1103,19 @@ mod tests { (dir_worktree.read(cx).id(), file_worktree.read(cx).id()) }); - let trusted_worktrees = init_trust_global(worktree_store, cx); + let trusted_worktrees = init_trust_global(worktree_store.clone(), cx); - let can_trust_file = - trusted_worktrees.update(cx, |store, cx| store.can_trust(file_worktree_id, cx)); + let can_trust_file = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, file_worktree_id, cx) + }); assert!( !can_trust_file, "single-file worktree should be restricted initially" ); - let can_trust_directory = - trusted_worktrees.update(cx, |store, cx| store.can_trust(dir_worktree_id, cx)); + let can_trust_directory = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, dir_worktree_id, cx) + }); assert!( !can_trust_directory, "directory worktree should be restricted initially" @@ -1061,16 +1123,18 @@ mod tests { trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::Worktree(dir_worktree_id)]), - None, cx, ); }); - let can_trust_dir = - trusted_worktrees.update(cx, |store, cx| store.can_trust(dir_worktree_id, cx)); - let can_trust_file_after = - trusted_worktrees.update(cx, |store, cx| store.can_trust(file_worktree_id, cx)); + let can_trust_dir = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, dir_worktree_id, cx) + }); + let can_trust_file_after = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, file_worktree_id, cx) + }); assert!(can_trust_dir, "directory worktree should be trusted"); assert!( can_trust_file_after, @@ -1112,17 +1176,19 @@ mod tests { (dir_worktree.read(cx).id(), file_worktree.read(cx).id()) }); - let trusted_worktrees = init_trust_global(worktree_store, cx); + let trusted_worktrees = init_trust_global(worktree_store.clone(), cx); - let can_trust_file = - trusted_worktrees.update(cx, |store, cx| store.can_trust(file_worktree_id, cx)); + let can_trust_file = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, file_worktree_id, cx) + }); assert!( !can_trust_file, "single-file worktree should be restricted initially" ); - let can_trust_directory = - trusted_worktrees.update(cx, |store, cx| store.can_trust(dir_worktree_id, cx)); + let can_trust_directory = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, dir_worktree_id, cx) + }); assert!( !can_trust_directory, "directory worktree should be restricted initially" @@ -1130,16 +1196,18 @@ mod tests { trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::AbsPath(PathBuf::from(path!("/project")))]), - None, cx, ); }); - let can_trust_dir = - trusted_worktrees.update(cx, |store, cx| store.can_trust(dir_worktree_id, cx)); - let can_trust_file_after = - trusted_worktrees.update(cx, |store, cx| store.can_trust(file_worktree_id, cx)); + let can_trust_dir = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, dir_worktree_id, cx) + }); + let can_trust_file_after = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, file_worktree_id, cx) + }); assert!( can_trust_dir, "directory worktree should be trusted after its parent is trusted" @@ -1182,25 +1250,27 @@ mod tests { }); assert_eq!(worktree_ids.len(), 2); - let trusted_worktrees = init_trust_global(worktree_store, cx); + let trusted_worktrees = init_trust_global(worktree_store.clone(), cx); for &worktree_id in &worktree_ids { - let can_trust = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!(!can_trust, "worktree should be restricted initially"); } trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::AbsPath(PathBuf::from(path!("/root")))]), - None, cx, ); }); for &worktree_id in &worktree_ids { - let can_trust = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!( can_trust, "worktree should be trusted after parent path trust" @@ -1263,8 +1333,9 @@ mod tests { .detach(); for &worktree_id in &worktree_ids { - let can_trust = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!(!can_trust, "worktree should be restricted initially"); } @@ -1280,8 +1351,9 @@ mod tests { }); for &worktree_id in &worktree_ids { - let can_trust = - trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!( can_trust, "worktree {worktree_id:?} should be trusted after auto_trust_all" @@ -1341,19 +1413,23 @@ mod tests { }) .detach(); - let can_trust = trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!(!can_trust, "should be restricted initially"); assert_eq!(events.borrow().len(), 1); events.borrow_mut().clear(); trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::Worktree(worktree_id)]), - None, cx, ); }); - let can_trust = trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!(can_trust, "should be trusted after trust()"); assert_eq!(events.borrow().len(), 1); assert!(matches!( @@ -1364,12 +1440,14 @@ mod tests { trusted_worktrees.update(cx, |store, cx| { store.restrict( + worktree_store.downgrade(), HashSet::from_iter([PathTrust::Worktree(worktree_id)]), - None, cx, ); }); - let can_trust = trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!(!can_trust, "should be restricted after restrict()"); assert_eq!(events.borrow().len(), 1); assert!(matches!( @@ -1385,12 +1463,14 @@ mod tests { trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::Worktree(worktree_id)]), - None, cx, ); }); - let can_trust = trusted_worktrees.update(cx, |store, cx| store.can_trust(worktree_id, cx)); + let can_trust = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, worktree_id, cx) + }); assert!(can_trust, "should be trusted again after second trust()"); assert_eq!(events.borrow().len(), 1); assert!(matches!( @@ -1438,27 +1518,27 @@ mod tests { let local_worktree = worktree_ids[0]; let _remote_worktree = worktree_ids[1]; - let trusted_worktrees = init_trust_global(worktree_store, cx); - - let host_a: Option = None; + let trusted_worktrees = init_trust_global(worktree_store.clone(), cx); - let can_trust_local = - trusted_worktrees.update(cx, |store, cx| store.can_trust(local_worktree, cx)); + let can_trust_local = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, local_worktree, cx) + }); assert!(!can_trust_local, "local worktree restricted on host_a"); trusted_worktrees.update(cx, |store, cx| { store.trust( + &worktree_store, HashSet::from_iter([PathTrust::Worktree(local_worktree)]), - host_a.clone(), cx, ); }); - let can_trust_local_after = - trusted_worktrees.update(cx, |store, cx| store.can_trust(local_worktree, cx)); + let can_trust_local_after = trusted_worktrees.update(cx, |store, cx| { + store.can_trust(&worktree_store, local_worktree, cx) + }); assert!( can_trust_local_after, - "local worktree should be trusted on host_a" + "local worktree should be trusted on local host" ); } } diff --git a/crates/remote_server/src/headless_project.rs b/crates/remote_server/src/headless_project.rs index c83cc6aa34402a082fe104d64a8cb47f460704b8..7fa08de96ec76854878277ae6b2ac4a24be1eef8 100644 --- a/crates/remote_server/src/headless_project.rs +++ b/crates/remote_server/src/headless_project.rs @@ -1,4 +1,5 @@ use anyhow::{Context as _, Result, anyhow}; +use client::ProjectId; use collections::HashSet; use language::File; use lsp::LanguageServerId; @@ -104,7 +105,7 @@ impl HeadlessProject { project::trusted_worktrees::track_worktree_trust( worktree_store.clone(), None::, - Some((session.clone(), REMOTE_SERVER_PROJECT_ID)), + Some((session.clone(), ProjectId(REMOTE_SERVER_PROJECT_ID))), None, cx, ); @@ -611,22 +612,23 @@ impl HeadlessProject { } pub async fn handle_trust_worktrees( - _: Entity, + this: Entity, envelope: TypedEnvelope, mut cx: AsyncApp, ) -> Result { let trusted_worktrees = cx .update(|cx| TrustedWorktrees::try_get_global(cx))? .context("missing trusted worktrees")?; + let worktree_store = this.read_with(&cx, |project, _| project.worktree_store.clone())?; trusted_worktrees.update(&mut cx, |trusted_worktrees, cx| { trusted_worktrees.trust( + &worktree_store, envelope .payload .trusted_paths .into_iter() .filter_map(PathTrust::from_proto) .collect(), - None, cx, ); })?; @@ -634,13 +636,15 @@ impl HeadlessProject { } pub async fn handle_restrict_worktrees( - _: Entity, + this: Entity, envelope: TypedEnvelope, mut cx: AsyncApp, ) -> Result { let trusted_worktrees = cx .update(|cx| TrustedWorktrees::try_get_global(cx))? .context("missing trusted worktrees")?; + let worktree_store = + this.read_with(&cx, |project, _| project.worktree_store.downgrade())?; trusted_worktrees.update(&mut cx, |trusted_worktrees, cx| { let restricted_paths = envelope .payload @@ -649,7 +653,7 @@ impl HeadlessProject { .map(WorktreeId::from_proto) .map(PathTrust::Worktree) .collect::>(); - trusted_worktrees.restrict(restricted_paths, None, cx); + trusted_worktrees.restrict(worktree_store, restricted_paths, cx); })?; Ok(proto::Ack {}) } diff --git a/crates/remote_server/src/unix.rs b/crates/remote_server/src/unix.rs index 449b8491ece2494dacf8bfb1fa89aeeb8f6a81ac..a81859a685e1ce34eccc0a0d4809b9914743f79f 100644 --- a/crates/remote_server/src/unix.rs +++ b/crates/remote_server/src/unix.rs @@ -1,7 +1,7 @@ use crate::HeadlessProject; use crate::headless_project::HeadlessAppState; use anyhow::{Context as _, Result, anyhow}; -use client::ProxySettings; +use client::{ProjectId, ProxySettings}; use collections::HashMap; use project::trusted_worktrees; use util::ResultExt; @@ -419,7 +419,7 @@ pub fn execute_run( log::info!("gpui app started, initializing server"); let session = start_server(listeners, log_rx, cx, is_wsl_interop); - trusted_worktrees::init(HashMap::default(), Some((session.clone(), REMOTE_SERVER_PROJECT_ID)), None, cx); + trusted_worktrees::init(HashMap::default(), Some((session.clone(), ProjectId(REMOTE_SERVER_PROJECT_ID))), None, cx); GitHostingProviderRegistry::set_global(git_hosting_provider_registry, cx); git_hosting_providers::init(cx); diff --git a/crates/rpc/src/proto_client.rs b/crates/rpc/src/proto_client.rs index 3850ff5820e6d73289b5714d6b880ecb584bf8d9..052964c0853f4e2262d37dbc913ad4b8d4a7ff93 100644 --- a/crates/rpc/src/proto_client.rs +++ b/crates/rpc/src/proto_client.rs @@ -20,7 +20,7 @@ use std::{ time::Duration, }; -#[derive(Clone)] +#[derive(Debug, Clone)] pub struct AnyProtoClient(Arc); type RequestIds = Arc< @@ -45,6 +45,15 @@ struct State { request_ids: RequestIds, } +impl std::fmt::Debug for State { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("State") + .field("next_lsp_request_id", &self.next_lsp_request_id) + .field("request_ids", &self.request_ids) + .finish_non_exhaustive() + } +} + pub trait ProtoClient: Send + Sync { fn request( &self, diff --git a/crates/workspace/src/persistence.rs b/crates/workspace/src/persistence.rs index cb37a86f25d5c0f22787f1e4a5b6fc3ad2d46893..3f6fc9262377a689f7dc190f5dc517c81e81db37 100644 --- a/crates/workspace/src/persistence.rs +++ b/crates/workspace/src/persistence.rs @@ -15,11 +15,10 @@ use db::{ sqlez::{connection::Connection, domain::Domain}, sqlez_macros::sql, }; -use gpui::{Axis, Bounds, Entity, Task, WindowBounds, WindowId, point, size}; +use gpui::{Axis, Bounds, Task, WindowBounds, WindowId, point, size}; use project::{ debugger::breakpoint_store::{BreakpointState, SourceBreakpoint}, - trusted_worktrees::{PathTrust, RemoteHostLocation, find_worktree_in_store}, - worktree_store::WorktreeStore, + trusted_worktrees::{DbTrustedPaths, RemoteHostLocation}, }; use language::{LanguageName, Toolchain, ToolchainScope}; @@ -1888,18 +1887,12 @@ VALUES {placeholders};"# Ok(()) } - pub fn fetch_trusted_worktrees( - &self, - worktree_store: Option>, - host: Option, - cx: &App, - ) -> Result, HashSet>> { + pub fn fetch_trusted_worktrees(&self) -> Result { let trusted_worktrees = DB.trusted_worktrees()?; Ok(trusted_worktrees .into_iter() .filter_map(|(abs_path, user_name, host_name)| { let db_host = match (user_name, host_name) { - (_, None) => None, (None, Some(host_name)) => Some(RemoteHostLocation { user_name: None, host_identifier: SharedString::new(host_name), @@ -1908,24 +1901,14 @@ VALUES {placeholders};"# user_name: Some(SharedString::new(user_name)), host_identifier: SharedString::new(host_name), }), + _ => None, }; - - let abs_path = abs_path?; - Some(if db_host != host { - (db_host, PathTrust::AbsPath(abs_path)) - } else if let Some(worktree_store) = &worktree_store { - find_worktree_in_store(worktree_store.read(cx), &abs_path, cx) - .map(PathTrust::Worktree) - .map(|trusted_worktree| (host.clone(), trusted_worktree)) - .unwrap_or_else(|| (db_host.clone(), PathTrust::AbsPath(abs_path))) - } else { - (db_host, PathTrust::AbsPath(abs_path)) - }) + Some((db_host, abs_path?)) }) - .fold(HashMap::default(), |mut acc, (remote_host, path_trust)| { + .fold(HashMap::default(), |mut acc, (remote_host, abs_path)| { acc.entry(remote_host) .or_insert_with(HashSet::default) - .insert(path_trust); + .insert(abs_path); acc })) } diff --git a/crates/workspace/src/security_modal.rs b/crates/workspace/src/security_modal.rs index bb1482d7cce2a9849a78a9512598e389a6e5eea0..664aa891550cecdd602d54bfca579d04e03f33dc 100644 --- a/crates/workspace/src/security_modal.rs +++ b/crates/workspace/src/security_modal.rs @@ -267,7 +267,9 @@ impl SecurityModal { } fn trust_and_dismiss(&mut self, cx: &mut Context) { - if let Some(trusted_worktrees) = TrustedWorktrees::try_get_global(cx) { + if let Some((trusted_worktrees, worktree_store)) = + TrustedWorktrees::try_get_global(cx).zip(self.worktree_store.upgrade()) + { trusted_worktrees.update(cx, |trusted_worktrees, cx| { let mut paths_to_trust = self .restricted_paths @@ -288,7 +290,7 @@ impl SecurityModal { }, )); } - trusted_worktrees.trust(paths_to_trust, self.remote_host.clone(), cx); + trusted_worktrees.trust(&worktree_store, paths_to_trust, cx); }); } @@ -305,7 +307,7 @@ impl SecurityModal { if let Some(worktree_store) = self.worktree_store.upgrade() { let new_restricted_worktrees = trusted_worktrees .read(cx) - .restricted_worktrees(worktree_store.read(cx), cx) + .restricted_worktrees(&worktree_store, cx) .into_iter() .filter_map(|(worktree_id, abs_path)| { let worktree = worktree_store.read(cx).worktree_for_id(worktree_id, cx)?; diff --git a/crates/workspace/src/workspace.rs b/crates/workspace/src/workspace.rs index 80d41a0779794dc49b0e5041bb8733f9816f3b65..f25fb25b7bb0221df7b67820f15bbbf3e86ab9bb 100644 --- a/crates/workspace/src/workspace.rs +++ b/crates/workspace/src/workspace.rs @@ -80,7 +80,7 @@ use project::{ debugger::{breakpoint_store::BreakpointStoreEvent, session::ThreadStatus}, project_settings::ProjectSettings, toolchain_store::ToolchainStoreEvent, - trusted_worktrees::{TrustedWorktrees, TrustedWorktreesEvent}, + trusted_worktrees::{RemoteHostLocation, TrustedWorktrees, TrustedWorktreesEvent}, }; use remote::{ RemoteClientDelegate, RemoteConnection, RemoteConnectionOptions, @@ -1188,7 +1188,6 @@ pub struct Workspace { _observe_current_user: Task>, _schedule_serialize_workspace: Option>, _schedule_serialize_ssh_paths: Option>, - _schedule_serialize_worktree_trust: Task<()>, pane_history_timestamp: Arc, bounds: Bounds, pub centered_layout: bool, @@ -1235,23 +1234,26 @@ impl Workspace { cx: &mut Context, ) -> Self { if let Some(trusted_worktrees) = TrustedWorktrees::try_get_global(cx) { - cx.subscribe(&trusted_worktrees, |workspace, worktrees_store, e, cx| { + cx.subscribe(&trusted_worktrees, |_, worktrees_store, e, cx| { if let TrustedWorktreesEvent::Trusted(..) = e { // Do not persist auto trusted worktrees if !ProjectSettings::get_global(cx).session.trust_all_worktrees { - let new_trusted_worktrees = - worktrees_store.update(cx, |worktrees_store, cx| { - worktrees_store.trusted_paths_for_serialization(cx) - }); - let timeout = cx.background_executor().timer(SERIALIZATION_THROTTLE_TIME); - workspace._schedule_serialize_worktree_trust = - cx.background_spawn(async move { - timeout.await; - persistence::DB - .save_trusted_worktrees(new_trusted_worktrees) - .await - .log_err(); - }); + worktrees_store.update(cx, |worktrees_store, cx| { + worktrees_store.schedule_serialization( + cx, + |new_trusted_worktrees, cx| { + let timeout = + cx.background_executor().timer(SERIALIZATION_THROTTLE_TIME); + cx.background_spawn(async move { + timeout.await; + persistence::DB + .save_trusted_worktrees(new_trusted_worktrees) + .await + .log_err(); + }) + }, + ) + }); } } }) @@ -1282,7 +1284,11 @@ impl Workspace { project::Event::WorktreeUpdatedEntries(worktree_id, _) => { if let Some(trusted_worktrees) = TrustedWorktrees::try_get_global(cx) { trusted_worktrees.update(cx, |trusted_worktrees, cx| { - trusted_worktrees.can_trust(*worktree_id, cx); + trusted_worktrees.can_trust( + &this.project().read(cx).worktree_store(), + *worktree_id, + cx, + ); }); } } @@ -1294,7 +1300,11 @@ impl Workspace { project::Event::WorktreeAdded(worktree_id) => { if let Some(trusted_worktrees) = TrustedWorktrees::try_get_global(cx) { trusted_worktrees.update(cx, |trusted_worktrees, cx| { - trusted_worktrees.can_trust(*worktree_id, cx); + trusted_worktrees.can_trust( + &this.project().read(cx).worktree_store(), + *worktree_id, + cx, + ); }); } this.update_worktree_data(window, cx); @@ -1583,7 +1593,6 @@ impl Workspace { _apply_leader_updates, _schedule_serialize_workspace: None, _schedule_serialize_ssh_paths: None, - _schedule_serialize_worktree_trust: Task::ready(()), leader_updates_tx, _subscriptions: subscriptions, pane_history_timestamp, @@ -6542,7 +6551,9 @@ impl Workspace { .unwrap_or(false); if has_restricted_worktrees { let project = self.project().read(cx); - let remote_host = project.remote_connection_options(cx); + let remote_host = project + .remote_connection_options(cx) + .map(RemoteHostLocation::from); let worktree_store = project.worktree_store().downgrade(); self.toggle_modal(window, cx, |_, cx| { SecurityModal::new(worktree_store, remote_host, cx) diff --git a/crates/zed/src/main.rs b/crates/zed/src/main.rs index 2db32547da75ae284a1592572414cb57b1ec97d1..2b3a055105d660fe096e53d887945e1de17de896 100644 --- a/crates/zed/src/main.rs +++ b/crates/zed/src/main.rs @@ -406,14 +406,14 @@ pub fn main() { }); app.run(move |cx| { - let trusted_paths = match workspace::WORKSPACE_DB.fetch_trusted_worktrees(None, None, cx) { + let db_trusted_paths = match workspace::WORKSPACE_DB.fetch_trusted_worktrees() { Ok(trusted_paths) => trusted_paths, Err(e) => { log::error!("Failed to do initial trusted worktrees fetch: {e:#}"); HashMap::default() } }; - trusted_worktrees::init(trusted_paths, None, None, cx); + trusted_worktrees::init(db_trusted_paths, None, None, cx); menu::init(); zed_actions::init();