|  | 
|  | 1 | +//! APIs for interacting with the scheduler and with processes, | 
|  | 2 | +//! corresponding to <linux/sched.h> and related header files. | 
|  | 3 | +#![allow(improper_ctypes)] | 
|  | 4 | + | 
|  | 5 | +use core::ptr; | 
|  | 6 | + | 
|  | 7 | +use crate::bindings; | 
|  | 8 | +use crate::rcu; | 
|  | 9 | + | 
|  | 10 | +extern "C" { | 
|  | 11 | +    fn task_lock_helper(p: *mut bindings::task_struct); | 
|  | 12 | +    fn task_unlock_helper(p: *mut bindings::task_struct); | 
|  | 13 | +    fn next_task_helper(p: *mut bindings::task_struct) -> *mut bindings::task_struct; | 
|  | 14 | +} | 
|  | 15 | + | 
|  | 16 | +/// Represents a `struct task_struct *`. | 
|  | 17 | +pub struct TaskStruct<'a>(&'a mut bindings::task_struct); | 
|  | 18 | + | 
|  | 19 | +impl TaskStruct<'_> { | 
|  | 20 | +    /// Returns the threadgroup ID (what userspace calls the process ID). | 
|  | 21 | +    pub fn tgid(&self) -> i32 { | 
|  | 22 | +        self.0.tgid | 
|  | 23 | +    } | 
|  | 24 | + | 
|  | 25 | +    /// Returns the command name / process title. This is a short name, | 
|  | 26 | +    /// typically the base name of the command, and does not have the | 
|  | 27 | +    /// full path or arguments. It's a fixed-sized set of bytes, but by | 
|  | 28 | +    /// convention it's interpreted as NUL-terminated. | 
|  | 29 | +    pub fn comm(&mut self) -> [u8; bindings::TASK_COMM_LEN as usize] { | 
|  | 30 | +        let mut result = [0u8; bindings::TASK_COMM_LEN as usize]; | 
|  | 31 | +        unsafe { | 
|  | 32 | +            task_lock_helper(self.0); | 
|  | 33 | +        } | 
|  | 34 | +        // if only char were unsigned char | 
|  | 35 | +        for (src, dst) in self.0.comm.iter().zip(result.iter_mut()) { | 
|  | 36 | +            if *src == 0 { | 
|  | 37 | +                break; | 
|  | 38 | +            } | 
|  | 39 | +            *dst = *src as _; | 
|  | 40 | +        } | 
|  | 41 | +        unsafe { | 
|  | 42 | +            task_unlock_helper(self.0); | 
|  | 43 | +        } | 
|  | 44 | +        result | 
|  | 45 | +    } | 
|  | 46 | +} | 
|  | 47 | + | 
|  | 48 | +/// Iterate over every process on the system. Returns only processes, | 
|  | 49 | +/// i.e., thread group leaders. | 
|  | 50 | +/// | 
|  | 51 | +/// ``` | 
|  | 52 | +/// let g = rcu::RcuReadGuard::new(); | 
|  | 53 | +/// for p in each_process(&g) { | 
|  | 54 | +///     println!("{:?}", p.comm()); | 
|  | 55 | +/// } | 
|  | 56 | +/// ``` | 
|  | 57 | +struct EachProcess<'g> { | 
|  | 58 | +    p: *mut bindings::task_struct, | 
|  | 59 | +    _g: &'g rcu::RcuReadGuard, | 
|  | 60 | +} | 
|  | 61 | + | 
|  | 62 | +pub fn each_process(g: &rcu::RcuReadGuard) -> impl Iterator<Item = TaskStruct> { | 
|  | 63 | +    // unsafe is bogus here because we don't read it | 
|  | 64 | +    // https://github.com/rust-lang/rust/issues/74843 | 
|  | 65 | +    EachProcess { | 
|  | 66 | +        p: unsafe { &mut bindings::init_task }, | 
|  | 67 | +        _g: g, | 
|  | 68 | +    } | 
|  | 69 | +} | 
|  | 70 | + | 
|  | 71 | +impl<'g> Iterator for EachProcess<'g> { | 
|  | 72 | +    type Item = TaskStruct<'g>; | 
|  | 73 | + | 
|  | 74 | +    fn next(&mut self) -> Option<TaskStruct<'g>> { | 
|  | 75 | +        // Safety: | 
|  | 76 | +        // - oldp is valid if not null, because it is either &init_task | 
|  | 77 | +        //   (a static location) or updated by this function. | 
|  | 78 | +        // - next_task calls rcu_dereference internally, which is safe | 
|  | 79 | +        //   because we hold self._g. | 
|  | 80 | +        // - The returned reference has lifetime 'g, which is valid | 
|  | 81 | +        //   because self._g lives at least that long. | 
|  | 82 | +        let oldp = unsafe { self.p.as_mut()? }; | 
|  | 83 | +        self.p = unsafe { next_task_helper(self.p) }; | 
|  | 84 | +        if self.p == unsafe { &mut bindings::init_task } { | 
|  | 85 | +            self.p = ptr::null_mut(); | 
|  | 86 | +        } | 
|  | 87 | +        Some(TaskStruct(oldp)) | 
|  | 88 | +    } | 
|  | 89 | +} | 
0 commit comments