auto_update_helper.rs

  1#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
  2
  3#[cfg(target_os = "windows")]
  4mod dialog;
  5#[cfg(target_os = "windows")]
  6mod updater;
  7
  8#[cfg(target_os = "windows")]
  9fn main() {
 10    if let Err(e) = windows_impl::run() {
 11        log::error!("Error: Zed update failed, {:?}", e);
 12        windows_impl::show_error(format!("Error: {:?}", e));
 13    }
 14}
 15
 16#[cfg(not(target_os = "windows"))]
 17fn main() {}
 18
 19#[cfg(target_os = "windows")]
 20mod windows_impl {
 21    use std::{borrow::Cow, path::Path};
 22
 23    use super::dialog::create_dialog_window;
 24    use super::updater::perform_update;
 25    use anyhow::{Context as _, Result};
 26    use windows::{
 27        Win32::{
 28            Foundation::{HWND, LPARAM, WPARAM},
 29            UI::WindowsAndMessaging::{
 30                DispatchMessageW, GetMessageW, MB_ICONERROR, MB_SYSTEMMODAL, MSG, MessageBoxW,
 31                PostMessageW, WM_USER,
 32            },
 33        },
 34        core::HSTRING,
 35    };
 36
 37    pub(crate) const WM_JOB_UPDATED: u32 = WM_USER + 1;
 38    pub(crate) const WM_TERMINATE: u32 = WM_USER + 2;
 39
 40    #[derive(Debug, Default)]
 41    struct Args {
 42        launch: bool,
 43    }
 44
 45    pub(crate) fn run() -> Result<()> {
 46        let helper_dir = std::env::current_exe()?
 47            .parent()
 48            .context("No parent directory")?
 49            .to_path_buf();
 50        init_log(&helper_dir)?;
 51        let app_dir = helper_dir
 52            .parent()
 53            .context("No parent directory")?
 54            .to_path_buf();
 55
 56        log::info!("======= Starting Zed update =======");
 57        let (tx, rx) = std::sync::mpsc::channel();
 58        let hwnd = create_dialog_window(rx)?.0 as isize;
 59        let args = parse_args(std::env::args().skip(1));
 60        std::thread::spawn(move || {
 61            let result = perform_update(app_dir.as_path(), Some(hwnd), args.launch);
 62            tx.send(result).ok();
 63            unsafe { PostMessageW(Some(HWND(hwnd as _)), WM_TERMINATE, WPARAM(0), LPARAM(0)) }.ok();
 64        });
 65        unsafe {
 66            let mut message = MSG::default();
 67            while GetMessageW(&mut message, None, 0, 0).as_bool() {
 68                DispatchMessageW(&message);
 69            }
 70        }
 71        Ok(())
 72    }
 73
 74    fn init_log(helper_dir: &Path) -> Result<()> {
 75        simplelog::WriteLogger::init(
 76            simplelog::LevelFilter::Info,
 77            simplelog::Config::default(),
 78            std::fs::File::options()
 79                .append(true)
 80                .create(true)
 81                .open(helper_dir.join("auto_update_helper.log"))?,
 82        )?;
 83        Ok(())
 84    }
 85
 86    fn parse_args(input: impl IntoIterator<Item = String>) -> Args {
 87        let mut args: Args = Args { launch: true };
 88
 89        let mut input = input.into_iter();
 90        if let Some(arg) = input.next() {
 91            let launch_arg;
 92
 93            if arg == "--launch" {
 94                launch_arg = input.next().map(Cow::Owned);
 95            } else if let Some(rest) = arg.strip_prefix("--launch=") {
 96                launch_arg = Some(Cow::Borrowed(rest));
 97            } else {
 98                launch_arg = None;
 99            }
100
101            if launch_arg.as_deref() == Some("false") {
102                args.launch = false;
103            }
104        }
105
106        args
107    }
108
109    pub(crate) fn show_error(mut content: String) {
110        if content.len() > 600 {
111            content.truncate(600);
112            content.push_str("...\n");
113        }
114        let _ = unsafe {
115            MessageBoxW(
116                None,
117                &HSTRING::from(content),
118                windows::core::w!("Error: Zed update failed."),
119                MB_ICONERROR | MB_SYSTEMMODAL,
120            )
121        };
122    }
123
124    #[cfg(test)]
125    mod tests {
126        use crate::windows_impl::parse_args;
127
128        #[test]
129        fn test_parse_args() {
130            // launch can be specified via two separate arguments
131            assert!(parse_args(["--launch".into(), "true".into()]).launch);
132            assert!(!parse_args(["--launch".into(), "false".into()]).launch);
133
134            // launch can be specified via one single argument
135            assert!(parse_args(["--launch=true".into()]).launch);
136            assert!(!parse_args(["--launch=false".into()]).launch);
137
138            // launch defaults to true on no arguments
139            assert!(parse_args([]).launch);
140
141            // launch defaults to true on invalid arguments
142            assert!(parse_args(["--launch".into()]).launch);
143            assert!(parse_args(["--launch=".into()]).launch);
144            assert!(parse_args(["--launch=invalid".into()]).launch);
145        }
146    }
147}