From b614ffd603f4a276fd060659e14e5794bb26381f Mon Sep 17 00:00:00 2001
From: Alessandro Decina <alessandro.d@gmail.com>
Date: Sat, 25 Mar 2023 17:30:01 +1100
Subject: [PATCH] aya: make it possible to use set_global() with slices of
 Pod(s)

---
 aya/src/bpf.rs  | 56 ++++++++++++++++++++++++++++++++++++++-----------
 aya/src/util.rs | 16 ++++++++++++--
 2 files changed, 58 insertions(+), 14 deletions(-)

diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs
index 719f1195..e94b2f6c 100644
--- a/aya/src/bpf.rs
+++ b/aya/src/bpf.rs
@@ -38,7 +38,7 @@ use crate::{
         is_btf_func_supported, is_btf_supported, is_btf_type_tag_supported, is_perf_link_supported,
         is_prog_name_supported, retry_with_verifier_logs,
     },
-    util::{bytes_of, possible_cpus, VerifierLog, POSSIBLE_CPUS},
+    util::{bytes_of, bytes_of_slice, possible_cpus, VerifierLog, POSSIBLE_CPUS},
 };
 
 pub(crate) const BPF_OBJ_NAME_LEN: usize = 16;
@@ -210,15 +210,17 @@ impl<'a> BpfLoader<'a> {
         self
     }
 
-    /// Sets the value of a global variable
+    /// Sets the value of a global variable.
+    ///
+    /// From Rust eBPF, a global variable can be defined as follows:
     ///
-    /// From Rust eBPF, a global variable would be constructed as follows:
     /// ```no_run
     /// #[no_mangle]
     /// static VERSION: i32 = 0;
     /// ```
-    /// Then it would be accessed with `core::ptr::read_volatile` inside
-    /// functions:
+    ///
+    /// Then it can be accessed using `core::ptr::read_volatile`:
+    ///
     /// ```no_run
     /// # #[no_mangle]
     /// # static VERSION: i32 = 0;
@@ -226,10 +228,12 @@ impl<'a> BpfLoader<'a> {
     /// let version = core::ptr::read_volatile(&VERSION);
     /// # }
     /// ```
-    /// If using a struct, ensure that it is `#[repr(C)]` to ensure the size will
-    /// match that of the corresponding ELF symbol.
     ///
-    /// From C eBPF, you would annotate a variable as `volatile const`
+    /// The type of a global variable must be `Pod` (plain old data), for instance `u8`, `u32` and
+    /// all other primitive types. You may use custom types as well, but you must ensure that those
+    /// types are `#[repr(C)]` and only contain other `Pod` types.
+    ///
+    /// From C eBPF, you would annotate a global variable as `volatile const`.
     ///
     /// # Example
     ///
@@ -238,14 +242,17 @@ impl<'a> BpfLoader<'a> {
     ///
     /// let bpf = BpfLoader::new()
     ///     .set_global("VERSION", &2)
+    ///     .set_global("PIDS", &[1234u16, 5678])
     ///     .load_file("file.o")?;
     /// # Ok::<(), aya::BpfError>(())
     /// ```
     ///
-    pub fn set_global<V: Pod>(&mut self, name: &'a str, value: &'a V) -> &mut BpfLoader<'a> {
-        // Safety: value is POD
-        let data = unsafe { bytes_of(value) };
-        self.globals.insert(name, data);
+    pub fn set_global<T: Into<GlobalData<'a>>>(
+        &mut self,
+        name: &'a str,
+        value: T,
+    ) -> &mut BpfLoader<'a> {
+        self.globals.insert(name, value.into().bytes);
         self
     }
 
@@ -897,3 +904,28 @@ fn load_btf(raw_btf: Vec<u8>) -> Result<RawFd, BtfError> {
         }
     }
 }
+
+/// Global data that can be exported to eBPF programs before they are loaded.
+///
+/// Valid global data includes `Pod` types and slices of `Pod` types. See also
+/// [BpfLoader::set_global].
+pub struct GlobalData<'a> {
+    bytes: &'a [u8],
+}
+
+impl<'a, T: Pod> From<&'a [T]> for GlobalData<'a> {
+    fn from(s: &'a [T]) -> Self {
+        GlobalData {
+            bytes: bytes_of_slice(s),
+        }
+    }
+}
+
+impl<'a, T: Pod> From<&'a T> for GlobalData<'a> {
+    fn from(v: &'a T) -> Self {
+        GlobalData {
+            // Safety: v is Pod
+            bytes: unsafe { bytes_of(v) },
+        }
+    }
+}
diff --git a/aya/src/util.rs b/aya/src/util.rs
index 34a82681..9584e820 100644
--- a/aya/src/util.rs
+++ b/aya/src/util.rs
@@ -8,7 +8,10 @@ use std::{
     str::FromStr,
 };
 
-use crate::generated::{TC_H_MAJ_MASK, TC_H_MIN_MASK};
+use crate::{
+    generated::{TC_H_MAJ_MASK, TC_H_MIN_MASK},
+    Pod,
+};
 
 use libc::{if_nametoindex, sysconf, _SC_PAGESIZE};
 
@@ -150,11 +153,20 @@ pub(crate) fn page_size() -> usize {
 }
 
 // bytes_of converts a <T> to a byte slice
-pub(crate) unsafe fn bytes_of<T>(val: &T) -> &[u8] {
+pub(crate) unsafe fn bytes_of<T: Pod>(val: &T) -> &[u8] {
     let size = mem::size_of::<T>();
     slice::from_raw_parts(slice::from_ref(val).as_ptr().cast(), size)
 }
 
+pub(crate) fn bytes_of_slice<T: Pod>(val: &[T]) -> &[u8] {
+    let size = val.len().wrapping_mul(mem::size_of::<T>());
+    // Safety:
+    // Any alignment is allowed.
+    // The size is determined in this function.
+    // The Pod trait ensures the type is valid to cast to bytes.
+    unsafe { slice::from_raw_parts(val.as_ptr().cast(), size) }
+}
+
 const MIN_LOG_BUF_SIZE: usize = 1024 * 10;
 const MAX_LOG_BUF_SIZE: usize = (std::u32::MAX >> 8) as usize;