aya: Add `XskMap::unset`

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<T: Borrow<MapData>> XskMap<T> {
 }
 
 impl<T: BorrowMut<MapData>> XskMap<T> {
+    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<T: BorrowMut<MapData>> XskMap<T> {
     /// 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)
+        })
     }
 }

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]

xtask/public-api/aya.txt

@@ -1084,6 +1084,7 @@ impl<T: core::borrow::Borrow<aya::maps::MapData>> aya::maps::XskMap<T>
 pub fn aya::maps::XskMap<T>::pin<P: core::convert::AsRef<std::path::Path>>(self, path: P) -> core::result::Result<(), aya::pin::PinError>
 impl<T: core::borrow::BorrowMut<aya::maps::MapData>> aya::maps::XskMap<T>
 pub fn aya::maps::XskMap<T>::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<T>::unset(&mut self, index: u32) -> core::result::Result<(), aya::maps::MapError>
 impl core::convert::TryFrom<aya::maps::Map> for aya::maps::XskMap<aya::maps::MapData>
 pub type aya::maps::XskMap<aya::maps::MapData>::Error = aya::maps::MapError
 pub fn aya::maps::XskMap<aya::maps::MapData>::try_from(map: aya::maps::Map) -> core::result::Result<Self, Self::Error>
@@ -2423,6 +2424,7 @@ impl<T: core::borrow::Borrow<aya::maps::MapData>> aya::maps::XskMap<T>
 pub fn aya::maps::XskMap<T>::pin<P: core::convert::AsRef<std::path::Path>>(self, path: P) -> core::result::Result<(), aya::pin::PinError>
 impl<T: core::borrow::BorrowMut<aya::maps::MapData>> aya::maps::XskMap<T>
 pub fn aya::maps::XskMap<T>::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<T>::unset(&mut self, index: u32) -> core::result::Result<(), aya::maps::MapError>
 impl core::convert::TryFrom<aya::maps::Map> for aya::maps::XskMap<aya::maps::MapData>
 pub type aya::maps::XskMap<aya::maps::MapData>::Error = aya::maps::MapError
 pub fn aya::maps::XskMap<aya::maps::MapData>::try_from(map: aya::maps::Map) -> core::result::Result<Self, Self::Error>