From 04bfb45a970a4e8c7c590c50785b1012c5fe4ee6 Mon Sep 17 00:00:00 2001
From: Andrew Barchuk <andrew@raindev.io>
Date: Mon, 8 Apr 2024 13:43:32 +0200
Subject: [PATCH] Fix local host detection for remotes with user  (#755)

---
 src/config.rs | 65 +++++++++++++++++++++++++++++++++++++++++++++------
 src/main.rs   |  4 ++--
 2 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 6ce001be..abbf8cd1 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -22,7 +22,7 @@ use which_crate::which;
 use super::utils::editor;
 use crate::command::CommandExt;
 use crate::sudo::SudoKind;
-use crate::utils::{hostname, string_prepend_str};
+use crate::utils::string_prepend_str;
 use tracing::{debug, error};
 
 pub static EXAMPLE_CONFIG: &str = include_str!("../config.example.toml");
@@ -1458,15 +1458,17 @@ impl Config {
     #[cfg(target_os = "linux")]
     str_value!(linux, emerge_update_flags);
 
-    pub fn should_execute_remote(&self, remote: &str) -> bool {
-        if let Ok(hostname) = hostname() {
-            if remote == hostname {
+    pub fn should_execute_remote(&self, hostname: Result<String>, remote: &str) -> bool {
+        let remote_host = remote.split_once('@').map_or(remote, |(_, host)| host);
+
+        if let Ok(hostname) = hostname {
+            if remote_host == hostname {
                 return false;
             }
         }
 
-        if let Some(limit) = self.opt.remote_host_limit.as_ref() {
-            return limit.is_match(remote);
+        if let Some(limit) = &self.opt.remote_host_limit.as_ref() {
+            return limit.is_match(remote_host);
         }
 
         true
@@ -1523,7 +1525,9 @@ impl Config {
 
 #[cfg(test)]
 mod test {
-    use crate::config::ConfigFile;
+
+    use crate::config::*;
+    use color_eyre::eyre::eyre;
 
     /// Test the default configuration in `config.example.toml` is valid.
     #[test]
@@ -1532,4 +1536,51 @@ mod test {
 
         assert!(toml::from_str::<ConfigFile>(str).is_ok());
     }
+
+    fn config() -> Config {
+        Config {
+            opt: CommandLineArgs::parse_from::<_, String>([]),
+            config_file: ConfigFile::default(),
+            allowed_steps: Vec::new(),
+        }
+    }
+
+    #[test]
+    fn test_should_execute_remote_different_hostname() {
+        assert!(config().should_execute_remote(Ok("hostname".to_string()), "remote_hostname"))
+    }
+
+    #[test]
+    fn test_should_execute_remote_different_hostname_with_user() {
+        assert!(config().should_execute_remote(Ok("hostname".to_string()), "user@remote_hostname"))
+    }
+
+    #[test]
+    fn test_should_execute_remote_unknown_hostname() {
+        assert!(config().should_execute_remote(Err(eyre!("failed to get hostname")), "remote_hostname"))
+    }
+
+    #[test]
+    fn test_should_not_execute_remote_same_hostname() {
+        assert!(!config().should_execute_remote(Ok("hostname".to_string()), "hostname"))
+    }
+
+    #[test]
+    fn test_should_not_execute_remote_same_hostname_with_user() {
+        assert!(!config().should_execute_remote(Ok("hostname".to_string()), "user@hostname"))
+    }
+
+    #[test]
+    fn test_should_execute_remote_matching_limit() {
+        let mut config = config();
+        config.opt = CommandLineArgs::parse_from(["topgrade", "--remote-host-limit", "remote_hostname"]);
+        assert!(config.should_execute_remote(Ok("hostname".to_string()), "user@remote_hostname"))
+    }
+
+    #[test]
+    fn test_should_not_execute_remote_not_matching_limit() {
+        let mut config = config();
+        config.opt = CommandLineArgs::parse_from(["topgrade", "--remote-host-limit", "other_hostname"]);
+        assert!(!config.should_execute_remote(Ok("hostname".to_string()), "user@remote_hostname"))
+    }
 }
diff --git a/src/main.rs b/src/main.rs
index bbee9980..4c12d5ac 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -27,7 +27,7 @@ use self::error::Upgraded;
 use self::steps::{remote::*, *};
 use self::terminal::*;
 
-use self::utils::{install_color_eyre, install_tracing, update_tracing};
+use self::utils::{hostname, install_color_eyre, install_tracing, update_tracing};
 
 mod breaking_changes;
 mod command;
@@ -182,7 +182,7 @@ fn run() -> Result<()> {
     }
 
     if let Some(topgrades) = config.remote_topgrades() {
-        for remote_topgrade in topgrades.iter().filter(|t| config.should_execute_remote(t)) {
+        for remote_topgrade in topgrades.iter().filter(|t| config.should_execute_remote(hostname(), t)) {
             runner.execute(Step::Remotes, format!("Remote ({remote_topgrade})"), || {
                 ssh::ssh_step(&ctx, remote_topgrade)
             })?;