Skip to main content

zsh/
zselect.rs

1//! Select/poll builtin module - port of Modules/zselect.c
2//!
3//! Provides zselect builtin for select/poll system calls on file descriptors.
4
5use std::collections::HashMap;
6use std::os::unix::io::RawFd;
7
8/// Which type of event to monitor
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub enum SelectMode {
11    Read,
12    Write,
13    Error,
14}
15
16impl SelectMode {
17    pub fn flag_char(&self) -> char {
18        match self {
19            SelectMode::Read => 'r',
20            SelectMode::Write => 'w',
21            SelectMode::Error => 'e',
22        }
23    }
24
25    pub fn from_char(c: char) -> Option<Self> {
26        match c {
27            'r' => Some(SelectMode::Read),
28            'w' => Some(SelectMode::Write),
29            'e' => Some(SelectMode::Error),
30            _ => None,
31        }
32    }
33}
34
35/// Options for zselect builtin
36#[derive(Debug, Default)]
37pub struct ZselectOptions {
38    pub array_name: Option<String>,
39    pub hash_name: Option<String>,
40    pub timeout_hundredths: Option<i64>,
41    pub fds: Vec<(RawFd, SelectMode)>,
42}
43
44/// Result of select operation
45#[derive(Debug)]
46pub struct SelectResult {
47    pub ready_fds: Vec<(RawFd, SelectMode)>,
48    pub as_array: Vec<String>,
49    pub as_hash: HashMap<String, String>,
50}
51
52/// Perform select/poll on file descriptors
53#[cfg(unix)]
54pub fn zselect(options: &ZselectOptions) -> Result<SelectResult, String> {
55    use std::collections::HashSet;
56
57    if options.fds.is_empty() {
58        return Ok(SelectResult {
59            ready_fds: Vec::new(),
60            as_array: Vec::new(),
61            as_hash: HashMap::new(),
62        });
63    }
64
65    let mut read_fds: HashSet<RawFd> = HashSet::new();
66    let mut write_fds: HashSet<RawFd> = HashSet::new();
67    let mut error_fds: HashSet<RawFd> = HashSet::new();
68
69    let mut max_fd: RawFd = 0;
70
71    for (fd, mode) in &options.fds {
72        max_fd = max_fd.max(*fd);
73        match mode {
74            SelectMode::Read => {
75                read_fds.insert(*fd);
76            }
77            SelectMode::Write => {
78                write_fds.insert(*fd);
79            }
80            SelectMode::Error => {
81                error_fds.insert(*fd);
82            }
83        }
84    }
85
86    let mut poll_fds: Vec<libc::pollfd> = Vec::new();
87
88    for (fd, mode) in &options.fds {
89        let events = match mode {
90            SelectMode::Read => libc::POLLIN,
91            SelectMode::Write => libc::POLLOUT,
92            SelectMode::Error => libc::POLLERR | libc::POLLPRI,
93        };
94
95        if let Some(existing) = poll_fds.iter_mut().find(|p| p.fd == *fd) {
96            existing.events |= events;
97        } else {
98            poll_fds.push(libc::pollfd {
99                fd: *fd,
100                events,
101                revents: 0,
102            });
103        }
104    }
105
106    let timeout_ms = options
107        .timeout_hundredths
108        .map(|t| (t * 10) as libc::c_int)
109        .unwrap_or(-1);
110
111    let result = loop {
112        let ret = unsafe {
113            libc::poll(
114                poll_fds.as_mut_ptr(),
115                poll_fds.len() as libc::nfds_t,
116                timeout_ms,
117            )
118        };
119
120        if ret < 0 {
121            let err = std::io::Error::last_os_error();
122            if err.kind() == std::io::ErrorKind::Interrupted {
123                continue;
124            }
125            return Err(format!("error on select: {}", err));
126        }
127
128        break ret;
129    };
130
131    if result == 0 {
132        return Ok(SelectResult {
133            ready_fds: Vec::new(),
134            as_array: Vec::new(),
135            as_hash: HashMap::new(),
136        });
137    }
138
139    let mut ready_fds = Vec::new();
140    let mut fd_modes: HashMap<RawFd, String> = HashMap::new();
141
142    for pfd in &poll_fds {
143        if pfd.revents != 0 {
144            if pfd.revents & libc::POLLIN != 0 && read_fds.contains(&pfd.fd) {
145                ready_fds.push((pfd.fd, SelectMode::Read));
146                fd_modes.entry(pfd.fd).or_insert_with(String::new).push('r');
147            }
148            if pfd.revents & libc::POLLOUT != 0 && write_fds.contains(&pfd.fd) {
149                ready_fds.push((pfd.fd, SelectMode::Write));
150                fd_modes.entry(pfd.fd).or_insert_with(String::new).push('w');
151            }
152            if (pfd.revents & (libc::POLLERR | libc::POLLPRI) != 0) && error_fds.contains(&pfd.fd) {
153                ready_fds.push((pfd.fd, SelectMode::Error));
154                fd_modes.entry(pfd.fd).or_insert_with(String::new).push('e');
155            }
156        }
157    }
158
159    let as_hash: HashMap<String, String> = fd_modes
160        .iter()
161        .map(|(fd, modes)| (fd.to_string(), modes.clone()))
162        .collect();
163
164    let mut as_array = Vec::new();
165    let mut current_mode: Option<SelectMode> = None;
166
167    for (fd, mode) in &ready_fds {
168        if current_mode != Some(*mode) {
169            as_array.push(format!("-{}", mode.flag_char()));
170            current_mode = Some(*mode);
171        }
172        as_array.push(fd.to_string());
173    }
174
175    Ok(SelectResult {
176        ready_fds,
177        as_array,
178        as_hash,
179    })
180}
181
182#[cfg(not(unix))]
183pub fn zselect(_options: &ZselectOptions) -> Result<SelectResult, String> {
184    Err("your system does not implement the select system call".to_string())
185}
186
187/// Parse zselect arguments
188pub fn parse_zselect_args(args: &[&str]) -> Result<ZselectOptions, String> {
189    let mut options = ZselectOptions::default();
190    let mut current_mode = SelectMode::Read;
191    let mut i = 0;
192
193    while i < args.len() {
194        let arg = args[i];
195
196        if arg.starts_with('-') && arg.len() > 1 {
197            let chars: Vec<char> = arg[1..].chars().collect();
198            let mut j = 0;
199
200            while j < chars.len() {
201                match chars[j] {
202                    'a' => {
203                        let name = if j + 1 < chars.len() {
204                            chars[j + 1..].iter().collect::<String>()
205                        } else if i + 1 < args.len() {
206                            i += 1;
207                            args[i].to_string()
208                        } else {
209                            return Err("argument expected after -a".to_string());
210                        };
211
212                        if name
213                            .chars()
214                            .next()
215                            .map(|c| c.is_ascii_digit())
216                            .unwrap_or(false)
217                        {
218                            return Err(format!("invalid array name: {}", name));
219                        }
220                        options.array_name = Some(name);
221                        break;
222                    }
223                    'A' => {
224                        let name = if j + 1 < chars.len() {
225                            chars[j + 1..].iter().collect::<String>()
226                        } else if i + 1 < args.len() {
227                            i += 1;
228                            args[i].to_string()
229                        } else {
230                            return Err("argument expected after -A".to_string());
231                        };
232
233                        if name
234                            .chars()
235                            .next()
236                            .map(|c| c.is_ascii_digit())
237                            .unwrap_or(false)
238                        {
239                            return Err(format!("invalid array name: {}", name));
240                        }
241                        options.hash_name = Some(name);
242                        break;
243                    }
244                    'r' => current_mode = SelectMode::Read,
245                    'w' => current_mode = SelectMode::Write,
246                    'e' => current_mode = SelectMode::Error,
247                    't' => {
248                        let timeout_str = if j + 1 < chars.len() {
249                            chars[j + 1..].iter().collect::<String>()
250                        } else if i + 1 < args.len() {
251                            i += 1;
252                            args[i].to_string()
253                        } else {
254                            return Err("argument expected after -t".to_string());
255                        };
256
257                        let timeout: i64 = timeout_str
258                            .parse()
259                            .map_err(|_| format!("number expected after -t: {}", timeout_str))?;
260                        options.timeout_hundredths = Some(timeout);
261                        break;
262                    }
263                    c if c.is_ascii_digit() => {
264                        let fd_str: String = chars[j..].iter().collect();
265                        let fd: RawFd = fd_str
266                            .parse()
267                            .map_err(|_| format!("expecting file descriptor: {}", fd_str))?;
268                        options.fds.push((fd, current_mode));
269                        break;
270                    }
271                    c => {
272                        return Err(format!("unknown option: -{}", c));
273                    }
274                }
275                j += 1;
276            }
277        } else if arg.chars().all(|c| c.is_ascii_digit()) {
278            let fd: RawFd = arg
279                .parse()
280                .map_err(|_| format!("expecting file descriptor: {}", arg))?;
281            options.fds.push((fd, current_mode));
282        } else {
283            return Err(format!("expecting file descriptor: {}", arg));
284        }
285
286        i += 1;
287    }
288
289    Ok(options)
290}
291
292/// Execute zselect builtin
293pub fn builtin_zselect(args: &[&str]) -> (i32, Vec<String>, HashMap<String, String>) {
294    let options = match parse_zselect_args(args) {
295        Ok(opts) => opts,
296        Err(e) => {
297            eprintln!("zselect: {}", e);
298            return (1, Vec::new(), HashMap::new());
299        }
300    };
301
302    match zselect(&options) {
303        Ok(result) => {
304            if result.ready_fds.is_empty() {
305                (1, Vec::new(), HashMap::new())
306            } else {
307                (0, result.as_array, result.as_hash)
308            }
309        }
310        Err(e) => {
311            eprintln!("zselect: {}", e);
312            (1, Vec::new(), HashMap::new())
313        }
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320
321    #[test]
322    fn test_select_mode_char() {
323        assert_eq!(SelectMode::Read.flag_char(), 'r');
324        assert_eq!(SelectMode::Write.flag_char(), 'w');
325        assert_eq!(SelectMode::Error.flag_char(), 'e');
326    }
327
328    #[test]
329    fn test_select_mode_from_char() {
330        assert_eq!(SelectMode::from_char('r'), Some(SelectMode::Read));
331        assert_eq!(SelectMode::from_char('w'), Some(SelectMode::Write));
332        assert_eq!(SelectMode::from_char('e'), Some(SelectMode::Error));
333        assert_eq!(SelectMode::from_char('x'), None);
334    }
335
336    #[test]
337    fn test_parse_basic_args() {
338        let args = vec!["-r", "0", "-w", "1"];
339        let options = parse_zselect_args(&args).unwrap();
340
341        assert_eq!(options.fds.len(), 2);
342        assert!(options.fds.contains(&(0, SelectMode::Read)));
343        assert!(options.fds.contains(&(1, SelectMode::Write)));
344    }
345
346    #[test]
347    fn test_parse_timeout() {
348        let args = vec!["-t", "100", "-r", "0"];
349        let options = parse_zselect_args(&args).unwrap();
350
351        assert_eq!(options.timeout_hundredths, Some(100));
352    }
353
354    #[test]
355    fn test_parse_combined_args() {
356        let args = vec!["-r0", "-w1"];
357        let options = parse_zselect_args(&args).unwrap();
358
359        assert_eq!(options.fds.len(), 2);
360    }
361
362    #[test]
363    fn test_parse_array_name() {
364        let args = vec!["-a", "myarray", "-r", "0"];
365        let options = parse_zselect_args(&args).unwrap();
366
367        assert_eq!(options.array_name, Some("myarray".to_string()));
368    }
369
370    #[test]
371    fn test_parse_hash_name() {
372        let args = vec!["-A", "myhash", "-r", "0"];
373        let options = parse_zselect_args(&args).unwrap();
374
375        assert_eq!(options.hash_name, Some("myhash".to_string()));
376    }
377
378    #[test]
379    fn test_parse_invalid_fd() {
380        let args = vec!["-r", "abc"];
381        let result = parse_zselect_args(&args);
382        assert!(result.is_err());
383    }
384
385    #[test]
386    fn test_zselect_empty() {
387        let options = ZselectOptions::default();
388        let result = zselect(&options).unwrap();
389        assert!(result.ready_fds.is_empty());
390    }
391}