From 50a0563bb46c4eb7807af149debaf9a422973877 Mon Sep 17 00:00:00 2001
From: Roey Darwish Dror <roey.ghost@gmail.com>
Date: Mon, 3 Jun 2019 09:41:25 +0300
Subject: [PATCH] Implement respawn after upgrade in Windows

---
 README.md          |  3 +--
 src/error.rs       | 13 +++++++++++++
 src/main.rs        | 22 +++++++++++++++++++++-
 src/self_update.rs | 26 +++++++++++++++++---------
 4 files changed, 52 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index a6f14808..9a518cdc 100644
--- a/README.md
+++ b/README.md
@@ -43,8 +43,7 @@ distribution which ships the latest version of Rust, such as Arch Linux.
 ## Usage
 Just run `topgrade`. It will run the following steps:
 
-* Try to self-upgrade if compiled with this feature. On Unix systems Topgrade will also respawn
-  itself if it was upgraded
+* Try to self-upgrade if compiled with this feature. Topgrade will respawn itself if it was upgraded.
 * **Linux**: Run the system package manager:
   * **Arch based**: Run [yay](https://github.com/Jguer/yay) or fall back to pacman
   * **Redhat based**: Run `yum upgrade` (or `dnf` if present)
diff --git a/src/error.rs b/src/error.rs
index acc75255..5aca9b23 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -44,6 +44,10 @@ pub enum ErrorKind {
 
     #[fail(display = "A step should be skipped")]
     SkipStep,
+
+    #[cfg(all(windows, feature = "self-update"))]
+    #[fail(display = "Topgrade Upgraded")]
+    Upgraded(ExitStatus),
 }
 
 impl Fail for Error {
@@ -66,6 +70,15 @@ impl Error {
     pub fn kind(&self) -> ErrorKind {
         *self.inner.get_context()
     }
+
+    #[cfg(all(windows, feature = "self-update"))]
+    pub fn upgraded(&self) -> bool {
+        if let ErrorKind::Upgraded(_) = self.kind() {
+            true
+        } else {
+            false
+        }
+    }
 }
 
 impl From<ErrorKind> for Error {
diff --git a/src/main.rs b/src/main.rs
index 18b086a8..854aeb14 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -88,7 +88,20 @@ fn run() -> Result<(), Error> {
     {
         openssl_probe::init_ssl_cert_env_vars();
         if !run_type.dry() && env::var("TOPGRADE_NO_SELF_UPGRADE").is_err() {
-            if let Err(e) = self_update::self_update() {
+            let result = self_update::self_update();
+
+            #[cfg(windows)]
+            {
+                let upgraded = match &result {
+                    Ok(()) => false,
+                    Err(e) => e.upgraded(),
+                };
+                if upgraded {
+                    return result;
+                }
+            }
+
+            if let Err(e) = result {
                 print_warning(format!("Self update error: {}", e));
                 if let Some(cause) = e.cause() {
                     print_warning(format!("Caused by: {}", cause));
@@ -466,6 +479,13 @@ fn main() {
             exit(0);
         }
         Err(error) => {
+            #[cfg(all(windows, feature = "self-update"))]
+            {
+                if let ErrorKind::Upgraded(status) = error.kind() {
+                    exit(status.code().unwrap());
+                }
+            }
+
             let should_print = match error.kind() {
                 ErrorKind::StepFailed => false,
                 ErrorKind::Retry => error
diff --git a/src/self_update.rs b/src/self_update.rs
index 5a780108..328c6d39 100644
--- a/src/self_update.rs
+++ b/src/self_update.rs
@@ -3,16 +3,13 @@ use super::terminal::*;
 use failure::ResultExt;
 use self_update_crate;
 use self_update_crate::backends::github::{GitHubUpdateStatus, Update};
-#[cfg(unix)]
 use std::env;
 #[cfg(unix)]
 use std::os::unix::process::CommandExt;
-#[cfg(unix)]
 use std::process::Command;
 
 pub fn self_update() -> Result<(), Error> {
     print_separator("Self update");
-    #[cfg(unix)]
     let current_exe = env::current_exe();
 
     let target = self_update_crate::get_target().context(ErrorKind::SelfUpdate)?;
@@ -38,15 +35,26 @@ pub fn self_update() -> Result<(), Error> {
         println!("Topgrade is up-to-date");
     }
 
-    #[cfg(unix)]
     {
         if result.updated() {
             print_warning("Respawning...");
-            let err = Command::new(current_exe.context(ErrorKind::SelfUpdate)?)
-                .args(env::args().skip(1))
-                .env("TOPGRADE_NO_SELF_UPGRADE", "")
-                .exec();
-            Err(err).context(ErrorKind::SelfUpdate)?
+            let mut command = Command::new(current_exe.context(ErrorKind::SelfUpdate)?);
+            command.args(env::args().skip(1)).env("TOPGRADE_NO_SELF_UPGRADE", "");
+
+            #[cfg(unix)]
+            {
+                let err = command.exec();
+                Err(err).context(ErrorKind::SelfUpdate)?
+            }
+
+            #[cfg(windows)]
+            {
+                let status = command
+                    .spawn()
+                    .and_then(|mut c| c.wait())
+                    .context(ErrorKind::SelfUpdate)?;
+                Err(ErrorKind::Upgraded(status))?
+            }
         }
     }