dope-kqueue 0.2.3

#![cfg(any(
    target_os = "macos",
    target_os = "freebsd",
    target_os = "netbsd",
    target_os = "openbsd",
    target_os = "dragonfly"
))]
#![allow(clippy::missing_safety_doc)]

mod driver_config;
mod errno;

pub mod driver;

use std::future::Future;
use std::io;
use std::os::fd::{AsRawFd, RawFd};
use std::pin::Pin;
use std::task::{Context, Poll};

use dope_core::Fd;
use dope_core::addr::SockAddr;
use dope_core::driver::DriverOps;
use dope_core::submit::{RecvMsgPacket, SendMsgPacket, SubmitOps};

pub use driver::{
    CurrentSlot, Driver, ProvidedBuf, ProvidedRing, Request, Token, current, current_driver_mut,
    set_current,
};
pub use driver_config::DriverConfig;
pub use errno::Errno;

pub fn new_driver(cfg: DriverConfig) -> io::Result<Driver> {
    Driver::new(cfg)
}

#[derive(Clone, Copy, Debug, Default)]
pub struct KqueueBackend;

impl dope_core::driver::Backend for KqueueBackend {
    type Driver = driver::Driver;
    type Config = DriverConfig;
    type RecvMsgOut = driver::RecvMsgOut;
    type ProvidedBuf = driver::ProvidedBuf;

    fn new_driver(cfg: Self::Config) -> io::Result<Self::Driver> {
        Driver::new(cfg)
    }
}

thread_local! {
    static AUX: std::cell::Cell<Option<dope_core::driver::ExternalDispatch<KqueueBackend>>> =
        const { std::cell::Cell::new(None) };
}

impl dope_core::driver::CompletionBackend for KqueueBackend {
    type Token = driver::Token;

    fn current_slot() -> Option<dope_core::driver::CurrentSlot<Self>> {
        driver::current()
    }

    fn set_slot(slot: Option<dope_core::driver::CurrentSlot<Self>>) {
        driver::set_current(slot);
    }

    fn install_external(
        aux: Option<dope_core::driver::ExternalDispatch<Self>>,
    ) -> Option<dope_core::driver::ExternalDispatch<Self>> {
        AUX.with(|c| c.replace(aux))
    }

    fn current_external() -> Option<dope_core::driver::ExternalDispatch<Self>> {
        AUX.with(|c| c.get())
    }

    fn drive_with<R: dope_core::driver::CompletionRouter<Self>>(
        driver: &mut Self::Driver,
        router: &mut R,
    ) -> io::Result<bool> {
        driver.drive_with(router)
    }

    fn dispatch_completions<R: dope_core::driver::CompletionRouter<Self>>(
        driver: &mut Self::Driver,
        router: &mut R,
    ) {
        driver.dispatch_completions(router)
    }
}

impl SubmitOps for &mut Driver {
    fn submit_accept(
        self,
        fd: Fd,
    ) -> impl Future<Output = io::Result<(RawFd, SockAddr)>> + 'static {
        let raw = fd.as_raw_fd();
        driver::Sock::set_nonblocking(raw);
        AcceptFuture::new(raw)
    }

    fn submit_read<B>(self, fd: Fd, buf: B) -> impl Future<Output = (io::Result<u32>, B)> + 'static
    where
        B: AsMut<[u8]> + Unpin + 'static,
    {
        let raw = fd.as_raw_fd();
        async move {
            let mut buf = buf;
            let slice = buf.as_mut();
            let n = unsafe { libc::read(raw, slice.as_mut_ptr().cast(), slice.len()) };
            (Errno::syscall_result(n), buf)
        }
    }

    fn submit_send_raw<B>(
        self,
        fd: Fd,
        buf: B,
        len: usize,
    ) -> impl Future<Output = (io::Result<u32>, B)> + 'static
    where
        B: AsRef<[u8]> + Unpin + 'static,
    {
        let raw = fd.as_raw_fd();
        async move {
            let bytes = buf.as_ref();
            let n = len.min(bytes.len());
            let r = unsafe { libc::send(raw, bytes.as_ptr().cast(), n, 0) };
            (Errno::syscall_result(r), buf)
        }
    }

    fn submit_sendmsg_raw<P>(
        self,
        fd: Fd,
        packet: P,
    ) -> impl Future<Output = (io::Result<u32>, P)> + 'static
    where
        P: SendMsgPacket + Unpin + 'static,
    {
        let req = self.try_submit_sendmsg_raw(fd, packet);
        Request::await_or_err(req)
    }

    fn submit_recvmsg_raw<P>(
        self,
        fd: Fd,
        packet: P,
    ) -> impl Future<Output = (io::Result<u32>, P)> + 'static
    where
        P: RecvMsgPacket + Unpin + 'static,
    {
        let raw = fd.as_raw_fd();
        async move {
            let mut packet = packet;
            let msg = packet.msg_mut_ptr();
            let n = unsafe { libc::recvmsg(raw, msg, 0) };
            (Errno::syscall_result(n), packet)
        }
    }
}

enum AcceptState {
    Idle,
    Waiting(Request<()>),
    Done,
}

pub struct AcceptFuture {
    fd: RawFd,
    state: AcceptState,
}

impl AcceptFuture {
    fn new(fd: RawFd) -> Self {
        Self {
            fd,
            state: AcceptState::Idle,
        }
    }

    fn try_accept(&mut self) -> Option<io::Result<(RawFd, SockAddr)>> {
        let mut addr = SockAddr::empty();
        let r = unsafe { libc::accept(self.fd, addr.addr_mut_ptr(), addr.addr_len_ptr()) };
        if r >= 0 {
            return Some(Ok((r as RawFd, addr)));
        }
        let err = io::Error::last_os_error();
        if matches!(err.raw_os_error(), Some(libc::EAGAIN)) {
            return None;
        }
        Some(Err(err))
    }
}

impl Future for AcceptFuture {
    type Output = io::Result<(RawFd, SockAddr)>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let me = self.get_mut();
        loop {
            match &mut me.state {
                AcceptState::Done => panic!("AcceptFuture polled after completion"),
                AcceptState::Idle => {
                    if let Some(out) = me.try_accept() {
                        me.state = AcceptState::Done;
                        return Poll::Ready(out);
                    }
                    let driver = driver::current_driver_mut();
                    let req = driver.register_read_oneshot(me.fd);
                    me.state = AcceptState::Waiting(req);
                }
                AcceptState::Waiting(req) => match Pin::new(req).poll(cx) {
                    Poll::Pending => {
                        return Poll::Pending;
                    }
                    Poll::Ready((res, _flags, _)) => {
                        res?;
                        me.state = AcceptState::Idle;
                    }
                },
            }
        }
    }
}