From 73a34e1571a606124c7e89ecf71ff001508643dc Mon Sep 17 00:00:00 2001 From: arctic-alpaca <67190338+arctic-alpaca@users.noreply.github.com> Date: Thu, 7 Dec 2023 15:44:56 +0100 Subject: [PATCH] aya: Add `XskMap::unset` --- aya/src/maps/xdp/xsk_map.rs | 42 ++++++++++++++++++++------ test/integration-test/src/tests/xdp.rs | 23 +++++++++++--- xtask/public-api/aya.txt | 2 ++ 3 files changed, 54 insertions(+), 13 deletions(-) diff --git a/aya/src/maps/xdp/xsk_map.rs b/aya/src/maps/xdp/xsk_map.rs index f38b97fc..f23a3f19 100644 --- a/aya/src/maps/xdp/xsk_map.rs +++ b/aya/src/maps/xdp/xsk_map.rs @@ -2,12 +2,12 @@ use std::{ borrow::{Borrow, BorrowMut}, - os::fd::{AsFd, AsRawFd, RawFd}, + os::fd::{AsFd, AsRawFd, BorrowedFd, RawFd}, }; use crate::{ maps::{MapData, MapError, check_bounds, check_kv_size}, - sys::{SyscallError, bpf_map_update_elem}, + sys::{SyscallError, bpf_map_delete_elem, bpf_map_update_elem}, }; /// An array of AF_XDP sockets. @@ -57,6 +57,16 @@ impl> XskMap { } impl> XskMap { + fn with_fd( + &mut self, + index: u32, + f: impl FnOnce(BorrowedFd<'_>) -> Result<(), SyscallError>, + ) -> Result<(), MapError> { + let data = self.inner.borrow_mut(); + check_bounds(data, index)?; + f(data.fd().as_fd()).map_err(Into::into) + } + /// Sets the `AF_XDP` socket at a given index. /// /// When redirecting a packet, the `AF_XDP` socket at `index` will recieve the packet. Note @@ -68,14 +78,28 @@ impl> XskMap { /// Returns [`MapError::OutOfBounds`] if `index` is out of bounds, [`MapError::SyscallError`] /// if `bpf_map_update_elem` fails. pub fn set(&mut self, index: u32, socket_fd: impl AsRawFd, flags: u64) -> Result<(), MapError> { - let data = self.inner.borrow_mut(); - check_bounds(data, index)?; - let fd = data.fd().as_fd(); - bpf_map_update_elem(fd, Some(&index), &socket_fd.as_raw_fd(), flags) - .map_err(|io_error| SyscallError { - call: "bpf_map_update_elem", + self.with_fd(index, |fd| { + bpf_map_update_elem(fd, Some(&index), &socket_fd.as_raw_fd(), flags).map_err( + |io_error| SyscallError { + call: "bpf_map_update_elem", + io_error, + }, + ) + }) + } + + /// Un-sets the `AF_XDP` socket at a given index. + /// + /// # Errors + /// + /// Returns [`MapError::OutOfBounds`] if `index` is out of bounds, [`MapError::SyscallError`] + /// if `bpf_map_delete_elem` fails. + pub fn unset(&mut self, index: u32) -> Result<(), MapError> { + self.with_fd(index, |fd| { + bpf_map_delete_elem(fd, &index).map_err(|io_error| SyscallError { + call: "bpf_map_delete_elem", io_error, }) - .map_err(Into::into) + }) } } diff --git a/test/integration-test/src/tests/xdp.rs b/test/integration-test/src/tests/xdp.rs index 6e0687a2..af921032 100644 --- a/test/integration-test/src/tests/xdp.rs +++ b/test/integration-test/src/tests/xdp.rs @@ -26,12 +26,14 @@ fn af_xdp() { xdp.load().unwrap(); xdp.attach("lo", XdpFlags::default()).unwrap(); + const SIZE: usize = 2 * 4096; + // So this needs to be page aligned. Pages are 4k on all mainstream architectures except for // Apple Silicon which uses 16k pages. So let's align on that for tests to run natively there. #[repr(C, align(16384))] - struct PageAligned([u8; 4096]); + struct PageAligned([u8; SIZE]); - let mut alloc = Box::new(PageAligned([0; 4096])); + let mut alloc = Box::new(PageAligned([0; SIZE])); let umem = { let PageAligned(mem) = alloc.as_mut(); let mem = mem.as_mut().into(); @@ -57,10 +59,12 @@ fn af_xdp() { socks.set(0, rx.as_raw_fd(), 0).unwrap(); let frame = umem.frame(BufIdx(0)).unwrap(); + let frame1 = umem.frame(BufIdx(1)).unwrap(); - // Produce a frame to be filled by the kernel - let mut writer = fq_cq.fill(1); + // Produce two frames to be filled by the kernel + let mut writer = fq_cq.fill(2); writer.insert_once(frame.offset); + writer.insert_once(frame1.offset); writer.commit(); let sock = UdpSocket::bind("127.0.0.1:0").unwrap(); @@ -81,6 +85,17 @@ fn af_xdp() { assert_eq!(&udp[0..2], port.to_be_bytes().as_slice()); // Source assert_eq!(&udp[2..4], 1777u16.to_be_bytes().as_slice()); // Dest assert_eq!(payload, b"hello AF_XDP"); + + assert_eq!(rx.available(), 1); + // Removes socket from map, no more packets will be redirected. + socks.unset(0).unwrap(); + assert_eq!(rx.available(), 1); + sock.send_to(b"hello AF_XDP", "127.0.0.1:1777").unwrap(); + assert_eq!(rx.available(), 1); + // Adds socket to map again, packets will be redirected again. + socks.set(0, rx.as_raw_fd(), 0).unwrap(); + sock.send_to(b"hello AF_XDP", "127.0.0.1:1777").unwrap(); + assert_eq!(rx.available(), 2); } #[test] diff --git a/xtask/public-api/aya.txt b/xtask/public-api/aya.txt index 390001fb..6ed2a959 100644 --- a/xtask/public-api/aya.txt +++ b/xtask/public-api/aya.txt @@ -1084,6 +1084,7 @@ impl> aya::maps::XskMap pub fn aya::maps::XskMap::pin>(self, path: P) -> core::result::Result<(), aya::pin::PinError> impl> aya::maps::XskMap pub fn aya::maps::XskMap::set(&mut self, index: u32, socket_fd: impl std::os::fd::raw::AsRawFd, flags: u64) -> core::result::Result<(), aya::maps::MapError> +pub fn aya::maps::XskMap::unset(&mut self, index: u32) -> core::result::Result<(), aya::maps::MapError> impl core::convert::TryFrom for aya::maps::XskMap pub type aya::maps::XskMap::Error = aya::maps::MapError pub fn aya::maps::XskMap::try_from(map: aya::maps::Map) -> core::result::Result @@ -2423,6 +2424,7 @@ impl> aya::maps::XskMap pub fn aya::maps::XskMap::pin>(self, path: P) -> core::result::Result<(), aya::pin::PinError> impl> aya::maps::XskMap pub fn aya::maps::XskMap::set(&mut self, index: u32, socket_fd: impl std::os::fd::raw::AsRawFd, flags: u64) -> core::result::Result<(), aya::maps::MapError> +pub fn aya::maps::XskMap::unset(&mut self, index: u32) -> core::result::Result<(), aya::maps::MapError> impl core::convert::TryFrom for aya::maps::XskMap pub type aya::maps::XskMap::Error = aya::maps::MapError pub fn aya::maps::XskMap::try_from(map: aya::maps::Map) -> core::result::Result