ntex/ntex-net/src/rt_polling/connect.rs

use std::os::fd::{AsRawFd, RawFd};
use std::{cell::RefCell, io, rc::Rc, task::Poll};

use ntex_neon::driver::{DriverApi, Event, Handler};
use ntex_neon::{syscall, Runtime};
use ntex_util::channel::oneshot::Sender;
use slab::Slab;
use socket2::SockAddr;

#[derive(Clone)]
pub(crate) struct ConnectOps(Rc<ConnectOpsInner>);

#[derive(Debug)]
enum Change {
    Event(Event),
    Error(io::Error),
}

struct ConnectOpsBatcher {
    inner: Rc<ConnectOpsInner>,
}

struct Item {
    fd: RawFd,
    tag: &'static str,
    sender: Sender<io::Result<()>>,
}

struct ConnectOpsInner {
    api: DriverApi,
    connects: RefCell<Slab<Item>>,
}

impl ConnectOps {
    pub(crate) fn current() -> Self {
        Runtime::value(|rt| {
            let mut inner = None;
            rt.driver().register(|api| {
                let ops = Rc::new(ConnectOpsInner {
                    api,
                    connects: RefCell::new(Slab::new()),
                });
                inner = Some(ops.clone());
                Box::new(ConnectOpsBatcher { inner: ops })
            });

            ConnectOps(inner.unwrap())
        })
    }

    pub(crate) fn connect(
        &self,
        tag: &'static str,
        fd: RawFd,
        addr: SockAddr,
        sender: Sender<io::Result<()>>,
    ) -> io::Result<usize> {
        let result = syscall!(break libc::connect(fd, addr.as_ptr(), addr.len()));

        if let Poll::Ready(res) = result {
            res?;
        }

        let item = Item { tag, fd, sender };
        let id = self.0.connects.borrow_mut().insert(item);

        self.0
            .api
            .attach(tag, fd, id as u32, Some(Event::writable(0)));
        Ok(id)
    }
}

impl Handler for ConnectOpsBatcher {
    fn event(&mut self, id: usize, event: Event) {
        log::debug!("connect-fd is readable {:?}", id);

        let mut connects = self.inner.connects.borrow_mut();

        if connects.contains(id) {
            let item = connects.remove(id);
            if event.writable {
                let mut err: libc::c_int = 0;
                let mut err_len = std::mem::size_of::<libc::c_int>() as libc::socklen_t;

                let res = syscall!(libc::getsockopt(
                    item.fd.as_raw_fd(),
                    libc::SOL_SOCKET,
                    libc::SO_ERROR,
                    &mut err as *mut _ as *mut _,
                    &mut err_len
                ));

                let res = if err == 0 {
                    res.map(|_| ())
                } else {
                    Err(io::Error::from_raw_os_error(err))
                };

                self.inner.api.detach(item.tag, item.fd, id as u32);
                let _ = item.sender.send(res);
            }
        }
    }

    fn error(&mut self, id: usize, err: io::Error) {
        let mut connects = self.inner.connects.borrow_mut();

        if connects.contains(id) {
            let item = connects.remove(id);
            let _ = item.sender.send(Err(err));
            self.inner.api.detach(item.tag, item.fd, id as u32);
        }
    }
}