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

use super::AsRawFd;

use std::io;
use core::mem::MaybeUninit;
use core::{time, cmp, ptr};

pub type RawFd = std::os::unix::io::RawFd;
///Limit on file descriptors in select's set
pub const FD_LIMIT: usize = libc::FD_SETSIZE as usize;

impl<T: std::os::unix::io::AsRawFd> AsRawFd for T {
    #[inline(always)]
    fn as_raw_fd(&self) -> RawFd {
        std::os::unix::io::AsRawFd::as_raw_fd(self)
    }
}

#[derive(Clone, Copy)]
pub struct FdSet {
    max_fd: RawFd,
    count: usize,
    inner: MaybeUninit<libc::fd_set>,
}

impl FdSet {
    #[inline]
    pub fn new() -> Self {
        let mut inner = MaybeUninit::uninit();
        unsafe {
            libc::FD_ZERO(inner.as_mut_ptr());
        }

        Self {
            max_fd: 0,
            count: 0,
            inner,
        }
    }

    #[inline]
    pub fn add<T: AsRawFd>(&mut self, source: &T) {
        self.count += 1;
        self.max_fd = cmp::max(source.as_raw_fd(), self.max_fd);
        unsafe {
            libc::FD_SET(source.as_raw_fd(), self.inner.as_mut_ptr())
        }
    }

    #[inline]
    pub fn is_present<T: AsRawFd>(&self, source: &T) -> bool {
        unsafe {
            libc::FD_ISSET(source.as_raw_fd(), self.inner.as_ptr() as _)
        }
    }

    #[inline]
    pub fn clear(&mut self) {
        self.count = 0;
        self.max_fd = 0;
        unsafe {
            libc::FD_ZERO(self.inner.as_mut_ptr());
        }
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.count
    }
}

pub fn select(read: &mut FdSet, write: &mut FdSet, except: &mut FdSet) -> io::Result<usize> {
    let nfds = cmp::max(cmp::max(read.max_fd, write.max_fd), except.max_fd) + 1;

    let read_fd = match read.max_fd {
        0 => ptr::null_mut(),
        _ => read.inner.as_mut_ptr(),
    };

    let write_fd = match write.max_fd {
        0 => ptr::null_mut(),
        _ => write.inner.as_mut_ptr(),
    };

    let except_fd = match except.max_fd {
        0 => ptr::null_mut(),
        _ => except.inner.as_mut_ptr(),
    };

    let result = unsafe {
        libc::select(nfds, read_fd, write_fd, except_fd, ptr::null_mut())
    };

    match result {
        -1 => Err(io::Error::last_os_error()),
        _ => Ok(result as usize)
    }
}

pub fn select_timeout(read: &mut FdSet, write: &mut FdSet, except: &mut FdSet, timeout: time::Duration) -> io::Result<usize> {
    use core::convert::TryInto;

    let mut timeout = libc::timeval {
        tv_sec: match timeout.as_secs().try_into() {
            Ok(secs) => secs,
            Err(_) => return Err(io::Error::new(io::ErrorKind::InvalidInput, "Duration as seconds doesn't fit timeval")),
        },
        tv_usec: timeout.subsec_micros() as _,
    };

    let nfds = cmp::max(cmp::max(read.max_fd, write.max_fd), except.max_fd) + 1;

    let read_fd = match read.max_fd {
        0 => ptr::null_mut(),
        _ => read.inner.as_mut_ptr(),
    };

    let write_fd = match write.max_fd {
        0 => ptr::null_mut(),
        _ => write.inner.as_mut_ptr(),
    };

    let except_fd = match except.max_fd {
        0 => ptr::null_mut(),
        _ => except.inner.as_mut_ptr(),
    };

    let result = unsafe {
        libc::select(nfds, read_fd, write_fd, except_fd, &mut timeout)
    };

    match result {
        -1 => Err(io::Error::last_os_error()),
        _ => Ok(result as usize)
    }
}