//! Support for file descriptor passing with `LISTEN_FD`
//!
//! The entrypoints of this module are [`FileDescriptor::receive()`] and [`FileDescriptor::receive_with_names()`]. (And their **`unsafe`** `*_no_dup` variants.)
//!
//! ```no_run
//! # main() -> Result<(), Box<dyn std::error::Error>> {
//! let received_fd: Vec<FileDescriptor> =  FileDescriptor::receive_with_names(true)
//!     .into_iter()
//!     .filter_map(|(name, res)| {
//!         if name == "<expected_filename>" {
//!             Some(res)
//!         } else {
//!             None
//!         }
//!     })
//!     .collect()?
//! // Use the received file descriptors.
//! # Ok(())
//! # }
//! ```
use std::ffi::{OsStr, OsString};
use std::fs::File;
use std::net::{TcpListener, TcpStream, UdpSocket};
use std::os::fd::RawFd;
use std::os::unix::net::{UnixDatagram, UnixListener, UnixStream};
use std::{env, process};

use error::{DupError, GetFdsError, ReceiveError, ReceiveNameError};
use rustix::fd::{AsRawFd, BorrowedFd, FromRawFd, OwnedFd};
use rustix::fs::FileType;
use rustix::io::FdFlags;
use rustix::net::SocketType;

pub mod error;

const SD_LISTEN_FDS_START: RawFd = 3;
const PID_VAR: &str = "LISTEN_PID";
const FD_NUMBER_VAR: &str = "LISTEN_FDS";
const FD_NAMES_VAR: &str = "LISTEN_FDNAMES";

/// File Descriptor passed by systemd.
///
/// They should truly be safe to use from rust code.
#[derive(Debug)]
pub enum FileDescriptor {
    /// The file descriptor is a [`File`].
    ///
    /// If this is an FD provided by a .socket unit it corresponds to `ListenSpecial=`.
    File(File),
    /// The file descriptor is a directory.
    Directory(OwnedFd),
    /// The file descriptor is a FIFO
    ///
    /// If this is an FD provided by a .socket unit it corresponds to `ListenFIFO=`
    Fifo(File),
    /// The file descriptor is a TCP socket listening for incoming connexions.
    ///
    /// If this is an FD provided by a .socket unit it corresponds to `ListenStream=` with an IP address.
    TcpListener(TcpListener),
    /// The file descriptor is a TCP socket that doesn't listen for incoming connexions.
    TcpStream(TcpStream),
    /// The file descriptor is an UDP .
    ///
    /// If this is an FD provided by a .socket unit it corresponds to `ListenDatagram=` with an IP address.
    UdpSocket(UdpSocket),
    /// The file descriptor is another inet socket
    ///
    /// You should figure out what exactly it is before you use it.
    InetOther(OwnedFd),
    /// The file descriptor is a Unix Stream Socket listening for incoming connexions.
    ///
    /// If this is an FD provided by a .socket unit it corresponds to `ListenStream=` with a path or starting with `@` (abstract socket).
    UnixListener(UnixListener),
    /// The file descriptor is a Unix Stream Socket.
    UnixStream(UnixStream),
    /// The file descriptor is a Unix Datagram Socket.
    ///
    /// If this is an FD provided by a .socket unit it corresponds to `ListenDatagram=` with a path or starting with `@` (abstract socket).
    UnixDatagram(UnixDatagram),
    /// The file descriptor is a Unix SEQPACKET Socket.
    ///
    /// If this is an FD provided by a .socket unit it corresponds to `ListenSequentialPacket=`
    UnixSequentialPaquet(OwnedFd),
    /// The file descriptor is another type of Unix Socket.
    UnixOther(OwnedFd),
    /// The file descriptor is some other socket type.
    ///
    /// Probably a Netlink family socket if from a .socket unit but you should check yourself.
    SocketOther(OwnedFd),
    /// The file descriptor is a message queue.
    ///
    /// If this is an FD provided by a .socket unit it corresponds to `ListenMessageQueue=`
    MessageQueue(OwnedFd),
    /// The file descriptor is something else.
    Other(OwnedFd),
}

impl FileDescriptor {
    /// Get any file descriptor passed by systemd or anything implementing the `LISTEN_FD` protocol.
    ///
    /// This isn't necessarily limited to File descriptor of listening sockets, IPCs or FIFOs but also anything that is in the file descriptor store.
    ///
    /// The file descriptors are duplicated using [`fcntl_dupfd_cloexec`](rustix::fs::fcntl_dupfd_cloexec) so they can safely be used from rust and will not be propagated to children process automatically.
    ///
    /// # Errors
    ///
    /// This function will fail if no file descriptor colud be received which might not actually be an error. See [`ReceiveError`] for details.
    pub fn receive(
        unset_env: bool,
    ) -> Result<impl IntoIterator<Item = Result<Self, DupError>>, ReceiveError> {
        let span = tracing::info_span!("receive", f_unset_env = ?unset_env);
        let _enter = span.enter();

        let fds = Self::inner_receive(unset_env)?;

        match Self::from_fds(fds) {
            Ok(fds) => Ok(fds),
            Err(error) => Err(ReceiveError::GetFds(error)),
        }
    }

    /// Get any file descriptor passed by systemd or anything implementing the `LISTEN_FD` protocol.
    ///
    /// This isn't necessarily limited to File descriptor of listening sockets, IPCs or FIFOs but also anything that is in the file descriptor store.
    ///
    /// The file descriptors are taken directly as [`OwnedFd`]s instead of being duplicated. In order to limit unsoundness this function therefore always unset the environment.
    ///
    /// # Errors
    ///
    /// This function will fail if no file descriptor colud be received which might not actually be an error. See [`ReceiveError`] for details.
    ///
    /// # Safety
    ///
    /// This function is safe if (and only if) the received file descriptors weren't already taken into a owned rust struct.
    /// (In short it needs to follow the safety constraints of [`OwnedFd::from_raw_fd`](std::os::fd::FromRawFd)).
    /// The simplest way to insure that it is so is to only use functions from this crate to get these file descriptors.
    pub unsafe fn receive_no_dup() -> Result<impl IntoIterator<Item = Self>, ReceiveError> {
        let span = tracing::info_span!("receive_no_dup");
        let _enter = span.enter();

        let fds = Self::inner_receive(true)?;

        Self::from_fds_no_dup(fds).map_err(ReceiveError::GetFds)
    }

    fn inner_receive(unset_env: bool) -> Result<usize, ReceiveError> {
        let pid = env::var_os(PID_VAR).ok_or(ReceiveError::NoListenPID)?;
        let fds = env::var_os(FD_NUMBER_VAR).ok_or(ReceiveError::NoListenFD)?;
        tracing::trace!("{PID_VAR} = {pid:?}; {FD_NUMBER_VAR} = {fds:?}");

        if unset_env {
            env::remove_var(PID_VAR);
            env::remove_var(FD_NUMBER_VAR);
            env::remove_var(FD_NAMES_VAR);
        }

        let pid = pid
            .into_string()
            .map_err(ReceiveError::NotUnicodeListenPID)?
            .parse::<u32>()
            .map_err(ReceiveError::ListenPIDParse)?;
        let fds = fds
            .into_string()
            .map_err(ReceiveError::NotUnicodeListenFD)?
            .parse::<usize>()
            .map_err(ReceiveError::ListenFDParse)?;

        let current_pid = process::id();
        if current_pid != pid {
            return Err(ReceiveError::PidMismatch {
                expected: pid,
                found: current_pid,
            });
        }

        Ok(fds)
    }

    /// Get any file descriptor passed using `LISTEN_FD` and their names.
    ///
    /// This isn't necessarily limited to File descriptor of listening sockets, IPCs or FIFOs but also anything that is in the file descriptor store.
    ///
    /// The file descriptors are duplicated using [`fcntl_dupfd_cloexec`](rustix::fs::fcntl_dupfd_cloexec) so they can safely be used from rust and will not be propagated to children process automatically.
    ///
    /// # Errors
    ///
    /// This function will fail if no file descriptors could be obtained or the names associated with them couldn't be obtained. See [`ReceiveNameError`] for details.
    pub fn receive_with_names(
        unset_env: bool,
    ) -> Result<impl IntoIterator<Item = (OsString, Result<Self, DupError>)>, ReceiveNameError>
    {
        let span = tracing::info_span!("receive_with_names", f_unset_env = ?unset_env);
        let _enter = span.enter();
        let ret = Self::get_names()?
            .into_iter()
            .zip(Self::receive(unset_env)?);
        Ok(ret)
    }

    /// Get any file descriptor passed using `LISTEN_FD` and their names.
    ///
    /// This isn't necessarily limited to File descriptor of listening sockets, IPCs or FIFOs but also anything that is in the file descriptor store.
    ///
    /// The file descriptors are taken directly as [`OwnedFd`]s instead of being duplicated. In order to limit unsoundness this function therefore always unset the environment.
    ///
    /// # Errors
    ///
    /// This function will fail if no file descriptors could be obtained or the names associated with them couldn't be obtained. See [`ReceiveNameError`] for details.
    ///
    /// # Safety
    ///
    /// This function is safe if (and only if) the received file descriptors weren't already taken into a owned rust struct.
    /// (In short it needs to follow the safety constraints of [`OwnedFd::from_raw_fd`](std::os::fd::FromRawFd)).
    /// The simplest way to insure that it is so is to only use functions from this crate to get these file descriptors.
    pub unsafe fn receive_with_names_no_dup(
    ) -> Result<impl IntoIterator<Item = (OsString, Self)>, ReceiveNameError> {
        let span = tracing::info_span!("receive_with_names_no_dups");
        let _enter = span.enter();
        let ret = Self::get_names()?.into_iter().zip(Self::receive_no_dup()?);
        Ok(ret)
    }

    fn get_names() -> Result<impl IntoIterator<Item = OsString>, ReceiveNameError> {
        let fd_names = env::var_os(FD_NAMES_VAR).ok_or(ReceiveNameError::NoListenFDName)?;
        tracing::trace!("{FD_NAMES_VAR} = {fd_names:?}");
        Ok(fd_names
            .as_encoded_bytes()
            .split(|b| *b == b':')
            .map(|name| {
                // SAFETY:
                // - Each `word` only contains content that originated from `OsStr::as_encoded_bytes`
                // - Only split with ASCII colon which is a non-empty UTF-8 substring
                unsafe { OsStr::from_encoded_bytes_unchecked(name) }
            })
            .map(OsStr::to_os_string)
            .collect::<Vec<_>>())
    }

    fn from_fd_offset(fd_offset: usize) -> Result<Self, DupError> {
        SD_LISTEN_FDS_START
        .checked_add(fd_offset as RawFd)
        .map(|fd| {
            // SAFETY: The file descriptor won't be closed by the time we duplicate it.
            let fd = unsafe { BorrowedFd::borrow_raw(fd) };
            rustix::fs::fcntl_dupfd_cloexec(fd, 0).map_err(DupError::from)
        })
        .expect("Already checked against overflow.")
        .map(Self::from_fd)
    }

    fn from_fds(
        num_fds: usize,
    ) -> Result<impl IntoIterator<Item = Result<Self, DupError>>, GetFdsError> {
        if SD_LISTEN_FDS_START.checked_add(num_fds as RawFd).is_none() {
            return Err(GetFdsError::TooManyFDs(num_fds));
        }

        let ret = (0..num_fds).map(Self::from_fd_offset);
        Ok(ret)
    }

    unsafe fn from_fds_no_dup(
        num_fds: usize,
    ) -> Result<impl IntoIterator<Item = Self>, GetFdsError> {
        if SD_LISTEN_FDS_START.checked_add(num_fds as RawFd).is_none() {
            return Err(GetFdsError::TooManyFDs(num_fds));
        }

        let ret = (0..num_fds).map(|fd_offset| {
            SD_LISTEN_FDS_START
                .checked_add(fd_offset as RawFd)
                .map(|fd| {
                    let fd = unsafe { OwnedFd::from_raw_fd(fd) };
                    let flags = rustix::fs::fcntl_getfd(&fd).unwrap();
                    rustix::fs::fcntl_setfd(&fd, flags.union(FdFlags::CLOEXEC)).unwrap();
                    fd
                })
                .map(Self::from_fd)
                .expect("Already checked against overflow.")
        });
        Ok(ret)
    }

    fn from_fd(fd: OwnedFd) -> Self {
        let stat = rustix::fs::fstat(&fd)
            .expect("This should only ever be called on a valid file descriptor.");
        let file_type = FileType::from_raw_mode(stat.st_mode);
        match file_type {
            FileType::RegularFile => Self::File(fd.into()),
            FileType::Directory => Self::Directory(fd),
            FileType::Fifo => Self::Fifo(fd.into()),
            FileType::Socket => match rustix::net::getsockname(&fd).unwrap() {
                rustix::net::SocketAddrAny::V4(_) | rustix::net::SocketAddrAny::V6(_) => {
                    match rustix::net::sockopt::get_socket_type(&fd).unwrap() {
                        SocketType::DGRAM => Self::UdpSocket(fd.into()),
                        SocketType::STREAM => {
                            if rustix::net::sockopt::get_socket_acceptconn(&fd).unwrap_or_default()
                            {
                                Self::TcpListener(fd.into())
                            } else {
                                Self::TcpStream(fd.into())
                            }
                        }
                        _ => Self::InetOther(fd),
                    }
                }
                rustix::net::SocketAddrAny::Unix(_) => {
                    match rustix::net::sockopt::get_socket_type(&fd).unwrap() {
                        SocketType::DGRAM => Self::UnixDatagram(fd.into()),
                        SocketType::STREAM => {
                            if rustix::net::sockopt::get_socket_acceptconn(&fd).unwrap_or_default()
                            {
                                Self::UnixListener(fd.into())
                            } else {
                                Self::UnixStream(fd.into())
                            }
                        }
                        SocketType::SEQPACKET => Self::UnixSequentialPaquet(fd),
                        _ => Self::UnixOther(fd),
                    }
                }
                _ => Self::SocketOther(fd),
            },
            _ => {
                // `rustix` does not enable us to test if a raw fd is a mq, so we must drop to libc here.
                // SAFETY: `mq_getattr` is specified to return -1 when passed a fd which is not a mq.
                //         Furthermore, we ignore `attr` and rely only on the return value.
                let mut attr = std::mem::MaybeUninit::<libc::mq_attr>::uninit();
                let res = unsafe { libc::mq_getattr(fd.as_raw_fd(), attr.as_mut_ptr()) };
                if res == 0 {
                    Self::MessageQueue(fd)
                } else {
                    Self::Other(fd)
                }
            }
        }
    }
}