From 60766e6d7c9f8e7398fd4c1b310294f8d944d85c Mon Sep 17 00:00:00 2001
From: Teo Voinea <58236992+tevoinea@users.noreply.github.com>
Date: Thu, 21 Sep 2023 13:31:45 -0400
Subject: [PATCH] Tevoinea/add version checking in local tasks (#3517)

* Compare task version to service version

* Swallow output when looking for appropriate name
---
 .../onefuzz-task/src/check_for_update.rs      | 78 +++++++++++++++++++
 src/agent/onefuzz-task/src/main.rs            | 14 +++-
 2 files changed, 90 insertions(+), 2 deletions(-)
 create mode 100644 src/agent/onefuzz-task/src/check_for_update.rs

diff --git a/src/agent/onefuzz-task/src/check_for_update.rs b/src/agent/onefuzz-task/src/check_for_update.rs
new file mode 100644
index 0000000000..51c0178158
--- /dev/null
+++ b/src/agent/onefuzz-task/src/check_for_update.rs
@@ -0,0 +1,78 @@
+use std::process::Stdio;
+
+use anyhow::Result;
+use serde_json::Value;
+
+pub fn run(onefuzz_built_version: &str) -> Result<()> {
+    // Find onefuzz cli
+    let common_names = ["onefuzz", "onefuzz.exe", "onefuzz.cmd"];
+    let mut valid_commands: Vec<_> = common_names
+        .into_iter()
+        .map(|name| {
+            (
+                name,
+                std::process::Command::new(name)
+                    .stderr(Stdio::null())
+                    .stdout(Stdio::null())
+                    .arg("-h")
+                    .spawn(),
+            )
+        })
+        .filter_map(|(name, child)| child.ok().map(|c| (name, c)))
+        .collect();
+
+    if valid_commands.is_empty() {
+        bail!(
+            "Could not find any of the following common names for the onefuzz-cli: {:?}",
+            common_names
+        );
+    }
+
+    let (name, child) = valid_commands
+        .first_mut()
+        .expect("Expected valid_commands to not be empty");
+
+    info!("Found the onefuzz cli at: {}", name);
+
+    // We just used this to check if it exists, we'll invoke it again later
+    let _ = child.kill();
+
+    // Run onefuzz info get
+    let output = std::process::Command::new(&name)
+        .args(["info", "get"])
+        .output()?;
+
+    if !output.status.success() {
+        bail!(
+            "Failed to run command `{} info get`. stderr: {:?}, stdout: {:?}",
+            name,
+            String::from_utf8(output.stderr),
+            String::from_utf8(output.stdout)
+        )
+    }
+
+    let stdout = String::from_utf8(output.stdout)?;
+    let info: Value = serde_json::from_str(&stdout)?;
+
+    if let Some(onefuzz_service_version) = info["versions"]["onefuzz"]["version"].as_str() {
+        if onefuzz_service_version == onefuzz_built_version {
+            println!("You are up to date!");
+        } else {
+            println!(
+                "Version mismatch. onefuzz-task version: {} | onefuzz service version: {}",
+                onefuzz_built_version, onefuzz_service_version
+            );
+            println!(
+                "To update, please run the following command: {} tools get .",
+                name
+            );
+            println!("Then extract the onefuzz-task binary from the appropriate OS folder");
+        }
+        return Ok(());
+    }
+
+    bail!(
+        "Failed to get onefuzz service version from cli response: {}",
+        stdout
+    )
+}
diff --git a/src/agent/onefuzz-task/src/main.rs b/src/agent/onefuzz-task/src/main.rs
index 77fd7a59ea..d230f92ff5 100644
--- a/src/agent/onefuzz-task/src/main.rs
+++ b/src/agent/onefuzz-task/src/main.rs
@@ -11,8 +11,10 @@ extern crate onefuzz;
 
 use anyhow::Result;
 use clap::{ArgMatches, Command};
+
 use std::io::{stdout, Write};
 
+mod check_for_update;
 mod local;
 mod managed;
 mod tasks;
@@ -20,12 +22,15 @@ mod tasks;
 const LICENSE_CMD: &str = "licenses";
 const LOCAL_CMD: &str = "local";
 const MANAGED_CMD: &str = "managed";
+const CHECK_FOR_UPDATE: &str = "check_for_update";
+
+const ONEFUZZ_BUILT_VERSION: &str = env!("ONEFUZZ_VERSION");
 
 fn main() -> Result<()> {
     let built_version = format!(
         "{} onefuzz:{} git:{}",
         crate_version!(),
-        env!("ONEFUZZ_VERSION"),
+        ONEFUZZ_BUILT_VERSION,
         env!("GIT_VERSION")
     );
 
@@ -33,7 +38,11 @@ fn main() -> Result<()> {
         .version(built_version)
         .subcommand(managed::cmd::args(MANAGED_CMD))
         .subcommand(local::cmd::args(LOCAL_CMD))
-        .subcommand(Command::new(LICENSE_CMD).about("display third-party licenses"));
+        .subcommand(Command::new(LICENSE_CMD).about("display third-party licenses"))
+        .subcommand(
+            Command::new(CHECK_FOR_UPDATE)
+                .about("compares the version of onefuzz-task with the onefuzz service"),
+        );
 
     let matches = app.get_matches();
 
@@ -55,6 +64,7 @@ async fn run(args: ArgMatches) -> Result<()> {
         Some((LICENSE_CMD, _)) => licenses(),
         Some((LOCAL_CMD, sub)) => local::cmd::run(sub.to_owned()).await,
         Some((MANAGED_CMD, sub)) => managed::cmd::run(sub).await,
+        Some((CHECK_FOR_UPDATE, _)) => check_for_update::run(ONEFUZZ_BUILT_VERSION),
         _ => anyhow::bail!("No command provided. Run with 'help' to see available commands."),
     }
 }