Merge pull request #449 from dezgeg/env

pgrep/pidwait/pkill: Implement --env

Author: cakebaker

src/uu/pgrep/src/process.rs

@@ -491,6 +491,19 @@ impl ProcessInformation {
         self.cached_thread_ids = Some(Rc::clone(&result));
         Rc::clone(&result)
     }
+
+    pub fn env_vars(&self) -> Result<HashMap<String, String>, io::Error> {
+        let content = fs::read_to_string(format!("/proc/{}/environ", self.pid))?;
+
+        let mut env_vars = HashMap::new();
+        for entry in content.split('\0') {
+            if let Some((key, value)) = entry.split_once('=') {
+                env_vars.insert(key.to_string(), value.to_string());
+            }
+        }
+
+        Ok(env_vars)
+    }
 }
 impl TryFrom<DirEntry> for ProcessInformation {
     type Error = io::Error;
@@ -692,4 +705,16 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    #[cfg(target_os = "linux")]
+    fn test_environ() {
+        let pid_entry = ProcessInformation::current_process_info().unwrap();
+        let env_vars = pid_entry.env_vars().unwrap();
+
+        assert_eq!(
+            *env_vars.get("HOME").unwrap(),
+            std::env::var("HOME").unwrap()
+        );
+    }
 }

src/uu/pgrep/src/process_matcher.rs

@@ -47,6 +47,7 @@ pub struct Settings {
     pub pgroup: Option<HashSet<u64>>,
     pub session: Option<HashSet<u64>>,
     pub cgroup: Option<HashSet<String>>,
+    pub env: Option<HashSet<String>>,
     pub threads: bool,
 
     pub pidfile: Option<String>,
@@ -111,6 +112,9 @@ pub fn get_match_settings(matches: &ArgMatches) -> UResult<Settings> {
         cgroup: matches
             .get_many::<String>("cgroup")
             .map(|groups| groups.cloned().collect()),
+        env: matches
+            .get_many::<String>("env")
+            .map(|env_vars| env_vars.cloned().collect()),
         threads: false,
         pidfile: matches.get_one::<String>("pidfile").cloned(),
         logpidfile: matches.get_flag("logpidfile"),
@@ -129,6 +133,7 @@ pub fn get_match_settings(matches: &ArgMatches) -> UResult<Settings> {
         && settings.pgroup.is_none()
         && settings.session.is_none()
         && settings.cgroup.is_none()
+        && settings.env.is_none()
         && !settings.require_handler
         && settings.pidfile.is_none()
         && pattern.is_empty()
@@ -279,6 +284,22 @@ fn collect_matched_pids(settings: &Settings) -> UResult<Vec<ProcessInformation>>
                 pid.cgroup_v2_path().unwrap_or("/".to_string()),
             );
 
+            let env_matched = match &settings.env {
+                Some(env_filters) => {
+                    let env_vars = pid.env_vars().unwrap_or_default();
+                    env_filters.iter().any(|filter| {
+                        if let Some((key, expected_value)) = filter.split_once('=') {
+                            // Match specific key=value pair
+                            env_vars.get(key) == Some(&expected_value.to_string())
+                        } else {
+                            // Match key existence only
+                            env_vars.contains_key(filter)
+                        }
+                    })
+                }
+                None => true,
+            };
+
             let ids_matched = any_matches(&settings.uid, pid.uid().unwrap())
                 && any_matches(&settings.euid, pid.euid().unwrap())
                 && any_matches(&settings.gid, pid.gid().unwrap());
@@ -311,6 +332,7 @@ fn collect_matched_pids(settings: &Settings) -> UResult<Vec<ProcessInformation>>
                 && pgroup_matched
                 && session_matched
                 && cgroup_matched
+                && env_matched
                 && ids_matched
                 && handler_matched
                 && pidfile_matched)
@@ -530,6 +552,7 @@ pub fn clap_args(pattern_help: &'static str, enable_v_flag: bool) -> Vec<Arg> {
         arg!(-r --runstates <state>    "match runstates [D,S,Z,...]"),
         arg!(-A --"ignore-ancestors"   "exclude our ancestors from results"),
         arg!(--cgroup <grp>            "match by cgroup v2 names").value_delimiter(','),
+        arg!(--env <"name[=val],...">      "match on environment variable").value_delimiter(','),
         // arg!(--ns <PID>                "match the processes that belong to the same namespace as <pid>"),
         // arg!(--nslist <ns>             "list which namespaces will be considered for the --ns option.")
         //     .value_delimiter(',')

src/uu/pidwait/src/pidwait.rs

@@ -48,7 +48,7 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
 }
 
 pub fn uu_app() -> Command {
-    Command::new(env!("CARGO_PKG_NAME"))
+    Command::new(uucore::util_name())
         .version(crate_version!())
         .about(ABOUT)
         .override_usage(format_usage(USAGE))

tests/by-util/test_pgrep.rs

@@ -685,3 +685,38 @@ fn test_ignore_ancestors() {
         .stdout_does_not_match(&Regex::new(&format!("(?m)^{our_pid}$")).unwrap())
         .stdout_does_not_match(&Regex::new("(?m)^1$").unwrap());
 }
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_env_nonexistent() {
+    new_ucmd!().arg("--env=NONEXISTENT").fails().code_is(1);
+}
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_env_nonmatching_value() {
+    new_ucmd!()
+        .arg("--env=PATH=not_a_valid_PATH")
+        .fails()
+        .code_is(1);
+}
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_env_key_match() {
+    new_ucmd!().arg("--env=PATH").succeeds();
+}
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_env_key_value_match() {
+    let home = std::env::var("HOME").unwrap();
+    new_ucmd!().arg(format!("--env=HOME={}", home)).succeeds();
+}
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_env_multiple_filters() {
+    // Multiple filters use OR logic
+    new_ucmd!().arg("--env=PATH,NONEXISTENT").succeeds();
+}