1use collections::{HashMap, HashSet};
2use gpui::{App, Entity, SharedString};
3use std::path::PathBuf;
4
5use db::{
6 query,
7 sqlez::{domain::Domain, thread_safe_connection::ThreadSafeConnection},
8 sqlez_macros::sql,
9};
10
11use crate::{
12 trusted_worktrees::{PathTrust, RemoteHostLocation, find_worktree_in_store},
13 worktree_store::WorktreeStore,
14};
15
16// https://www.sqlite.org/limits.html
17// > <..> the maximum value of a host parameter number is SQLITE_MAX_VARIABLE_NUMBER,
18// > which defaults to <..> 32766 for SQLite versions after 3.32.0.
19#[allow(unused)]
20const MAX_QUERY_PLACEHOLDERS: usize = 32000;
21
22#[allow(unused)]
23pub struct ProjectDb(ThreadSafeConnection);
24
25impl Domain for ProjectDb {
26 const NAME: &str = stringify!(ProjectDb);
27
28 const MIGRATIONS: &[&str] = &[sql!(
29 CREATE TABLE IF NOT EXISTS trusted_worktrees (
30 trust_id INTEGER PRIMARY KEY AUTOINCREMENT,
31 absolute_path TEXT,
32 user_name TEXT,
33 host_name TEXT
34 ) STRICT;
35 )];
36}
37
38db::static_connection!(PROJECT_DB, ProjectDb, []);
39
40impl ProjectDb {
41 pub(crate) async fn save_trusted_worktrees(
42 &self,
43 trusted_worktrees: HashMap<Option<RemoteHostLocation>, HashSet<PathBuf>>,
44 trusted_workspaces: HashSet<Option<RemoteHostLocation>>,
45 ) -> anyhow::Result<()> {
46 use anyhow::Context as _;
47 use db::sqlez::statement::Statement;
48 use itertools::Itertools as _;
49
50 PROJECT_DB
51 .clear_trusted_worktrees()
52 .await
53 .context("clearing previous trust state")?;
54
55 let trusted_worktrees = trusted_worktrees
56 .into_iter()
57 .flat_map(|(host, abs_paths)| {
58 abs_paths
59 .into_iter()
60 .map(move |abs_path| (Some(abs_path), host.clone()))
61 })
62 .chain(trusted_workspaces.into_iter().map(|host| (None, host)))
63 .collect::<Vec<_>>();
64 let mut first_worktree;
65 let mut last_worktree = 0_usize;
66 for (count, placeholders) in std::iter::once("(?, ?, ?)")
67 .cycle()
68 .take(trusted_worktrees.len())
69 .chunks(MAX_QUERY_PLACEHOLDERS / 3)
70 .into_iter()
71 .map(|chunk| {
72 let mut count = 0;
73 let placeholders = chunk
74 .inspect(|_| {
75 count += 1;
76 })
77 .join(", ");
78 (count, placeholders)
79 })
80 .collect::<Vec<_>>()
81 {
82 first_worktree = last_worktree;
83 last_worktree = last_worktree + count;
84 let query = format!(
85 r#"INSERT INTO trusted_worktrees(absolute_path, user_name, host_name)
86VALUES {placeholders};"#
87 );
88
89 let trusted_worktrees = trusted_worktrees[first_worktree..last_worktree].to_vec();
90 self.write(move |conn| {
91 let mut statement = Statement::prepare(conn, query)?;
92 let mut next_index = 1;
93 for (abs_path, host) in trusted_worktrees {
94 let abs_path = abs_path.as_ref().map(|abs_path| abs_path.to_string_lossy());
95 next_index = statement.bind(
96 &abs_path.as_ref().map(|abs_path| abs_path.as_ref()),
97 next_index,
98 )?;
99 next_index = statement.bind(
100 &host
101 .as_ref()
102 .and_then(|host| Some(host.user_name.as_ref()?.as_str())),
103 next_index,
104 )?;
105 next_index = statement.bind(
106 &host.as_ref().map(|host| host.host_identifier.as_str()),
107 next_index,
108 )?;
109 }
110 statement.exec()
111 })
112 .await
113 .context("inserting new trusted state")?;
114 }
115 Ok(())
116 }
117
118 pub(crate) fn fetch_trusted_worktrees(
119 &self,
120 worktree_store: Option<Entity<WorktreeStore>>,
121 host: Option<RemoteHostLocation>,
122 cx: &App,
123 ) -> anyhow::Result<HashMap<Option<RemoteHostLocation>, HashSet<PathTrust>>> {
124 let trusted_worktrees = PROJECT_DB.trusted_worktrees()?;
125 Ok(trusted_worktrees
126 .into_iter()
127 .map(|(abs_path, user_name, host_name)| {
128 let db_host = match (user_name, host_name) {
129 (_, None) => None,
130 (None, Some(host_name)) => Some(RemoteHostLocation {
131 user_name: None,
132 host_identifier: SharedString::new(host_name),
133 }),
134 (Some(user_name), Some(host_name)) => Some(RemoteHostLocation {
135 user_name: Some(SharedString::new(user_name)),
136 host_identifier: SharedString::new(host_name),
137 }),
138 };
139
140 match abs_path {
141 Some(abs_path) => {
142 if db_host != host {
143 (db_host, PathTrust::AbsPath(abs_path))
144 } else if let Some(worktree_store) = &worktree_store {
145 find_worktree_in_store(worktree_store.read(cx), &abs_path, cx)
146 .map(PathTrust::Worktree)
147 .map(|trusted_worktree| (host.clone(), trusted_worktree))
148 .unwrap_or_else(|| (db_host.clone(), PathTrust::AbsPath(abs_path)))
149 } else {
150 (db_host, PathTrust::AbsPath(abs_path))
151 }
152 }
153 None => (db_host, PathTrust::Workspace),
154 }
155 })
156 .fold(HashMap::default(), |mut acc, (remote_host, path_trust)| {
157 acc.entry(remote_host)
158 .or_insert_with(HashSet::default)
159 .insert(path_trust);
160 acc
161 }))
162 }
163
164 query! {
165 fn trusted_worktrees() -> Result<Vec<(Option<PathBuf>, Option<String>, Option<String>)>> {
166 SELECT absolute_path, user_name, host_name
167 FROM trusted_worktrees
168 }
169 }
170
171 query! {
172 pub async fn clear_trusted_worktrees() -> Result<()> {
173 DELETE FROM trusted_worktrees
174 }
175 }
176}
177
178#[cfg(test)]
179mod tests {
180 use std::path::PathBuf;
181
182 use collections::{HashMap, HashSet};
183 use gpui::{SharedString, TestAppContext};
184 use serde_json::json;
185 use settings::SettingsStore;
186 use smol::lock::Mutex;
187 use util::path;
188
189 use crate::{
190 FakeFs, Project,
191 persistence::PROJECT_DB,
192 trusted_worktrees::{PathTrust, RemoteHostLocation},
193 };
194
195 static TEST_LOCK: Mutex<()> = Mutex::new(());
196
197 #[gpui::test]
198 async fn test_save_and_fetch_trusted_worktrees(cx: &mut TestAppContext) {
199 cx.executor().allow_parking();
200 let _guard = TEST_LOCK.lock().await;
201 PROJECT_DB.clear_trusted_worktrees().await.unwrap();
202 cx.update(|cx| {
203 if cx.try_global::<SettingsStore>().is_none() {
204 let settings = SettingsStore::test(cx);
205 cx.set_global(settings);
206 }
207 });
208
209 let fs = FakeFs::new(cx.executor());
210 fs.insert_tree(
211 path!("/"),
212 json!({
213 "project_a": { "main.rs": "" },
214 "project_b": { "lib.rs": "" }
215 }),
216 )
217 .await;
218
219 let project = Project::test(
220 fs,
221 [path!("/project_a").as_ref(), path!("/project_b").as_ref()],
222 cx,
223 )
224 .await;
225 let worktree_store = project.read_with(cx, |p, _| p.worktree_store());
226
227 let mut trusted_paths: HashMap<Option<RemoteHostLocation>, HashSet<PathBuf>> =
228 HashMap::default();
229 trusted_paths.insert(
230 None,
231 HashSet::from_iter([
232 PathBuf::from(path!("/project_a")),
233 PathBuf::from(path!("/project_b")),
234 ]),
235 );
236
237 PROJECT_DB
238 .save_trusted_worktrees(trusted_paths, HashSet::default())
239 .await
240 .unwrap();
241
242 let fetched = cx.update(|cx| {
243 PROJECT_DB.fetch_trusted_worktrees(Some(worktree_store.clone()), None, cx)
244 });
245 let fetched = fetched.unwrap();
246
247 let local_trust = fetched.get(&None).expect("should have local host entry");
248 assert_eq!(local_trust.len(), 2);
249 assert!(
250 local_trust
251 .iter()
252 .all(|p| matches!(p, PathTrust::Worktree(_)))
253 );
254
255 let fetched_no_store = cx
256 .update(|cx| PROJECT_DB.fetch_trusted_worktrees(None, None, cx))
257 .unwrap();
258 let local_trust_no_store = fetched_no_store
259 .get(&None)
260 .expect("should have local host entry");
261 assert_eq!(local_trust_no_store.len(), 2);
262 assert!(
263 local_trust_no_store
264 .iter()
265 .all(|p| matches!(p, PathTrust::AbsPath(_)))
266 );
267 }
268
269 #[gpui::test]
270 async fn test_save_and_fetch_workspace_trust(cx: &mut TestAppContext) {
271 cx.executor().allow_parking();
272 let _guard = TEST_LOCK.lock().await;
273 PROJECT_DB.clear_trusted_worktrees().await.unwrap();
274 cx.update(|cx| {
275 if cx.try_global::<SettingsStore>().is_none() {
276 let settings = SettingsStore::test(cx);
277 cx.set_global(settings);
278 }
279 });
280
281 let fs = FakeFs::new(cx.executor());
282 fs.insert_tree(path!("/root"), json!({ "main.rs": "" }))
283 .await;
284
285 let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
286 let worktree_store = project.read_with(cx, |p, _| p.worktree_store());
287
288 let trusted_workspaces = HashSet::from_iter([None]);
289 PROJECT_DB
290 .save_trusted_worktrees(HashMap::default(), trusted_workspaces)
291 .await
292 .unwrap();
293
294 let fetched = cx.update(|cx| {
295 PROJECT_DB.fetch_trusted_worktrees(Some(worktree_store.clone()), None, cx)
296 });
297 let fetched = fetched.unwrap();
298
299 let local_trust = fetched.get(&None).expect("should have local host entry");
300 assert!(local_trust.contains(&PathTrust::Workspace));
301
302 let fetched_no_store = cx
303 .update(|cx| PROJECT_DB.fetch_trusted_worktrees(None, None, cx))
304 .unwrap();
305 let local_trust_no_store = fetched_no_store
306 .get(&None)
307 .expect("should have local host entry");
308 assert!(local_trust_no_store.contains(&PathTrust::Workspace));
309 }
310
311 #[gpui::test]
312 async fn test_save_and_fetch_remote_host_trust(cx: &mut TestAppContext) {
313 cx.executor().allow_parking();
314 let _guard = TEST_LOCK.lock().await;
315 PROJECT_DB.clear_trusted_worktrees().await.unwrap();
316 cx.update(|cx| {
317 if cx.try_global::<SettingsStore>().is_none() {
318 let settings = SettingsStore::test(cx);
319 cx.set_global(settings);
320 }
321 });
322
323 let fs = FakeFs::new(cx.executor());
324 fs.insert_tree(path!("/root"), json!({ "main.rs": "" }))
325 .await;
326
327 let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
328 let worktree_store = project.read_with(cx, |p, _| p.worktree_store());
329
330 let remote_host = Some(RemoteHostLocation {
331 user_name: Some(SharedString::from("testuser")),
332 host_identifier: SharedString::from("remote.example.com"),
333 });
334
335 let mut trusted_paths: HashMap<Option<RemoteHostLocation>, HashSet<PathBuf>> =
336 HashMap::default();
337 trusted_paths.insert(
338 remote_host.clone(),
339 HashSet::from_iter([PathBuf::from("/home/testuser/project")]),
340 );
341
342 PROJECT_DB
343 .save_trusted_worktrees(trusted_paths, HashSet::default())
344 .await
345 .unwrap();
346
347 let fetched = cx.update(|cx| {
348 PROJECT_DB.fetch_trusted_worktrees(Some(worktree_store.clone()), None, cx)
349 });
350 let fetched = fetched.unwrap();
351
352 let remote_trust = fetched
353 .get(&remote_host)
354 .expect("should have remote host entry");
355 assert_eq!(remote_trust.len(), 1);
356 assert!(remote_trust
357 .iter()
358 .any(|p| matches!(p, PathTrust::AbsPath(path) if path == &PathBuf::from("/home/testuser/project"))));
359
360 let fetched_no_store = cx
361 .update(|cx| PROJECT_DB.fetch_trusted_worktrees(None, None, cx))
362 .unwrap();
363 let remote_trust_no_store = fetched_no_store
364 .get(&remote_host)
365 .expect("should have remote host entry");
366 assert_eq!(remote_trust_no_store.len(), 1);
367 assert!(remote_trust_no_store
368 .iter()
369 .any(|p| matches!(p, PathTrust::AbsPath(path) if path == &PathBuf::from("/home/testuser/project"))));
370 }
371
372 #[gpui::test]
373 async fn test_clear_trusted_worktrees(cx: &mut TestAppContext) {
374 cx.executor().allow_parking();
375 let _guard = TEST_LOCK.lock().await;
376 PROJECT_DB.clear_trusted_worktrees().await.unwrap();
377 cx.update(|cx| {
378 if cx.try_global::<SettingsStore>().is_none() {
379 let settings = SettingsStore::test(cx);
380 cx.set_global(settings);
381 }
382 });
383
384 let fs = FakeFs::new(cx.executor());
385 fs.insert_tree(path!("/root"), json!({ "main.rs": "" }))
386 .await;
387
388 let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
389 let worktree_store = project.read_with(cx, |p, _| p.worktree_store());
390
391 let trusted_workspaces = HashSet::from_iter([None]);
392 PROJECT_DB
393 .save_trusted_worktrees(HashMap::default(), trusted_workspaces)
394 .await
395 .unwrap();
396
397 PROJECT_DB.clear_trusted_worktrees().await.unwrap();
398
399 let fetched = cx.update(|cx| {
400 PROJECT_DB.fetch_trusted_worktrees(Some(worktree_store.clone()), None, cx)
401 });
402 let fetched = fetched.unwrap();
403
404 assert!(fetched.is_empty(), "should be empty after clear");
405
406 let fetched_no_store = cx
407 .update(|cx| PROJECT_DB.fetch_trusted_worktrees(None, None, cx))
408 .unwrap();
409 assert!(fetched_no_store.is_empty(), "should be empty after clear");
410 }
411}