From a72977f7d4b2b12191cfcdffc0b07fb6f6d7881e Mon Sep 17 00:00:00 2001
From: yangsheng6810 <styang@fastmail.com>
Date: Sun, 6 Dec 2020 15:15:20 -0600
Subject: [PATCH] feat(username): Try harder to guess if inside ssh (#1964)

* feat: Try harder to guess if inside ssh

* Add test for SSH_CLIENT

* Update documentation on checking ssh connection

* Update docs/config/README.md

Co-authored-by: Thomas O'Donnell <andytom@users.noreply.github.com>

Co-authored-by: Thomas O'Donnell <andytom@users.noreply.github.com>
---
 docs/config/README.md   |  6 ++++++
 src/modules/username.rs | 35 +++++++++++++++++++++++++++++++++--
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/docs/config/README.md b/docs/config/README.md
index 8b80e4d1..59151dd5 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -2399,6 +2399,12 @@ The module will be shown if any of the following conditions are met:
 - The user is currently connected as an SSH session
 - The variable `show_always` is set to true
 
+::: tip
+SSH connection is detected by checking environment variables
+`SSH_CONNECTION`, `SSH_CLIENT`, and `SSH_TTY`. If your SSH host does not set up
+these variables, one workaround is to set one of them with a dummy value.
+:::
+
 ### Options
 
 | Option        | Default                 | Description                           |
diff --git a/src/modules/username.rs b/src/modules/username.rs
index ff9f280f..d2732559 100644
--- a/src/modules/username.rs
+++ b/src/modules/username.rs
@@ -13,7 +13,6 @@ use crate::utils;
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     let user = context.get_env("USER");
     let logname = context.get_env("LOGNAME");
-    let ssh_connection = context.get_env("SSH_CONNECTION");
 
     const ROOT_UID: Option<u32> = Some(0);
     let user_uid = get_uid();
@@ -21,7 +20,8 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     let mut module = context.new_module("username");
     let config: UsernameConfig = UsernameConfig::try_load(module.config);
 
-    if user != logname || ssh_connection.is_some() || user_uid == ROOT_UID || config.show_always {
+    if user != logname || is_ssh_connection(&context) || user_uid == ROOT_UID || config.show_always
+    {
         let username = user?;
         let parsed = StringFormatter::new(config.format).and_then(|formatter| {
             formatter
@@ -55,6 +55,13 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     }
 }
 
+fn is_ssh_connection(context: &Context) -> bool {
+    let ssh_env: Vec<&str> = vec!["SSH_CONNECTION", "SSH_CLIENT", "SSH_TTY"];
+    ssh_env
+        .into_iter()
+        .any(|env| context.get_env(env).is_some())
+}
+
 fn get_uid() -> Option<u32> {
     utils::exec_cmd("id", &["-u"])?
         .stdout
@@ -129,6 +136,30 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn ssh_connection_tty() -> io::Result<()> {
+        let actual = ModuleRenderer::new("username")
+            .env("USER", "astronaut")
+            .env("SSH_TTY", "/dev/pts/0")
+            .collect();
+        let expected = Some(format!("{} in ", Color::Yellow.bold().paint("astronaut")));
+
+        assert_eq!(expected, actual);
+        Ok(())
+    }
+
+    #[test]
+    fn ssh_connection_client() -> io::Result<()> {
+        let actual = ModuleRenderer::new("username")
+            .env("USER", "astronaut")
+            .env("SSH_CLIENT", "192.168.0.101 39323 22")
+            .collect();
+        let expected = Some(format!("{} in ", Color::Yellow.bold().paint("astronaut")));
+
+        assert_eq!(expected, actual);
+        Ok(())
+    }
+
     #[test]
     fn show_always() -> io::Result<()> {
         let actual = ModuleRenderer::new("username")