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; 2]> = LazyCell::new(|| {
213    [
214        Job {
215            apply: Box::new(|_| {
216                std::thread::sleep(Duration::from_millis(1000));
217                if let Ok(config) = std::env::var("ZED_AUTO_UPDATE") {
218                    match config.as_str() {
219                        "err1" => Err(std::io::Error::other("Simulated error")).context("Anyhow!"),
220                        "err2" => Ok(()),
221                        _ => panic!("Unknown ZED_AUTO_UPDATE value: {}", config),
222                    }
223                } else {
224                    Ok(())
225                }
226            }),
227            rollback: Box::new(|_| {
228                unsafe { std::env::set_var("ZED_AUTO_UPDATE_RB", "rollback1") };
229                Ok(())
230            }),
231        },
232        Job {
233            apply: Box::new(|_| {
234                std::thread::sleep(Duration::from_millis(1000));
235                if let Ok(config) = std::env::var("ZED_AUTO_UPDATE") {
236                    match config.as_str() {
237                        "err1" => Ok(()),
238                        "err2" => Err(std::io::Error::other("Simulated error")).context("Anyhow!"),
239                        _ => panic!("Unknown ZED_AUTO_UPDATE value: {}", config),
240                    }
241                } else {
242                    Ok(())
243                }
244            }),
245            rollback: Box::new(|_| Ok(())),
246        },
247    ]
248});
249
250pub(crate) fn perform_update(app_dir: &Path, hwnd: Option<isize>, launch: bool) -> Result<()> {
251    let hwnd = hwnd.map(|ptr| HWND(ptr as _));
252
253    let mut last_successful_job = None;
254    'outer: for (i, job) in JOBS.iter().enumerate() {
255        let start = Instant::now();
256        loop {
257            if start.elapsed().as_secs() > 2 {
258                log::error!("Timed out, rolling back");
259                break 'outer;
260            }
261            match (job.apply)(app_dir) {
262                Ok(_) => {
263                    last_successful_job = Some(i);
264                    unsafe { PostMessageW(hwnd, WM_JOB_UPDATED, WPARAM(0), LPARAM(0))? };
265                    break;
266                }
267                Err(err) => {
268                    // Check if it's a "not found" error
269                    let io_err = err.downcast_ref::<std::io::Error>().unwrap();
270                    if io_err.kind() == std::io::ErrorKind::NotFound {
271                        log::warn!("File or folder not found.");
272                        last_successful_job = Some(i);
273                        unsafe { PostMessageW(hwnd, WM_JOB_UPDATED, WPARAM(0), LPARAM(0))? };
274                        break;
275                    }
276
277                    log::error!("Operation failed: {} ({:?})", err, io_err.kind());
278                    std::thread::sleep(Duration::from_millis(50));
279                }
280            }
281        }
282    }
283
284    if last_successful_job
285        .map(|job| job != JOBS.len() - 1)
286        .unwrap_or(true)
287    {
288        let Some(last_successful_job) = last_successful_job else {
289            anyhow::bail!("Autoupdate failed, nothing to rollback");
290        };
291
292        for job in (0..=last_successful_job).rev() {
293            let job = &JOBS[job];
294            if let Err(e) = (job.rollback)(app_dir) {
295                anyhow::bail!(
296                    "Job rollback failed, the app might be left in an inconsistent state: ({:?})",
297                    e
298                );
299            }
300        }
301
302        anyhow::bail!("Autoupdate failed, rollback successful");
303    }
304
305    if launch {
306        #[allow(clippy::disallowed_methods, reason = "doesn't run in the main binary")]
307        let _ = std::process::Command::new(app_dir.join("Zed.exe")).spawn();
308    }
309    log::info!("Update completed successfully");
310    Ok(())
311}
312
313#[cfg(test)]
314mod test {
315    use super::perform_update;
316
317    #[test]
318    fn test_perform_update() {
319        let app_dir = std::path::Path::new("C:/");
320        assert!(perform_update(app_dir, None, false).is_ok());
321
322        // Simulate a timeout
323        unsafe { std::env::set_var("ZED_AUTO_UPDATE", "err1") };
324        let ret = perform_update(app_dir, None, false);
325        assert!(
326            ret.is_err_and(|e| e.to_string().as_str() == "Autoupdate failed, nothing to rollback")
327        );
328    }
329
330    #[test]
331    fn test_perform_update_rollback() {
332        let app_dir = std::path::Path::new("C:/");
333        assert!(perform_update(app_dir, None, false).is_ok());
334
335        // Simulate a timeout
336        unsafe { std::env::set_var("ZED_AUTO_UPDATE", "err2") };
337        let ret = perform_update(app_dir, None, false);
338        assert!(
339            ret.is_err_and(|e| e.to_string().as_str() == "Autoupdate failed, rollback successful")
340        );
341        assert!(std::env::var("ZED_AUTO_UPDATE_RB").is_ok_and(|e| e == "rollback1"));
342    }
343}