From 0d1a2f917c60b784b614c554dbbf1cec40b136ba Mon Sep 17 00:00:00 2001
From: Guillaume Raffin <theelectronwill@gmail.com>
Date: Sun, 14 May 2023 18:54:13 +0200
Subject: [PATCH] aya+ebpf: Implement read+write methods for PerfEventArray

This allow to read _and_ write a PerfEventArray, from userspace _and_ kernel.
---
 aya/src/maps/perf/perf_event_array.rs         | 25 +++++++++
 bpf/aya-bpf/src/maps/perf/perf_event_array.rs | 55 ++++++++++++++++---
 2 files changed, 73 insertions(+), 7 deletions(-)

diff --git a/aya/src/maps/perf/perf_event_array.rs b/aya/src/maps/perf/perf_event_array.rs
index c241a37b0..f3774d19b 100644
--- a/aya/src/maps/perf/perf_event_array.rs
+++ b/aya/src/maps/perf/perf_event_array.rs
@@ -8,10 +8,12 @@ use std::{
     sync::Arc,
 };
 
+use aya_obj::generated::BPF_ANY;
 use bytes::BytesMut;
 
 use crate::{
     maps::{
+        check_bounds,
         perf::{Events, PerfBuffer, PerfBufferError},
         MapData, MapError,
     },
@@ -194,4 +196,27 @@ impl<T: BorrowMut<MapData> + Borrow<MapData>> PerfEventArray<T> {
             _map: self.map.clone(),
         })
     }
+
+    /// Inserts a perf_event file descriptor at the given index.
+    ///
+    /// ## Errors
+    ///
+    /// Returns [`MapError::OutOfBounds`] if `index` is out of bounds, [`MapError::SyscallError`]
+    /// if `bpf_map_update_elem` fails.
+    pub fn set(&mut self, index: u32, value: i32) -> Result<(), MapError> {
+        let data: &MapData = self.map.deref().borrow();
+        check_bounds(data, index)?;
+        let fd = data.fd_or_err()?;
+
+        // only BPF_ANY or BPF_EXIST are allowed, and for arrays they do the same thing (the elements always exist)
+        let flags = BPF_ANY as u64;
+
+        bpf_map_update_elem(fd, Some(&index), &value, flags).map_err(|(_, io_error)| {
+            MapError::SyscallError {
+                call: "bpf_map_update_elem".to_owned(),
+                io_error,
+            }
+        })?;
+        Ok(())
+    }
 }
diff --git a/bpf/aya-bpf/src/maps/perf/perf_event_array.rs b/bpf/aya-bpf/src/maps/perf/perf_event_array.rs
index c881a885b..ce5659d2e 100644
--- a/bpf/aya-bpf/src/maps/perf/perf_event_array.rs
+++ b/bpf/aya-bpf/src/maps/perf/perf_event_array.rs
@@ -1,12 +1,19 @@
-use core::{cell::UnsafeCell, marker::PhantomData, mem};
+use core::{
+    cell::UnsafeCell,
+    marker::PhantomData,
+    mem::{self, MaybeUninit},
+};
+
+use aya_bpf_bindings::{bindings::bpf_perf_event_value, helpers::bpf_perf_event_read_value};
 
 use crate::{
-    bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_PERF_EVENT_ARRAY, BPF_F_CURRENT_CPU},
+    bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_PERF_EVENT_ARRAY, BPF_F_CURRENT_CPU, BPF_F_INDEX_MASK},
     helpers::bpf_perf_event_output,
     maps::PinningType,
     BpfContext,
 };
 
+/// A map of type `BPF_MAP_TYPE_PERF_EVENT_ARRAY`.
 #[repr(transparent)]
 pub struct PerfEventArray<T> {
     def: UnsafeCell<bpf_map_def>,
@@ -50,20 +57,54 @@ impl<T> PerfEventArray<T> {
         }
     }
 
-    pub fn output<C: BpfContext>(&self, ctx: &C, data: &T, flags: u32) {
-        self.output_at_index(ctx, BPF_F_CURRENT_CPU as u32, data, flags)
+    pub fn output_current_cpu<C: BpfContext>(&self, ctx: &C, data: &T) -> Result<(), i64> {
+        self.output(ctx, data, BPF_F_CURRENT_CPU)
+    }
+
+    pub fn output_at_index<C: BpfContext>(&self, ctx: &C, data: &T, index: u32) -> Result<(), i64> {
+        self.output(ctx, data, (index as u64) & BPF_F_INDEX_MASK)
     }
 
-    pub fn output_at_index<C: BpfContext>(&self, ctx: &C, index: u32, data: &T, flags: u32) {
-        let flags = (flags as u64) << 32 | index as u64;
+    fn output<C: BpfContext>(&self, ctx: &C, data: &T, flags: u64) -> Result<(), i64> {
         unsafe {
-            bpf_perf_event_output(
+            let ret = bpf_perf_event_output(
                 ctx.as_ptr(),
                 self.def.get() as *mut _,
                 flags,
                 data as *const _ as *mut _,
                 mem::size_of::<T>() as u64,
             );
+            if ret == 0 {
+                Ok(())
+            } else {
+                Err(ret)
+            }
+        }
+    }
+
+    pub fn read_current_cpu(&self) -> Result<bpf_perf_event_value, i64> {
+        self.read(BPF_F_CURRENT_CPU)
+    }
+
+    pub fn read_at_index(&self, index: u32) -> Result<bpf_perf_event_value, i64> {
+        self.read((index as u64) & BPF_F_INDEX_MASK)
+    }
+
+    fn read(&self, flags: u64) -> Result<bpf_perf_event_value, i64> {
+        let mut buf = MaybeUninit::<bpf_perf_event_value>::uninit();
+        unsafe {
+            // According to the Linux manual, `bpf_perf_event_read_value` is preferred over `bpf_perf_event_read`.
+            let ret = bpf_perf_event_read_value(
+                self.def.get() as *mut _,
+                flags,
+                buf.as_mut_ptr(),
+                mem::size_of::<bpf_perf_event_value>() as u32,
+            );
+            if ret == 0 {
+                Ok(buf.assume_init())
+            } else {
+                Err(ret)
+            }
         }
     }
 }