From 152bc2ec949ac1de43b4c566eae7635a830c6dff Mon Sep 17 00:00:00 2001
From: 0xMRTT <0xMRTT@tuta.io>
Date: Fri, 24 Jun 2022 23:40:08 +0200
Subject: [PATCH] Add yarn support. Fix #958

---
 src/config.rs     | 17 ++++++++++
 src/main.rs       |  1 +
 src/steps/node.rs | 83 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 101 insertions(+)

diff --git a/src/config.rs b/src/config.rs
index 9eec74c7..1e0b238e 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -166,6 +166,13 @@ pub struct Windows {
     enable_winget: Option<bool>,
 }
 
+#[derive(Deserialize, Default, Debug)]
+#[serde(deny_unknown_fields)]
+#[allow(clippy::upper_case_acronyms)]
+pub struct Yarn {
+    use_sudo: Option<bool>,
+}
+
 #[derive(Deserialize, Default, Debug)]
 #[serde(deny_unknown_fields)]
 #[allow(clippy::upper_case_acronyms)]
@@ -268,6 +275,7 @@ pub struct ConfigFile {
     git: Option<Git>,
     windows: Option<Windows>,
     npm: Option<NPM>,
+    yarn: Option<Yarn>,
     vim: Option<Vim>,
     firmware: Option<Firmware>,
     vagrant: Option<Vagrant>,
@@ -839,6 +847,15 @@ impl Config {
             .and_then(|npm| npm.use_sudo)
             .unwrap_or(false)
     }
+    #[cfg(target_os = "linux")]
+    pub fn yarn_use_sudo(&self) -> bool {
+        self.config_file
+            .yarn
+            .as_ref()
+            .and_then(|yarn| yarn.use_sudo)
+            .unwrap_or(false)
+    }
+    
 
     #[cfg(target_os = "linux")]
     pub fn firmware_upgrade(&self) -> bool {
diff --git a/src/main.rs b/src/main.rs
index 88d0f605..e967f8af 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -335,6 +335,7 @@ fn run() -> Result<()> {
     runner.execute(Step::Vim, "voom", || vim::run_voom(&base_dirs, run_type))?;
     runner.execute(Step::Kakoune, "Kakoune", || kakoune::upgrade_kak_plug(&ctx))?;
     runner.execute(Step::Node, "npm", || node::run_npm_upgrade(&ctx))?;
+    runner.execute(Step::Node, "yarn", || node::run_yarn_upgrade(&ctx))?;
     runner.execute(Step::Containers, "Containers", || containers::run_containers(&ctx))?;
     runner.execute(Step::Deno, "deno", || node::deno_upgrade(&ctx))?;
     runner.execute(Step::Composer, "composer", || generic::run_composer_update(&ctx))?;
diff --git a/src/steps/node.rs b/src/steps/node.rs
index 2541e7cd..c27b4b45 100644
--- a/src/steps/node.rs
+++ b/src/steps/node.rs
@@ -88,6 +88,59 @@ impl NPM {
     }
 }
 
+struct Yarn {
+    command: PathBuf,
+    yarn: Option<PathBuf>,
+}
+
+impl Yarn {
+    fn new(command: PathBuf) -> Self {
+        Self {
+            command,
+            yarn: require("yarn").ok(),
+        }
+    }
+
+    #[cfg(target_os = "linux")]
+    fn root(&self) -> Result<PathBuf> {
+        let args = ["global", "dir"];
+        Command::new(&self.command)
+            .args(args)
+            .check_output()
+            .map(|s| PathBuf::from(s.trim()))
+    }
+
+    fn upgrade(&self, run_type: RunType, use_sudo: bool) -> Result<()> {
+        print_separator("Yarn Package Manager");
+        let args = ["global", "upgrade"];
+        
+        if use_sudo {
+            run_type
+                .execute("sudo")
+                .arg(self.yarn.as_ref().unwrap_or(&self.command))
+                .args(args)
+                .check_run()?;
+        } else {
+            run_type.execute(&self.command).args(args).check_run()?;
+        }
+
+        Ok(())
+    }
+
+    #[cfg(target_os = "linux")]
+    pub fn should_use_sudo(&self) -> Result<bool> {
+        let yarn_root = self.root()?;
+        if !yarn_root.exists() {
+            return Err(SkipStep(format!("NPM root at {} doesn't exist", yarn_root.display(),)).into());
+        }
+
+        let metadata = std::fs::metadata(&yarn_root)?;
+        let uid = Uid::effective();
+
+        Ok(metadata.uid() != uid.as_raw() && metadata.uid() == 0)
+    }
+}
+
 #[cfg(target_os = "linux")]
 fn should_use_sudo(npm: &NPM, ctx: &ExecutionContext) -> Result<bool> {
     if npm.should_use_sudo()? {
@@ -102,6 +155,20 @@ fn should_use_sudo(npm: &NPM, ctx: &ExecutionContext) -> Result<bool> {
     }
 }
 
+#[cfg(target_os = "linux")]
+fn should_use_sudo_yarn(yarn: &Yarn, ctx: &ExecutionContext) -> Result<bool> {
+    if yarn.should_use_sudo()? {
+        if ctx.config().yarn_use_sudo() {
+            Ok(true)
+        } else {
+            Err(SkipStep("NPM root is owned by another user which is not the current user. Set use_sudo = true under the NPM section in your configuration to run NPM as sudo".to_string())
+                .into())
+        }
+    } else {
+        Ok(false)
+    }
+}
+
 pub fn run_npm_upgrade(ctx: &ExecutionContext) -> Result<()> {
     let npm = require("npm").map(NPM::new)?;
 
@@ -116,6 +183,22 @@ pub fn run_npm_upgrade(ctx: &ExecutionContext) -> Result<()> {
     }
 }
 
+pub fn run_yarn_upgrade(ctx: &ExecutionContext) -> Result<()> {
+    let yarn = require("yarn").map(Yarn::new)?;
+
+    #[cfg(target_os = "linux")]
+    {
+        yarn.upgrade(ctx.run_type(), should_use_sudo_yarn(&yarn, ctx)?)
+    }
+
+    #[cfg(not(target_os = "linux"))]
+    {
+        yarn.upgrade(ctx.run_type(), false)
+    }
+}
+
+
+
 pub fn deno_upgrade(ctx: &ExecutionContext) -> Result<()> {
     let deno = require("deno")?;
     let deno_dir = ctx.base_dirs().home_dir().join(".deno");