From 36420d929734beb7486cc5d14944fc7cf8e9b62a Mon Sep 17 00:00:00 2001
From: Adam Preuss <a_preuss@apple.com>
Date: Fri, 1 Sep 2023 14:34:21 +0000
Subject: [PATCH] aya: support loading a map by fd

This adds support to loading maps by fd similarly to the way programs
can be loaded by fd.
---
 aya/src/maps/mod.rs      | 44 ++++++++++++++++++++++++++++++++++++++--
 aya/src/sys/bpf.rs       | 15 ++++++++++++++
 xtask/public-api/aya.txt |  1 +
 3 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/aya/src/maps/mod.rs b/aya/src/maps/mod.rs
index a934d361..ceab7307 100644
--- a/aya/src/maps/mod.rs
+++ b/aya/src/maps/mod.rs
@@ -68,8 +68,9 @@ use crate::{
     obj::{self, parse_map_info, BpfSectionKind},
     pin::PinError,
     sys::{
-        bpf_create_map, bpf_get_object, bpf_map_freeze, bpf_map_get_info_by_fd,
-        bpf_map_get_next_key, bpf_map_update_elem_ptr, bpf_pin_object, SyscallError,
+        bpf_create_map, bpf_get_object, bpf_map_freeze, bpf_map_get_fd_by_id,
+        bpf_map_get_info_by_fd, bpf_map_get_next_key, bpf_map_update_elem_ptr, bpf_pin_object,
+        SyscallError,
     },
     util::{nr_cpus, KernelVersion},
     PinningType, Pod,
@@ -648,6 +649,13 @@ impl MapData {
         })
     }
 
+    /// Loads a map from a map id.
+    pub fn from_id(id: u32) -> Result<Self, MapError> {
+        bpf_map_get_fd_by_id(id)
+            .map_err(MapError::from)
+            .and_then(Self::from_fd)
+    }
+
     /// Loads a map from a file descriptor.
     ///
     /// If loading from a BPF Filesystem (bpffs) you should use [`Map::from_pin`](crate::maps::MapData::from_pin).
@@ -935,6 +943,38 @@ mod tests {
         })
     }
 
+    #[test]
+    fn test_from_map_id() {
+        override_syscall(|call| match call {
+            Syscall::Bpf {
+                cmd: bpf_cmd::BPF_MAP_GET_FD_BY_ID,
+                attr,
+            } => {
+                assert_eq!(
+                    unsafe { attr.__bindgen_anon_6.__bindgen_anon_1.map_id },
+                    1234
+                );
+                Ok(42)
+            }
+            Syscall::Bpf {
+                cmd: bpf_cmd::BPF_OBJ_GET_INFO_BY_FD,
+                attr,
+            } => {
+                assert_eq!(unsafe { attr.info.bpf_fd }, 42);
+                Ok(0)
+            }
+            _ => Err((-1, io::Error::from_raw_os_error(EFAULT))),
+        });
+
+        assert_matches!(
+            MapData::from_id(1234),
+            Ok(MapData {
+                obj: _,
+                fd,
+            }) => assert_eq!(fd.as_fd().as_raw_fd(), 42)
+        );
+    }
+
     #[test]
     fn test_create() {
         override_syscall(|call| match call {
diff --git a/aya/src/sys/bpf.rs b/aya/src/sys/bpf.rs
index dd42120a..511f0163 100644
--- a/aya/src/sys/bpf.rs
+++ b/aya/src/sys/bpf.rs
@@ -550,6 +550,21 @@ pub(crate) fn bpf_prog_get_info_by_fd(
     })
 }
 
+pub(crate) fn bpf_map_get_fd_by_id(map_id: u32) -> Result<OwnedFd, SyscallError> {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+
+    attr.__bindgen_anon_6.__bindgen_anon_1.map_id = map_id;
+
+    // SAFETY: BPF_MAP_GET_FD_BY_ID returns a new file descriptor.
+    unsafe { fd_sys_bpf(bpf_cmd::BPF_MAP_GET_FD_BY_ID, &mut attr) }.map_err(|(code, io_error)| {
+        assert_eq!(code, -1);
+        SyscallError {
+            call: "bpf_map_get_fd_by_id",
+            io_error,
+        }
+    })
+}
+
 pub(crate) fn bpf_map_get_info_by_fd(fd: BorrowedFd<'_>) -> Result<bpf_map_info, SyscallError> {
     bpf_obj_get_info_by_fd(fd, |_| {})
 }
diff --git a/xtask/public-api/aya.txt b/xtask/public-api/aya.txt
index 28a82b88..3650d1f3 100644
--- a/xtask/public-api/aya.txt
+++ b/xtask/public-api/aya.txt
@@ -1686,6 +1686,7 @@ impl aya::maps::MapData
 pub fn aya::maps::MapData::create(obj: aya_obj::maps::Map, name: &str, btf_fd: core::option::Option<std::os::fd::owned::BorrowedFd<'_>>) -> core::result::Result<Self, aya::maps::MapError>
 pub fn aya::maps::MapData::fd(&self) -> &aya::maps::MapFd
 pub fn aya::maps::MapData::from_fd(fd: std::os::fd::owned::OwnedFd) -> core::result::Result<Self, aya::maps::MapError>
+pub fn aya::maps::MapData::from_id(id: u32) -> core::result::Result<Self, aya::maps::MapError>
 pub fn aya::maps::MapData::from_pin<P: core::convert::AsRef<std::path::Path>>(path: P) -> core::result::Result<Self, aya::maps::MapError>
 pub fn aya::maps::MapData::pin<P: core::convert::AsRef<std::path::Path>>(&mut self, path: P) -> core::result::Result<(), aya::pin::PinError>
 impl core::fmt::Debug for aya::maps::MapData