From 5458de0760361f26a10695a2102cda8fe1d50b9f Mon Sep 17 00:00:00 2001
From: jinlong <jinlong@tencent.com>
Date: Tue, 30 Jul 2024 20:22:10 +0800
Subject: [PATCH] feat: Add `set` for `Array<T>`

Signed-off-by: jinlong <jinlong@tencent.com>
---
 ebpf/aya-ebpf/src/maps/array.rs          | 21 ++++++++++++++++-
 test/integration-ebpf/src/map_test.rs    |  6 +++++
 test/integration-test/src/tests/smoke.rs | 30 +++++++++++++++++++++---
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/ebpf/aya-ebpf/src/maps/array.rs b/ebpf/aya-ebpf/src/maps/array.rs
index a062ed5b..5e644c54 100644
--- a/ebpf/aya-ebpf/src/maps/array.rs
+++ b/ebpf/aya-ebpf/src/maps/array.rs
@@ -1,6 +1,7 @@
 use core::{cell::UnsafeCell, marker::PhantomData, mem, ptr::NonNull};
 
-use aya_ebpf_cty::c_void;
+use aya_ebpf_bindings::helpers::bpf_map_update_elem;
+use aya_ebpf_cty::{c_long, c_void};
 
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_ARRAY},
@@ -71,4 +72,22 @@ impl<T> Array<T> {
         );
         NonNull::new(ptr as *mut T)
     }
+
+    /// Sets the value of the element at the given index.
+    pub fn set(&self, index: u32, value: &T, flags: u64) -> Result<(), c_long> {
+        insert(self.def.get(), &index, value, flags)
+    }
+}
+
+#[inline]
+fn insert<K, V>(def: *mut bpf_map_def, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
+    let ret = unsafe {
+        bpf_map_update_elem(
+            def as *mut _,
+            key as *const _ as *const _,
+            value as *const _ as *const _,
+            flags,
+        )
+    };
+    (ret == 0).then_some(()).ok_or(ret)
 }
diff --git a/test/integration-ebpf/src/map_test.rs b/test/integration-ebpf/src/map_test.rs
index e3fc28bc..6b62a91c 100644
--- a/test/integration-ebpf/src/map_test.rs
+++ b/test/integration-ebpf/src/map_test.rs
@@ -31,6 +31,12 @@ pub fn simple_prog(_ctx: SkBuffContext) -> i64 {
     0
 }
 
+#[xdp]
+pub fn foo_map_insert(_ctx: XdpContext) -> u32 {
+    FOO.set(0, &1234, 0).ok();
+    xdp_action::XDP_PASS
+}
+
 #[cfg(not(test))]
 #[panic_handler]
 fn panic(_info: &core::panic::PanicInfo) -> ! {
diff --git a/test/integration-test/src/tests/smoke.rs b/test/integration-test/src/tests/smoke.rs
index 0d57aea9..f49a1725 100644
--- a/test/integration-test/src/tests/smoke.rs
+++ b/test/integration-test/src/tests/smoke.rs
@@ -1,7 +1,7 @@
+use std::net::UdpSocket;
+
 use aya::{
-    programs::{Extension, TracePoint, Xdp, XdpFlags},
-    util::KernelVersion,
-    Ebpf, EbpfLoader,
+    maps::Array, programs::{Extension, TracePoint, Xdp, XdpFlags}, util::KernelVersion, Ebpf, EbpfLoader
 };
 use test_log::test;
 
@@ -69,3 +69,27 @@ fn extension() {
         .load(pass.fd().unwrap().try_clone().unwrap(), "xdp_pass")
         .unwrap();
 }
+
+#[test]
+fn map_set() {
+    let _netns = NetNsGuard::new();
+
+    let mut bpf = Ebpf::load(crate::MAP_TEST).unwrap();
+    let dispatcher: &mut Xdp = bpf
+        .program_mut("foo_map_insert")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    dispatcher.load().unwrap();
+    dispatcher.attach("lo", XdpFlags::default()).unwrap();
+    let map: Array<&aya::maps::MapData, u32> = Array::try_from(bpf.map("FOO").unwrap()).unwrap();
+
+    let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
+    socket.connect("127.0.0.1:12345").unwrap();
+    socket.send(&[0; 1]).unwrap();
+    if let Ok(val) = map.get(&0, 0) {
+        assert_eq!(val, 1234);
+    } else {
+        panic!("Key not found");
+    }
+}
\ No newline at end of file