use std::io;
use std::os::fd::RawFd;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Ready {
pub udp: bool,
pub tun: bool,
}
pub struct Epoll {
epoll_fd: RawFd,
udp_fd: RawFd,
tun_fd: RawFd,
}
pub fn set_nonblocking(fd: RawFd) -> io::Result<()> {
let rc = unsafe { libc::fcntl(fd, libc::F_SETFL, libc::O_NONBLOCK) };
if rc < 0 {
Err(io::Error::last_os_error())
} else {
Ok(())
}
}
pub fn read_fd(fd: RawFd, buf: &mut [u8]) -> io::Result<usize> {
let n = unsafe { libc::read(fd, buf.as_mut_ptr().cast(), buf.len()) };
if n < 0 {
return Err(io::Error::last_os_error());
}
Ok(usize::try_from(n).expect("non-negative read return fits usize"))
}
pub fn write_fd(fd: RawFd, buf: &[u8]) -> io::Result<usize> {
let n = unsafe { libc::write(fd, buf.as_ptr().cast(), buf.len()) };
if n < 0 {
return Err(io::Error::last_os_error());
}
Ok(usize::try_from(n).expect("non-negative write return fits usize"))
}
impl Epoll {
pub fn create() -> io::Result<Self> {
let epoll_fd = unsafe { libc::epoll_create1(libc::EPOLL_CLOEXEC) };
if epoll_fd < 0 {
return Err(io::Error::last_os_error());
}
Ok(Self {
epoll_fd,
udp_fd: -1,
tun_fd: -1,
})
}
pub fn add(&self, fd: RawFd, token: u64) -> io::Result<()> {
set_nonblocking(fd)?;
let mut ev = libc::epoll_event {
events: libc::EPOLLIN as u32,
u64: token,
};
let rc = unsafe { libc::epoll_ctl(self.epoll_fd, libc::EPOLL_CTL_ADD, fd, &raw mut ev) };
if rc < 0 {
return Err(io::Error::last_os_error());
}
Ok(())
}
pub fn remove(&self, fd: RawFd) -> io::Result<()> {
let rc = unsafe {
libc::epoll_ctl(self.epoll_fd, libc::EPOLL_CTL_DEL, fd, std::ptr::null_mut())
};
if rc < 0 {
return Err(io::Error::last_os_error());
}
Ok(())
}
pub fn wait_tokens(&self, timeout_ms: i32, out: &mut Vec<u64>) -> io::Result<()> {
out.clear();
let mut events = [libc::epoll_event { events: 0, u64: 0 }; 64];
let nfds = unsafe {
libc::epoll_wait(
self.epoll_fd,
events.as_mut_ptr(),
libc::c_int::try_from(events.len()).expect("event array len fits c_int"),
timeout_ms,
)
};
if nfds < 0 {
let e = io::Error::last_os_error();
if e.raw_os_error() == Some(libc::EINTR) {
return Ok(());
}
return Err(e);
}
for ev in &events[..usize::try_from(nfds).expect("non-negative nfds fits usize")] {
out.push(ev.u64);
}
Ok(())
}
pub fn new(udp_fd: RawFd, tun_fd: RawFd) -> io::Result<Self> {
let mut this = Self::create()?;
this.udp_fd = udp_fd;
this.tun_fd = tun_fd;
this.add(udp_fd, udp_fd as u64)?;
this.add(tun_fd, tun_fd as u64)?;
Ok(this)
}
pub fn wait(&self, timeout_ms: i32) -> io::Result<Ready> {
let mut events = [libc::epoll_event { events: 0, u64: 0 }; 2];
let nfds = unsafe {
libc::epoll_wait(
self.epoll_fd,
events.as_mut_ptr(),
libc::c_int::try_from(events.len()).expect("event array len fits c_int"),
timeout_ms,
)
};
if nfds < 0 {
let e = io::Error::last_os_error();
if e.raw_os_error() == Some(libc::EINTR) {
return Ok(Ready::default());
}
return Err(e);
}
let mut ready = Ready::default();
for ev in &events[..usize::try_from(nfds).expect("non-negative nfds fits usize")] {
let fd = ev.u64 as RawFd;
if fd == self.udp_fd {
ready.udp = true;
} else if fd == self.tun_fd {
ready.tun = true;
}
}
Ok(ready)
}
}
impl Drop for Epoll {
fn drop(&mut self) {
unsafe { libc::close(self.epoll_fd) };
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::UdpSocket;
use std::os::fd::AsRawFd;
#[test]
fn wait_reports_each_fd_independently() {
let udp_rx = UdpSocket::bind("127.0.0.1:0").unwrap();
let udp_tx = UdpSocket::bind("127.0.0.1:0").unwrap();
udp_tx.connect(udp_rx.local_addr().unwrap()).unwrap();
let tun_rx = UdpSocket::bind("127.0.0.1:0").unwrap();
let tun_tx = UdpSocket::bind("127.0.0.1:0").unwrap();
tun_tx.connect(tun_rx.local_addr().unwrap()).unwrap();
let ep = Epoll::new(udp_rx.as_raw_fd(), tun_rx.as_raw_fd()).unwrap();
assert_eq!(ep.wait(10).unwrap(), Ready::default());
udp_tx.send(b"ping").unwrap();
let after_udp = ep.wait(200).unwrap();
assert!(after_udp.udp, "udp must be ready after a datagram arrives");
assert!(
!after_udp.tun,
"tun must not be reported when only udp was fed"
);
udp_rx.recv(&mut [0u8; 8]).unwrap();
tun_tx.send(b"frame").unwrap();
let after_tun = ep.wait(200).unwrap();
assert!(after_tun.tun, "tun must be ready after a frame arrives");
assert!(
!after_tun.udp,
"udp must not be reported when only tun was fed"
);
}
#[test]
fn wait_times_out_with_nothing_ready() {
let a = UdpSocket::bind("127.0.0.1:0").unwrap();
let b = UdpSocket::bind("127.0.0.1:0").unwrap();
let ep = Epoll::new(a.as_raw_fd(), b.as_raw_fd()).unwrap();
assert_eq!(
ep.wait(5).unwrap(),
Ready::default(),
"idle fds must report nothing ready"
);
}
#[test]
fn test_read_write_fd_and_error_paths() {
let sock1 = UdpSocket::bind("127.0.0.1:0").unwrap();
let sock2 = UdpSocket::bind("127.0.0.1:0").unwrap();
sock1.connect(sock2.local_addr().unwrap()).unwrap();
let written = write_fd(sock1.as_raw_fd(), b"hello").unwrap();
assert_eq!(written, 5);
let mut buf = [0u8; 16];
let read = read_fd(sock2.as_raw_fd(), &mut buf).unwrap();
assert_eq!(read, 5);
assert_eq!(&buf[..5], b"hello");
assert!(set_nonblocking(-1).is_err());
assert!(read_fd(-1, &mut buf).is_err());
assert!(write_fd(-1, b"fail").is_err());
}
#[test]
fn epoll_dynamic_add_remove_and_wait_tokens() {
let epoll = Epoll::create().expect("create epoll");
let mut fds = [0; 2];
unsafe {
assert_eq!(libc::pipe(fds.as_mut_ptr()), 0);
}
let (rx, tx) = (fds[0], fds[1]);
epoll.add(rx, 42).expect("add rx fd");
let mut ready = Vec::new();
epoll.wait_tokens(0, &mut ready).expect("wait");
assert!(ready.is_empty());
let written = unsafe { libc::write(tx, [1u8].as_ptr().cast(), 1) };
assert_eq!(written, 1);
epoll.wait_tokens(50, &mut ready).expect("wait");
assert_eq!(ready, vec![42]);
epoll.remove(rx).expect("remove rx fd");
ready.clear();
epoll.wait_tokens(0, &mut ready).expect("wait");
assert!(ready.is_empty());
unsafe {
libc::close(rx);
libc::close(tx);
}
}
}