// Copyright 2020-2021 Ian Jackson and contributors to Otter // SPDX-License-Identifier: AGPL-3.0-or-later // There is NO WARRANTY. use crate::prelude::*; pub trait Timed { fn set_deadline(&mut self, deadline: Option); fn set_timeout(&mut self, timeout: Option); } pub trait TimedRead : Timed + Read { } pub trait TimedWrite: Timed + Write { } use nix::fcntl::{fcntl, OFlag, FcntlArg}; use nix::Error as NE; use mio::Token; pub struct TimedFd { fd: Fd, poll: mio::Poll, events: mio::event::Events, deadline: Option, rw: PhantomData, } pub trait TimedFdReadWrite { const INTEREST: mio::Interest; } pub type TimedFdReader = TimedFd; pub type TimedFdWriter = TimedFd; #[derive(Debug,Copy,Clone)] pub struct TimedFdRead; impl TimedFdReadWrite for TimedFdRead { const INTEREST : mio::Interest = mio::Interest::READABLE; } #[derive(Debug,Copy,Clone)] pub struct TimedFdWrite; impl TimedFdReadWrite for TimedFdWrite { const INTEREST : mio::Interest = mio::Interest::WRITABLE; } pub struct Fd(RawFd); impl Fd { pub fn from_raw_fd(fd: RawFd) -> Self { Fd(fd) } fn extract_raw_fd(&mut self) -> RawFd { mem::replace(&mut self.0, -1) } } impl IntoRawFd for Fd { fn into_raw_fd(mut self) -> RawFd { self.extract_raw_fd() } } impl AsRawFd for Fd { fn as_raw_fd(&self) -> RawFd { self.0 } } impl TimedFd where RW: TimedFdReadWrite { /// Takes ownership of the fd /// /// Will change the fd's open-file to nonblocking. #[throws(io::Error)] pub fn new(fd: F) -> TimedFd where F: IntoRawFd { Self::from_fd( Fd::from_raw_fd( fd.into_raw_fd() ))? } /// Takes ownership of the fd /// /// Will change the fd's open-file to nonblocking. #[throws(io::Error)] fn from_fd(fd: Fd) -> Self { fcntl(fd.as_raw_fd(), FcntlArg::F_SETFL(OFlag::O_NONBLOCK)) .map_err(|e| io::Error::from(e))?; let poll = mio::Poll::new()?; poll.registry().register( &mut mio::unix::SourceFd(&fd.as_raw_fd()), Token(0), RW::INTEREST, )?; let events = mio::event::Events::with_capacity(1); TimedFd { fd, poll, events, deadline: None, rw: PhantomData } } } impl Timed for TimedFd where RW: TimedFdReadWrite { fn set_deadline(&mut self, deadline: Option) { self.deadline = deadline; } fn set_timeout(&mut self, timeout: Option) { self.set_deadline(timeout.map(|timeout|{ Instant::now() + timeout })); } } impl TimedFd where RW: TimedFdReadWrite { #[throws(io::Error)] fn rw(&mut self, mut f: F) -> O where F: FnMut(i32) -> Result { 'again: loop { for event in &self.events { if event.token() == Token(0) { match f(self.fd.as_raw_fd()) { Ok(got) => { break 'again got }, Err(NE::EINTR) => continue 'again, Err(NE::EAGAIN) => break, Err(ne) => throw!(ne), } } } let timeout = if let Some(deadline) = self.deadline { let now = Instant::now(); if now >= deadline { throw!(io::ErrorKind::TimedOut) } Some(deadline - now) } else { None }; loop { match self.poll.poll(&mut self.events, timeout) { Err(e) if e.kind() == ErrorKind::Interrupted => continue, Err(e) => throw!(e), Ok(()) => break, } } if self.events.is_empty() { throw!(io::ErrorKind::TimedOut) } } } } impl Read for TimedFd { #[throws(io::Error)] fn read(&mut self, buf: &mut [u8]) -> usize { self.rw(|fd| unistd::read(fd, buf))? } } impl Write for TimedFd { #[throws(io::Error)] fn write(&mut self, buf: &[u8]) -> usize { self.rw(|fd| unistd::write(fd, buf))? } #[throws(io::Error)] fn flush(&mut self) { } } impl Drop for Fd { fn drop(&mut self) { let fd = self.extract_raw_fd(); if fd >= 2 { let _ = nix::unistd::close(fd); } } }