Merge pull request #385 from dezgeg/pidof-c

pidof: Support --check-root option

Author: sylvestre

src/uu/pgrep/src/process.rs

@@ -4,6 +4,7 @@
 // file that was distributed with this source code.
 
 use regex::Regex;
+use std::fs::read_link;
 use std::hash::Hash;
 use std::sync::LazyLock;
 use std::{
@@ -250,6 +251,17 @@ impl ProcessInformation {
         })
     }
 
+    pub fn current_process_info() -> Result<ProcessInformation, io::Error> {
+        use std::str::FromStr;
+
+        #[cfg(target_os = "linux")]
+        let pid = uucore::process::getpid();
+        #[cfg(not(target_os = "linux"))]
+        let pid = 0; // dummy
+
+        ProcessInformation::try_new(PathBuf::from_str(&format!("/proc/{}", pid)).unwrap())
+    }
+
     pub fn proc_status(&self) -> &str {
         &self.inner_status
     }
@@ -366,6 +378,11 @@ impl ProcessInformation {
         self.get_uid_or_gid_field("Gid", 1)
     }
 
+    // Root directory of the process (which can be changed by chroot)
+    pub fn root(&mut self) -> Result<PathBuf, io::Error> {
+        read_link(format!("/proc/{}/root", self.pid))
+    }
+
     /// Fetch run state from [ProcessInformation::cached_stat]
     ///
     /// - [The /proc Filesystem: Table 1-4](https://docs.kernel.org/filesystems/proc.html#id10)
@@ -495,7 +512,9 @@ pub fn walk_threads() -> impl Iterator<Item = ProcessInformation> {
 mod tests {
     use super::*;
     #[cfg(target_os = "linux")]
-    use std::{collections::HashSet, str::FromStr};
+    use std::collections::HashSet;
+    #[cfg(target_os = "linux")]
+    use uucore::process::getpid;
 
     #[test]
     fn test_run_state_conversion() {
@@ -515,41 +534,19 @@ mod tests {
         assert!(RunState::try_from("Rg").is_err());
     }
 
-    #[cfg(target_os = "linux")]
-    fn current_pid() -> usize {
-        // Direct read link of /proc/self.
-        // It's result must be current programs pid.
-        fs::read_link("/proc/self")
-            .unwrap()
-            .to_str()
-            .unwrap()
-            .parse::<usize>()
-            .unwrap()
-    }
-
-    #[cfg(target_os = "linux")]
-    fn current_process_info() -> ProcessInformation {
-        ProcessInformation::try_new(PathBuf::from_str(&format!("/proc/{}", current_pid())).unwrap())
-            .unwrap()
-    }
-
     #[test]
     #[cfg(target_os = "linux")]
     fn test_walk_pid() {
-        let current_pid = current_pid();
-
-        let find = walk_process().find(|it| it.pid == current_pid);
+        let find = walk_process().find(|it| it.pid == getpid() as usize);
 
         assert!(find.is_some());
     }
 
     #[test]
     #[cfg(target_os = "linux")]
     fn test_pid_entry() {
-        let current_pid = current_pid();
-
-        let pid_entry = current_process_info();
-        let mut result = WalkDir::new(format!("/proc/{}/fd", current_pid))
+        let pid_entry = ProcessInformation::current_process_info().unwrap();
+        let mut result = WalkDir::new(format!("/proc/{}/fd", getpid()))
             .into_iter()
             .flatten()
             .map(DirEntry::into_path)
@@ -569,7 +566,7 @@ mod tests {
     fn test_thread_ids() {
         let main_tid = unsafe { uucore::libc::gettid() };
         std::thread::spawn(move || {
-            let mut pid_entry = current_process_info();
+            let mut pid_entry = ProcessInformation::current_process_info().unwrap();
             let thread_ids = pid_entry.thread_ids();
 
             assert!(thread_ids.contains(&(main_tid as usize)));
@@ -599,7 +596,7 @@ mod tests {
     #[test]
     #[cfg(target_os = "linux")]
     fn test_ids() {
-        let mut pid_entry = current_process_info();
+        let mut pid_entry = ProcessInformation::current_process_info().unwrap();
         assert_eq!(
             pid_entry.ppid().unwrap(),
             unsafe { uucore::libc::getppid() } as u64
@@ -615,10 +612,17 @@ mod tests {
     #[test]
     #[cfg(target_os = "linux")]
     fn test_uid_gid() {
-        let mut pid_entry = current_process_info();
+        let mut pid_entry = ProcessInformation::current_process_info().unwrap();
         assert_eq!(pid_entry.uid().unwrap(), uucore::process::getuid());
         assert_eq!(pid_entry.euid().unwrap(), uucore::process::geteuid());
         assert_eq!(pid_entry.gid().unwrap(), uucore::process::getgid());
         assert_eq!(pid_entry.egid().unwrap(), uucore::process::getegid());
     }
+
+    #[test]
+    #[cfg(target_os = "linux")]
+    fn test_root() {
+        let mut pid_entry = ProcessInformation::current_process_info().unwrap();
+        assert_eq!(pid_entry.root().unwrap(), PathBuf::from("/"));
+    }
 }

src/uu/pidof/Cargo.toml

@@ -13,7 +13,7 @@ categories = ["command-line-utilities"]
 
 
 [dependencies]
-uucore = { workspace = true }
+uucore = { workspace = true, features = ["process"] }
 clap = { workspace = true }
 uu_pgrep = { path = "../pgrep" }
 

src/uu/pidof/src/pidof.rs

@@ -7,6 +7,8 @@ use std::path::PathBuf;
 
 use clap::{crate_version, Arg, ArgAction, ArgMatches, Command};
 use uu_pgrep::process::{walk_process, ProcessInformation};
+#[cfg(unix)]
+use uucore::process::geteuid;
 use uucore::{error::UResult, format_usage, help_about, help_usage};
 
 const ABOUT: &str = help_about!("pidof.md");
@@ -94,21 +96,35 @@ fn collect_matched_pids(matches: &ArgMatches) -> Vec<usize> {
         .copied()
         .collect::<Vec<_>>();
 
+    // Original pidof silently ignores the check-root option if the user is not root.
+    #[cfg(unix)]
+    let check_root = matches.get_flag("check-root") && geteuid() == 0;
+    #[cfg(not(unix))]
+    let check_root = false;
+    let our_root = ProcessInformation::current_process_info()
+        .unwrap()
+        .root()
+        .unwrap();
+
     program_names
         .into_iter()
         .flat_map(|program| {
             let mut processed = Vec::new();
             for mut process in collected.clone() {
-                let contains =
-                    match_process_name(&mut process, &program, with_workers, match_scripts);
-                let should_omit = arg_omit_pid.contains(&process.pid);
-
-                if contains && !should_omit {
-                    if matches.get_flag("t") {
-                        processed.extend_from_slice(&process.thread_ids());
-                    } else {
-                        processed.push(process.pid);
-                    }
+                if !match_process_name(&mut process, &program, with_workers, match_scripts) {
+                    continue;
+                }
+                if arg_omit_pid.contains(&process.pid) {
+                    continue;
+                }
+                if check_root && process.root().unwrap() != our_root {
+                    continue;
+                }
+
+                if matches.get_flag("t") {
+                    processed.extend_from_slice(&process.thread_ids());
+                } else {
+                    processed.push(process.pid);
                 }
             }
 
@@ -140,12 +156,13 @@ pub fn uu_app() -> Command {
                 .index(1)
                 .action(ArgAction::Append),
         )
-        // .arg(
-        //     Arg::new("c")
-        //         .short('c')
-        //         .help("Return PIDs with the same root directory")
-        //         .action(ArgAction::SetTrue),
-        // )
+        .arg(
+            Arg::new("check-root")
+                .short('c')
+                .long("check-root")
+                .help("Only return PIDs with the same root directory")
+                .action(ArgAction::SetTrue),
+        )
         .arg(
             Arg::new("S")
                 .short('S')

tests/by-util/test_pidof.rs

@@ -46,6 +46,16 @@ fn test_quiet() {
         .no_output();
 }
 
+#[test]
+#[cfg(target_os = "linux")]
+fn test_check_root_accepted() {
+    new_ucmd!()
+        .arg("-w")
+        .arg("--check-root")
+        .arg("kthreadd")
+        .succeeds();
+}
+
 #[test]
 #[cfg(target_os = "linux")]
 fn test_single_shot() {