feat: add a client auto-update feature

Signed-off-by: ljedrz <[email protected]>
ProvableHQ · May 5, 2022 · 7051e3b · 7051e3b
1 parent 2bd4e17
commit 7051e3b
Show file tree

Hide file tree

Showing 3 changed files with 187 additions and 13 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -35,6 +35,7 @@ prometheus = [ "snarkos-metrics/prometheus", "snarkos-network/prometheus" ]
 rpc = [ "snarkos-rpc" ]
 task-metrics = [ "snarkos-environment/task-metrics" ]
 test = [ "snarkos-metrics/test", "snarkos-network/test" ]
+auto-update = [ "reqwest", "serde_json", "tokio/process" ]
 
 [dependencies.aleo-std]
 version = "0.1.12"
@@ -59,7 +60,7 @@ version = "0.8"
 version = "1"
 
 [dependencies.self_update]
-version = "0.29"
+version = "0.30"
 features = [
   "archive-zip",
   "compression-zip-bzip2",
@@ -96,6 +97,15 @@ version = "0.8.0"
 version = "3.1"
 features = [ "derive" ]
 
+[dependencies.reqwest]
+version = "0.11"
+features = [ "json", "stream" ]
+optional = true
+
+[dependencies.serde_json]
+version = "1"
+optional = true
+
 [dependencies.thiserror]
 version = "1.0"
 

diff --git a/snarkos/node.rs b/snarkos/node.rs
@@ -209,14 +209,54 @@ impl Node {
             }
         }
 
+        #[cfg(feature = "auto-update")]
+        initialize_auto_updater(server).await;
+
         // Note: Do not move this. The pending await must be here otherwise
         // other snarkOS commands will not exit.
+        #[cfg(not(feature = "auto-update"))]
         std::future::pending::<()>().await;
 
         Ok(())
     }
 }
 
+#[cfg(feature = "auto-update")]
+async fn initialize_auto_updater<N: Network, E: Environment>(server: Server<N, E>) {
+    let mut auto_updater = crate::AutoUpdater::new().await.expect("Auto-updater failed");
+
+    loop {
+        match auto_updater.check_for_updates().await {
+            Ok(true) => {
+                warn!("[auto-updater]: You are using an outdated version of the snarkOS client! Automatically updating...");
+                if let Err(e) = auto_updater.update_local_repo().await {
+                    error!("[auto-updater]: Couldn't automatically update the local snarkOS repo: {}", e);
+                    continue;
+                }
+
+                server.shut_down().await;
+
+                if let Err(e) = auto_updater.rebuild_local_repo().await {
+                    error!("[auto-updater]: Couldn't automatically rebuild the local snarkOS repo: {}", e);
+                }
+
+                if let Err(e) = auto_updater.restart() {
+                    error!("[auto-updater]: Couldn't automatically restart the node: {}", e);
+                }
+
+                break;
+            }
+            Ok(false) => {
+                debug!("[auto-updater]: You are using an up-to-date version of the snarkOS client.");
+            }
+            Err(e) => {
+                error!("[auto-updater]: Couldn't check the snarkOS repo for updates: {}", e);
+            }
+        }
+        tokio::time::sleep(std::time::Duration::from_secs(crate::AUTO_UPDATE_INTERVAL_SECS)).await;
+    }
+}
+
 pub fn initialize_logger(verbosity: u8, log_sender: Option<mpsc::Sender<Vec<u8>>>) {
     match verbosity {
         0 => std::env::set_var("RUST_LOG", "info"),

diff --git a/snarkos/updater.rs b/snarkos/updater.rs
@@ -14,21 +14,145 @@
 // You should have received a copy of the GNU General Public License
 // along with the snarkOS library. If not, see <https://www.gnu.org/licenses/>.
 
+#[cfg(feature = "auto-update")]
+use anyhow::anyhow;
 use colored::Colorize;
 use self_update::{backends::github, version::bump_is_greater, Status};
+#[cfg(feature = "auto-update")]
+use std::{collections::HashMap, str};
+#[cfg(feature = "auto-update")]
+use tokio::process::Command;
+
+const SNARKOS_BIN_NAME: &str = "snarkos";
+const SNARKOS_REPO_NAME: &str = "snarkOS";
+const SNARKOS_REPO_OWNER: &str = "AleoHQ";
+
+#[cfg(feature = "auto-update")]
+const SNARKOS_CURRENT_BRANCH: &str = "testnet3";
+#[cfg(feature = "auto-update")]
+pub(crate) const AUTO_UPDATE_INTERVAL_SECS: u64 = 60 * 15; // 15 minutes
+
+#[cfg(feature = "auto-update")]
+pub struct AutoUpdater {
+    reqwest_client: reqwest::Client,
+    pub latest_sha: String,
+}
 
 pub struct Updater;
 
-impl Updater {
-    const SNARKOS_BIN_NAME: &'static str = "snarkos";
-    const SNARKOS_REPO_NAME: &'static str = "snarkOS";
-    const SNARKOS_REPO_OWNER: &'static str = "AleoHQ";
+#[cfg(feature = "auto-update")]
+impl AutoUpdater {
+    pub async fn get_build_sha() -> anyhow::Result<String> {
+        let local_repo_path = env!("CARGO_MANIFEST_DIR");
+        let local_repo_sha_bytes = Command::new("git")
+            .args(&["rev-parse", "HEAD"])
+            .current_dir(local_repo_path)
+            .output()
+            .await?
+            .stdout;
+        let local_repo_sha = str::from_utf8(&local_repo_sha_bytes)?.trim_end();
+
+        debug!("[auto-updater]: The local repo SHA is {}", local_repo_sha);
+
+        Ok(local_repo_sha.into())
+    }
+
+    pub async fn new() -> anyhow::Result<Self> {
+        let reqwest_client = reqwest::Client::builder().user_agent("curl").build()?;
+        let latest_sha = Self::get_build_sha().await?;
+
+        Ok(Self {
+            reqwest_client,
+            latest_sha,
+        })
+    }
+
+    pub async fn check_for_updates(&mut self) -> anyhow::Result<bool> {
+        let check_endpoint = format!(
+            "https://api.github.com/repos/{}/{}/commits/{}",
+            SNARKOS_REPO_OWNER, SNARKOS_REPO_NAME, SNARKOS_CURRENT_BRANCH
+        );
+        let req = self.reqwest_client.get(check_endpoint).build()?;
+        let remote_repo_sha = self
+            .reqwest_client
+            .execute(req)
+            .await?
+            .json::<HashMap<String, serde_json::Value>>()
+            .await?
+            .remove("sha")
+            .and_then(|v| v.as_str().map(|s| s.to_owned()))
+            .ok_or_else(|| anyhow!("GitHub's API didn't return the latest SHA"))?;
+
+        debug!("[auto-updater]: The remote repo SHA is {}", remote_repo_sha);
+
+        if self.latest_sha != remote_repo_sha {
+            self.latest_sha = remote_repo_sha;
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+
+    pub async fn update_local_repo(&self) -> anyhow::Result<()> {
+        let local_repo_path = env!("CARGO_MANIFEST_DIR");
+        let repo_addr = format!("https://github.com/{}/{}", SNARKOS_REPO_OWNER, SNARKOS_REPO_NAME);
+
+        Command::new("git")
+            .args(&["pull", &repo_addr, SNARKOS_CURRENT_BRANCH])
+            .current_dir(local_repo_path)
+            .output()
+            .await?;
+
+        debug!("[auto-updater]: Updated the local repo");
+
+        Ok(())
+    }
 
+    pub async fn rebuild_local_repo(&self) -> anyhow::Result<()> {
+        let local_repo_path = env!("CARGO_MANIFEST_DIR");
+
+        info!("[auto-updater]: Rebuilding the local repo...");
+
+        Command::new("cargo")
+            .args(&["build", "--release"])
+            .current_dir(local_repo_path)
+            .output()
+            .await?;
+
+        info!("[auto-updater]: Rebuilt the local repo; your snarkOS client is now up to date");
+
+        Ok(())
+    }
+
+    #[cfg(target_family = "unix")]
+    pub fn restart(&self) -> anyhow::Result<()> {
+        use std::{env, os::unix::process::CommandExt, path::PathBuf};
+
+        let mut binary_path: PathBuf = env!("CARGO_MANIFEST_DIR").parse().unwrap();
+        binary_path.push("target");
+        binary_path.push("release");
+
+        let original_args = env::args()
+            .skip(1)
+            .flat_map(|arg| arg.split('=').map(|s| s.to_owned()).collect::<Vec<_>>())
+            .collect::<Vec<String>>();
+
+        std::process::Command::new("cargo")
+            .args(&["run", "--release", "--"])
+            .args(&original_args)
+            .current_dir(binary_path)
+            .exec();
+
+        Ok(())
+    }
+}
+
+impl Updater {
     /// Show all available releases for `snarkos`.
     pub fn show_available_releases() -> Result<String, UpdaterError> {
         let releases = github::ReleaseList::configure()
-            .repo_owner(Self::SNARKOS_REPO_OWNER)
-            .repo_name(Self::SNARKOS_REPO_NAME)
+            .repo_owner(SNARKOS_REPO_OWNER)
+            .repo_name(SNARKOS_REPO_NAME)
             .build()?
             .fetch()?;
 
@@ -44,9 +168,9 @@ impl Updater {
         let mut update_builder = github::Update::configure();
 
         update_builder
-            .repo_owner(Self::SNARKOS_REPO_OWNER)
-            .repo_name(Self::SNARKOS_REPO_NAME)
-            .bin_name(Self::SNARKOS_BIN_NAME)
+            .repo_owner(SNARKOS_REPO_OWNER)
+            .repo_name(SNARKOS_REPO_NAME)
+            .bin_name(SNARKOS_BIN_NAME)
             .current_version(env!("CARGO_PKG_VERSION"))
             .show_download_progress(show_output)
             .no_confirm(true)
@@ -63,9 +187,9 @@ impl Updater {
     /// Check if there is an available update for `aleo` and return the newest release.
     pub fn update_available() -> Result<String, UpdaterError> {
         let updater = github::Update::configure()
-            .repo_owner(Self::SNARKOS_REPO_OWNER)
-            .repo_name(Self::SNARKOS_REPO_NAME)
-            .bin_name(Self::SNARKOS_BIN_NAME)
+            .repo_owner(SNARKOS_REPO_OWNER)
+            .repo_name(SNARKOS_REPO_NAME)
+            .bin_name(SNARKOS_BIN_NAME)
             .current_version(env!("CARGO_PKG_VERSION"))
             .build()?;