1use std::{
2 ffi::OsStr,
3 os::windows::ffi::OsStrExt,
4 path::Path,
5 sync::LazyLock,
6 time::{Duration, Instant},
7};
8
9use anyhow::{Context as _, Result};
10use windows::{
11 Win32::{
12 Foundation::{HWND, LPARAM, WPARAM},
13 System::RestartManager::{
14 CCH_RM_SESSION_KEY, RmEndSession, RmGetList, RmRegisterResources, RmShutdown,
15 RmStartSession,
16 },
17 UI::WindowsAndMessaging::PostMessageW,
18 },
19 core::{PCWSTR, PWSTR},
20};
21
22use crate::windows_impl::WM_JOB_UPDATED;
23
24pub(crate) struct Job {
25 pub apply: Box<dyn Fn(&Path) -> Result<()> + Send + Sync>,
26 pub rollback: Box<dyn Fn(&Path) -> Result<()> + Send + Sync>,
27}
28
29impl Job {
30 pub fn mkdir(name: &'static Path) -> Self {
31 Job {
32 apply: Box::new(move |app_dir| {
33 let dir = app_dir.join(name);
34 std::fs::create_dir_all(&dir)
35 .context(format!("Failed to create directory {}", dir.display()))
36 }),
37 rollback: Box::new(move |app_dir| {
38 let dir = app_dir.join(name);
39 std::fs::remove_dir_all(&dir)
40 .context(format!("Failed to remove directory {}", dir.display()))
41 }),
42 }
43 }
44
45 pub fn mkdir_if_exists(name: &'static Path, check: &'static Path) -> Self {
46 Job {
47 apply: Box::new(move |app_dir| {
48 let dir = app_dir.join(name);
49 let check = app_dir.join(check);
50
51 if check.exists() {
52 std::fs::create_dir_all(&dir)
53 .context(format!("Failed to create directory {}", dir.display()))?
54 }
55 Ok(())
56 }),
57 rollback: Box::new(move |app_dir| {
58 let dir = app_dir.join(name);
59
60 if dir.exists() {
61 std::fs::remove_dir_all(&dir)
62 .context(format!("Failed to remove directory {}", dir.display()))?
63 }
64
65 Ok(())
66 }),
67 }
68 }
69
70 pub fn move_file(filename: &'static Path, new_filename: &'static Path) -> Self {
71 Job {
72 apply: Box::new(move |app_dir| {
73 let old_file = app_dir.join(filename);
74 let new_file = app_dir.join(new_filename);
75 log::info!(
76 "Moving file: {}->{}",
77 old_file.display(),
78 new_file.display()
79 );
80
81 std::fs::rename(&old_file, new_file)
82 .context(format!("Failed to move file {}", old_file.display()))
83 }),
84 rollback: Box::new(move |app_dir| {
85 let old_file = app_dir.join(filename);
86 let new_file = app_dir.join(new_filename);
87 log::info!(
88 "Rolling back file move: {}->{}",
89 old_file.display(),
90 new_file.display()
91 );
92
93 std::fs::rename(&new_file, &old_file).context(format!(
94 "Failed to rollback file move {}->{}",
95 new_file.display(),
96 old_file.display()
97 ))
98 }),
99 }
100 }
101
102 pub fn move_if_exists(filename: &'static Path, new_filename: &'static Path) -> Self {
103 Job {
104 apply: Box::new(move |app_dir| {
105 let old_file = app_dir.join(filename);
106 let new_file = app_dir.join(new_filename);
107
108 if old_file.exists() {
109 log::info!(
110 "Moving file: {}->{}",
111 old_file.display(),
112 new_file.display()
113 );
114
115 std::fs::rename(&old_file, new_file)
116 .context(format!("Failed to move file {}", old_file.display()))?;
117 }
118
119 Ok(())
120 }),
121 rollback: Box::new(move |app_dir| {
122 let old_file = app_dir.join(filename);
123 let new_file = app_dir.join(new_filename);
124
125 if new_file.exists() {
126 log::info!(
127 "Rolling back file move: {}->{}",
128 old_file.display(),
129 new_file.display()
130 );
131
132 std::fs::rename(&new_file, &old_file).context(format!(
133 "Failed to rollback file move {}->{}",
134 new_file.display(),
135 old_file.display()
136 ))?
137 }
138
139 Ok(())
140 }),
141 }
142 }
143
144 pub fn rmdir_nofail(filename: &'static Path) -> Self {
145 Job {
146 apply: Box::new(move |app_dir| {
147 let filename = app_dir.join(filename);
148 log::info!("Removing file: {}", filename.display());
149 if let Err(e) = std::fs::remove_dir_all(&filename) {
150 log::warn!("Failed to remove directory: {}", e);
151 }
152
153 Ok(())
154 }),
155 rollback: Box::new(move |app_dir| {
156 let filename = app_dir.join(filename);
157 anyhow::bail!(
158 "Delete operations cannot be rolled back, file: {}",
159 filename.display()
160 )
161 }),
162 }
163 }
164}
165
166#[cfg(not(test))]
167pub(crate) static JOBS: LazyLock<[Job; 22]> = LazyLock::new(|| {
168 fn p(value: &str) -> &Path {
169 Path::new(value)
170 }
171 [
172 // Move old files
173 // Not deleting because installing new files can fail
174 Job::mkdir(p("old")),
175 Job::move_file(p("Zed.exe"), p("old\\Zed.exe")),
176 Job::mkdir(p("old\\bin")),
177 Job::move_file(p("bin\\Zed.exe"), p("old\\bin\\Zed.exe")),
178 Job::move_file(p("bin\\zed"), p("old\\bin\\zed")),
179 //
180 // TODO: remove after a few weeks once everyone is on the new version and this file never exists
181 Job::move_if_exists(p("OpenConsole.exe"), p("old\\OpenConsole.exe")),
182 Job::mkdir(p("old\\x64")),
183 Job::mkdir(p("old\\arm64")),
184 Job::move_if_exists(p("x64\\OpenConsole.exe"), p("old\\x64\\OpenConsole.exe")),
185 Job::move_if_exists(
186 p("arm64\\OpenConsole.exe"),
187 p("old\\arm64\\OpenConsole.exe"),
188 ),
189 //
190 Job::move_file(p("conpty.dll"), p("old\\conpty.dll")),
191 // Copy new files
192 Job::move_file(p("install\\Zed.exe"), p("Zed.exe")),
193 Job::move_file(p("install\\bin\\Zed.exe"), p("bin\\Zed.exe")),
194 Job::move_file(p("install\\bin\\zed"), p("bin\\zed")),
195 //
196 Job::mkdir_if_exists(p("x64"), p("install\\x64")),
197 Job::mkdir_if_exists(p("arm64"), p("install\\arm64")),
198 Job::move_if_exists(
199 p("install\\x64\\OpenConsole.exe"),
200 p("x64\\OpenConsole.exe"),
201 ),
202 Job::move_if_exists(
203 p("install\\arm64\\OpenConsole.exe"),
204 p("arm64\\OpenConsole.exe"),
205 ),
206 //
207 Job::move_file(p("install\\conpty.dll"), p("conpty.dll")),
208 // Cleanup installer and updates folder
209 Job::rmdir_nofail(p("updates")),
210 Job::rmdir_nofail(p("install")),
211 // Cleanup old installation
212 Job::rmdir_nofail(p("old")),
213 ]
214});
215
216#[cfg(test)]
217pub(crate) static JOBS: LazyLock<[Job; 9]> = LazyLock::new(|| {
218 fn p(value: &str) -> &Path {
219 Path::new(value)
220 }
221 [
222 Job {
223 apply: Box::new(|_| {
224 std::thread::sleep(Duration::from_millis(1000));
225 if let Ok(config) = std::env::var("ZED_AUTO_UPDATE") {
226 match config.as_str() {
227 "err1" => Err(std::io::Error::other("Simulated error")).context("Anyhow!"),
228 "err2" => Ok(()),
229 _ => panic!("Unknown ZED_AUTO_UPDATE value: {}", config),
230 }
231 } else {
232 Ok(())
233 }
234 }),
235 rollback: Box::new(|_| {
236 unsafe { std::env::set_var("ZED_AUTO_UPDATE_RB", "rollback1") };
237 Ok(())
238 }),
239 },
240 Job::mkdir(p("test1")),
241 Job::mkdir_if_exists(p("test_exists"), p("test1")),
242 Job::mkdir_if_exists(p("test_missing"), p("dont")),
243 Job {
244 apply: Box::new(|folder| {
245 std::fs::write(folder.join("test1/test"), "test")?;
246 Ok(())
247 }),
248 rollback: Box::new(|folder| {
249 std::fs::remove_file(folder.join("test1/test"))?;
250 Ok(())
251 }),
252 },
253 Job::move_file(p("test1/test"), p("test1/moved")),
254 Job::move_if_exists(p("test1/test"), p("test1/noop")),
255 Job {
256 apply: Box::new(|_| {
257 std::thread::sleep(Duration::from_millis(1000));
258 if let Ok(config) = std::env::var("ZED_AUTO_UPDATE") {
259 match config.as_str() {
260 "err1" => Ok(()),
261 "err2" => Err(std::io::Error::other("Simulated error")).context("Anyhow!"),
262 _ => panic!("Unknown ZED_AUTO_UPDATE value: {}", config),
263 }
264 } else {
265 Ok(())
266 }
267 }),
268 rollback: Box::new(|_| Ok(())),
269 },
270 Job::rmdir_nofail(p("test1/nofolder")),
271 ]
272});
273
274/// Attempts to use Windows Restart Manager to release file handles held by other processes
275/// (e.g., Explorer.exe) on the files we need to move during the update.
276///
277/// This is a best-effort operation - if it fails, we'll still try the update and rely on
278/// the retry logic.
279fn release_file_handles(app_dir: &Path) -> Result<()> {
280 // Files that commonly get locked by Explorer or other processes
281 let files_to_release = [
282 app_dir.join("Zed.exe"),
283 app_dir.join("bin\\Zed.exe"),
284 app_dir.join("bin\\zed"),
285 app_dir.join("conpty.dll"),
286 ];
287
288 log::info!("Attempting to release file handles using Restart Manager...");
289
290 let mut session: u32 = 0;
291 let mut session_key = [0u16; CCH_RM_SESSION_KEY as usize + 1];
292
293 // Start a Restart Manager session
294 let err = unsafe {
295 RmStartSession(
296 &mut session,
297 Some(0),
298 PWSTR::from_raw(session_key.as_mut_ptr()),
299 )
300 };
301 if err.is_err() {
302 anyhow::bail!("RmStartSession failed: {err:?}");
303 }
304
305 // Ensure we end the session when done
306 let _session_guard = scopeguard::guard(session, |s| {
307 let _ = unsafe { RmEndSession(s) };
308 });
309
310 // Convert paths to wide strings for Windows API
311 let wide_paths: Vec<Vec<u16>> = files_to_release
312 .iter()
313 .filter(|p| p.exists())
314 .map(|p| {
315 OsStr::new(p)
316 .encode_wide()
317 .chain(std::iter::once(0))
318 .collect()
319 })
320 .collect();
321
322 if wide_paths.is_empty() {
323 log::info!("No files to release handles for");
324 return Ok(());
325 }
326
327 let pcwstr_paths: Vec<PCWSTR> = wide_paths
328 .iter()
329 .map(|p| PCWSTR::from_raw(p.as_ptr()))
330 .collect();
331
332 // Register the files we want to modify
333 let err = unsafe { RmRegisterResources(session, Some(&pcwstr_paths), None, None) };
334 if err.is_err() {
335 anyhow::bail!("RmRegisterResources failed: {err:?}");
336 }
337
338 // Check if any processes are using these files
339 let mut needed: u32 = 0;
340 let mut count: u32 = 0;
341 let mut reboot_reasons: u32 = 0;
342 let _ = unsafe { RmGetList(session, &mut needed, &mut count, None, &mut reboot_reasons) };
343
344 if needed == 0 {
345 log::info!("No processes are holding handles to the files");
346 return Ok(());
347 }
348
349 log::info!(
350 "{} process(es) are holding handles to the files, requesting release...",
351 needed
352 );
353
354 // Request processes to release their handles
355 // RmShutdown with flags=0 asks applications to release handles gracefully
356 // For Explorer, this typically releases icon cache handles without closing Explorer
357 let err = unsafe { RmShutdown(session, 0, None) };
358 if err.is_err() {
359 anyhow::bail!("RmShutdown failed: {:?}", err);
360 }
361
362 log::info!("Successfully requested handle release");
363 Ok(())
364}
365
366pub(crate) fn perform_update(app_dir: &Path, hwnd: Option<isize>, launch: bool) -> Result<()> {
367 let hwnd = hwnd.map(|ptr| HWND(ptr as _));
368
369 // Try to release file handles before starting the update
370 if let Err(e) = release_file_handles(app_dir) {
371 log::warn!("Restart Manager failed (will continue anyway): {}", e);
372 }
373
374 let mut last_successful_job = None;
375 'outer: for (i, job) in JOBS.iter().enumerate() {
376 let start = Instant::now();
377 loop {
378 if start.elapsed().as_secs() > 2 {
379 log::error!("Timed out, rolling back");
380 break 'outer;
381 }
382 match (job.apply)(app_dir) {
383 Ok(_) => {
384 last_successful_job = Some(i);
385 unsafe { PostMessageW(hwnd, WM_JOB_UPDATED, WPARAM(0), LPARAM(0))? };
386 break;
387 }
388 Err(err) => match err.downcast_ref::<std::io::Error>() {
389 Some(io_err) => match io_err.kind() {
390 std::io::ErrorKind::NotFound => {
391 log::error!("Operation failed with file not found, aborting: {}", err);
392 break 'outer;
393 }
394 _ => {
395 log::error!("Operation failed (retrying): {}", err);
396 std::thread::sleep(Duration::from_millis(50));
397 }
398 },
399 None => {
400 log::error!("Operation failed with unexpected error, aborting: {}", err);
401 break 'outer;
402 }
403 },
404 }
405 }
406 }
407
408 if last_successful_job
409 .map(|job| job != JOBS.len() - 1)
410 .unwrap_or(true)
411 {
412 let Some(last_successful_job) = last_successful_job else {
413 anyhow::bail!("Autoupdate failed, nothing to rollback");
414 };
415
416 for job in (0..=last_successful_job).rev() {
417 let job = &JOBS[job];
418 if let Err(e) = (job.rollback)(app_dir) {
419 anyhow::bail!(
420 "Job rollback failed, the app might be left in an inconsistent state: ({:?})",
421 e
422 );
423 }
424 }
425
426 anyhow::bail!("Autoupdate failed, rollback successful");
427 }
428
429 if launch {
430 #[allow(clippy::disallowed_methods, reason = "doesn't run in the main binary")]
431 let _ = std::process::Command::new(app_dir.join("Zed.exe")).spawn();
432 }
433 log::info!("Update completed successfully");
434 Ok(())
435}
436
437#[cfg(test)]
438mod test {
439 use super::perform_update;
440
441 #[test]
442 fn test_perform_update() {
443 let app_dir = tempfile::tempdir().unwrap();
444 let app_dir = app_dir.path();
445 assert!(perform_update(app_dir, None, false).is_ok());
446
447 let app_dir = tempfile::tempdir().unwrap();
448 let app_dir = app_dir.path();
449 // Simulate a timeout
450 unsafe { std::env::set_var("ZED_AUTO_UPDATE", "err1") };
451 let ret = perform_update(app_dir, None, false);
452 assert!(
453 ret.is_err_and(|e| e.to_string().as_str() == "Autoupdate failed, nothing to rollback")
454 );
455
456 let app_dir = tempfile::tempdir().unwrap();
457 let app_dir = app_dir.path();
458 // Simulate a timeout
459 unsafe { std::env::set_var("ZED_AUTO_UPDATE", "err2") };
460 let ret = perform_update(app_dir, None, false);
461 assert!(
462 ret.is_err_and(|e| e.to_string().as_str() == "Autoupdate failed, rollback successful")
463 );
464 assert!(std::env::var("ZED_AUTO_UPDATE_RB").is_ok_and(|e| e == "rollback1"));
465 }
466}