lib: add register_buffers_* family of functions

io-uring-test/src/main.rs

@@ -76,6 +76,7 @@ fn test<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
 
     // register
     tests::register::test_register_files_sparse(&mut ring, &test)?;
+    tests::register::test_register_buffers(&mut ring, &test)?;
     tests::register_buf_ring::test_register_buf_ring(&mut ring, &test)?;
 
     // fs

io-uring-test/src/tests/register.rs

@@ -1,3 +1,5 @@
+use std::io;
+
 use crate::Test;
 use io_uring::{cqueue, opcode, squeue, IoUring};
 
@@ -68,3 +70,92 @@ pub fn test_register_files_sparse<S: squeue::EntryMarker, C: cqueue::EntryMarker
 
     Ok(())
 }
+
+pub fn test_register_buffers<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
+    ring: &mut IoUring<S, C>,
+    test: &Test,
+) -> anyhow::Result<()> {
+    test_register("register_buffers_sparse", ring, test, |ring| {
+        ring.submitter().register_buffers_sparse(4)
+    })?;
+    test_register("register_buffers", ring, test, |ring| {
+        let buf = [0xde, 0xed, 0xbe, 0xef];
+
+        unsafe {
+            ring.submitter().register_buffers(&[libc::iovec {
+                iov_base: buf.as_ptr() as _,
+                iov_len: buf.len(),
+            }])
+        }
+    })?;
+    test_register("register_buffers_tags", ring, test, |ring| {
+        let buf = [0xde, 0xed, 0xbe, 0xef];
+
+        unsafe {
+            ring.submitter().register_buffers_tags(
+                &[libc::iovec {
+                    iov_base: buf.as_ptr() as _,
+                    iov_len: buf.len(),
+                }],
+                &[0],
+            )
+        }
+    })?;
+
+    return Ok(());
+
+    fn test_register<
+        S: squeue::EntryMarker,
+        C: cqueue::EntryMarker,
+        F: FnMut(&mut IoUring<S, C>) -> io::Result<()>,
+    >(
+        name: &str,
+        ring: &mut IoUring<S, C>,
+        test: &Test,
+        mut register: F,
+    ) -> anyhow::Result<()> {
+        // register_files_sparse was introduced in kernel 5.19, as was the opcode for UringCmd16.
+        // So require the UringCmd16 to avoid running this test on earlier kernels.
+        require!(
+            test;
+            test.probe.is_supported(opcode::UringCmd16::CODE);
+        );
+
+        println!("test {name}");
+
+        if let Ok(()) = ring.submitter().unregister_buffers() {
+            return Err(anyhow::anyhow!(
+                "unregister_buffers should fail if not buffer table has been setup"
+            ));
+        }
+
+        if let Err(e) = register(ring) {
+            return Err(anyhow::anyhow!("{name} failed: {}", e));
+        }
+
+        // See that same call again, with any value, will fail because a direct table cannot be built
+        // over an existing one.
+
+        if let Ok(()) = register(ring) {
+            return Err(anyhow::anyhow!(
+                "{name} should not have succeeded twice in a row"
+            ));
+        }
+
+        // See that the direct table can be removed.
+
+        if let Err(e) = ring.submitter().unregister_buffers() {
+            return Err(anyhow::anyhow!("unregister_buffers failed: {}", e));
+        }
+
+        // See that a second attempt to remove the direct table would fail.
+
+        if let Ok(()) = ring.submitter().unregister_buffers() {
+            return Err(anyhow::anyhow!(
+                "unregister_buffers should not have succeeded twice in a row"
+            ));
+        }
+
+        Ok(())
+    }
+}

src/submit.rs

@@ -179,6 +179,106 @@ impl<'a> Submitter<'a> {
         .map(drop)
     }
 
+    /// Update a range of fixed buffers starting at `off`.
+    ///
+    /// This is required to use buffers registered using
+    /// [`register_buffers_sparse`](Self::register_buffers_sparse),
+    /// although it can be also be used with [`register_buffers`](Self::register_buffers).
+    ///
+    /// See [`register_buffers_tags`](Self::register_buffers_tags)
+    /// for more information about resource tagging.
+    ///
+    /// # Safety
+    ///
+    /// This function is unsafe because improper use may lead to memory problems.
+    /// For example, a use-after-free may occur if `iov_base` contains a pointer freed
+    /// before unregistering the buffer through [`unregister_buffers`](Self::unregister_buffers)
+    /// or [`register_buffers_update_tag`](Self::register_buffers_update_tag).
+    pub unsafe fn register_buffers_update_tag(
+        &self,
+        off: u32,
+        bufs: &[libc::iovec],
+        tags: &[u64],
+    ) -> io::Result<()> {
+        let rr = sys::io_uring_rsrc_update2 {
+            nr: bufs.len().min(tags.len()) as _,
+            data: bufs.as_ptr() as _,
+            tags: tags.as_ptr() as _,
+            offset: off,
+            ..Default::default()
+        };
+        let rr = cast_ptr::<sys::io_uring_rsrc_update2>(&rr);
+        execute(
+            self.fd.as_raw_fd(),
+            sys::IORING_REGISTER_BUFFERS_UPDATE,
+            rr as *const _,
+            std::mem::size_of::<sys::io_uring_rsrc_update2>() as _,
+        )
+        .map(drop)
+    }
+
+    /// Variant of [`register_buffers`](Self::register_buffers)
+    /// with resource tagging.
+    ///
+    /// `tags` should be the same length as `bufs` and contain the
+    /// tag value corresponding to the buffer at the same index.
+    ///
+    /// If a tag is zero, then tagging for this particular resource
+    /// (a buffer in this case) is disabled. Otherwise, after the
+    /// resource had been unregistered and it's not used anymore,
+    /// a CQE will be posted with `user_data` set to the specified
+    /// tag and all other fields zeroed.
+    ///
+    /// # Safety
+    ///
+    /// This function is unsafe because improper use may lead to memory problems.
+    /// For example, a use-after-free may occur if `iov_base` contains a pointer freed
+    /// before unregistering the buffer through [`unregister_buffers`](Self::unregister_buffers)
+    /// or [`register_buffers_update_tag`](Self::register_buffers_update_tag).
+    pub unsafe fn register_buffers_tags(
+        &self,
+        bufs: &[libc::iovec],
+        tags: &[u64],
+    ) -> io::Result<()> {
+        let rr = sys::io_uring_rsrc_register {
+            nr: bufs.len().min(tags.len()) as _,
+            data: bufs.as_ptr() as _,
+            tags: tags.as_ptr() as _,
+            ..Default::default()
+        };
+        let rr = cast_ptr::<sys::io_uring_rsrc_register>(&rr);
+        execute(
+            self.fd.as_raw_fd(),
+            sys::IORING_REGISTER_BUFFERS2,
+            rr as *const _,
+            std::mem::size_of::<sys::io_uring_rsrc_register>() as _,
+        )
+        .map(drop)
+    }
+
+    /// Registers an empty table of nr fixed buffers buffers.
+    ///
+    /// These must be updated before use, using eg.
+    /// [`register_buffers_update_tag`](Self::register_buffers_update_tag).
+    ///
+    /// See [`register_buffers`](Self::register_buffers)
+    /// for more information about fixed buffers.
+    pub fn register_buffers_sparse(&self, nr: u32) -> io::Result<()> {
+        let rr = sys::io_uring_rsrc_register {
+            nr,
+            flags: sys::IORING_RSRC_REGISTER_SPARSE,
+            ..Default::default()
+        };
+        let rr = cast_ptr::<sys::io_uring_rsrc_register>(&rr);
+        execute(
+            self.fd.as_raw_fd(),
+            sys::IORING_REGISTER_BUFFERS2,
+            rr as *const _,
+            std::mem::size_of::<sys::io_uring_rsrc_register>() as _,
+        )
+        .map(drop)
+    }
+
     /// Registers an empty file table of nr_files number of file descriptors. The sparse variant is
     /// available in kernels 5.19 and later.
     ///