magi-code 0.77.1

use super::HerdrTransport;
#[cfg(unix)]
use super::IPC_TIMEOUT;
use std::{io, path::Path};

#[cfg(unix)]
use std::time::{Duration, Instant};

#[derive(Debug, Clone, Copy)]
pub(super) struct UnixSocketTransport;

impl HerdrTransport for UnixSocketTransport {
    fn write_line(&self, socket_path: &Path, line: &str) -> io::Result<()> {
        write_line_to_socket(socket_path, line)
    }
}

#[cfg(unix)]
pub(super) fn write_line_to_socket(socket_path: &Path, line: &str) -> io::Result<()> {
    use std::{
        io::{ErrorKind, Write},
        os::{
            fd::{AsRawFd, FromRawFd, OwnedFd},
            unix::net::UnixStream,
        },
    };

    let (address, address_len) = unix_socket_address(socket_path)?;
    let raw_fd = unsafe {
        // SAFETY: The arguments are the standard stream Unix-domain socket constants. On success
        // the returned descriptor is adopted exactly once immediately below.
        libc::socket(libc::AF_UNIX, libc::SOCK_STREAM, 0)
    };
    if raw_fd < 0 {
        return Err(io::Error::last_os_error());
    }
    let owned_fd = unsafe {
        // SAFETY: `raw_fd` is a newly created, valid descriptor and no other owner exists yet.
        OwnedFd::from_raw_fd(raw_fd)
    };
    let mut stream = UnixStream::from(owned_fd);
    stream.set_nonblocking(true)?;
    // Keep the kernel-level write timeout as a second guard. Failure to install it is a transport
    // failure, not a reason to continue with an unbounded socket operation.
    stream.set_write_timeout(Some(IPC_TIMEOUT))?;
    let deadline = Instant::now()
        .checked_add(IPC_TIMEOUT)
        .ok_or_else(|| io::Error::new(ErrorKind::TimedOut, "HERDR socket deadline overflow"))?;

    let result = unsafe {
        // SAFETY: `address` is initialized, remains alive for this call, and `address_len`
        // covers only its initialized Unix-socket address bytes.
        libc::connect(
            stream.as_raw_fd(),
            (&address as *const libc::sockaddr_un).cast::<libc::sockaddr>(),
            address_len,
        )
    };
    let error = (result != 0).then(io::Error::last_os_error);
    match classify_connect_result(
        result,
        error.as_ref().and_then(|error| error.raw_os_error()),
    ) {
        ConnectDisposition::Connected => {}
        ConnectDisposition::Pending => {
            // A nonblocking connect is issued once. Even EINTR is resolved by polling this
            // descriptor and reading SO_ERROR; retrying connect can change the operation state.
            wait_for_unix_socket(stream.as_raw_fd(), libc::POLLOUT, deadline)?;
            if let Some(error) = stream.take_error()? {
                return Err(error);
            }
        }
        ConnectDisposition::Failed => return Err(error.unwrap_or_else(io::Error::last_os_error)),
    }

    let bytes = line.as_bytes();
    let mut written = 0;
    while written < bytes.len() {
        if Instant::now() >= deadline {
            return Err(io::Error::new(
                ErrorKind::TimedOut,
                "HERDR socket write timed out",
            ));
        }
        match stream.write(&bytes[written..]) {
            Ok(0) => {
                return Err(io::Error::new(
                    ErrorKind::WriteZero,
                    "HERDR socket closed during write",
                ));
            }
            Ok(count) => written += count,
            Err(error) if error.kind() == ErrorKind::Interrupted => {}
            Err(error) if error.kind() == ErrorKind::WouldBlock => {
                wait_for_unix_socket(stream.as_raw_fd(), libc::POLLOUT, deadline)?;
            }
            Err(error) => return Err(error),
        }
    }
    Ok(())
}

#[cfg(unix)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ConnectDisposition {
    Connected,
    Pending,
    Failed,
}

#[cfg(unix)]
pub(super) fn classify_connect_result(result: i32, raw_error: Option<i32>) -> ConnectDisposition {
    if result == 0 {
        ConnectDisposition::Connected
    } else if raw_error.is_some_and(is_pending_connect_error) {
        ConnectDisposition::Pending
    } else {
        ConnectDisposition::Failed
    }
}

#[cfg(unix)]
fn is_pending_connect_error(code: i32) -> bool {
    code == libc::EINTR
        || code == libc::EINPROGRESS
        || code == libc::EALREADY
        || code == libc::EAGAIN
        || code == libc::EWOULDBLOCK
}

#[cfg(unix)]
fn unix_socket_address(socket_path: &Path) -> io::Result<(libc::sockaddr_un, libc::socklen_t)> {
    use std::{mem, os::unix::ffi::OsStrExt, ptr};

    let bytes = socket_path.as_os_str().as_bytes();
    if bytes.is_empty() || bytes.contains(&0) {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "HERDR socket path must be a non-empty pathname without NUL bytes",
        ));
    }
    // SAFETY: all-zero is a valid representation for `sockaddr_un`; the family and pathname are
    // initialized before the value is passed to libc.
    let mut address: libc::sockaddr_un = unsafe { mem::zeroed() };
    address.sun_family = libc::AF_UNIX as libc::sa_family_t;
    if bytes.len() >= address.sun_path.len() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "HERDR socket path exceeds the platform Unix-socket limit",
        ));
    }
    // SAFETY: the source and destination are valid, non-overlapping buffers and the destination
    // has at least `bytes.len()` bytes; the zeroed tail supplies the terminator.
    unsafe {
        ptr::copy_nonoverlapping(
            bytes.as_ptr(),
            address.sun_path.as_mut_ptr().cast::<u8>(),
            bytes.len(),
        );
    }
    let address_len =
        (mem::offset_of!(libc::sockaddr_un, sun_path) + bytes.len() + 1) as libc::socklen_t;
    Ok((address, address_len))
}

#[cfg(unix)]
pub(super) fn poll_timeout_ms(remaining: Duration) -> i32 {
    remaining
        .as_nanos()
        .div_ceil(1_000_000)
        .clamp(1, i32::MAX as u128) as i32
}

#[cfg(unix)]
fn wait_for_unix_socket(
    fd: std::os::fd::RawFd,
    events: libc::c_short,
    deadline: Instant,
) -> io::Result<()> {
    loop {
        let Some(remaining) = deadline.checked_duration_since(Instant::now()) else {
            return Err(io::Error::new(
                io::ErrorKind::TimedOut,
                "HERDR socket operation timed out",
            ));
        };
        let timeout_ms = poll_timeout_ms(remaining);
        let mut poll_fd = libc::pollfd {
            fd,
            events,
            revents: 0,
        };
        let result = unsafe {
            // SAFETY: `poll_fd` is a valid, exclusively borrowed one-element poll array.
            libc::poll(&mut poll_fd, 1, timeout_ms)
        };
        if result > 0 {
            if poll_fd.revents & libc::POLLNVAL != 0 {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "HERDR socket descriptor became invalid",
                ));
            }
            if poll_fd.revents & (events | libc::POLLERR | libc::POLLHUP) != 0 {
                return Ok(());
            }
        } else if result == 0 {
            return Err(io::Error::new(
                io::ErrorKind::TimedOut,
                "HERDR socket operation timed out",
            ));
        } else {
            let error = io::Error::last_os_error();
            if error.raw_os_error() != Some(libc::EINTR) {
                return Err(error);
            }
        }
    }
}

#[cfg(not(unix))]
pub(super) fn write_line_to_socket(_socket_path: &Path, _line: &str) -> io::Result<()> {
    Ok(())
}