wireshift-uring 0.1.1

#![allow(clippy::unnecessary_wraps)]
use io_uring::{opcode, squeue, IoUring};
use wireshift_core::backend::BackendSubmission;
use wireshift_core::op::OpDescriptor;
use wireshift_core::{Error, Result};

use crate::file_registry::FileRegistryTable;
use crate::helpers::link_flag_between;
use crate::types::{Inflight, InflightData};
use crate::uring_sys;

mod file_ops;
mod io_ops;
mod net_ops;

use file_ops::{build_close_fixed, build_fsync, build_openat, build_openat_direct, build_statx};
use io_ops::{
    build_read, build_read_fixed, build_read_gpu, build_read_vectored, build_write,
    build_write_vectored,
};
use net_ops::{build_accept, build_connect, build_recv, build_send};

/// An SQE entry prepared for submission.
pub struct PreparedEntry {
    /// The `io_uring` SQE.
    pub entry: squeue::Entry,
    /// Data that must stay pinned while this SQE is in flight.
    pub pinned: InflightData,
    /// Optional flags to link this SQE to the next one.
    pub link_to_next: Option<squeue::Flags>,
}

/// Prepares a submission for `io_uring` and returns an inflight tracking object.
pub fn prepare_submission(
    ring: &mut IoUring,
    submission: BackendSubmission,
    fixed_files: &FileRegistryTable,
) -> std::result::Result<Inflight, (BackendSubmission, Error)> {
    let id = submission.id;
    match build_entries(&submission.descriptor, id, fixed_files) {
        Ok(entries_data) => {
            let space_needed = entries_data.len();
            let submission_queue = ring.submission();
            let available = submission_queue.capacity() - submission_queue.len();
            drop(submission_queue);
            if available < space_needed {
                return Err((
                    submission,
                    Error::submission(
                        "io_uring submission queue is full or does not have enough capacity for the chain",
                        "drain completions or increase queue_depth in RingConfig",
                    ),
                ));
            }

            let mut pinned_list = Vec::with_capacity(space_needed);
            for mut prepared in entries_data {
                if let Some(link_flag) = prepared.link_to_next {
                    prepared.entry = prepared.entry.flags(link_flag);
                }
                if let Err(error) = uring_sys::push_entry(ring, &prepared.entry) {
                    return Err((submission, error));
                }
                pinned_list.push(prepared.pinned);
            }

            let expected_cqes = space_needed;
            let pinned = if space_needed == 1 {
                match pinned_list.into_iter().next() {
                    Some(p) => p,
                    None => InflightData::None,
                }
            } else {
                InflightData::Linked(pinned_list)
            };

            Ok(Inflight {
                id,
                descriptor: submission.descriptor,
                pinned,
                expected_cqes,
                results: Vec::with_capacity(expected_cqes),
            })
        }
        Err(error) => Err((submission, error)),
    }
}

/// Builds a list of native `io_uring` SQEs for a given operation descriptor.
pub fn build_entries(
    descriptor: &OpDescriptor,
    id: u64,
    fixed_files: &FileRegistryTable,
) -> Result<Vec<PreparedEntry>> {
    // Reserve top 2 bits for internal use (Wakeup, Cancel, etc.)
    if id & (0b11 << 62) != 0 {
        return Err(Error::validation(
            format!("submission ID 0x{id:x} uses reserved bits"),
            "ensure submission IDs are under 62 bits",
        ));
    }

    match descriptor {
        OpDescriptor::Read { .. } => build_read(descriptor, id, fixed_files),
        OpDescriptor::ReadFixed { .. } => build_read_fixed(descriptor, id, fixed_files),
        OpDescriptor::ReadGpu { .. } => build_read_gpu(descriptor, id, fixed_files),
        OpDescriptor::Write { .. } => build_write(descriptor, id, fixed_files),
        OpDescriptor::ReadVectored { .. } => build_read_vectored(descriptor, id, fixed_files),
        OpDescriptor::WriteVectored { .. } => build_write_vectored(descriptor, id, fixed_files),
        OpDescriptor::OpenAt { .. } => build_openat(descriptor, id, fixed_files),
        OpDescriptor::OpenAtDirect { .. } => build_openat_direct(descriptor, id, fixed_files),
        OpDescriptor::Statx { .. } => build_statx(descriptor, id, fixed_files),
        OpDescriptor::Connect { .. } => build_connect(descriptor, id, fixed_files),
        OpDescriptor::Accept { .. } => build_accept(descriptor, id, fixed_files),
        OpDescriptor::Send { .. } => build_send(descriptor, id, fixed_files),
        OpDescriptor::Recv { .. } => build_recv(descriptor, id, fixed_files),
        OpDescriptor::Fsync { .. } => build_fsync(descriptor, id, fixed_files),
        OpDescriptor::CloseFixed { .. } => build_close_fixed(descriptor, id, fixed_files),
        OpDescriptor::Cancel { .. } => build_cancel(descriptor, id, fixed_files),
        OpDescriptor::Nop => build_nop(descriptor, id, fixed_files),
        OpDescriptor::Linked { .. } => build_linked(descriptor, id, fixed_files),
        // Splice and madvise fall through to nop for the uring SQE builder.
        // The actual splice/madvise syscalls are handled by the fallback path.
        OpDescriptor::Splice { .. } | OpDescriptor::Madvise { .. } => {
            build_nop(descriptor, id, fixed_files)
        }
        _ => build_nop(descriptor, id, fixed_files),
    }
}

fn build_cancel(
    descriptor: &OpDescriptor,
    id: u64,
    _fixed_files: &FileRegistryTable,
) -> Result<Vec<PreparedEntry>> {
    let OpDescriptor::Cancel { target } = descriptor else {
        unreachable!()
    };
    let entry = opcode::AsyncCancel::new(*target).build().user_data(id);
    Ok(vec![PreparedEntry {
        entry,
        pinned: InflightData::None,
        link_to_next: None,
    }])
}

fn build_nop(
    _descriptor: &OpDescriptor,
    id: u64,
    _fixed_files: &FileRegistryTable,
) -> Result<Vec<PreparedEntry>> {
    let entry = opcode::Nop::new().build().user_data(id);
    Ok(vec![PreparedEntry {
        entry,
        pinned: InflightData::None,
        link_to_next: None,
    }])
}

fn build_linked(
    descriptor: &OpDescriptor,
    id: u64,
    fixed_files: &FileRegistryTable,
) -> Result<Vec<PreparedEntry>> {
    let OpDescriptor::Linked { descriptors } = descriptor else {
        unreachable!()
    };
    let mut all = Vec::with_capacity(descriptors.len());
    for (index, inner) in descriptors.iter().enumerate() {
        let mut res = build_entries(inner, id, fixed_files)?;
        if index + 1 < descriptors.len() {
            let next = &descriptors[index + 1];
            let last = res.last_mut().ok_or_else(|| {
                Error::completion(
                    "linked descriptor expanded into zero SQEs",
                    "ensure every linked operation contributes at least one native SQE",
                )
            })?;
            last.link_to_next = Some(link_flag_between(inner, next));
        }
        all.append(&mut res);
    }
    Ok(all)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;
    use std::fs::File;
    use std::os::fd::AsRawFd;

    use crate::helpers::sequential_hint_bit;
    use io_uring::{opcode, squeue};
    use wireshift_core::buffer::BufferPool;
    use wireshift_core::op::OpDescriptor;

    #[test]
    fn build_entries_uses_fixed_file_slots_for_registered_reads() {
        let file = File::open("/dev/null").expect("open /dev/null");
        let fd = file.as_raw_fd();
        let buffer = BufferPool::new(4096, 1)
            .expect("pool")
            .acquire()
            .expect("buffer")
            .into_submitted();
        let descriptor = OpDescriptor::Read {
            file,
            offset: 0,
            buffer,
            len: None,
        };
        let registry = FileRegistryTable {
            slots_by_fd: HashMap::from([(fd, 7)]),
        };

        let entries = build_entries(&descriptor, 11, &registry).expect("entries");
        let entry_debug = format!("{:?}", entries[0].entry);

        assert_eq!(entries.len(), 1);
        assert!(
            entry_debug.contains(&format!("flags: {}", squeue::Flags::FIXED_FILE.bits())),
            "registered files must use IOSQE_FIXED_FILE"
        );
    }

    #[test]
    fn build_entries_links_fadvise_after_sequential_openat_direct() {
        let descriptor = OpDescriptor::OpenAtDirect {
            dir: None,
            path: "/tmp/wireshift-sequential".into(),
            flags: sequential_hint_bit(),
            read: true,
            write: false,
            create: false,
            truncate: false,
            slot: 3,
        };

        let entries =
            build_entries(&descriptor, 21, &FileRegistryTable::default()).expect("entries");
        let entry_debug = format!("{:?}", entries[1].entry);

        assert_eq!(entries.len(), 2);
        assert_eq!(entries[0].link_to_next, Some(squeue::Flags::IO_LINK));
        assert!(
            entry_debug.contains(&format!("op_code: {}", opcode::Fadvise::CODE)),
            "sequential hint must emit a native Fadvise SQE"
        );
        assert!(
            entry_debug.contains(&format!("flags: {}", squeue::Flags::FIXED_FILE.bits())),
            "sequential fadvise must target the direct descriptor slot"
        );
    }

    #[test]
    fn duplicate_fixed_fd_registration_is_rejected() {
        use crate::file_registry::duplicate_fd;
        let duplicate = duplicate_fd(&[10, 11, 10]);
        assert_eq!(duplicate, Some(10));
    }

    #[test]
    fn build_entries_links_read_fixed_to_close_with_hard_link() {
        let buffer = BufferPool::new(4096, 1)
            .expect("pool")
            .acquire()
            .expect("buffer")
            .into_submitted();
        let descriptor = OpDescriptor::Linked {
            descriptors: vec![
                OpDescriptor::ReadFixed {
                    slot: 1,
                    offset: 0,
                    buffer,
                },
                OpDescriptor::CloseFixed { slot: 1 },
            ],
        };

        let entries =
            build_entries(&descriptor, 42, &FileRegistryTable::default()).expect("entries");

        assert_eq!(entries.len(), 2);
        assert_eq!(entries[0].link_to_next, Some(squeue::Flags::IO_HARDLINK));
        assert_eq!(entries[1].link_to_next, None);
    }
}