diff --git a/crates/auto_update/src/auto_update.rs b/crates/auto_update/src/auto_update.rs index 8f1a76c974..0d86651c00 100644 --- a/crates/auto_update/src/auto_update.rs +++ b/crates/auto_update/src/auto_update.rs @@ -1492,7 +1492,9 @@ mod tests { ); let will_restart = cx.expect_restart(); cx.update(|cx| cx.restart()); - let path = will_restart.await.unwrap().unwrap(); + let (path, arguments) = will_restart.await.unwrap(); + assert!(arguments.is_empty()); + let path = path.unwrap(); assert_eq!(path, tmp_dir.path().join("zed")); assert_eq!(std::fs::read_to_string(path).unwrap(), ""); } diff --git a/crates/auto_update_helper/src/auto_update_helper.rs b/crates/auto_update_helper/src/auto_update_helper.rs index 21ead701b2..e8f5237ce7 100644 --- a/crates/auto_update_helper/src/auto_update_helper.rs +++ b/crates/auto_update_helper/src/auto_update_helper.rs @@ -18,7 +18,10 @@ fn main() {} #[cfg(target_os = "windows")] mod windows_impl { - use std::{borrow::Cow, path::Path}; + use std::{ + ffi::{OsStr, OsString}, + path::Path, + }; use super::dialog::create_dialog_window; use super::updater::perform_update; @@ -37,9 +40,10 @@ mod windows_impl { pub(crate) const WM_JOB_UPDATED: u32 = WM_USER + 1; pub(crate) const WM_TERMINATE: u32 = WM_USER + 2; - #[derive(Debug, Default)] + #[derive(Debug, PartialEq, Eq)] struct Args { launch: bool, + launch_arguments: Vec, } pub(crate) fn run() -> Result<()> { @@ -56,9 +60,14 @@ mod windows_impl { log::info!("======= Starting Zed update ======="); let (tx, rx) = std::sync::mpsc::channel(); let hwnd = create_dialog_window(rx)?.0 as isize; - let args = parse_args(std::env::args().skip(1)); + let args = parse_args(std::env::args_os().skip(1)); std::thread::spawn(move || { - let result = perform_update(app_dir.as_path(), Some(hwnd), args.launch); + let result = perform_update( + app_dir.as_path(), + Some(hwnd), + args.launch, + &args.launch_arguments, + ); tx.send(result).ok(); unsafe { PostMessageW(Some(HWND(hwnd as _)), WM_TERMINATE, WPARAM(0), LPARAM(0)) }.ok(); }); @@ -83,24 +92,23 @@ mod windows_impl { Ok(()) } - fn parse_args(input: impl IntoIterator) -> Args { - let mut args: Args = Args { launch: true }; + fn parse_args(input: impl IntoIterator) -> Args { + let mut args = Args { + launch: true, + launch_arguments: Vec::new(), + }; let mut input = input.into_iter(); if let Some(arg) = input.next() { - let launch_arg; - - if arg == "--launch" { - launch_arg = input.next().map(Cow::Owned); - } else if let Some(rest) = arg.strip_prefix("--launch=") { - launch_arg = Some(Cow::Borrowed(rest)); + if arg == OsStr::new("--launch") { + args.launch = input.next().as_deref() != Some(OsStr::new("false")); + } else if let Some(launch) = arg.to_str().and_then(|arg| arg.strip_prefix("--launch=")) + { + args.launch = launch != "false"; } else { - launch_arg = None; - } - - if launch_arg.as_deref() == Some("false") { - args.launch = false; + args.launch_arguments.push(arg); } + args.launch_arguments.extend(input); } args @@ -123,6 +131,8 @@ mod windows_impl { #[cfg(test)] mod tests { + use std::ffi::OsString; + use crate::windows_impl::parse_args; #[test] @@ -143,5 +153,17 @@ mod windows_impl { assert!(parse_args(["--launch=".into()]).launch); assert!(parse_args(["--launch=invalid".into()]).launch); } + + #[test] + fn test_parse_args_preserves_launch_arguments() { + let launch_arguments = vec![ + OsString::from("--user-data-dir"), + OsString::from(r"C:\Zed Data"), + ]; + assert_eq!( + parse_args(launch_arguments.clone()).launch_arguments, + launch_arguments + ); + } } } diff --git a/crates/auto_update_helper/src/updater.rs b/crates/auto_update_helper/src/updater.rs index 7821c908c4..d8dcb7311b 100644 --- a/crates/auto_update_helper/src/updater.rs +++ b/crates/auto_update_helper/src/updater.rs @@ -1,5 +1,5 @@ use std::{ - ffi::OsStr, + ffi::{OsStr, OsString}, os::windows::ffi::OsStrExt, path::Path, sync::LazyLock, @@ -363,7 +363,19 @@ fn release_file_handles(app_dir: &Path) -> Result<()> { Ok(()) } -pub(crate) fn perform_update(app_dir: &Path, hwnd: Option, launch: bool) -> Result<()> { +#[allow(clippy::disallowed_methods, reason = "doesn't run in the main binary")] +fn zed_launch_command(app_dir: &Path, launch_arguments: &[OsString]) -> std::process::Command { + let mut command = std::process::Command::new(app_dir.join("Zed.exe")); + command.args(launch_arguments); + command +} + +pub(crate) fn perform_update( + app_dir: &Path, + hwnd: Option, + launch: bool, + launch_arguments: &[OsString], +) -> Result<()> { let hwnd = hwnd.map(|ptr| HWND(ptr as _)); // Try to release file handles before starting the update @@ -428,7 +440,9 @@ pub(crate) fn perform_update(app_dir: &Path, hwnd: Option, launch: bool) if launch { #[allow(clippy::disallowed_methods, reason = "doesn't run in the main binary")] - let _ = std::process::Command::new(app_dir.join("Zed.exe")).spawn(); + let _child = zed_launch_command(app_dir, launch_arguments) + .spawn() + .context("Failed to launch Zed after update")?; } log::info!("Update completed successfully"); Ok(()) @@ -436,19 +450,42 @@ pub(crate) fn perform_update(app_dir: &Path, hwnd: Option, launch: bool) #[cfg(test)] mod test { - use super::perform_update; + use std::{ffi::OsString, path::Path}; + + use super::{perform_update, zed_launch_command}; + + #[test] + fn test_zed_launch_command_preserves_arguments() { + let arguments = vec![ + OsString::from("--user-data-dir"), + OsString::from(r"C:\Zed Data"), + ]; + let command = zed_launch_command(Path::new(r"C:\Program Files\Zed"), &arguments); + + assert_eq!( + command.get_program(), + Path::new(r"C:\Program Files\Zed\Zed.exe").as_os_str() + ); + assert_eq!( + command.get_args().collect::>(), + arguments + .iter() + .map(OsString::as_os_str) + .collect::>() + ); + } #[test] fn test_perform_update() { let app_dir = tempfile::tempdir().unwrap(); let app_dir = app_dir.path(); - assert!(perform_update(app_dir, None, false).is_ok()); + assert!(perform_update(app_dir, None, false, &[]).is_ok()); let app_dir = tempfile::tempdir().unwrap(); let app_dir = app_dir.path(); // Simulate a timeout unsafe { std::env::set_var("ZED_AUTO_UPDATE", "err1") }; - let ret = perform_update(app_dir, None, false); + let ret = perform_update(app_dir, None, false, &[]); assert!( ret.is_err_and(|e| e.to_string().as_str() == "Autoupdate failed, nothing to rollback") ); @@ -457,7 +494,7 @@ mod test { let app_dir = app_dir.path(); // Simulate a timeout unsafe { std::env::set_var("ZED_AUTO_UPDATE", "err2") }; - let ret = perform_update(app_dir, None, false); + let ret = perform_update(app_dir, None, false, &[]); assert!( ret.is_err_and(|e| e.to_string().as_str() == "Autoupdate failed, rollback successful") ); diff --git a/crates/gpui/src/app.rs b/crates/gpui/src/app.rs index dedebbe585..452143676a 100644 --- a/crates/gpui/src/app.rs +++ b/crates/gpui/src/app.rs @@ -2,6 +2,7 @@ use scheduler::Instant; use std::{ any::{TypeId, type_name}, cell::{BorrowMutError, Cell, Ref, RefCell, RefMut}, + ffi::OsString, marker::PhantomData, mem, ops::{Deref, DerefMut}, @@ -206,6 +207,12 @@ impl Application { self } + /// Configures arguments to pass when restarting the application. + pub fn with_restart_arguments(self, arguments: Vec) -> Self { + self.0.borrow_mut().restart_arguments = arguments; + self + } + /// Sets the HTTP client for the application. pub fn with_http_client(self, http_client: Arc) -> Self { let mut context_lock = self.0.borrow_mut(); @@ -729,6 +736,7 @@ pub struct App { pub(crate) pending_notifications: FxHashSet, pub(crate) pending_global_notifications: TypeIdHashSet, pub(crate) restart_path: Option, + pub(crate) restart_arguments: Vec, pub(crate) layout_id_buffer: Vec, // We recycle this memory across layout requests. pub(crate) propagate_event: bool, pub(crate) prompt_builder: Option, @@ -834,6 +842,7 @@ impl App { quit_observers: SubscriberSet::new(), restart_observers: SubscriberSet::new(), restart_path: None, + restart_arguments: Vec::new(), window_closed_observers: SubscriberSet::new(), layout_id_buffer: Default::default(), propagate_event: true, @@ -1587,7 +1596,10 @@ impl App { self.restart_observers .clone() .retain(&(), |observer| observer(self)); - self.platform.restart(self.restart_path.take()) + self.platform.restart( + self.restart_path.take(), + std::mem::take(&mut self.restart_arguments), + ) } /// Sets the path to use when restarting the application. @@ -3050,7 +3062,10 @@ impl<'a, T> Drop for GpuiBorrow<'a, T> { #[cfg(test)] mod test { - use std::{cell::RefCell, rc::Rc}; + use std::{cell::RefCell, ffi::OsString, path::PathBuf, rc::Rc}; + + #[cfg(unix)] + use std::os::unix::ffi::OsStringExt; use crate::{AppContext, TestAppContext}; @@ -3084,4 +3099,26 @@ mod test { assert_eq!(*observation_count.borrow(), 2); } + + #[gpui::test] + async fn test_restart_preserves_path_and_arguments(cx: &mut TestAppContext) { + #[cfg(unix)] + let user_data_dir = OsString::from_vec(b"/tmp/zed data/\xff".to_vec()); + #[cfg(not(unix))] + let user_data_dir = OsString::from("C:\\zed data"); + let arguments = vec![OsString::from("--user-data-dir"), user_data_dir]; + let restart_path = PathBuf::from("updated-zed"); + let _application = + super::Application(cx.app.clone()).with_restart_arguments(arguments.clone()); + let restart = cx.expect_restart(); + + cx.update(|cx| { + cx.set_restart_path(restart_path.clone()); + cx.restart(); + }); + + let (path, restart_arguments) = restart.await.expect("restart was not requested"); + assert_eq!(path, Some(restart_path)); + assert_eq!(restart_arguments, arguments); + } } diff --git a/crates/gpui/src/app/test_context.rs b/crates/gpui/src/app/test_context.rs index c80c259345..1dceee09d7 100644 --- a/crates/gpui/src/app/test_context.rs +++ b/crates/gpui/src/app/test_context.rs @@ -404,7 +404,7 @@ impl TestAppContext { } /// Returns true if there's an alert dialog open. - pub fn expect_restart(&self) -> oneshot::Receiver> { + pub fn expect_restart(&self) -> oneshot::Receiver<(Option, Vec)> { let (tx, rx) = futures::channel::oneshot::channel(); self.test_platform.expect_restart.borrow_mut().replace(tx); rx diff --git a/crates/gpui/src/platform.rs b/crates/gpui/src/platform.rs index 7445765293..51c24da048 100644 --- a/crates/gpui/src/platform.rs +++ b/crates/gpui/src/platform.rs @@ -64,6 +64,7 @@ use std::io::Cursor; use std::ops; use std::time::Duration; use std::{ + ffi::OsString, fmt::{self, Debug}, ops::Range, path::{Path, PathBuf}, @@ -129,7 +130,7 @@ pub trait Platform: 'static { fn run(&self, on_finish_launching: Box); fn quit(&self); - fn restart(&self, binary_path: Option); + fn restart(&self, binary_path: Option, arguments: Vec); fn activate(&self, ignoring_other_apps: bool); fn hide(&self); fn hide_other_apps(&self); diff --git a/crates/gpui/src/platform/test/platform.rs b/crates/gpui/src/platform/test/platform.rs index 42edcd6cba..3982190f5c 100644 --- a/crates/gpui/src/platform/test/platform.rs +++ b/crates/gpui/src/platform/test/platform.rs @@ -35,7 +35,8 @@ pub(crate) struct TestPlatform { pub opened_url: RefCell>, pub(crate) system_notifications: RefCell, pub text_system: Arc, - pub expect_restart: RefCell>>>, + pub expect_restart: + RefCell, Vec)>>>, headless_renderer_factory: Option Option>>>, weak: Weak, } @@ -340,9 +341,9 @@ impl Platform for TestPlatform { fn quit(&self) {} - fn restart(&self, path: Option) { + fn restart(&self, path: Option, arguments: Vec) { if let Some(tx) = self.expect_restart.take() { - tx.send(path).unwrap(); + tx.send((path, arguments)).unwrap(); } } diff --git a/crates/gpui/src/platform/visual_test.rs b/crates/gpui/src/platform/visual_test.rs index 41f93ae7fc..3ff561b26a 100644 --- a/crates/gpui/src/platform/visual_test.rs +++ b/crates/gpui/src/platform/visual_test.rs @@ -83,7 +83,7 @@ impl Platform for VisualTestPlatform { fn quit(&self) {} - fn restart(&self, _binary_path: Option) {} + fn restart(&self, _binary_path: Option, _arguments: Vec) {} fn activate(&self, _ignoring_other_apps: bool) {} diff --git a/crates/gpui_linux/src/linux/platform.rs b/crates/gpui_linux/src/linux/platform.rs index 33307034ae..7e8ec2224e 100644 --- a/crates/gpui_linux/src/linux/platform.rs +++ b/crates/gpui_linux/src/linux/platform.rs @@ -285,7 +285,7 @@ impl Platform for LinuxPlatform

{ self.inner.compositor_name() } - fn restart(&self, binary_path: Option) { + fn restart(&self, binary_path: Option, arguments: Vec) { use std::os::unix::process::CommandExt as _; // get the process id of the current process @@ -312,7 +312,9 @@ impl Platform for LinuxPlatform

{ sleep 0.1 done - "$1" + app_path="$1" + shift + "$app_path" "$@" "#; #[allow( @@ -325,6 +327,7 @@ impl Platform for LinuxPlatform

{ .arg(script) .arg(&app_pid) .arg(&app_path) + .args(arguments) .process_group(0) .spawn(); diff --git a/crates/gpui_macos/src/platform.rs b/crates/gpui_macos/src/platform.rs index c9bd307600..3ec54e8452 100644 --- a/crates/gpui_macos/src/platform.rs +++ b/crates/gpui_macos/src/platform.rs @@ -537,7 +537,7 @@ impl Platform for MacPlatform { } } - fn restart(&self, binary_path: Option) { + fn restart(&self, binary_path: Option, arguments: Vec) { use std::os::unix::process::CommandExt as _; let app_pid = std::process::id().to_string(); @@ -557,7 +557,13 @@ impl Platform for MacPlatform { while kill -0 $0 2> /dev/null; do sleep 0.1 done - open "$1" + app_path="$1" + shift + if (($# > 0)); then + open "$app_path" --args "$@" + else + open "$app_path" + fi "#; #[allow( @@ -569,6 +575,7 @@ impl Platform for MacPlatform { .arg(script) .arg(app_pid) .arg(app_path) + .args(arguments) .process_group(0) .spawn(); diff --git a/crates/gpui_web/src/platform.rs b/crates/gpui_web/src/platform.rs index 65e092cac8..6e6f9c0c77 100644 --- a/crates/gpui_web/src/platform.rs +++ b/crates/gpui_web/src/platform.rs @@ -307,7 +307,7 @@ impl Platform for WebPlatform { log::warn!("WebPlatform::quit called, but quitting is not supported in the browser ."); } - fn restart(&self, _binary_path: Option) {} + fn restart(&self, _binary_path: Option, _arguments: Vec) {} fn activate(&self, _ignoring_other_apps: bool) {} diff --git a/crates/gpui_windows/src/platform.rs b/crates/gpui_windows/src/platform.rs index d3bb150407..0c0f18a0ee 100644 --- a/crates/gpui_windows/src/platform.rs +++ b/crates/gpui_windows/src/platform.rs @@ -1,6 +1,7 @@ use std::{ cell::{Cell, RefCell}, - ffi::OsStr, + ffi::{OsStr, OsString}, + os::windows::ffi::{OsStrExt as _, OsStringExt as _}, path::{Path, PathBuf}, rc::{Rc, Weak}, sync::{ @@ -371,6 +372,39 @@ fn translate_accelerator(msg: &MSG) -> Option<()> { (result.0 == 0).then_some(()) } +fn encode_restart_arguments(arguments: &[OsString]) -> OsString { + // `Start-Process` accepts a single native command line, so quote each argument according to + // the Windows argv parsing rules before passing the complete string through the environment. + let mut encoded = Vec::new(); + + for (index, argument) in arguments.iter().enumerate() { + if index > 0 { + encoded.push(b' ' as u16); + } + encoded.push(b'"' as u16); + + let mut backslash_count = 0; + for code_unit in argument.encode_wide() { + if code_unit == b'\\' as u16 { + backslash_count += 1; + } else { + if code_unit == b'"' as u16 { + encoded.extend(std::iter::repeat_n(b'\\' as u16, backslash_count * 2 + 1)); + } else { + encoded.extend(std::iter::repeat_n(b'\\' as u16, backslash_count)); + } + backslash_count = 0; + encoded.push(code_unit); + } + } + + encoded.extend(std::iter::repeat_n(b'\\' as u16, backslash_count * 2)); + encoded.push(b'"' as u16); + } + + OsString::from_wide(&encoded) +} + impl Platform for WindowsPlatform { fn background_executor(&self) -> BackgroundExecutor { self.background_executor.clone() @@ -436,28 +470,33 @@ impl Platform for WindowsPlatform { .detach(); } - fn restart(&self, binary_path: Option) { + fn restart(&self, binary_path: Option, arguments: Vec) { let pid = std::process::id(); let Some(app_path) = binary_path.or(self.app_path().log_err()) else { return; }; - let script = format!( - r#" - $pidToWaitFor = {} - $exePath = "{}" + let script = r#" + $pidToWaitFor = $env:ZED_RESTART_PID + $exePath = $env:ZED_RESTART_EXECUTABLE + $argumentList = $env:ZED_RESTART_ARGUMENTS - while ($true) {{ + [Environment]::SetEnvironmentVariable("ZED_RESTART_PID", $null) + [Environment]::SetEnvironmentVariable("ZED_RESTART_EXECUTABLE", $null) + [Environment]::SetEnvironmentVariable("ZED_RESTART_ARGUMENTS", $null) + + while ($true) { $process = Get-Process -Id $pidToWaitFor -ErrorAction SilentlyContinue - if (-not $process) {{ - Start-Process -FilePath $exePath + if (-not $process) { + if ([string]::IsNullOrEmpty($argumentList)) { + Start-Process -FilePath $exePath + } else { + Start-Process -FilePath $exePath -ArgumentList $argumentList + } break - }} + } Start-Sleep -Seconds 0.1 - }} - "#, - pid, - app_path.display(), - ); + } + "#; // Defer spawning to the foreground executor so it runs after the // current `AppCell` borrow is released. On Windows, `Command::spawn()` @@ -466,14 +505,19 @@ impl Platform for WindowsPlatform { // borrow of the `AppCell` ending up with a double borrow panic self.foreground_executor .spawn(async move { + let mut command = new_std_command(get_windows_system_shell()); + let arguments = encode_restart_arguments(&arguments); + command + .arg("-command") + .arg(script) + .env("ZED_RESTART_PID", pid.to_string()) + .env("ZED_RESTART_EXECUTABLE", app_path) + .env("ZED_RESTART_ARGUMENTS", arguments); #[allow( clippy::disallowed_methods, reason = "We are restarting ourselves, using std command thus is fine" )] - let restart_process = new_std_command(get_windows_system_shell()) - .arg("-command") - .arg(script) - .spawn(); + let restart_process = command.spawn(); match restart_process { Ok(_) => unsafe { PostQuitMessage(0) }, @@ -1493,9 +1537,29 @@ unsafe extern "system" fn window_procedure( #[cfg(test)] mod tests { + use std::ffi::{OsStr, OsString}; + use crate::{read_from_clipboard, write_to_clipboard}; use gpui::ClipboardItem; + use super::encode_restart_arguments; + + #[test] + fn test_encode_restart_arguments() { + assert_eq!(encode_restart_arguments(&[]), OsStr::new("")); + assert_eq!( + encode_restart_arguments(&[ + OsString::from("--user-data-dir"), + OsString::from(r"C:\Zed Data"), + ]), + OsStr::new(r#""--user-data-dir" "C:\Zed Data""#) + ); + assert_eq!( + encode_restart_arguments(&[OsString::from(r"C:\")]), + OsStr::new(r#""C:\\""#) + ); + } + #[test] fn test_clipboard() { let item = ClipboardItem::new_string("你好,我是张小白".to_string()); diff --git a/crates/zed/src/main.rs b/crates/zed/src/main.rs index 04d2b78671..b8eba3bf97 100644 --- a/crates/zed/src/main.rs +++ b/crates/zed/src/main.rs @@ -268,10 +268,15 @@ fn main() { return; } - // Set custom data directory. - if let Some(dir) = &args.user_data_dir { - paths::set_custom_data_dir(dir); - } + let restart_arguments = if let Some(directory) = args.user_data_dir.as_deref() { + let directory = paths::set_custom_data_dir(directory); + vec![ + std::ffi::OsString::from("--user-data-dir"), + directory.as_os_str().to_owned(), + ] + } else { + Vec::new() + }; #[cfg(target_os = "windows")] match util::get_zed_cli_path() { @@ -340,7 +345,9 @@ fn main() { #[cfg(windows)] check_for_conpty_dll(); - let app = build_application().with_assets(Assets); + let app = build_application() + .with_assets(Assets) + .with_restart_arguments(restart_arguments); let app_db = db::AppDatabase::new(); let system_id = app.background_executor().spawn(system_id());