From 2a795d54a22f725c07d7ce78a4bd53cc6bd2ce86 Mon Sep 17 00:00:00 2001
From: Ryan Seipp <rseipp@ryanseipp.com>
Date: Sat, 7 Sep 2024 10:13:51 -0400
Subject: [PATCH] Add `io_uring_buf_ring` struct

---
 src/io_uring.rs            | 44 +++++++++++++++++++++++
 tests/io_uring/register.rs | 73 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 114 insertions(+), 3 deletions(-)

diff --git a/src/io_uring.rs b/src/io_uring.rs
index e1b13d3c6..fd71c1efc 100644
--- a/src/io_uring.rs
+++ b/src/io_uring.rs
@@ -1433,6 +1433,39 @@ pub struct io_uring_buf {
     pub resv: u16,
 }
 
+#[allow(missing_docs)]
+#[repr(C)]
+#[derive(Debug, Copy, Clone, Default)]
+pub struct buf_ring_tail_struct {
+    pub resv1: u64,
+    pub resv2: u32,
+    pub resv3: u16,
+    pub tail: u16,
+}
+
+#[allow(missing_docs)]
+#[repr(C)]
+#[derive(Debug, Default)]
+pub struct buf_ring_bufs_struct {
+    pub bufs: sys::__IncompleteArrayField<io_uring_buf>,
+}
+
+#[allow(missing_docs)]
+#[repr(C)]
+#[derive(Debug, Default)]
+pub struct tail_or_bufs_struct {
+    pub tail: sys::__BindgenUnionField<buf_ring_tail_struct>,
+    pub bufs: sys::__BindgenUnionField<buf_ring_bufs_struct>,
+    pub union_field: [u64; 2],
+}
+
+#[allow(missing_docs)]
+#[repr(C)]
+#[derive(Debug, Default)]
+pub struct io_uring_buf_ring {
+    pub tail_or_bufs: tail_or_bufs_struct,
+}
+
 impl Default for ioprio_union {
     #[inline]
     fn default() -> Self {
@@ -1595,4 +1628,15 @@ fn io_uring_layouts() {
     check_struct!(io_uring_buf_reg, ring_addr, ring_entries, bgid, pad, resv);
     check_struct!(io_uring_buf, addr, len, bid, resv);
     check_struct!(io_uring_sync_cancel_reg, addr, fd, flags, timeout, pad);
+
+    check_renamed_type!(tail_or_bufs_struct, io_uring_buf_ring__bindgen_ty_1);
+    check_renamed_type!(
+        buf_ring_tail_struct,
+        io_uring_buf_ring__bindgen_ty_1__bindgen_ty_1
+    );
+    check_renamed_type!(
+        buf_ring_bufs_struct,
+        io_uring_buf_ring__bindgen_ty_1__bindgen_ty_2
+    );
+    check_struct_renamed_field!(io_uring_buf_ring, tail_or_bufs, __bindgen_anon_1);
 }
diff --git a/tests/io_uring/register.rs b/tests/io_uring/register.rs
index 929241e76..ea0896d5d 100644
--- a/tests/io_uring/register.rs
+++ b/tests/io_uring/register.rs
@@ -1,10 +1,14 @@
+use std::ptr;
+
 use libc::c_void;
 use rustix::fd::{AsFd, AsRawFd, BorrowedFd};
-use rustix::io::Result;
+use rustix::io::{Errno, Result};
 use rustix::io_uring::{
-    io_uring_params, io_uring_register_with, io_uring_rsrc_update, io_uring_setup,
-    IoringFeatureFlags, IoringRegisterFlags, IoringRegisterOp,
+    io_uring_buf, io_uring_buf_reg, io_uring_buf_ring, io_uring_params, io_uring_register_with,
+    io_uring_rsrc_update, io_uring_setup, IoringFeatureFlags, IoringRegisterFlags,
+    IoringRegisterOp,
 };
+use rustix::mm::{MapFlags, ProtFlags};
 
 fn do_register<FD>(
     fd: FD,
@@ -87,6 +91,19 @@ where
     Ok(())
 }
 
+fn register_buf_ring<FD>(fd: FD, reg: &io_uring_buf_reg) -> Result<()>
+where
+    FD: AsFd,
+{
+    do_register(
+        fd,
+        false,
+        IoringRegisterOp::RegisterPbufRing,
+        reg as *const io_uring_buf_reg as *const c_void,
+        1,
+    )
+}
+
 #[test]
 fn test_io_uring_register_with() {
     let mut params = io_uring_params::default();
@@ -104,3 +121,53 @@ fn test_io_uring_register_with() {
     unregister_ring(ring_fd).unwrap();
     register_result.unwrap();
 }
+
+#[test]
+fn io_uring_buf_ring_can_be_registered() {
+    const ENTRIES: usize = 8;
+    const BGID: u16 = 42;
+
+    let mut params = io_uring_params::default();
+    let ring_fd = io_uring_setup(4, &mut params).unwrap();
+
+    // Test that the kernel version supports IORING_REGISTER_PBUF_RING. If it doesn't, the kernel
+    // will return EINVAL. Not setting a `ring_addr` on `io_uring_buf_reg` will return `EFAULT`.
+    if let Err(e) = register_buf_ring(ring_fd.as_fd(), &io_uring_buf_reg::default()) {
+        if e == Errno::INVAL {
+            // Skip the test, as the current kernel version doesn't support what we need to test.
+            return;
+        }
+    }
+
+    let buf_ring_size = ENTRIES * std::mem::size_of::<io_uring_buf>();
+
+    let br_ptr = unsafe {
+        rustix::mm::mmap_anonymous(
+            ptr::null_mut(),
+            buf_ring_size,
+            ProtFlags::READ | ProtFlags::WRITE,
+            MapFlags::PRIVATE,
+        )
+    }
+    .unwrap() as *mut io_uring_buf_ring;
+
+    let br = unsafe { br_ptr.as_mut() }.expect("A valid io_uring_buf_ring struct");
+
+    let reg = io_uring_buf_reg {
+        ring_addr: br_ptr as u64,
+        ring_entries: ENTRIES as u32,
+        bgid: BGID,
+        pad: 0,
+        resv: [0u64; 3],
+    };
+
+    assert_eq!(register_buf_ring(ring_fd, &reg), Ok(()));
+
+    let tail = unsafe { br.tail_or_bufs.tail.as_mut() };
+    tail.tail = 0;
+    let bufs = unsafe { br.tail_or_bufs.bufs.as_mut().bufs.as_mut_slice(ENTRIES) };
+
+    assert_eq!(bufs[0].bid, 0);
+    bufs[7].bid = 7;
+    assert_eq!(bufs[7].bid, 7);
+}