1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
/// A limit for a GNU/Linux specific limitable property.
///
/// Any given limit always contain a _soft_ and a _hard_ limit.
///
/// A soft or hard limited whose value is `None` here means there is no actual limit, i.e. the value
/// found in `/proc/<pid>/limits` is `unlimited`.
#[derive(Debug, Default, Eq, PartialEq)]
pub struct Limit {
    pub soft: Option<u32>,
    pub hard: Option<u32>,
}

/// A structure containing all possible properties that can be limited by a GNU/Linux operating
/// system.
#[derive(Debug, Default, Eq, PartialEq)]
pub struct Limits {
    pub max_cpu_time: Limit,
    pub max_file_size: Limit,
    pub max_data_size: Limit,
    pub max_stack_size: Limit,
    pub max_core_file_size: Limit,
    pub max_resident_set: Limit,
    pub max_processes: Limit,
    pub max_open_files: Limit,
    pub max_locked_memory: Limit,
    pub max_address_space: Limit,
    pub max_file_locks: Limit,
    pub max_pending_signals: Limit,
    pub max_msgqueue_size: Limit,
    pub max_nice_priority: Limit,
    pub max_realtime_priority: Limit,
    pub max_realtime_timeout: Limit,
}

impl Limits {
    /// Set properties on a `Limit` structure, as read from strings.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use limits_rs::{Limit, Limits};
    ///
    /// // Create a new limits structure
    /// let mut limits = Limits::default();
    ///
    /// // Trying to set a non-existing property should do nothing
    /// limits.set_property_from_strings("Does_not_exist", "123", "456");
    /// assert_eq!(limits, Limits::default());
    ///
    /// // Let's set a limit for a existing property and assert that the limit is actually stored in
    /// // the structure
    /// limits.set_property_from_strings("Max file locks", "123", "456");
    /// assert_eq!(limits.max_file_locks, Limit { soft: Some(123), hard: Some(456) })
    ///
    /// ```
    pub fn set_property_from_strings(&mut self, name: &str, soft_string: &str, hard_string: &str) {
        use std::str::FromStr;

        let lower_case = name.to_lowercase();

        let soft = if soft_string == "unlimited" {
            None
        } else {
            u32::from_str(soft_string).ok()
        };

        let hard = if hard_string == "unlimited" {
            None
        } else {
            u32::from_str(hard_string).ok()
        };

        let new_limit = Limit { soft, hard };

        match lower_case.as_str() {
            "max cpu time" => self.max_cpu_time = new_limit,
            "max file_size" => self.max_file_size = new_limit,
            "max data size" => self.max_data_size = new_limit,
            "max stack size" => self.max_stack_size = new_limit,
            "max core file size" => self.max_core_file_size = new_limit,
            "max resident set" => self.max_resident_set = new_limit,
            "max processes" => self.max_processes = new_limit,
            "max open files" => self.max_open_files = new_limit,
            "max locked memory" => self.max_locked_memory = new_limit,
            "max address space" => self.max_address_space = new_limit,
            "max file locks" => self.max_file_locks = new_limit,
            "max pending signals" => self.max_pending_signals = new_limit,
            "max msgqueue size" => self.max_msgqueue_size = new_limit,
            "max nice priority" => self.max_nice_priority = new_limit,
            "max realtime priority" => self.max_realtime_priority = new_limit,
            "max realtime timeout" => self.max_realtime_timeout = new_limit,
            _ => (),
        }
    }
}

/// Get the limits for a specific process identifier.
///
/// Along `get_own_limits`, this method provides the core functionality of this crate.
///
/// # Examples
///
/// ```rust
/// use limits_rs::get_pid_limits;
///
/// // Let's check what the CPU time hard limit is for process `1`.
/// let limits = get_pid_limits(1).unwrap();
/// let max_cpu_time_hard_limit = limits.max_cpu_time.hard;
/// ```
pub fn get_pid_limits(pid: u32) -> Result<Limits, crate::Error> {
    // Rad the limits file for the process, and put all the lines into an iterator.
    let file_path = format!("/proc/{}/limits", pid);
    let file = std::fs::File::open(&file_path)
        .map_err(|io_error| crate::Error::ProcFileNotFound(file_path, io_error))?;
    let reader = std::io::BufReader::new(file);

    get_limits_from_reader(reader)
}

/// Read limits from any type that implements the `std::io::BufRead` crate, such as
/// `std::io::BufReader` or `std::io::Cursor`.
fn get_limits_from_reader<T>(reader: T) -> Result<Limits, crate::Error>
where
    T: std::io::BufRead,
{
    let mut limits = Limits::default();
    let mut lines = std::io::BufRead::lines(reader).filter_map(Result::ok);

    // Skip first line, which always contains the table header.
    lines.next();

    for line in lines {
        // Separate the name of the property from the rest of the table, which is padded to 27
        // characters, i.e. the soft limits always start at character 27.
        let (property, values) = line.split_at(26);
        let property = property.trim();
        let values: Vec<&str> = values.split_whitespace().collect();
        limits.set_property_from_strings(property, values[0], values[1]);
    }

    Ok(limits)
}

#[cfg(test)]
mod tests {
    use crate::{Limit, Limits};

    #[test]
    fn test_own_limits_does_not_panic() {
        crate::get_own_limits().unwrap();
    }

    #[test]
    fn test_pid_limits_does_not_panic() {
        crate::get_pid_limits(1).unwrap();
    }

    #[test]
    fn test_proc_file_not_found() {
        let error = format!("{:?}", super::get_pid_limits(std::u32::MAX).unwrap_err());
        let expected_error = String::from(
            r#"ProcFileNotFound("/proc/4294967295/limits", Os { code: 2, kind: NotFound, message: "No such file or directory" })"#,
        );

        assert_eq!(error, expected_error);
    }

    #[test]
    fn test_from_empty_string() {
        let reader = std::io::Cursor::new("");
        let limits = super::get_limits_from_reader(reader).unwrap();

        let expected_limits = Limits::default();

        assert_eq!(limits, expected_limits);
    }

    #[test]
    fn test_from_correct_string() {
        let reader = std::io::Cursor::new(
            r#"Limit                     Soft Limit           Hard Limit           Units
Max cpu time              unlimited            unlimited            seconds
Max file size             unlimited            unlimited            bytes
Max data size             unlimited            unlimited            bytes
Max stack size            8388608              unlimited            bytes
Max core file size        unlimited            unlimited            bytes
Max resident set          unlimited            unlimited            bytes
Max processes             62935                62935                processes
Max open files            1024                 524288               files
Max locked memory         65536                65536                bytes
Max address space         unlimited            unlimited            bytes
Max file locks            unlimited            unlimited            locks
Max pending signals       62935                62935                signals
Max msgqueue size         819200               819200               bytes
Max nice priority         0                    0
Max realtime priority     99                   99
Max realtime timeout      unlimited            unlimited            us"#,
        );
        let limits = super::get_limits_from_reader(reader).unwrap();

        let expected_limits = Limits {
            max_cpu_time: Default::default(),
            max_file_size: Default::default(),
            max_data_size: Default::default(),
            max_stack_size: Limit {
                soft: Some(8388608),
                hard: None,
            },
            max_core_file_size: Default::default(),
            max_resident_set: Default::default(),
            max_processes: Limit {
                soft: Some(62935),
                hard: Some(62935),
            },
            max_open_files: Limit {
                soft: Some(1024),
                hard: Some(524288),
            },
            max_locked_memory: Limit {
                soft: Some(65536),
                hard: Some(65536),
            },
            max_address_space: Default::default(),
            max_file_locks: Default::default(),
            max_pending_signals: Limit {
                soft: Some(62935),
                hard: Some(62935),
            },
            max_msgqueue_size: Limit {
                soft: Some(819200),
                hard: Some(819200),
            },
            max_nice_priority: Limit {
                soft: Some(0),
                hard: Some(0),
            },
            max_realtime_priority: Limit {
                soft: Some(99),
                hard: Some(99),
            },
            max_realtime_timeout: Default::default(),
        };

        assert_eq!(limits, expected_limits);
    }
}