From 862d12a904f935119430289fa31ec549909a430c Mon Sep 17 00:00:00 2001
From: Ylarod <me@ylarod.cn>
Date: Wed, 5 Apr 2023 11:47:35 +0800
Subject: [PATCH] ksud: add /data/adb/ksu/bin to PATH (#343)

---
 userspace/ksud/src/defs.rs  |  6 ++++--
 userspace/ksud/src/event.rs | 18 +++++++++++++++++-
 userspace/ksud/src/ksu.rs   | 18 +++++++++++++++++-
 3 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/userspace/ksud/src/defs.rs b/userspace/ksud/src/defs.rs
index 17843abf..39899c57 100644
--- a/userspace/ksud/src/defs.rs
+++ b/userspace/ksud/src/defs.rs
@@ -1,11 +1,13 @@
 use const_format::concatcp;
 
 pub const ADB_DIR: &str = "/data/adb/";
+pub const WORKING_DIR: &str = concatcp!(ADB_DIR, "ksu/");
+pub const BINARY_DIR: &str = concatcp!(WORKING_DIR, "bin/");
 
 pub const DAEMON_PATH: &str = concatcp!(ADB_DIR, "ksud");
 
-pub const WORKING_DIR: &str = concatcp!(ADB_DIR, "ksu/");
-pub const BINARY_DIR: &str = concatcp!(WORKING_DIR, "bin/");
+#[cfg(target_os = "android")]
+pub const DAEMON_LINK_PATH: &str = concatcp!(BINARY_DIR, "ksud");
 
 pub const MODULE_DIR: &str = concatcp!(ADB_DIR, "modules/");
 pub const MODULE_IMG: &str = concatcp!(WORKING_DIR, "modules.img");
diff --git a/userspace/ksud/src/event.rs b/userspace/ksud/src/event.rs
index 4f066306..e2a29574 100644
--- a/userspace/ksud/src/event.rs
+++ b/userspace/ksud/src/event.rs
@@ -248,5 +248,21 @@ pub fn install() -> Result<()> {
     std::fs::copy("/proc/self/exe", defs::DAEMON_PATH)?;
 
     // install binary assets
-    assets::ensure_binaries().with_context(|| "Failed to extract assets")
+    assets::ensure_binaries().with_context(|| "Failed to extract assets")?;
+
+    #[cfg(target_os = "android")]
+    link_ksud_to_bin()?;
+
+    Ok(())
+}
+
+#[cfg(target_os = "android")]
+fn link_ksud_to_bin() -> Result<()> {
+    use std::path::PathBuf;
+    let ksu_bin = PathBuf::from(defs::DAEMON_PATH);
+    let ksu_bin_link = PathBuf::from(defs::DAEMON_LINK_PATH);
+    if ksu_bin.exists() && !ksu_bin_link.exists() {
+        std::os::unix::fs::symlink(&ksu_bin, &ksu_bin_link)?;
+    }
+    Ok(())
 }
diff --git a/userspace/ksud/src/ksu.rs b/userspace/ksud/src/ksu.rs
index 0afc417d..2bf5fcbf 100644
--- a/userspace/ksud/src/ksu.rs
+++ b/userspace/ksud/src/ksu.rs
@@ -3,8 +3,10 @@ use anyhow::{Ok, Result};
 #[cfg(unix)]
 use anyhow::ensure;
 use getopts::Options;
+use std::env;
 #[cfg(unix)]
 use std::os::unix::process::CommandExt;
+use std::path::PathBuf;
 use std::{ffi::CStr, process::Command};
 
 use crate::{
@@ -84,7 +86,7 @@ pub fn root_shell() -> Result<()> {
         "COMMAND",
     );
     opts.optflag("h", "help", "display this help message and exit");
-    opts.optflag("l", "login", "force run in the global mount namespace");
+    opts.optflag("l", "login", "pretend the shell to be a login shell");
     opts.optflag(
         "p",
         "preserve-environment",
@@ -208,6 +210,10 @@ pub fn root_shell() -> Result<()> {
         }
     }
 
+    // add /data/adb/ksu/bin to PATH
+    #[cfg(any(target_os = "linux", target_os = "android"))]
+    add_path_to_env(defs::BINARY_DIR)?;
+
     // escape from the current cgroup and become session leader
     // WARNING!!! This cause some root shell hang forever!
     // command = command.process_group(0);
@@ -233,6 +239,16 @@ pub fn root_shell() -> Result<()> {
     Err(command.exec().into())
 }
 
+fn add_path_to_env(path: &str) -> Result<()> {
+    let mut paths =
+        env::var_os("PATH").map_or(Vec::new(), |val| env::split_paths(&val).collect::<Vec<_>>());
+    let new_path = PathBuf::from(path);
+    paths.push(new_path);
+    let new_path_env = env::join_paths(paths)?;
+    env::set_var("PATH", new_path_env);
+    Ok(())
+}
+
 pub fn get_version() -> i32 {
     let mut result: i32 = 0;
     #[cfg(any(target_os = "linux", target_os = "android"))]