From 93435fc85400aa036f3890c43c78c9c9eb4baa96 Mon Sep 17 00:00:00 2001
From: Mary <mary@mary.zone>
Date: Wed, 28 Jun 2023 15:04:46 +0200
Subject: [PATCH] aya: allow global value to be optional

This allow to not error out when a global symbol is missing from the object.
---
 aya-obj/src/obj.rs | 16 +++++++++++-----
 aya/src/bpf.rs     | 10 ++++++----
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/aya-obj/src/obj.rs b/aya-obj/src/obj.rs
index 0a7235fd..5ff4f1c3 100644
--- a/aya-obj/src/obj.rs
+++ b/aya-obj/src/obj.rs
@@ -597,7 +597,10 @@ impl Object {
     }
 
     /// Patches map data
-    pub fn patch_map_data(&mut self, globals: HashMap<&str, &[u8]>) -> Result<(), ParseError> {
+    pub fn patch_map_data(
+        &mut self,
+        globals: HashMap<&str, (&[u8], bool)>,
+    ) -> Result<(), ParseError> {
         let symbols: HashMap<String, &Symbol> = self
             .symbol_table
             .iter()
@@ -605,7 +608,7 @@ impl Object {
             .map(|(_, s)| (s.name.as_ref().unwrap().clone(), s))
             .collect();
 
-        for (name, data) in globals {
+        for (name, (data, must_exist)) in globals {
             if let Some(symbol) = symbols.get(name) {
                 if data.len() as u64 != symbol.size {
                     return Err(ParseError::InvalidGlobalData {
@@ -633,7 +636,7 @@ impl Object {
                     });
                 }
                 map.data_mut().splice(start..end, data.iter().cloned());
-            } else {
+            } else if must_exist {
                 return Err(ParseError::SymbolNotFound {
                     name: name.to_owned(),
                 });
@@ -2354,8 +2357,11 @@ mod tests {
         );
 
         let test_data: &[u8] = &[1, 2, 3];
-        obj.patch_map_data(HashMap::from([("my_config", test_data)]))
-            .unwrap();
+        obj.patch_map_data(HashMap::from([
+            ("my_config", (test_data, true)),
+            ("optional_variable", (test_data, false)),
+        ]))
+        .unwrap();
 
         let map = obj.maps.get(".rodata").unwrap();
         assert_eq!(test_data, map.data());
diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs
index 2b499c12..4a6df7a1 100644
--- a/aya/src/bpf.rs
+++ b/aya/src/bpf.rs
@@ -119,7 +119,7 @@ fn detect_features() -> Features {
 pub struct BpfLoader<'a> {
     btf: Option<Cow<'a, Btf>>,
     map_pin_path: Option<PathBuf>,
-    globals: HashMap<&'a str, &'a [u8]>,
+    globals: HashMap<&'a str, (&'a [u8], bool)>,
     max_entries: HashMap<&'a str, u32>,
     extensions: HashSet<&'a str>,
     verifier_log_level: VerifierLogLevel,
@@ -203,6 +203,7 @@ impl<'a> BpfLoader<'a> {
     }
 
     /// Sets the value of a global variable.
+    /// If the `must_exist` argument is `true`, [`BpfLoader::load`] will fail with [`ParseError::SymbolNotFound`] if the loaded object code does not contain the variable.
     ///
     /// From Rust eBPF, a global variable can be defined as follows:
     ///
@@ -233,8 +234,8 @@ impl<'a> BpfLoader<'a> {
     /// use aya::BpfLoader;
     ///
     /// let bpf = BpfLoader::new()
-    ///     .set_global("VERSION", &2)
-    ///     .set_global("PIDS", &[1234u16, 5678])
+    ///     .set_global("VERSION", &2, true)
+    ///     .set_global("PIDS", &[1234u16, 5678], true)
     ///     .load_file("file.o")?;
     /// # Ok::<(), aya::BpfError>(())
     /// ```
@@ -243,8 +244,9 @@ impl<'a> BpfLoader<'a> {
         &mut self,
         name: &'a str,
         value: T,
+        must_exist: bool,
     ) -> &mut BpfLoader<'a> {
-        self.globals.insert(name, value.into().bytes);
+        self.globals.insert(name, (value.into().bytes, must_exist));
         self
     }