updater.rs

  1use std::{
  2    cell::LazyCell,
  3    path::Path,
  4    time::{Duration, Instant},
  5};
  6
  7use anyhow::{Context as _, Result};
  8use windows::Win32::{
  9    Foundation::{HWND, LPARAM, WPARAM},
 10    UI::WindowsAndMessaging::PostMessageW,
 11};
 12
 13use crate::windows_impl::WM_JOB_UPDATED;
 14
 15pub(crate) struct Job {
 16    pub apply: Box<dyn Fn(&Path) -> Result<()>>,
 17    pub rollback: Box<dyn Fn(&Path) -> Result<()>>,
 18}
 19
 20impl Job {
 21    pub fn mkdir(name: &'static Path) -> Self {
 22        Job {
 23            apply: Box::new(move |app_dir| {
 24                let dir = app_dir.join(name);
 25                std::fs::create_dir_all(&dir)
 26                    .context(format!("Failed to create directory {}", dir.display()))
 27            }),
 28            rollback: Box::new(move |app_dir| {
 29                let dir = app_dir.join(name);
 30                std::fs::remove_dir_all(&dir)
 31                    .context(format!("Failed to remove directory {}", dir.display()))
 32            }),
 33        }
 34    }
 35
 36    pub fn mkdir_if_exists(name: &'static Path, check: &'static Path) -> Self {
 37        Job {
 38            apply: Box::new(move |app_dir| {
 39                let dir = app_dir.join(name);
 40                let check = app_dir.join(check);
 41
 42                if check.exists() {
 43                    std::fs::create_dir_all(&dir)
 44                        .context(format!("Failed to create directory {}", dir.display()))?
 45                }
 46                Ok(())
 47            }),
 48            rollback: Box::new(move |app_dir| {
 49                let dir = app_dir.join(name);
 50
 51                if dir.exists() {
 52                    std::fs::remove_dir_all(&dir)
 53                        .context(format!("Failed to remove directory {}", dir.display()))?
 54                }
 55
 56                Ok(())
 57            }),
 58        }
 59    }
 60
 61    pub fn move_file(filename: &'static Path, new_filename: &'static Path) -> Self {
 62        Job {
 63            apply: Box::new(move |app_dir| {
 64                let old_file = app_dir.join(filename);
 65                let new_file = app_dir.join(new_filename);
 66                log::info!(
 67                    "Moving file: {}->{}",
 68                    old_file.display(),
 69                    new_file.display()
 70                );
 71
 72                std::fs::rename(&old_file, new_file)
 73                    .context(format!("Failed to move file {}", old_file.display()))
 74            }),
 75            rollback: Box::new(move |app_dir| {
 76                let old_file = app_dir.join(filename);
 77                let new_file = app_dir.join(new_filename);
 78                log::info!(
 79                    "Rolling back file move: {}->{}",
 80                    old_file.display(),
 81                    new_file.display()
 82                );
 83
 84                std::fs::rename(&new_file, &old_file).context(format!(
 85                    "Failed to rollback file move {}->{}",
 86                    new_file.display(),
 87                    old_file.display()
 88                ))
 89            }),
 90        }
 91    }
 92
 93    pub fn move_if_exists(filename: &'static Path, new_filename: &'static Path) -> Self {
 94        Job {
 95            apply: Box::new(move |app_dir| {
 96                let old_file = app_dir.join(filename);
 97                let new_file = app_dir.join(new_filename);
 98
 99                if old_file.exists() {
100                    log::info!(
101                        "Moving file: {}->{}",
102                        old_file.display(),
103                        new_file.display()
104                    );
105
106                    std::fs::rename(&old_file, new_file)
107                        .context(format!("Failed to move file {}", old_file.display()))?;
108                }
109
110                Ok(())
111            }),
112            rollback: Box::new(move |app_dir| {
113                let old_file = app_dir.join(filename);
114                let new_file = app_dir.join(new_filename);
115
116                if new_file.exists() {
117                    log::info!(
118                        "Rolling back file move: {}->{}",
119                        old_file.display(),
120                        new_file.display()
121                    );
122
123                    std::fs::rename(&new_file, &old_file).context(format!(
124                        "Failed to rollback file move {}->{}",
125                        new_file.display(),
126                        old_file.display()
127                    ))?
128                }
129
130                Ok(())
131            }),
132        }
133    }
134
135    pub fn rmdir_nofail(filename: &'static Path) -> Self {
136        Job {
137            apply: Box::new(move |app_dir| {
138                let filename = app_dir.join(filename);
139                log::info!("Removing file: {}", filename.display());
140                if let Err(e) = std::fs::remove_dir_all(&filename) {
141                    log::warn!("Failed to remove directory: {}", e);
142                }
143
144                Ok(())
145            }),
146            rollback: Box::new(move |app_dir| {
147                let filename = app_dir.join(filename);
148                anyhow::bail!(
149                    "Delete operations cannot be rolled back, file: {}",
150                    filename.display()
151                )
152            }),
153        }
154    }
155}
156
157// app is single threaded
158#[cfg(not(test))]
159#[allow(clippy::declare_interior_mutable_const)]
160pub(crate) const JOBS: LazyCell<[Job; 22]> = LazyCell::new(|| {
161    fn p(value: &str) -> &Path {
162        Path::new(value)
163    }
164    [
165        // Move old files
166        // Not deleting because installing new files can fail
167        Job::mkdir(p("old")),
168        Job::move_file(p("Zed.exe"), p("old\\Zed.exe")),
169        Job::mkdir(p("old\\bin")),
170        Job::move_file(p("bin\\Zed.exe"), p("old\\bin\\Zed.exe")),
171        Job::move_file(p("bin\\zed"), p("old\\bin\\zed")),
172        //
173        // TODO: remove after a few weeks once everyone is on the new version and this file never exists
174        Job::move_if_exists(p("OpenConsole.exe"), p("old\\OpenConsole.exe")),
175        Job::mkdir(p("old\\x64")),
176        Job::mkdir(p("old\\arm64")),
177        Job::move_if_exists(p("x64\\OpenConsole.exe"), p("old\\x64\\OpenConsole.exe")),
178        Job::move_if_exists(
179            p("arm64\\OpenConsole.exe"),
180            p("old\\arm64\\OpenConsole.exe"),
181        ),
182        //
183        Job::move_file(p("conpty.dll"), p("old\\conpty.dll")),
184        // Copy new files
185        Job::move_file(p("install\\Zed.exe"), p("Zed.exe")),
186        Job::move_file(p("install\\bin\\Zed.exe"), p("bin\\Zed.exe")),
187        Job::move_file(p("install\\bin\\zed"), p("bin\\zed")),
188        //
189        Job::mkdir_if_exists(p("x64"), p("install\\x64")),
190        Job::mkdir_if_exists(p("arm64"), p("install\\arm64")),
191        Job::move_if_exists(
192            p("install\\x64\\OpenConsole.exe"),
193            p("x64\\OpenConsole.exe"),
194        ),
195        Job::move_if_exists(
196            p("install\\arm64\\OpenConsole.exe"),
197            p("arm64\\OpenConsole.exe"),
198        ),
199        //
200        Job::move_file(p("install\\conpty.dll"), p("conpty.dll")),
201        // Cleanup installer and updates folder
202        Job::rmdir_nofail(p("updates")),
203        Job::rmdir_nofail(p("install")),
204        // Cleanup old installation
205        Job::rmdir_nofail(p("old")),
206    ]
207});
208
209// app is single threaded
210#[cfg(test)]
211#[allow(clippy::declare_interior_mutable_const)]
212pub(crate) const JOBS: LazyCell<[Job; 9]> = LazyCell::new(|| {
213    fn p(value: &str) -> &Path {
214        Path::new(value)
215    }
216    [
217        Job {
218            apply: Box::new(|_| {
219                std::thread::sleep(Duration::from_millis(1000));
220                if let Ok(config) = std::env::var("ZED_AUTO_UPDATE") {
221                    match config.as_str() {
222                        "err1" => Err(std::io::Error::other("Simulated error")).context("Anyhow!"),
223                        "err2" => Ok(()),
224                        _ => panic!("Unknown ZED_AUTO_UPDATE value: {}", config),
225                    }
226                } else {
227                    Ok(())
228                }
229            }),
230            rollback: Box::new(|_| {
231                unsafe { std::env::set_var("ZED_AUTO_UPDATE_RB", "rollback1") };
232                Ok(())
233            }),
234        },
235        Job::mkdir(p("test1")),
236        Job::mkdir_if_exists(p("test_exists"), p("test1")),
237        Job::mkdir_if_exists(p("test_missing"), p("dont")),
238        Job {
239            apply: Box::new(|folder| {
240                std::fs::write(folder.join("test1/test"), "test")?;
241                Ok(())
242            }),
243            rollback: Box::new(|folder| {
244                std::fs::remove_file(folder.join("test1/test"))?;
245                Ok(())
246            }),
247        },
248        Job::move_file(p("test1/test"), p("test1/moved")),
249        Job::move_if_exists(p("test1/test"), p("test1/noop")),
250        Job {
251            apply: Box::new(|_| {
252                std::thread::sleep(Duration::from_millis(1000));
253                if let Ok(config) = std::env::var("ZED_AUTO_UPDATE") {
254                    match config.as_str() {
255                        "err1" => Ok(()),
256                        "err2" => Err(std::io::Error::other("Simulated error")).context("Anyhow!"),
257                        _ => panic!("Unknown ZED_AUTO_UPDATE value: {}", config),
258                    }
259                } else {
260                    Ok(())
261                }
262            }),
263            rollback: Box::new(|_| Ok(())),
264        },
265        Job::rmdir_nofail(p("test1/nofolder")),
266    ]
267});
268
269pub(crate) fn perform_update(app_dir: &Path, hwnd: Option<isize>, launch: bool) -> Result<()> {
270    let hwnd = hwnd.map(|ptr| HWND(ptr as _));
271
272    let mut last_successful_job = None;
273    'outer: for (i, job) in JOBS.iter().enumerate() {
274        let start = Instant::now();
275        loop {
276            if start.elapsed().as_secs() > 2 {
277                log::error!("Timed out, rolling back");
278                break 'outer;
279            }
280            match (job.apply)(app_dir) {
281                Ok(_) => {
282                    last_successful_job = Some(i);
283                    unsafe { PostMessageW(hwnd, WM_JOB_UPDATED, WPARAM(0), LPARAM(0))? };
284                    break;
285                }
286                Err(err) => {
287                    // Check if it's a "not found" error
288                    let io_err = err.downcast_ref::<std::io::Error>().unwrap();
289                    if io_err.kind() == std::io::ErrorKind::NotFound {
290                        log::warn!("File or folder not found.");
291                        last_successful_job = Some(i);
292                        unsafe { PostMessageW(hwnd, WM_JOB_UPDATED, WPARAM(0), LPARAM(0))? };
293                        break;
294                    }
295
296                    log::error!("Operation failed: {} ({:?})", err, io_err.kind());
297                    std::thread::sleep(Duration::from_millis(50));
298                }
299            }
300        }
301    }
302
303    if last_successful_job
304        .map(|job| job != JOBS.len() - 1)
305        .unwrap_or(true)
306    {
307        let Some(last_successful_job) = last_successful_job else {
308            anyhow::bail!("Autoupdate failed, nothing to rollback");
309        };
310
311        for job in (0..=last_successful_job).rev() {
312            let job = &JOBS[job];
313            if let Err(e) = (job.rollback)(app_dir) {
314                anyhow::bail!(
315                    "Job rollback failed, the app might be left in an inconsistent state: ({:?})",
316                    e
317                );
318            }
319        }
320
321        anyhow::bail!("Autoupdate failed, rollback successful");
322    }
323
324    if launch {
325        #[allow(clippy::disallowed_methods, reason = "doesn't run in the main binary")]
326        let _ = std::process::Command::new(app_dir.join("Zed.exe")).spawn();
327    }
328    log::info!("Update completed successfully");
329    Ok(())
330}
331
332#[cfg(test)]
333mod test {
334    use super::perform_update;
335
336    #[test]
337    fn test_perform_update() {
338        let app_dir = tempfile::tempdir().unwrap();
339        let app_dir = app_dir.path();
340        assert!(perform_update(app_dir, None, false).is_ok());
341
342        let app_dir = tempfile::tempdir().unwrap();
343        let app_dir = app_dir.path();
344        // Simulate a timeout
345        unsafe { std::env::set_var("ZED_AUTO_UPDATE", "err1") };
346        let ret = perform_update(app_dir, None, false);
347        assert!(
348            ret.is_err_and(|e| e.to_string().as_str() == "Autoupdate failed, nothing to rollback")
349        );
350
351        let app_dir = tempfile::tempdir().unwrap();
352        let app_dir = app_dir.path();
353        // Simulate a timeout
354        unsafe { std::env::set_var("ZED_AUTO_UPDATE", "err2") };
355        let ret = perform_update(app_dir, None, false);
356        assert!(
357            ret.is_err_and(|e| e.to_string().as_str() == "Autoupdate failed, rollback successful")
358        );
359        assert!(std::env::var("ZED_AUTO_UPDATE_RB").is_ok_and(|e| e == "rollback1"));
360    }
361}