Merge pull request #340 from dezgeg/pgroup-session

pgrep/pkill/pidwait: Support --session and --pgroup

Author: cakebaker

src/uu/pgrep/src/pgrep.rs

@@ -45,15 +45,15 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
                     // pgrep from procps-ng outputs the process name inside square brackets
                     // if /proc/<PID>/cmdline is empty
                     if it.cmdline.is_empty() {
-                        format!("{} [{}]", it.pid, it.clone().status().get("Name").unwrap())
+                        format!("{} [{}]", it.pid, it.clone().name().unwrap())
                     } else {
                         format!("{} {}", it.pid, it.cmdline)
                     }
                 })
                 .collect()
         } else if matches.get_flag("list-name") {
             pids.into_iter()
-                .map(|it| format!("{} {}", it.pid, it.clone().status().get("Name").unwrap()))
+                .map(|it| format!("{} {}", it.pid, it.clone().name().unwrap()))
                 .collect()
         } else {
             pids.into_iter().map(|it| format!("{}", it.pid)).collect()

src/uu/pgrep/src/process.rs

@@ -274,6 +274,21 @@ impl ProcessInformation {
         Rc::clone(&result)
     }
 
+    pub fn name(&mut self) -> Result<String, io::Error> {
+        self.status()
+            .get("Name")
+            .cloned()
+            .ok_or(io::ErrorKind::InvalidData.into())
+    }
+
+    fn get_numeric_stat_field(&mut self, index: usize) -> Result<u64, io::Error> {
+        self.stat()
+            .get(index)
+            .ok_or(io::ErrorKind::InvalidData)?
+            .parse::<u64>()
+            .map_err(|_| io::ErrorKind::InvalidData.into())
+    }
+
     /// Fetch start time from [ProcessInformation::cached_stat]
     ///
     /// - [The /proc Filesystem: Table 1-4](https://docs.kernel.org/filesystems/proc.html#id10)
@@ -284,12 +299,7 @@ impl ProcessInformation {
 
         // Kernel doc: https://docs.kernel.org/filesystems/proc.html#process-specific-subdirectories
         // Table 1-4
-        let time = self
-            .stat()
-            .get(21)
-            .ok_or(io::ErrorKind::InvalidData)?
-            .parse::<u64>()
-            .map_err(|_| io::ErrorKind::InvalidData)?;
+        let time = self.get_numeric_stat_field(21)?;
 
         self.cached_start_time = Some(time);
 
@@ -299,11 +309,19 @@ impl ProcessInformation {
     pub fn ppid(&mut self) -> Result<u64, io::Error> {
         // the PPID is the fourth field in /proc/<PID>/stat
         // (https://www.kernel.org/doc/html/latest/filesystems/proc.html#id10)
-        self.stat()
-            .get(3)
-            .ok_or(io::ErrorKind::InvalidData)?
-            .parse::<u64>()
-            .map_err(|_| io::ErrorKind::InvalidData.into())
+        self.get_numeric_stat_field(3)
+    }
+
+    pub fn pgid(&mut self) -> Result<u64, io::Error> {
+        // the process group ID is the fifth field in /proc/<PID>/stat
+        // (https://www.kernel.org/doc/html/latest/filesystems/proc.html#id10)
+        self.get_numeric_stat_field(4)
+    }
+
+    pub fn sid(&mut self) -> Result<u64, io::Error> {
+        // the session ID is the sixth field in /proc/<PID>/stat
+        // (https://www.kernel.org/doc/html/latest/filesystems/proc.html#id10)
+        self.get_numeric_stat_field(5)
     }
 
     fn get_uid_or_gid_field(&mut self, field: &str, index: usize) -> Result<u32, io::Error> {
@@ -475,6 +493,12 @@ mod tests {
             .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() {
@@ -490,11 +514,7 @@ mod tests {
     fn test_pid_entry() {
         let current_pid = current_pid();
 
-        let pid_entry = ProcessInformation::try_new(
-            PathBuf::from_str(&format!("/proc/{}", current_pid)).unwrap(),
-        )
-        .unwrap();
-
+        let pid_entry = current_process_info();
         let mut result = WalkDir::new(format!("/proc/{}/fd", current_pid))
             .into_iter()
             .flatten()
@@ -515,10 +535,7 @@ mod tests {
     fn test_thread_ids() {
         let main_tid = unsafe { uucore::libc::gettid() };
         std::thread::spawn(move || {
-            let mut pid_entry = ProcessInformation::try_new(
-                PathBuf::from_str(&format!("/proc/{}", current_pid())).unwrap(),
-            )
-            .unwrap();
+            let mut pid_entry = current_process_info();
             let thread_ids = pid_entry.thread_ids();
 
             assert!(thread_ids.contains(&(main_tid as usize)));
@@ -545,13 +562,26 @@ mod tests {
         assert!(stat_split(case)[1] == "sleep (2) .sh");
     }
 
+    #[test]
+    #[cfg(target_os = "linux")]
+    fn test_ids() {
+        let mut pid_entry = current_process_info();
+        assert_eq!(
+            pid_entry.ppid().unwrap(),
+            unsafe { uucore::libc::getppid() } as u64
+        );
+        assert_eq!(
+            pid_entry.pgid().unwrap(),
+            unsafe { uucore::libc::getpgid(0) } as u64
+        );
+        assert_eq!(pid_entry.sid().unwrap(), unsafe { uucore::libc::getsid(0) }
+            as u64);
+    }
+
     #[test]
     #[cfg(target_os = "linux")]
     fn test_uid_gid() {
-        let mut pid_entry = ProcessInformation::try_new(
-            PathBuf::from_str(&format!("/proc/{}", current_pid())).unwrap(),
-        )
-        .unwrap();
+        let mut pid_entry = current_process_info();
         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());

src/uu/pgrep/src/process_matcher.rs

@@ -11,6 +11,8 @@ use std::{collections::HashSet, io};
 use clap::{arg, Arg, ArgAction, ArgMatches};
 use regex::Regex;
 #[cfg(unix)]
+use uucore::libc::{getpgrp, getsid};
+#[cfg(unix)]
 use uucore::{
     display::Quotable,
     entries::{grp2gid, usr2uid},
@@ -40,6 +42,8 @@ pub struct Settings {
     pub uid: Option<HashSet<u32>>,
     pub euid: Option<HashSet<u32>>,
     pub gid: Option<HashSet<u32>>,
+    pub pgroup: Option<HashSet<u64>>,
+    pub session: Option<HashSet<u64>>,
 }
 
 pub fn get_match_settings(matches: &ArgMatches) -> UResult<Settings> {
@@ -76,6 +80,26 @@ pub fn get_match_settings(matches: &ArgMatches) -> UResult<Settings> {
         gid: matches
             .get_many::<u32>("group")
             .map(|ids| ids.cloned().collect()),
+        pgroup: matches.get_many::<u64>("pgroup").map(|xs| {
+            xs.map(|pg| {
+                if *pg == 0 {
+                    unsafe { getpgrp() as u64 }
+                } else {
+                    *pg
+                }
+            })
+            .collect()
+        }),
+        session: matches.get_many::<u64>("session").map(|xs| {
+            xs.map(|sid| {
+                if *sid == 0 {
+                    unsafe { getsid(0) as u64 }
+                } else {
+                    *sid
+                }
+            })
+            .collect()
+        }),
     };
 
     if !settings.newest
@@ -87,6 +111,8 @@ pub fn get_match_settings(matches: &ArgMatches) -> UResult<Settings> {
         && settings.uid.is_none()
         && settings.euid.is_none()
         && settings.gid.is_none()
+        && settings.pgroup.is_none()
+        && settings.session.is_none()
         && pattern.is_empty()
     {
         return Err(USimpleError::new(
@@ -182,12 +208,11 @@ fn collect_matched_pids(settings: &Settings) -> Vec<ProcessInformation> {
                 _ => true,
             };
 
-            let binding = pid.status();
-            let name = binding.get("Name").unwrap();
+            let name = pid.name().unwrap();
             let name = if settings.ignore_case {
                 name.to_lowercase()
             } else {
-                name.into()
+                name
             };
             let pattern_matched = {
                 let want = if settings.full {
@@ -207,6 +232,8 @@ fn collect_matched_pids(settings: &Settings) -> Vec<ProcessInformation> {
             let older_matched = pid.start_time().unwrap() >= arg_older;
 
             let parent_matched = any_matches(&settings.parent, pid.ppid().unwrap());
+            let pgroup_matched = any_matches(&settings.pgroup, pid.pgid().unwrap());
+            let session_matched = any_matches(&settings.session, pid.sid().unwrap());
 
             let ids_matched = any_matches(&settings.uid, pid.uid().unwrap())
                 && any_matches(&settings.euid, pid.euid().unwrap())
@@ -217,6 +244,8 @@ fn collect_matched_pids(settings: &Settings) -> Vec<ProcessInformation> {
                 && tty_matched
                 && older_matched
                 && parent_matched
+                && pgroup_matched
+                && session_matched
                 && ids_matched)
                 ^ settings.inverse
             {
@@ -290,6 +319,22 @@ pub fn grp2gid(_name: &str) -> io::Result<u32> {
     ))
 }
 
+/// # Safety
+///
+/// Dummy implementation for unsupported platforms.
+#[cfg(not(unix))]
+pub unsafe fn getpgrp() -> u32 {
+    panic!("unsupported on this platform");
+}
+
+/// # Safety
+///
+/// Dummy implementation for unsupported platforms.
+#[cfg(not(unix))]
+pub unsafe fn getsid(_pid: u32) -> u32 {
+    panic!("unsupported on this platform");
+}
+
 fn parse_uid_or_username(uid_or_username: &str) -> io::Result<u32> {
     uid_or_username
         .parse::<u32>()
@@ -315,9 +360,9 @@ pub fn clap_args(pattern_help: &'static str, enable_v_flag: bool) -> Vec<Arg> {
         arg!(-H --"require-handler"    "match only if signal handler is present"),
         arg!(-c --count                "count of matching processes"),
         arg!(-f --full                 "use full process name to match"),
-        // arg!(-g --pgroup <PGID>        "match listed process group IDs")
-        //     .value_delimiter(',')
-        //     .value_parser(clap::value_parser!(u64)),
+        arg!(-g --pgroup <PGID>        "match listed process group IDs")
+            .value_delimiter(',')
+            .value_parser(clap::value_parser!(u64)),
         arg!(-G --group <GID>          "match real group IDs")
             .value_delimiter(',')
             .value_parser(parse_gid_or_group_name),
@@ -331,9 +376,9 @@ pub fn clap_args(pattern_help: &'static str, enable_v_flag: bool) -> Vec<Arg> {
         arg!(-P --parent <PPID>        "match only child processes of the given parent")
             .value_delimiter(',')
             .value_parser(clap::value_parser!(u64)),
-        // arg!(-s --session <SID>        "match session IDs")
-        //     .value_delimiter(',')
-        //     .value_parser(clap::value_parser!(u64)),
+        arg!(-s --session <SID>        "match session IDs")
+            .value_delimiter(',')
+            .value_parser(clap::value_parser!(u64)),
         arg!(--signal <sig>            "signal to send (either number or name)")
             .default_value("SIGTERM"),
         arg!(-t --terminal <tty>       "match by controlling terminal").value_delimiter(','),

tests/by-util/test_pgrep.rs

@@ -424,11 +424,57 @@ fn test_current_user() {
 #[test]
 #[cfg(target_os = "linux")]
 fn test_does_not_match_current_process() {
-    let our_pid = std::process::id();
-    dbg!(&our_pid);
     new_ucmd!()
         .arg("-f")
         .arg("UNIQUE_STRING_THAT_DOES_NOT_MATCH_ANY_OTHER_PROCESS")
         .fails()
         .no_output();
 }
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_pgroup() {
+    let our_pid = std::process::id();
+    let our_pgroup = unsafe { uucore::libc::getpgid(0) };
+    new_ucmd!()
+        .arg("--pgroup")
+        .arg(our_pgroup.to_string())
+        .succeeds()
+        .stdout_contains(our_pid.to_string());
+
+    new_ucmd!()
+        .arg("--pgroup")
+        .arg("0")
+        .succeeds()
+        .stdout_contains(our_pid.to_string());
+}
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_nonexisting_pgroup() {
+    new_ucmd!().arg("--pgroup=9999999999").fails();
+}
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_session() {
+    let our_pid = std::process::id();
+    let our_sid = unsafe { uucore::libc::getsid(0) };
+    new_ucmd!()
+        .arg("--session")
+        .arg(our_sid.to_string())
+        .succeeds()
+        .stdout_contains(our_pid.to_string());
+
+    new_ucmd!()
+        .arg("--session")
+        .arg("0")
+        .succeeds()
+        .stdout_contains(our_pid.to_string());
+}
+
+#[test]
+#[cfg(target_os = "linux")]
+fn test_nonexisting_session() {
+    new_ucmd!().arg("--session=9999999999").fails();
+}