From a04d83651b37cd9548a1df514bca0228ad88d56b Mon Sep 17 00:00:00 2001
From: Mary <mary@mary.zone>
Date: Wed, 29 Mar 2023 15:06:40 +0200
Subject: [PATCH] aya: Allows sharing Map between different eBPFs

Allows to share maps between eBPFs without relying on pinned maps.
---
 aya/src/bpf.rs      | 37 +++++++++++++++++++++++++++++++++++--
 aya/src/maps/mod.rs | 22 ++++++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs
index 6e9dcdac..25e56fba 100644
--- a/aya/src/bpf.rs
+++ b/aya/src/bpf.rs
@@ -126,6 +126,7 @@ pub struct BpfLoader<'a> {
     btf: Option<Cow<'a, Btf>>,
     map_pin_path: Option<PathBuf>,
     globals: HashMap<&'a str, (&'a [u8], bool)>,
+    maps: HashMap<&'a str, &'a Map>,
     max_entries: HashMap<&'a str, u32>,
     extensions: HashSet<&'a str>,
     verifier_log_level: VerifierLogLevel,
@@ -160,6 +161,7 @@ impl<'a> BpfLoader<'a> {
             btf: Btf::from_sys_fs().ok().map(Cow::Owned),
             map_pin_path: None,
             globals: HashMap::new(),
+            maps: HashMap::new(),
             max_entries: HashMap::new(),
             extensions: HashSet::new(),
             verifier_log_level: VerifierLogLevel::default(),
@@ -283,6 +285,29 @@ impl<'a> BpfLoader<'a> {
         self
     }
 
+    /// Allows to share a map between multiple eBPF without pinning.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// use aya::BpfLoader;
+    /// use aya::maps::Map;
+    ///
+    /// let mut shared_bpf = BpfLoader::new()
+    ///     .load_file("shared.o")?;
+    ///
+    /// let shared_map = shared_bpf.take_map("shared_map").unwrap();
+    ///
+    /// let bpf = BpfLoader::new()
+    ///     .map("shared_map", &shared_map)
+    ///     .load_file("file.o")?;
+    /// # Ok::<(), aya::BpfError>(())
+    /// ```
+    pub fn map(&mut self, name: &'a str, map: &'a Map) -> &mut BpfLoader<'a> {
+        self.maps.insert(name, map);
+        self
+    }
+
     /// Set the max_entries for specified map.
     ///
     /// Overwrite the value of max_entries of the map that matches
@@ -383,6 +408,7 @@ impl<'a> BpfLoader<'a> {
             extensions,
             verifier_log_level,
             allow_unsupported_maps,
+            ..
         } = self;
         let mut obj = Object::parse(data)?;
         obj.patch_map_data(globals.clone())?;
@@ -429,7 +455,14 @@ impl<'a> BpfLoader<'a> {
             }
             let mut map = MapData {
                 obj,
-                fd: None,
+                fd: self.maps.get(name.as_str()).and_then(|x| {
+                    let map_data = (*x).map_data().clone();
+                    let map_fd = map_data.fd;
+
+                    std::mem::forget(map_data);
+
+                    map_fd
+                }),
                 pinned: false,
                 btf_fd,
             };
@@ -455,7 +488,7 @@ impl<'a> BpfLoader<'a> {
                         }
                     }
                 }
-                PinningType::None => map.create(&name)?,
+                PinningType::None => map.fd_or_err().or_else(|_| map.create(&name))?,
             };
             if !map.obj.data().is_empty() && map.obj.section_kind() != BpfSectionKind::Bss {
                 bpf_map_update_elem_ptr(fd, &0 as *const _, map.obj.data_mut().as_mut_ptr(), 0)
diff --git a/aya/src/maps/mod.rs b/aya/src/maps/mod.rs
index dbb94e35..6b1f0e05 100644
--- a/aya/src/maps/mod.rs
+++ b/aya/src/maps/mod.rs
@@ -274,6 +274,28 @@ pub enum Map {
 }
 
 impl Map {
+    /// Returns the low level map.
+    pub(crate) fn map_data(&self) -> &MapData {
+        match self {
+            Map::Array(map) => map,
+            Map::PerCpuArray(map) => map,
+            Map::ProgramArray(map) => map,
+            Map::HashMap(map) => map,
+            Map::LruHashMap(map) => map,
+            Map::PerCpuHashMap(map) => map,
+            Map::PerCpuLruHashMap(map) => map,
+            Map::PerfEventArray(map) => map,
+            Map::SockHash(map) => map,
+            Map::SockMap(map) => map,
+            Map::BloomFilter(map) => map,
+            Map::LpmTrie(map) => map,
+            Map::Stack(map) => map,
+            Map::StackTraceMap(map) => map,
+            Map::Queue(map) => map,
+            Map::Unsupported(map) => map,
+        }
+    }
+
     /// Returns the low level map type.
     fn map_type(&self) -> u32 {
         match self {