aya-{common,ebpf}: Add spin lock support

Cargo.toml

@@ -2,6 +2,7 @@
 members = [
     "aya",
     "aya-build",
+    "aya-common",
     "aya-log",
     "aya-log-common",
     "aya-log-parser",

aya-common/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+description = "Library shared across eBPF and user-space"
+documentation = "https://docs.rs/aya-common"
+keywords = ["bpf", "ebpf", "common"]
+name = "aya-common"
+version = "0.1.0"
+
+authors.workspace = true
+edition.workspace = true
+homepage.workspace = true
+license.workspace = true
+repository.workspace = true
+rust-version.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+aya-ebpf-bindings = { version = "^0.1.2", path = "../ebpf/aya-ebpf-bindings" }
+
+[features]
+user = ["aya-ebpf-bindings/user"]

aya-common/src/lib.rs

@@ -0,0 +1,9 @@
+#![cfg_attr(
+    target_arch = "bpf",
+    expect(unused_crate_dependencies, reason = "compiler_builtins")
+)]
+#![no_std]
+
+mod spin_lock;
+
+pub use spin_lock::SpinLock;

aya-common/src/spin_lock.rs

@@ -0,0 +1,2 @@
+/// A spin lock that can be used to protect shared data in eBPF maps
+pub type SpinLock = aya_ebpf_bindings::bindings::bpf_spin_lock;

ebpf/aya-ebpf-bindings/Cargo.toml

@@ -13,7 +13,11 @@ repository.workspace = true
 workspace = true
 
 [dependencies]
+aya = { path = "../../aya", version = "^0.13.1", optional = true }
 aya-ebpf-cty = { version = "^0.2.3", path = "../aya-ebpf-cty" }
 
 [build-dependencies]
 aya-build = { version = "^0.1.3", path = "../../aya-build" }
+
+[features]
+user = ["dep:aya"]

ebpf/aya-ebpf-bindings/src/lib.rs

@@ -118,4 +118,7 @@ pub mod bindings {
         pub id: ::aya_ebpf_cty::c_uint,
         pub pinning: ::aya_ebpf_cty::c_uint,
     }
+
+    #[cfg(feature = "user")]
+    unsafe impl aya::Pod for bpf_spin_lock {}
 }

ebpf/aya-ebpf/Cargo.toml

@@ -14,6 +14,7 @@ rust-version.workspace = true
 workspace = true
 
 [dependencies]
+aya-common = { version = "^0.1.0", path = "../../aya-common" }
 aya-ebpf-bindings = { version = "^0.1.2", path = "../aya-ebpf-bindings" }
 aya-ebpf-cty = { version = "^0.2.3", path = "../aya-ebpf-cty" }
 aya-ebpf-macros = { version = "^0.1.2", path = "../../aya-ebpf-macros" }

ebpf/aya-ebpf/src/lib.rs

@@ -45,6 +45,7 @@ pub mod btf_maps;
 pub mod helpers;
 pub mod maps;
 pub mod programs;
+pub mod spin_lock;
 
 use core::{
     mem::MaybeUninit,

ebpf/aya-ebpf/src/spin_lock.rs

@@ -0,0 +1,56 @@
+pub use aya_common::SpinLock;
+
+use crate::helpers;
+
+/// An RAII implementation of a scope of a spin lock. When this structure is
+/// dropped (falls out of scope), the lock will be unlocked.
+#[must_use = "if unused the spin lock will immediately unlock"]
+pub struct SpinLockGuard<'a> {
+    spin_lock: &'a mut SpinLock,
+}
+
+impl Drop for SpinLockGuard<'_> {
+    fn drop(&mut self) {
+        // SAFETY: Call to an eBPF helper. `self.spin_lock` is always
+        // initialized.
+        unsafe {
+            helpers::bpf_spin_unlock(core::ptr::from_mut(self.spin_lock));
+        }
+    }
+}
+
+mod sealed {
+    use super::{SpinLock, SpinLockGuard, helpers};
+
+    pub trait EbpfSpinLock {
+        fn lock(&mut self) -> SpinLockGuard<'_>;
+    }
+
+    impl EbpfSpinLock for SpinLock {
+        fn lock(&mut self) -> SpinLockGuard<'_> {
+            // SAFETY: Call to an eBPF helper. `self` is always initialized.
+            unsafe {
+                helpers::bpf_spin_lock(core::ptr::from_mut(self));
+            }
+            SpinLockGuard { spin_lock: self }
+        }
+    }
+}
+
+/// Extension trait for [`SpinLock`] exposing eBPF-only helpers. These helpers
+/// are not available in user-space.
+pub trait EbpfSpinLockExt: sealed::EbpfSpinLock {
+    fn lock(&mut self) -> SpinLockGuard<'_>;
+}
+
+impl<T> EbpfSpinLockExt for T
+where
+    T: sealed::EbpfSpinLock,
+{
+    /// Acquires a spin lock and returns a [`SpinLockGuard`]. The lock is
+    /// acquired as long as the guard is alive.
+    #[inline]
+    fn lock(&mut self) -> SpinLockGuard<'_> {
+        sealed::EbpfSpinLock::lock(self)
+    }
+}

test/integration-common/Cargo.toml

@@ -15,6 +15,7 @@ workspace = true
 
 [dependencies]
 aya = { path = "../../aya", optional = true }
+aya-common = { path = "../../aya-common" }
 
 [features]
 user = ["aya"]

test/integration-common/src/lib.rs

@@ -1,3 +1,7 @@
+#![cfg_attr(
+    target_arch = "bpf",
+    expect(unused_crate_dependencies, reason = "compiler_builtins")
+)]
 #![no_std]
 
 pub mod array {
@@ -76,6 +80,20 @@ pub mod ring_buf {
     unsafe impl aya::Pod for Registers {}
 }
 
+pub mod spin_lock {
+    use aya_common::SpinLock;
+
+    #[derive(Copy, Clone)]
+    #[repr(C)]
+    pub struct Counter {
+        pub count: u32,
+        pub spin_lock: SpinLock,
+    }
+
+    #[cfg(feature = "user")]
+    unsafe impl aya::Pod for Counter {}
+}
+
 pub mod strncmp {
     #[derive(Copy, Clone)]
     #[repr(C)]

test/integration-ebpf/Cargo.toml

@@ -84,6 +84,10 @@ path = "src/sk_storage.rs"
 name = "simple_prog"
 path = "src/simple_prog.rs"
 
+[[bin]]
+name = "spin_lock"
+path = "src/spin_lock.rs"
+
 [[bin]]
 name = "strncmp"
 path = "src/strncmp.rs"

test/integration-ebpf/src/spin_lock.rs

@@ -0,0 +1,32 @@
+#![no_std]
+#![no_main]
+#![expect(unused_crate_dependencies, reason = "used in other bins")]
+
+#[cfg(not(test))]
+extern crate ebpf_panic;
+
+use aya_ebpf::{
+    bindings::xdp_action,
+    btf_maps::Array,
+    macros::{btf_map, xdp},
+    programs::XdpContext,
+    spin_lock::EbpfSpinLock as _,
+};
+use integration_common::spin_lock::Counter;
+
+#[btf_map]
+static COUNTER: Array<Counter, 1> = Array::new();
+
+#[xdp]
+fn packet_counter(_ctx: XdpContext) -> u32 {
+    let Some(counter) = COUNTER.get_ptr_mut(0) else {
+        return xdp_action::XDP_PASS;
+    };
+    let counter = unsafe { &mut *counter };
+    {
+        let _guard = counter.spin_lock.lock();
+        counter.count = counter.count.saturating_add(1);
+    }
+
+    xdp_action::XDP_PASS
+}

test/integration-test/src/lib.rs

@@ -55,6 +55,7 @@ bpf_file!(
     RING_BUF => "ring_buf",
     SIMPLE_PROG => "simple_prog",
     SK_STORAGE => "sk_storage",
+    SPIN_LOCK => "spin_lock",
     STRNCMP => "strncmp",
     TCX => "tcx",
     TEST => "test",

test/integration-test/src/tests.rs

@@ -29,6 +29,7 @@ mod relocations;
 mod ring_buf;
 mod sk_storage;
 mod smoke;
+mod spin_lock;
 mod strncmp;
 mod tcx;
 mod uprobe_cookie;

test/integration-test/src/tests/spin_lock.rs

@@ -0,0 +1,51 @@
+use std::{net::UdpSocket, time::Duration};
+
+use aya::{
+    EbpfLoader,
+    maps::Array,
+    programs::{Xdp, XdpFlags},
+};
+use integration_common::spin_lock::Counter;
+
+use crate::utils::NetNsGuard;
+
+#[test_log::test]
+fn test_spin_lock() {
+    let _netns = NetNsGuard::new();
+
+    let mut ebpf = EbpfLoader::new().load(crate::SPIN_LOCK).unwrap();
+
+    let prog: &mut Xdp = ebpf
+        .program_mut("packet_counter")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog.load().unwrap();
+    prog.attach("lo", XdpFlags::default()).unwrap();
+
+    const PAYLOAD: &str = "hello counter";
+
+    let sock = UdpSocket::bind("127.0.0.1:0").unwrap();
+    let addr = sock.local_addr().unwrap();
+    sock.set_read_timeout(Some(Duration::from_secs(60)))
+        .unwrap();
+
+    let num_packets = 10;
+    for _ in 0..num_packets {
+        sock.send_to(PAYLOAD.as_bytes(), addr).unwrap();
+    }
+
+    // Read back the packets to ensure it went through the entire network stack,
+    // including the XDP program.
+    let mut buf = [0u8; PAYLOAD.len() + 1];
+    for _ in 0..num_packets {
+        let n = sock.recv(&mut buf).unwrap();
+        assert_eq!(n, PAYLOAD.len());
+        assert_eq!(&buf[..n], PAYLOAD.as_bytes());
+    }
+
+    let counter_map = ebpf.map("COUNTER").unwrap();
+    let counter_map = Array::<_, Counter>::try_from(counter_map).unwrap();
+    let Counter { count, .. } = counter_map.get(&0, 0).unwrap();
+    assert_eq!(count, num_packets);
+}

xtask/public-api/aya-common.txt

@@ -0,0 +1,2 @@
+pub mod aya_common
+pub type aya_common::SpinLock = aya_ebpf_bindings::x86_64::bindings::bpf_spin_lock

xtask/public-api/aya-ebpf.txt

@@ -3061,6 +3061,37 @@ impl<T> core::borrow::BorrowMut<T> for aya_ebpf::programs::xdp::XdpContext where
 pub fn aya_ebpf::programs::xdp::XdpContext::borrow_mut(&mut self) -> &mut T
 impl<T> core::convert::From<T> for aya_ebpf::programs::xdp::XdpContext
 pub fn aya_ebpf::programs::xdp::XdpContext::from(t: T) -> T
+pub mod aya_ebpf::spin_lock
+pub use aya_ebpf::spin_lock::SpinLock
+pub struct aya_ebpf::spin_lock::SpinLockGuard<'a>
+impl core::ops::drop::Drop for aya_ebpf::spin_lock::SpinLockGuard<'_>
+pub fn aya_ebpf::spin_lock::SpinLockGuard<'_>::drop(&mut self)
+impl<'a> core::marker::Freeze for aya_ebpf::spin_lock::SpinLockGuard<'a>
+impl<'a> core::marker::Send for aya_ebpf::spin_lock::SpinLockGuard<'a>
+impl<'a> core::marker::Sync for aya_ebpf::spin_lock::SpinLockGuard<'a>
+impl<'a> core::marker::Unpin for aya_ebpf::spin_lock::SpinLockGuard<'a>
+impl<'a> core::panic::unwind_safe::RefUnwindSafe for aya_ebpf::spin_lock::SpinLockGuard<'a>
+impl<'a> !core::panic::unwind_safe::UnwindSafe for aya_ebpf::spin_lock::SpinLockGuard<'a>
+impl<T, U> core::convert::Into<U> for aya_ebpf::spin_lock::SpinLockGuard<'a> where U: core::convert::From<T>
+pub fn aya_ebpf::spin_lock::SpinLockGuard<'a>::into(self) -> U
+impl<T, U> core::convert::TryFrom<U> for aya_ebpf::spin_lock::SpinLockGuard<'a> where U: core::convert::Into<T>
+pub type aya_ebpf::spin_lock::SpinLockGuard<'a>::Error = core::convert::Infallible
+pub fn aya_ebpf::spin_lock::SpinLockGuard<'a>::try_from(value: U) -> core::result::Result<T, <T as core::convert::TryFrom<U>>::Error>
+impl<T, U> core::convert::TryInto<U> for aya_ebpf::spin_lock::SpinLockGuard<'a> where U: core::convert::TryFrom<T>
+pub type aya_ebpf::spin_lock::SpinLockGuard<'a>::Error = <U as core::convert::TryFrom<T>>::Error
+pub fn aya_ebpf::spin_lock::SpinLockGuard<'a>::try_into(self) -> core::result::Result<U, <U as core::convert::TryFrom<T>>::Error>
+impl<T> core::any::Any for aya_ebpf::spin_lock::SpinLockGuard<'a> where T: 'static + ?core::marker::Sized
+pub fn aya_ebpf::spin_lock::SpinLockGuard<'a>::type_id(&self) -> core::any::TypeId
+impl<T> core::borrow::Borrow<T> for aya_ebpf::spin_lock::SpinLockGuard<'a> where T: ?core::marker::Sized
+pub fn aya_ebpf::spin_lock::SpinLockGuard<'a>::borrow(&self) -> &T
+impl<T> core::borrow::BorrowMut<T> for aya_ebpf::spin_lock::SpinLockGuard<'a> where T: ?core::marker::Sized
+pub fn aya_ebpf::spin_lock::SpinLockGuard<'a>::borrow_mut(&mut self) -> &mut T
+impl<T> core::convert::From<T> for aya_ebpf::spin_lock::SpinLockGuard<'a>
+pub fn aya_ebpf::spin_lock::SpinLockGuard<'a>::from(t: T) -> T
+pub trait aya_ebpf::spin_lock::EbpfSpinLock: aya_ebpf::spin_lock::sealed::Sealed
+pub fn aya_ebpf::spin_lock::EbpfSpinLock::lock(&mut self) -> aya_ebpf::spin_lock::SpinLockGuard<'_>
+impl aya_ebpf::spin_lock::EbpfSpinLock for aya_common::spin_lock::SpinLock
+pub fn aya_common::spin_lock::SpinLock::lock(&mut self) -> aya_ebpf::spin_lock::SpinLockGuard<'_>
 pub macro aya_ebpf::bpf_printk!
 #[repr(transparent)] pub struct aya_ebpf::EbpfGlobal<T>
 impl<T> aya_ebpf::EbpfGlobal<T> where T: core::marker::Copy