1use std::collections::HashMap;
6use std::os::unix::io::RawFd;
7
8#[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#[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#[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#[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
187pub 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
292pub 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}