dynix/src/daemon.rs

use std::{
    env, io,
    ops::Deref,
    os::fd::{AsFd, BorrowedFd, IntoRawFd, OwnedFd, RawFd},
    sync::{
        LazyLock,
        atomic::{AtomicUsize, Ordering},
    },
    time::Duration,
};

use iddqd::{BiHashMap, IdOrdMap};

use mio::{Events, Interest, Poll, Token, net::UnixListener, unix::SourceFd};

use rustix::{
    buffer::spare_capacity,
    fs::{FileType, Stat},
    net::SocketFlags,
    process::Uid,
};

use serde::{Deserialize, Serialize};
use serde_json::StreamDeserializer;

use crate::{
    daemon_tokfd::{FdInfo, FdKind},
    prelude::*,
};

use crate::{OwnedFdWithFlags, TokenFd};

pub static UID: LazyLock<Uid> = LazyLock::new(|| rustix::process::getuid());

pub static USER_SOCKET_DIR: LazyLock<&'static Path> = LazyLock::new(|| {
    let dir: Box<Path> = env::var_os("XDG_RUNTIME_DIR")
        .map(PathBuf::from)
        .unwrap_or_else(|| ["/", "run", "user", &UID.to_string()].into_iter().collect())
        .into_boxed_path();

    Box::leak(dir)
});

pub static TMPDIR: LazyLock<&'static Path> = LazyLock::new(|| {
    let dir: Box<Path> = env::temp_dir().into_boxed_path();

    Box::leak(dir)
});

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[derive(Deserialize, Serialize)]
#[serde(untagged)]
pub enum ConvenientAttrPath {
    Dotted(Box<str>),
    Split(Box<[Box<str>]>),
}

impl ConvenientAttrPath {
    pub fn clone_from_dotted(s: &str) -> Self {
        Self::Dotted(Box::from(s))
    }

    pub fn clone_from_split(s: &[&str]) -> Self {
        Self::from_str_iter(s.into_iter().map(Deref::deref))
    }

    pub fn from_str_iter<'i, I>(iter: I) -> Self
    where
        I: Iterator<Item = &'i str>,
    {
        let boxed = iter.map(|s| Box::from(s));
        Self::Split(Box::from_iter(boxed))
    }
}

impl<'i> FromIterator<&'i str> for ConvenientAttrPath {
    fn from_iter<I>(iter: I) -> Self
    where
        I: IntoIterator<Item = &'i str>,
    {
        Self::from_str_iter(iter.into_iter())
    }
}

impl FromIterator<Box<str>> for ConvenientAttrPath {
    fn from_iter<I>(iter: I) -> Self
    where
        I: IntoIterator<Item = Box<str>>,
    {
        Self::Split(Box::from_iter(iter))
    }
}

#[derive(Debug, Clone, PartialEq)]
#[derive(serde::Deserialize, serde::Serialize)]
#[serde(tag = "action", content = "args", rename_all = "snake_case")]
pub enum DaemonCmd {
    Append {
        name: ConvenientAttrPath,
        value: Box<str>,
    },
}

const TIMEOUT_NEVER: Option<Duration> = None;

static NEXT_TOKEN_NUMBER: AtomicUsize = AtomicUsize::new(1);
fn next_token() -> Token {
    let tok = NEXT_TOKEN_NUMBER.fetch_add(1, Ordering::SeqCst);

    // If the increment wrapped to 0, then we just increment it again.
    if tok == 0 {
        return Token(NEXT_TOKEN_NUMBER.fetch_add(1, Ordering::SeqCst));
    }

    Token(tok)
}

#[derive(Debug)]
pub struct Daemon {
    fd: OwnedFdWithFlags,
    path: Option<Box<Path>>,

    poller: Poll,

    //fd_error_buffer: CircularBuffer<ERROR_BUFFER_LEN, IoError>,
    fd_info: IdOrdMap<FdInfo>,

    // Bijective mapping of [`mio::Token`]s to [`RawFd`]s.
    tokfd: BiHashMap<TokenFd>,

    cmd_buffer: Vec<u8>,
    next_timeout: Option<Duration>,
}

/// `tokfd` handling.
impl Daemon {
    fn fd_error_pop(&mut self, fd: RawFd) -> Option<IoError> {
        self.fd_info
            .get_mut(&fd)
            .unwrap_or_else(|| {
                if let Ok(name) = FdInfo::guess_name(fd) {
                    panic!(
                        "tried to pop error for unknown fd {fd} ({})",
                        name.to_string_lossy(),
                    );
                }
                panic!("tried to pop error for unknown fd {fd}")
            })
            .error_buffer
            .pop_front()
    }

    fn fd_error_push(&mut self, fd: RawFd, error: IoError) -> Result<(), IoError> {
        self.fd_info
            .get_mut(&fd)
            .unwrap_or_else(|| panic!("tried to push error for unknown fd {fd}"))
            .error_buffer
            .try_push_back(error)
    }

    fn register(&mut self, fd: RawFd) -> Token {
        self.register_with_kind(fd, FdKind::Unknown)
    }

    fn register_with_kind(&mut self, fd: RawFd, kind: FdKind) -> Token {
        let token = next_token();

        self.fd_info.insert_unique(FdInfo::new(fd, kind)).unwrap();

        self.tokfd
            .insert_unique(TokenFd { token, fd })
            .unwrap_or_else(|e| todo!("{e}"));

        let mut source = SourceFd(&fd);
        self.poller
            .registry()
            .register(&mut source, token, Interest::READABLE)
            .unwrap_or_else(|e| unreachable!("registering {fd:?} with poller failed: {e}"));

        token
    }

    fn fd_for_token(&self, token: Token) -> Option<RawFd> {
        self.tokfd
            .get1(&token)
            .map(|TokenFd { fd, .. }| fd)
            .copied()
    }

    fn token_for_fd(&self, fd: RawFd) -> Option<Token> {
        self.tokfd
            .get2(&fd)
            .map(|TokenFd { token, .. }| token)
            .copied()
    }
}

impl Daemon {
    pub fn new(fd: OwnedFd, kind: FdKind, name: Option<Box<OsStr>>) -> Self {
        let mut fd_info: IdOrdMap<FdInfo> = Default::default();

        // Supposedly, the only possible ways this can fail are EMFILE, ENFILE, and ENOMEM.
        // If any of those are the case, we're screwed anyway.
        let poller = Poll::new().unwrap_or_else(|e| panic!("can't create new mio::Poll: {e}"));
        // Make sure we register the poller in `fd_info`, so we can keep track of its errors.
        fd_info
            .insert_unique(FdInfo::new(poller.as_raw_fd(), FdKind::Poller))
            .unwrap_or_else(|e| unreachable!("{e}"));

        let fd = OwnedFdWithFlags::new_with_fallback(fd);

        fd_info
            .insert_unique(FdInfo::new(fd.as_raw_fd(), kind))
            .unwrap_or_else(|e| unreachable!("{e}"));

        debug!("opened daemon to {:?} file descriptor {fd:?}", name);

        let path = match &name {
            Some(name) => Some(PathBuf::from(name).into_boxed_path()),
            None => None,
        };

        Self {
            fd,
            path,
            poller,
            //poll_error_buffer: Default::default(),
            fd_info,
            tokfd: Default::default(),
            //fd_error_buffer: Default::default(),
            cmd_buffer: Vec::with_capacity(1024),
            next_timeout: TIMEOUT_NEVER,
        }
    }

    pub fn from_unix_socket_path(path: &Path) -> Result<Self, IoError> {
        // We unconditionally unlink() `path` before binding, but completely ignore the result.
        let _ = rustix::fs::unlink(path);
        let listener = UnixListener::bind(path)
            .tap_err(|e| error!("failed to bind AF_UNIX socket at {}: {e}", path.display()))?;
        let listener_owned_fd = OwnedFd::from(listener);
        // FIXME: should we KEEP_ALIVE?
        rustix::net::sockopt::set_socket_keepalive(&listener_owned_fd, true).unwrap();
        let path: Box<OsStr> = path.to_path_buf().into_boxed_path().into_boxed_os_str();

        Ok(Self::new(listener_owned_fd, FdKind::Socket, Some(path)))
    }

    pub fn open_default_socket() -> Result<Self, IoError> {
        use IoErrorKind::*;
        let preferred = USER_SOCKET_DIR.join("dynix.sock").into_boxed_path();
        let constructed = match Self::from_unix_socket_path(&preferred) {
            Ok(v) => v,
            Err(e) if e.kind() == Unsupported => {
                //
                return Err(e);
            },
            Err(e) => {
                warn!(
                    "failed binding AF_UNIX socket at {}: {e}; trying elsewhere",
                    preferred.display()
                );

                let fallback = TMPDIR.join("dynix.sock").into_boxed_path();
                Self::from_unix_socket_path(&fallback).tap_err(|e| {
                    error!(
                        "failed binding AF_UNIX socket at {}: {e}",
                        fallback.display()
                    )
                })?
            },
        };

        Ok(constructed)
    }

    /// This panics if stdin cannot be opened.
    ///
    /// If you want to handle that error, use [`Daemon::from_raw_parts()`].
    pub fn from_stdin() -> Self {
        let stdin = io::stdin();
        let fd = stdin
            .as_fd()
            .try_clone_to_owned()
            .expect("dynix daemon could not open stdin; try a Unix socket?");

        Self::new(fd, FdKind::File, None)
    }

    //pub unsafe fn from_raw_parts(fd: OwnedFd) -> Self {
    //    Self {
    //        fd: OwnedFdWithFlags::new_with_fallback(fd),
    //    }
    //}

    pub fn fd(&self) -> BorrowedFd<'_> {
        self.fd.as_fd()
    }
}

/// Private helpers.
impl Daemon {
    fn read_cmd(&mut self, fd: &BorrowedFd) -> Result<(), IoError> {
        //let fd = match fd.as_ref() {
        //    Some(fd) => fd.as_fd(),
        //    None => self.fd.as_fd(),
        //};

        if self.cmd_buffer.len() == self.cmd_buffer.capacity() {
            self.cmd_buffer.reserve(1024);
        }

        let _count = rustix::io::read(fd, spare_capacity(&mut self.cmd_buffer))
            .tap_err(|e| error!("read() on daemon fd {fd:?} failed: {e}"))?;

        // The buffer might have existing data from the last read.
        let deserializer = serde_json::Deserializer::from_slice(&self.cmd_buffer);
        let stream: StreamDeserializer<_, DaemonCmd> = deserializer.into_iter();
        for cmd in stream {
            let cmd = match cmd {
                Ok(cmd) => cmd,
                Err(e) if e.is_eof() => {
                    self.next_timeout = Some(Duration::from_secs(4));
                    warn!("Didn't get a valid daemon command; giving the other side 4 seconds...");
                    return Ok(());
                },
                Err(e) => {
                    warn!("error deserializing command: {e}");
                    debug!("command buffer was: {}", self.cmd_buffer.as_bstr());
                    //warn!("error count for fd {fd:?}: {}", self.fd_error_buffer.len());
                    self.cmd_buffer.clear();
                    // Don't propagate the error unless we have too many.
                    self.fd_error_push(fd.as_raw_fd(), e.into()).tap_err(|e| {
                        error!("Accumulated too many errors for daemon fd {fd:?}: {e}")
                    })?;
                    //self.fd_error_buffer.try_push_back(e.into()).tap_err(|e| {
                    //    error!("Accumulated too many errors for daemon fd {fd:?}: {e}")
                    //})?;
                    return Ok(());
                },
            };
            debug!("got cmd: {cmd:?}");
        }

        self.cmd_buffer.clear();

        Ok(())
    }

    pub(crate) fn enter_loop(&mut self) -> Result<Option<()>, IoError> {
        let raw_fd = self.fd.as_raw_fd();
        if cfg!(debug_assertions) {
            assert!(
                self.fd_info.contains_key(&raw_fd),
                "we should know about daemon fd {raw_fd}",
            );
            assert!(
                self.fd_info.contains_key(&self.poller.as_raw_fd()),
                "we should know about poller fd {}",
                self.poller.as_raw_fd(),
            );
        }
        let mut daemon_source = SourceFd(&raw_fd);
        const DAEMON: Token = Token(0);
        self.tokfd
            .insert_unique(TokenFd {
                token: DAEMON,
                fd: raw_fd,
            })
            .unwrap();

        self.poller
            .registry()
            .register(&mut daemon_source, DAEMON, Interest::READABLE)
            .unwrap_or_else(|e| unreachable!("registering mio Poll for daemon fd {raw_fd:?}: {e}"));

        let mut events = Events::with_capacity(1024);

        loop {
            if let Some(timeout) = self.next_timeout {
                debug!(
                    "epoll_wait() with a timeout: {}",
                    humantime::format_duration(timeout),
                );
            }

            match self.poller.poll(&mut events, self.next_timeout.take()) {
                Ok(_) => {
                    if events.is_empty() {
                        warn!("timeout expired");
                        self.cmd_buffer.clear();
                    } else {
                        //let _ = self.poll_error_buffer.pop_front();
                        let _ = self.fd_error_pop(self.poller.as_raw_fd());
                    }
                },
                Err(e) => {
                    warn!("mio Poll::poll() error: {e}");
                    self.fd_error_push(self.poller.as_raw_fd(), e)
                        .tap_err(|e| {
                            error!("accumulated too many errors for mio::Poll::poll(): {e}")
                        })?;
                },
            }

            for event in &events {
                match event.token() {
                    DAEMON => {
                        let is_sock = rustix::fs::fstat(&self.fd)
                            .map(|Stat { st_mode, .. }| st_mode)
                            .map(FileType::from_raw_mode)
                            .map(FileType::is_socket)
                            .unwrap_or(false);

                        if !is_sock {
                            // SAFETY: oh boy: disjoint borrows with extra steps.
                            let file_fd = unsafe { BorrowedFd::borrow_raw(self.fd.as_raw_fd()) };
                            self.read_cmd(&file_fd).unwrap();
                            continue;
                        }

                        // Accept, first.
                        let flags = SocketFlags::NONBLOCK | SocketFlags::CLOEXEC;
                        let stream_fd = match rustix::net::accept_with(&self.fd, flags) {
                            Ok(stream) => {
                                debug!(
                                    "Accepted connection from socket {:?} as stream {:?}",
                                    self.fd, stream,
                                );
                                stream
                            },
                            Err(e) => {
                                error!("accept4 on daemon socket failed: {e}");
                                self.fd_error_push(self.fd.as_raw_fd(), e.into())
                                    .tap_err(|e| {
                                        error!(
                                            "Accumulated too many errors for daemon fd {:?}: {e}",
                                            self.fd
                                        )
                                    })?;
                                //self.fd_error_buffer.try_push_back(e.into()).tap_err(|e| {
                                //    error!(
                                //        "Accumulated too many errors for daemon fd {:?}: {e}",
                                //        self.fd,
                                //    )
                                //})?;
                                continue;
                            },
                        };

                        // Add this stream to our poll interest list.
                        // NOTE: `stream_fd` is now effectively `ManuallyDrop`.
                        let stream_fd = stream_fd.into_raw_fd();
                        let _token = self.register(stream_fd);

                        // Wait for the next poll to handle.
                    },
                    other_token => {
                        // This must be a stream fd.
                        let stream_fd = self.fd_for_token(other_token).unwrap_or_else(|| {
                            unreachable!("tried to get fd for no-existent token? {other_token:?}")
                        });

                        // SAFETY: oh boy.
                        let stream_fd = unsafe { BorrowedFd::borrow_raw(stream_fd) };
                        self.read_cmd(&stream_fd).unwrap();
                    },
                }
            }
        }
    }
}

impl Drop for Daemon {
    fn drop(&mut self) {
        //let probably_synthetic = self
        //    .path
        //    .to_str()
        //    .map(|p| p.starts_with("«") && p.ends_with("»"))
        //    .unwrap_or(false);
        //if probably_synthetic {
        //    return;
        //}
        if let Some(path) = self.path.as_deref() {
            let _ = rustix::fs::unlink(path);
        }
    }
}