diff --git a/src/uu/pgrep/src/process.rs b/src/uu/pgrep/src/process.rs index 50521b22..36010640 100644 --- a/src/uu/pgrep/src/process.rs +++ b/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 { + 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 { + 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 { 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,30 +534,10 @@ 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::() - .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()); } @@ -546,10 +545,8 @@ mod tests { #[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("/")); + } } diff --git a/src/uu/pidof/Cargo.toml b/src/uu/pidof/Cargo.toml index 8e7361a5..7bbf82be 100644 --- a/src/uu/pidof/Cargo.toml +++ b/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" } diff --git a/src/uu/pidof/src/pidof.rs b/src/uu/pidof/src/pidof.rs index 898ec1a5..cf9e6988 100644 --- a/src/uu/pidof/src/pidof.rs +++ b/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 { .copied() .collect::>(); + // 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') diff --git a/tests/by-util/test_pidof.rs b/tests/by-util/test_pidof.rs index c9083e01..44a314db 100644 --- a/tests/by-util/test_pidof.rs +++ b/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() {