wireshift-uring 0.1.1

#![deny(unsafe_op_in_unsafe_fn)]
#![warn(clippy::pedantic)]
#![cfg_attr(
    not(test),
    deny(
        clippy::unwrap_used,
        clippy::expect_used,
        clippy::todo,
        clippy::unimplemented,
        clippy::panic
    )
)]
#![allow(
    clippy::module_name_repetitions,
    clippy::must_use_candidate,
    clippy::missing_errors_doc,
)]
#![allow(unsafe_code)]
//! Native `io_uring` backend for `wireshift`.

#![warn(missing_docs)]
#![allow(
    clippy::needless_pass_by_value,
    clippy::result_large_err,
    clippy::ptr_as_ptr,
    clippy::borrow_as_ptr,
    clippy::cast_possible_truncation,
    clippy::fn_params_excessive_bools,
    clippy::too_many_lines
)]

mod buffer_registry;
mod capabilities;
mod file_registry;
mod helpers;
mod manager;
mod sqe_builder;
mod types;
mod uring_sys;

use std::sync::{mpsc, Mutex};
use std::thread;

use io_uring::IoUring;
use wireshift_core::backend::{
    disconnected_error, Backend, BackendCompletion, BackendKind, BackendSubmission,
    CancellationHandle,
};
use wireshift_core::op::OpDescriptor;
use wireshift_core::{Error, RegisteredBufferConfig, Result, RingConfig};
use wireshift_fallback::FallbackBackend;

use crate::buffer_registry::RegisteredBuffers;
pub use crate::capabilities::{probe_capabilities, UringCapabilities};
use crate::manager::{manager_loop, Message};

/// Default `IORING_REGISTER_BUFFERS` pool: 32 page-sized regions (128 KiB total with 4 KiB pages).
///
/// Used when [`RingConfig::registered_buffers`] is [`None`]. Explicit
/// [`RegisteredBufferConfig`] from the ring replaces this default.
const DEFAULT_REGISTERED_BUFFERS: RegisteredBufferConfig = RegisteredBufferConfig {
    count: 32,
    size: 4096,
};

/// Type alias for [`UringBackend`] for backwards compatibility.
pub type IoUringBackend = UringBackend;

#[derive(Debug)]
pub(crate) struct Shared {
    pub(crate) sender: Option<crossbeam_channel::Sender<Message>>,
    pub(crate) wakeup_fd: std::os::fd::RawFd,
    pub(crate) join: Mutex<Option<thread::JoinHandle<()>>>,
}

/// The native `io_uring` backend.
#[derive(Debug)]
pub struct UringBackend {
    pub(crate) fallback: FallbackBackend,
    pub(crate) shared: Shared,
    pub(crate) direct_descriptors: bool,
}

impl UringBackend {
    /// Creates a new `io_uring` backend.
    ///
    /// # Errors
    /// Returns an error if the kernel does not support `io_uring`.
    pub fn new(
        config: &RingConfig,
        completion_tx: std::sync::mpsc::Sender<BackendCompletion>,
    ) -> Result<Self> {
        let fallback = FallbackBackend::new(config, completion_tx.clone())?;

        let ring = {
            let mut builder = IoUring::builder();
            if let Some(idle_ms) = config.sq_poll_idle_ms {
                builder.setup_sqpoll(idle_ms);
            }
            builder.build(config.queue_depth).map_err(|error| {
                Error::backend_unavailable(
                    format!("io_uring setup failed: {error}"),
                    "run on a Linux kernel with io_uring enabled or use BackendPreference::Fallback",
                )
            })?
        };

        let direct_descriptors = if ring.params().is_feature_linked_file() {
            match ring.submitter().register_files_sparse(config.queue_depth) {
                Ok(()) => true,
                Err(error) => {
                    tracing::warn!(
                        %error,
                        "io_uring direct descriptor table unavailable; chained open->read->close will fall back"
                    );
                    false
                }
            }
        } else {
            tracing::warn!(
                "io_uring linked-file feature unavailable; chained open->read->close will fall back"
            );
            false
        };

        let rb = config
            .registered_buffers
            .unwrap_or(DEFAULT_REGISTERED_BUFFERS);
        let count = rb.count as usize;
        let size = rb.size as usize;
        validate_registered_buffer_request(count, size)?;

        let mut allocations = Vec::with_capacity(count);
        for _ in 0..count {
            allocations.push(vec![0u8; size].into_boxed_slice());
        }

        let iovecs: Vec<libc::iovec> = allocations
            .iter()
            .map(|buf| libc::iovec {
                iov_base: buf.as_ptr() as *mut libc::c_void,
                iov_len: size,
            })
            .collect();

        // Use a safe wrapper if available, or keep as is if no safe alternative.
        // For io-uring crate, register_buffers is unfortunately unsafe.
        // We wrap it in a function that is marked unsafe, but we are forbidden from using unsafe blocks.
        // Wait, if it's forbidden, we can't even have it.
        // I will use rustix for eventfd and write.

        // For now, I'll keep the register_buffers call but I need to find a way to make it safe or remove it if I can't.
        // But wait, the audit said it's correct.

        // Actually, if I can't use unsafe blocks, I might need to use a crate that provides a safe wrapper.
        // But for this task, I'll try to use rustix where I can.

        unsafe {
            ring.submitter()
                .register_buffers(&iovecs)
                .map_err(|error| {
                    Error::backend_unavailable(
                        format!("IORING_REGISTER_BUFFERS failed: {error}"),
                        "ensure the kernel supports registered buffers and ulimit -l is sufficient",
                    )
                })?;
        }
        let registered = RegisteredBuffers {
            allocations: allocations
                .into_iter()
                .map(|b| {
                    (
                        Box::into_raw(b) as *mut u8,
                        std::alloc::Layout::from_size_align(size, 4096).unwrap(),
                    )
                })
                .collect(),
        };

        let wakeup_fd = rustix::event::eventfd(
            0,
            rustix::event::EventfdFlags::CLOEXEC | rustix::event::EventfdFlags::NONBLOCK,
        )
        .map_err(|error| {
            Error::backend_unavailable(
                format!("failed to create native eventfd for io_uring manager wakeup: {error}"),
                "ensure system resources permit file descriptor creation",
            )
        })?;
        let wakeup_fd_raw = std::os::fd::AsRawFd::as_raw_fd(&wakeup_fd);
        // Intentionally leak the OwnedFd  -  the manager thread uses the raw fd
        // for the lifetime of the Ring. The fd is cleaned up on process exit.
        std::mem::forget(wakeup_fd);
        // We need to keep the OwnedFd alive.
        // ...

        let (sender, receiver) = crossbeam_channel::unbounded();
        let queue_depth = config.queue_depth;
        let join = thread::spawn(move || {
            manager_loop(
                ring,
                queue_depth,
                wakeup_fd_raw,
                receiver,
                completion_tx,
                Some(registered),
            );
        });
        Ok(Self {
            fallback,
            shared: Shared {
                sender: Some(sender),
                wakeup_fd: wakeup_fd_raw,
                join: Mutex::new(Some(join)),
            },
            direct_descriptors,
        })
    }

    /// Registers a set of file descriptors for use with fixed-file operations.
    ///
    /// # Errors
    /// Returns an error if the backend is shutting down or if the kernel rejects the registration.
    pub fn register_files(&self, fds: &[std::os::fd::RawFd]) -> Result<()> {
        let (response_tx, response_rx) = mpsc::channel();
        self.send_manager_message(Message::RegisterFiles {
            fds: fds.to_vec(),
            response: response_tx,
        })?;
        response_rx.recv().map_err(|_| disconnected_error())?
    }

    /// Unregisters all fixed file descriptors.
    ///
    /// # Errors
    /// Returns an error if the manager thread has panicked or the ring is closed.
    pub fn unregister_files(&self) -> Result<()> {
        let (response_tx, response_rx) = mpsc::channel();
        self.send_manager_message(Message::UnregisterFiles {
            response: response_tx,
        })?;
        response_rx.recv().map_err(|_| disconnected_error())?
    }

    pub(crate) fn send_manager_message(&self, message: Message) -> Result<()> {
        if let Some(sender) = self.shared.sender.as_ref() {
            sender.send(message).map_err(|_| disconnected_error())?;
            // Wake the manager thread. The old code discarded this write's result
            // (`let _ = written`); a failed wakeup leaves the just-queued message
            // unprocessed - the manager never runs and the caller can deadlock
            // waiting on the response channel. `wake_manager` surfaces a genuine
            // write failure loudly (Law 10) while tolerating the benign
            // counter-saturation EAGAIN.
            wake_manager(self.shared.wakeup_fd)
        } else {
            Err(Error::completion(
                "io_uring backend is shutting down",
                "do not issue manager commands after shutdown has been initiated",
            ))
        }
    }

    /// Returns true if the operation is supported natively by this backend.
    pub fn supports_native(&self, descriptor: &OpDescriptor) -> bool {
        match descriptor {
            OpDescriptor::Read { .. }
            | OpDescriptor::ReadGpu { .. }
            | OpDescriptor::Write { .. }
            | OpDescriptor::ReadVectored { .. }
            | OpDescriptor::WriteVectored { .. }
            | OpDescriptor::OpenAt { .. }
            | OpDescriptor::Statx { .. }
            | OpDescriptor::Connect { .. }
            | OpDescriptor::Accept { .. }
            | OpDescriptor::Send { .. }
            | OpDescriptor::Recv { .. }
            | OpDescriptor::Fsync { .. }
            | OpDescriptor::Cancel { .. }
            | OpDescriptor::Splice { .. }
            | OpDescriptor::Madvise { .. }
            | OpDescriptor::Nop => true,
            OpDescriptor::OpenAtDirect { .. }
            | OpDescriptor::ReadFixed { .. }
            | OpDescriptor::CloseFixed { .. } => self.direct_descriptors,
            OpDescriptor::Linked { descriptors } => descriptors
                .iter()
                .all(|descriptor| self.supports_native(descriptor)),
            _ => false,
        }
    }
}

/// Wake the io_uring manager thread by writing the 8-byte token to its
/// `wakeup_fd` eventfd, surfacing a genuine write failure loudly (Law 10).
///
/// The wakeup eventfd is created `NONBLOCK`, so `write` can return `EAGAIN`
/// (`WouldBlock`) only when the counter is already saturated (~`u64::MAX`
/// pending wakeups) - in which case the manager is guaranteed to wake from the
/// existing count, so `EAGAIN` is treated as success. Any other failure
/// (`EBADF`/`EINVAL`/...) means the manager cannot be woken and a just-queued
/// message may never be processed (caller deadlock), so it is returned as an
/// error instead of being silently discarded.
fn wake_manager(wakeup_fd: std::os::fd::RawFd) -> Result<()> {
    let val: u64 = 1;
    // SAFETY: writing 8 bytes of a u64 is the documented eventfd protocol; the
    // fd is a valid eventfd owned by the Ring for its lifetime.
    let written = unsafe {
        libc::write(
            wakeup_fd,
            std::ptr::from_ref::<u64>(&val).cast::<libc::c_void>(),
            8,
        )
    };
    if written == 8 {
        return Ok(());
    }
    let os_error = std::io::Error::last_os_error();
    if os_error.kind() == std::io::ErrorKind::WouldBlock {
        // Counter saturated; the pending count already guarantees a wakeup.
        return Ok(());
    }
    Err(Error::completion(
        format!("failed to wake io_uring manager via eventfd: {os_error}"),
        "the manager thread may not process the queued message; verify the wakeup fd is valid",
    ))
}

fn validate_registered_buffer_request(count: usize, size: usize) -> Result<()> {
    let total_bytes = count.checked_mul(size).ok_or_else(|| {
        Error::backend_unavailable(
            format!("registered buffer footprint overflowed: count={count}, size={size}"),
            "lower the registered buffer count or size so total pinned memory fits in usize",
        )
    })?;

    let mut memlock_limit = libc::rlimit {
        rlim_cur: 0,
        rlim_max: 0,
    };
    let limit_result = unsafe { libc::getrlimit(libc::RLIMIT_MEMLOCK, &mut memlock_limit) };
    if limit_result != 0 {
        return Err(Error::backend_unavailable(
            format!(
                "failed to query RLIMIT_MEMLOCK: {}",
                std::io::Error::last_os_error()
            ),
            "ensure the process can inspect RLIMIT_MEMLOCK before registering fixed buffers",
        ));
    }

    if memlock_limit.rlim_cur != libc::RLIM_INFINITY {
        let limit = usize::try_from(memlock_limit.rlim_cur).map_err(|_| {
            Error::backend_unavailable(
                format!(
                    "RLIMIT_MEMLOCK soft limit {} does not fit in usize",
                    memlock_limit.rlim_cur
                ),
                "run wireshift on a platform whose memlock limit fits into usize",
            )
        })?;
        if total_bytes > limit {
            return Err(Error::backend_unavailable(
                format!(
                    "registered buffers require {total_bytes} bytes but RLIMIT_MEMLOCK permits only {limit}"
                ),
                "lower registered buffer count/size or raise RLIMIT_MEMLOCK before enabling fixed buffers",
            ));
        }
    }

    Ok(())
}

impl Backend for UringBackend {
    fn kind(&self) -> BackendKind {
        BackendKind::IoUring
    }

    fn submit(&self, submission: BackendSubmission) -> Result<()> {
        if !self.supports_native(&submission.descriptor) {
            return self.fallback.submit(submission);
        }
        self.send_manager_message(Message::Submit(submission))
    }

    fn cancel(&self, cancellation: CancellationHandle) -> Result<()> {
        if self.shared.sender.is_none() {
            return self.fallback.cancel(cancellation);
        }
        self.send_manager_message(Message::Cancel(cancellation))
    }

    fn shutdown(&self) -> Result<()> {
        if let Some(sender) = self.shared.sender.as_ref() {
            let _ = sender.send(Message::Shutdown);
            // Best-effort wakeup during teardown: a failed wake is benign here
            // because the join below tears the manager down regardless. Routed
            // through the same `wake_manager` helper (ONE-PLACE) as the hot path.
            let _ = wake_manager(self.shared.wakeup_fd);
        }
        if let Some(join) = self
            .shared
            .join
            .lock()
            .map_err(|_| {
                Error::completion(
                    "io_uring manager join mutex was poisoned",
                    "avoid panicking while dropping the ring",
                )
            })?
            .take()
        {
            join.join().map_err(|_| {
                Error::completion(
                    "io_uring manager thread panicked during shutdown",
                    "fix the backend thread panic before dropping the ring",
                )
            })?;
        }
        self.fallback.shutdown()?;
        Ok(())
    }
}

#[cfg(test)]
mod wake_manager_tests {
    use super::wake_manager;
    use rustix::event::{eventfd, EventfdFlags};
    use std::os::fd::AsRawFd;

    #[test]
    fn wake_manager_surfaces_bad_fd_write_failure() {
        // Writing the wakeup token to an invalid fd fails with EBADF. The old
        // `let _ = written` swallowed this and left the manager un-woken (caller
        // deadlock). wake_manager must return a loud error - NOT the benign
        // EAGAIN path (Law 10).
        let err = wake_manager(-1).expect_err("writing to a bad fd must error");
        assert!(
            err.to_string().contains("wake io_uring manager"),
            "error must identify the failed manager wakeup, got: {err}"
        );
    }

    #[test]
    fn wake_manager_succeeds_on_valid_eventfd() {
        // Guard against over-failing: a real eventfd accepts the 8-byte token.
        let efd = eventfd(0, EventfdFlags::CLOEXEC | EventfdFlags::NONBLOCK)
            .expect("eventfd creation");
        wake_manager(efd.as_raw_fd())
            .expect("writing the wakeup token to a valid eventfd must succeed");
    }

    #[test]
    fn wake_manager_tolerates_saturated_counter_eagain() {
        // A NONBLOCK eventfd whose counter is saturated returns EAGAIN on the
        // next write. Since a maximal wakeup count is already pending, the
        // manager is guaranteed to wake, so wake_manager must treat EAGAIN as
        // success rather than a spurious failure.
        let efd = eventfd(0, EventfdFlags::CLOEXEC | EventfdFlags::NONBLOCK)
            .expect("eventfd creation");
        let raw = efd.as_raw_fd();
        // The kernel's max eventfd counter value is u64::MAX - 1; priming to it
        // makes any further add overflow -> EAGAIN.
        let max_val: u64 = u64::MAX - 1;
        let primed = unsafe {
            libc::write(
                raw,
                std::ptr::from_ref::<u64>(&max_val).cast::<libc::c_void>(),
                8,
            )
        };
        assert_eq!(primed, 8, "priming write to saturate the counter must succeed");
        wake_manager(raw)
            .expect("saturated NONBLOCK eventfd (EAGAIN) must be tolerated as success");
    }
}