zsh/
rlimits.rs

1//! Resource limits - port of Builtins/rlimits.c
2//!
3//! Provides `limit`, `ulimit`, and `unlimit` builtins for managing resource limits.
4
5#[cfg(unix)]
6use libc::{
7    getrlimit, rlimit, setrlimit, RLIMIT_AS, RLIMIT_CORE, RLIMIT_CPU, RLIMIT_DATA, RLIMIT_FSIZE,
8    RLIMIT_NOFILE, RLIMIT_STACK, RLIM_INFINITY,
9};
10
11/// Resource limit type
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum LimitType {
14    Memory,
15    Number,
16    Time,
17    Microseconds,
18    Unknown,
19}
20
21/// Resource information
22#[derive(Debug, Clone)]
23pub struct ResInfo {
24    pub res: i32,
25    pub name: &'static str,
26    pub limit_type: LimitType,
27    pub unit: u64,
28    pub opt: char,
29    pub descr: &'static str,
30}
31
32/// Known resource limits
33#[cfg(unix)]
34pub static KNOWN_RESOURCES: &[ResInfo] = &[
35    ResInfo {
36        res: RLIMIT_CPU as i32,
37        name: "cputime",
38        limit_type: LimitType::Time,
39        unit: 1,
40        opt: 't',
41        descr: "cpu time (seconds)",
42    },
43    ResInfo {
44        res: RLIMIT_FSIZE as i32,
45        name: "filesize",
46        limit_type: LimitType::Memory,
47        unit: 512,
48        opt: 'f',
49        descr: "file size (blocks)",
50    },
51    ResInfo {
52        res: RLIMIT_DATA as i32,
53        name: "datasize",
54        limit_type: LimitType::Memory,
55        unit: 1024,
56        opt: 'd',
57        descr: "data seg size (kbytes)",
58    },
59    ResInfo {
60        res: RLIMIT_STACK as i32,
61        name: "stacksize",
62        limit_type: LimitType::Memory,
63        unit: 1024,
64        opt: 's',
65        descr: "stack size (kbytes)",
66    },
67    ResInfo {
68        res: RLIMIT_CORE as i32,
69        name: "coredumpsize",
70        limit_type: LimitType::Memory,
71        unit: 512,
72        opt: 'c',
73        descr: "core file size (blocks)",
74    },
75    ResInfo {
76        res: RLIMIT_NOFILE as i32,
77        name: "descriptors",
78        limit_type: LimitType::Number,
79        unit: 1,
80        opt: 'n',
81        descr: "file descriptors",
82    },
83    ResInfo {
84        res: RLIMIT_AS as i32,
85        name: "addressspace",
86        limit_type: LimitType::Memory,
87        unit: 1024,
88        opt: 'v',
89        descr: "address space (kbytes)",
90    },
91];
92
93#[cfg(not(unix))]
94pub static KNOWN_RESOURCES: &[ResInfo] = &[];
95
96/// A resource limit value
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98pub enum LimitValue {
99    Unlimited,
100    Value(u64),
101}
102
103impl LimitValue {
104    #[cfg(unix)]
105    pub fn from_rlim(val: u64) -> Self {
106        if val == RLIM_INFINITY as u64 {
107            LimitValue::Unlimited
108        } else {
109            LimitValue::Value(val)
110        }
111    }
112
113    #[cfg(unix)]
114    pub fn to_rlim(&self) -> u64 {
115        match self {
116            LimitValue::Unlimited => RLIM_INFINITY as u64,
117            LimitValue::Value(v) => *v,
118        }
119    }
120
121    pub fn format(&self, info: Option<&ResInfo>) -> String {
122        match self {
123            LimitValue::Unlimited => "unlimited".to_string(),
124            LimitValue::Value(val) => {
125                if let Some(info) = info {
126                    match info.limit_type {
127                        LimitType::Time => {
128                            let hours = val / 3600;
129                            let mins = (val / 60) % 60;
130                            let secs = val % 60;
131                            format!("{}:{:02}:{:02}", hours, mins, secs)
132                        }
133                        LimitType::Microseconds => format!("{}us", val),
134                        LimitType::Memory => {
135                            if *val >= 1024 * 1024 {
136                                format!("{}MB", val / (1024 * 1024))
137                            } else {
138                                format!("{}kB", val / 1024)
139                            }
140                        }
141                        _ => format!("{}", val),
142                    }
143                } else {
144                    format!("{}", val)
145                }
146            }
147        }
148    }
149}
150
151impl std::fmt::Display for LimitValue {
152    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
153        match self {
154            LimitValue::Unlimited => write!(f, "unlimited"),
155            LimitValue::Value(v) => write!(f, "{}", v),
156        }
157    }
158}
159
160/// Resource limits manager
161#[derive(Debug, Default)]
162pub struct ResourceLimits {
163    #[cfg(unix)]
164    cached: std::collections::HashMap<i32, (LimitValue, LimitValue)>,
165}
166
167impl ResourceLimits {
168    pub fn new() -> Self {
169        Self {
170            #[cfg(unix)]
171            cached: std::collections::HashMap::new(),
172        }
173    }
174
175    /// Get resource info by name (prefix match)
176    pub fn find_by_name(&self, name: &str) -> Option<&'static ResInfo> {
177        let mut found: Option<&'static ResInfo> = None;
178        let mut ambiguous = false;
179
180        for info in KNOWN_RESOURCES {
181            if info.name.starts_with(name) {
182                if found.is_some() {
183                    ambiguous = true;
184                    break;
185                }
186                found = Some(info);
187            }
188        }
189
190        if ambiguous {
191            None
192        } else {
193            found
194        }
195    }
196
197    /// Get resource info by option character
198    pub fn find_by_opt(&self, opt: char) -> Option<&'static ResInfo> {
199        KNOWN_RESOURCES.iter().find(|info| info.opt == opt)
200    }
201
202    /// Get resource info by resource number
203    pub fn find_by_res(&self, res: i32) -> Option<&'static ResInfo> {
204        KNOWN_RESOURCES.iter().find(|info| info.res == res)
205    }
206
207    /// Get current limit (soft and hard)
208    #[cfg(unix)]
209    pub fn get(&self, res: i32) -> Result<(LimitValue, LimitValue), String> {
210        let mut rlim = rlimit {
211            rlim_cur: 0,
212            rlim_max: 0,
213        };
214
215        if unsafe { getrlimit(res as _, &mut rlim) } < 0 {
216            return Err(format!(
217                "can't read limit: {}",
218                std::io::Error::last_os_error()
219            ));
220        }
221
222        Ok((
223            LimitValue::from_rlim(rlim.rlim_cur),
224            LimitValue::from_rlim(rlim.rlim_max),
225        ))
226    }
227
228    #[cfg(not(unix))]
229    pub fn get(&self, _res: i32) -> Result<(LimitValue, LimitValue), String> {
230        Err("resource limits not supported on this platform".to_string())
231    }
232
233    /// Set a limit
234    #[cfg(unix)]
235    pub fn set(
236        &mut self,
237        res: i32,
238        soft: Option<LimitValue>,
239        hard: Option<LimitValue>,
240    ) -> Result<(), String> {
241        let (cur_soft, cur_hard) = self.get(res)?;
242
243        let new_soft = soft.unwrap_or(cur_soft);
244        let new_hard = hard.unwrap_or(cur_hard);
245
246        if let LimitValue::Value(s) = new_soft {
247            if let LimitValue::Value(h) = new_hard {
248                if s > h {
249                    return Err("soft limit exceeds hard limit".to_string());
250                }
251            }
252        }
253
254        let euid = unsafe { libc::geteuid() };
255        if euid != 0 {
256            if let (LimitValue::Value(new_h), LimitValue::Value(cur_h)) = (new_hard, cur_hard) {
257                if new_h > cur_h {
258                    return Err("can't raise hard limits".to_string());
259                }
260            }
261        }
262
263        let rlim = rlimit {
264            rlim_cur: new_soft.to_rlim(),
265            rlim_max: new_hard.to_rlim(),
266        };
267
268        if unsafe { setrlimit(res as _, &rlim) } < 0 {
269            return Err(format!(
270                "setrlimit failed: {}",
271                std::io::Error::last_os_error()
272            ));
273        }
274
275        self.cached.insert(res, (new_soft, new_hard));
276        Ok(())
277    }
278
279    #[cfg(not(unix))]
280    pub fn set(
281        &mut self,
282        _res: i32,
283        _soft: Option<LimitValue>,
284        _hard: Option<LimitValue>,
285    ) -> Result<(), String> {
286        Err("resource limits not supported on this platform".to_string())
287    }
288
289    /// Remove a limit (set to unlimited)
290    pub fn unlimit(&mut self, res: i32, hard: bool) -> Result<(), String> {
291        if hard {
292            self.set(
293                res,
294                Some(LimitValue::Unlimited),
295                Some(LimitValue::Unlimited),
296            )
297        } else {
298            let (_, cur_hard) = self.get(res)?;
299            self.set(res, Some(cur_hard), None)
300        }
301    }
302
303    /// List all limits
304    pub fn list_all(&self, hard: bool) -> Vec<(String, LimitValue)> {
305        let mut result = Vec::new();
306
307        for info in KNOWN_RESOURCES {
308            if let Ok((soft, hard_val)) = self.get(info.res) {
309                let val = if hard { hard_val } else { soft };
310                result.push((info.name.to_string(), val));
311            }
312        }
313
314        result
315    }
316}
317
318/// Parse a limit value string
319pub fn parse_limit_value(s: &str, info: Option<&ResInfo>) -> Result<LimitValue, String> {
320    if s == "unlimited" {
321        return Ok(LimitValue::Unlimited);
322    }
323
324    let info = info.ok_or("unknown resource type")?;
325
326    match info.limit_type {
327        LimitType::Time => {
328            if let Some(colon_pos) = s.find(':') {
329                let hours: u64 = s[..colon_pos].parse().map_err(|_| "invalid number")?;
330                let rest = &s[colon_pos + 1..];
331
332                let (mins, secs) = if let Some(colon2) = rest.find(':') {
333                    let m: u64 = rest[..colon2].parse().map_err(|_| "invalid number")?;
334                    let s: u64 = rest[colon2 + 1..].parse().map_err(|_| "invalid number")?;
335                    (m, s)
336                } else {
337                    let m: u64 = rest.parse().map_err(|_| "invalid number")?;
338                    (m, 0)
339                };
340
341                Ok(LimitValue::Value(hours * 3600 + mins * 60 + secs))
342            } else {
343                let s_lower = s.to_lowercase();
344                let (num_str, multiplier) = if s_lower.ends_with('h') {
345                    (&s[..s.len() - 1], 3600)
346                } else if s_lower.ends_with('m') {
347                    (&s[..s.len() - 1], 60)
348                } else {
349                    (s, 1)
350                };
351
352                let val: u64 = num_str.parse().map_err(|_| "invalid number")?;
353                Ok(LimitValue::Value(val * multiplier))
354            }
355        }
356        LimitType::Memory => {
357            let s_lower = s.to_lowercase();
358            let (num_str, multiplier) = if s_lower.ends_with('g') {
359                (&s[..s.len() - 1], 1024 * 1024 * 1024)
360            } else if s_lower.ends_with('m') {
361                (&s[..s.len() - 1], 1024 * 1024)
362            } else if s_lower.ends_with('k') {
363                (&s[..s.len() - 1], 1024)
364            } else {
365                (s, 1024)
366            };
367
368            let val: u64 = num_str.parse().map_err(|_| "invalid number")?;
369            Ok(LimitValue::Value(val * multiplier))
370        }
371        _ => {
372            let val: u64 = s.parse().map_err(|_| "limit must be a number")?;
373            Ok(LimitValue::Value(val))
374        }
375    }
376}
377
378/// Format a limit for display (limit builtin style)
379pub fn format_limit_display(name: &str, val: LimitValue, info: Option<&ResInfo>) -> String {
380    format!("{:<16}{}", name, val.format(info))
381}
382
383/// Format a limit for display (ulimit builtin style)
384pub fn format_ulimit_display(info: &ResInfo, val: LimitValue, show_header: bool) -> String {
385    let mut result = String::new();
386
387    if show_header {
388        result.push_str(&format!("-{}: {:<32}", info.opt, info.descr));
389    }
390
391    match val {
392        LimitValue::Unlimited => result.push_str("unlimited"),
393        LimitValue::Value(v) => {
394            let display_val = v / info.unit;
395            result.push_str(&format!("{}", display_val));
396        }
397    }
398
399    result
400}
401
402/// Execute the limit builtin
403pub fn builtin_limit(
404    args: &[&str],
405    limits: &mut ResourceLimits,
406    hard: bool,
407    set: bool,
408) -> (i32, String) {
409    let mut output = String::new();
410
411    if args.is_empty() {
412        for (name, val) in limits.list_all(hard) {
413            let info = limits.find_by_name(&name);
414            output.push_str(&format_limit_display(&name, val, info));
415            output.push('\n');
416        }
417        return (0, output);
418    }
419
420    let mut i = 0;
421    while i < args.len() {
422        let name = args[i];
423
424        if name.chars().all(|c| c.is_ascii_digit()) {
425            let res: i32 = match name.parse() {
426                Ok(n) => n,
427                Err(_) => return (1, "limit: invalid resource number\n".to_string()),
428            };
429
430            if i + 1 >= args.len() {
431                match limits.get(res) {
432                    Ok((soft, hard_val)) => {
433                        let val = if hard { hard_val } else { soft };
434                        output.push_str(&format!("{:<16}{}\n", res, val));
435                    }
436                    Err(e) => return (1, format!("limit: {}\n", e)),
437                }
438                i += 1;
439                continue;
440            }
441
442            let val_str = args[i + 1];
443            let val = match parse_limit_value(val_str, None) {
444                Ok(v) => v,
445                Err(e) => return (1, format!("limit: {}\n", e)),
446            };
447
448            if set {
449                let (soft, hard_opt) = if hard {
450                    (None, Some(val))
451                } else {
452                    (Some(val), None)
453                };
454
455                if let Err(e) = limits.set(res, soft, hard_opt) {
456                    return (1, format!("limit: {}\n", e));
457                }
458            }
459
460            i += 2;
461            continue;
462        }
463
464        let info = match limits.find_by_name(name) {
465            Some(info) => info,
466            None => return (1, format!("limit: no such resource: {}\n", name)),
467        };
468
469        if i + 1 >= args.len() {
470            match limits.get(info.res) {
471                Ok((soft, hard_val)) => {
472                    let val = if hard { hard_val } else { soft };
473                    output.push_str(&format_limit_display(info.name, val, Some(info)));
474                    output.push('\n');
475                }
476                Err(e) => return (1, format!("limit: {}\n", e)),
477            }
478            i += 1;
479            continue;
480        }
481
482        let val_str = args[i + 1];
483        let val = match parse_limit_value(val_str, Some(info)) {
484            Ok(v) => v,
485            Err(e) => return (1, format!("limit: {}\n", e)),
486        };
487
488        if set {
489            let (soft, hard_opt) = if hard {
490                (None, Some(val))
491            } else {
492                (Some(val), None)
493            };
494
495            if let Err(e) = limits.set(info.res, soft, hard_opt) {
496                return (1, format!("limit: {}\n", e));
497            }
498        }
499
500        i += 2;
501    }
502
503    (0, output)
504}
505
506/// Execute the ulimit builtin
507pub fn builtin_ulimit(
508    args: &[&str],
509    limits: &mut ResourceLimits,
510    hard: bool,
511    soft: bool,
512) -> (i32, String) {
513    let mut output = String::new();
514    let show_all = args.iter().any(|a| *a == "-a");
515
516    if show_all || args.is_empty() {
517        let use_hard = hard && !soft;
518
519        for info in KNOWN_RESOURCES {
520            if let Ok((s, h)) = limits.get(info.res) {
521                let val = if use_hard { h } else { s };
522                output.push_str(&format_ulimit_display(info, val, true));
523                output.push('\n');
524            }
525        }
526        return (0, output);
527    }
528
529    let mut i = 0;
530    let mut res = RLIMIT_FSIZE as i32;
531    let mut use_hard = hard && !soft;
532
533    while i < args.len() {
534        let arg = args[i];
535
536        if arg.starts_with('-') {
537            for c in arg[1..].chars() {
538                match c {
539                    'H' => use_hard = true,
540                    'S' => use_hard = false,
541                    'a' => {}
542                    _ => {
543                        if let Some(info) = limits.find_by_opt(c) {
544                            res = info.res;
545                        } else {
546                            return (1, format!("ulimit: bad option: -{}\n", c));
547                        }
548                    }
549                }
550            }
551            i += 1;
552            continue;
553        }
554
555        let info = limits.find_by_res(res);
556        let val = match parse_limit_value(arg, info) {
557            Ok(v) => v,
558            Err(e) => return (1, format!("ulimit: {}\n", e)),
559        };
560
561        let (soft_opt, hard_opt) = if use_hard {
562            (None, Some(val))
563        } else {
564            (Some(val), None)
565        };
566
567        if let Err(e) = limits.set(res, soft_opt, hard_opt) {
568            return (1, format!("ulimit: {}\n", e));
569        }
570
571        i += 1;
572    }
573
574    if let Some(info) = limits.find_by_res(res) {
575        if let Ok((s, h)) = limits.get(res) {
576            let val = if use_hard { h } else { s };
577            output.push_str(&format_ulimit_display(info, val, false));
578            output.push('\n');
579        }
580    }
581
582    (0, output)
583}
584
585/// Execute the unlimit builtin
586pub fn builtin_unlimit(args: &[&str], limits: &mut ResourceLimits, hard: bool) -> (i32, String) {
587    if args.is_empty() {
588        for info in KNOWN_RESOURCES {
589            if let Err(e) = limits.unlimit(info.res, hard) {
590                if hard {
591                    return (1, format!("unlimit: {}: {}\n", info.name, e));
592                }
593            }
594        }
595        return (0, String::new());
596    }
597
598    for name in args {
599        let info = match limits.find_by_name(name) {
600            Some(info) => info,
601            None => {
602                if name.chars().all(|c| c.is_ascii_digit()) {
603                    let res: i32 = match name.parse() {
604                        Ok(n) => n,
605                        Err(_) => return (1, "unlimit: invalid resource number\n".to_string()),
606                    };
607                    if let Err(e) = limits.unlimit(res, hard) {
608                        return (1, format!("unlimit: {}\n", e));
609                    }
610                    continue;
611                }
612                return (1, format!("unlimit: no such resource: {}\n", name));
613            }
614        };
615
616        if let Err(e) = limits.unlimit(info.res, hard) {
617            return (1, format!("unlimit: {}: {}\n", info.name, e));
618        }
619    }
620
621    (0, String::new())
622}
623
624#[cfg(test)]
625mod tests {
626    use super::*;
627
628    #[test]
629    fn test_limit_value_format() {
630        assert_eq!(LimitValue::Unlimited.format(None), "unlimited");
631        assert_eq!(LimitValue::Value(1234).format(None), "1234");
632    }
633
634    #[test]
635    fn test_parse_limit_unlimited() {
636        let info = &KNOWN_RESOURCES[0]; // cputime
637        assert_eq!(
638            parse_limit_value("unlimited", Some(info)).unwrap(),
639            LimitValue::Unlimited
640        );
641    }
642
643    #[test]
644    #[cfg(unix)]
645    fn test_parse_limit_time() {
646        let info = KNOWN_RESOURCES
647            .iter()
648            .find(|i| i.limit_type == LimitType::Time)
649            .unwrap();
650
651        assert_eq!(
652            parse_limit_value("60", Some(info)).unwrap(),
653            LimitValue::Value(60)
654        );
655        assert_eq!(
656            parse_limit_value("1h", Some(info)).unwrap(),
657            LimitValue::Value(3600)
658        );
659        assert_eq!(
660            parse_limit_value("5m", Some(info)).unwrap(),
661            LimitValue::Value(300)
662        );
663        assert_eq!(
664            parse_limit_value("1:30", Some(info)).unwrap(),
665            LimitValue::Value(3600 + 30 * 60)
666        );
667        assert_eq!(
668            parse_limit_value("1:30:45", Some(info)).unwrap(),
669            LimitValue::Value(3600 + 30 * 60 + 45)
670        );
671    }
672
673    #[test]
674    #[cfg(unix)]
675    fn test_parse_limit_memory() {
676        let info = KNOWN_RESOURCES
677            .iter()
678            .find(|i| i.limit_type == LimitType::Memory)
679            .unwrap();
680
681        assert_eq!(
682            parse_limit_value("100", Some(info)).unwrap(),
683            LimitValue::Value(100 * 1024)
684        );
685        assert_eq!(
686            parse_limit_value("100k", Some(info)).unwrap(),
687            LimitValue::Value(100 * 1024)
688        );
689        assert_eq!(
690            parse_limit_value("10M", Some(info)).unwrap(),
691            LimitValue::Value(10 * 1024 * 1024)
692        );
693        assert_eq!(
694            parse_limit_value("1G", Some(info)).unwrap(),
695            LimitValue::Value(1024 * 1024 * 1024)
696        );
697    }
698
699    #[test]
700    #[cfg(unix)]
701    fn test_find_resource() {
702        let limits = ResourceLimits::new();
703
704        assert!(limits.find_by_name("cpu").is_some());
705        assert!(limits.find_by_name("cputime").is_some());
706        assert!(limits.find_by_name("file").is_some());
707        assert!(limits.find_by_name("nonexistent").is_none());
708
709        assert!(limits.find_by_opt('t').is_some());
710        assert!(limits.find_by_opt('f').is_some());
711        assert!(limits.find_by_opt('z').is_none());
712    }
713
714    #[test]
715    #[cfg(unix)]
716    fn test_get_limits() {
717        let limits = ResourceLimits::new();
718
719        let result = limits.get(RLIMIT_NOFILE as i32);
720        assert!(result.is_ok());
721
722        let (soft, hard) = result.unwrap();
723        match soft {
724            LimitValue::Unlimited => {}
725            LimitValue::Value(v) => assert!(v > 0),
726        }
727        match hard {
728            LimitValue::Unlimited => {}
729            LimitValue::Value(v) => assert!(v > 0),
730        }
731    }
732
733    #[test]
734    #[cfg(unix)]
735    fn test_list_all() {
736        let limits = ResourceLimits::new();
737        let all = limits.list_all(false);
738
739        assert!(!all.is_empty());
740        assert!(all.iter().any(|(name, _)| name == "cputime"));
741        assert!(all.iter().any(|(name, _)| name == "filesize"));
742    }
743}
zsh/rlimits.rs

zsh/
rlimits.rs
