Embedding a skeleton / object

[atenart]

Embedding a skeleton (Skel) in an internal structure is quite useful, eg. for loading a program once and attaching it to functions on-demand later on (attach_kprobe). Up until 0.22 it was possible to either embed a Skel or an Object (taking it out of the Skel).

With the addition of commit c01e9b9 ("libbpf-cargo: Define maps as publicly accessible members") and follow-ups I'm now not sure how to have the same behavior, or if it is even possible. There are a few things preventing this:

1. A storage (MaybeUninit<libbpf_rs::OpenObject>) needs to be provided to the skeleton builder to hold the (open)object and obviously have to live long enough. The internal structure now needs to embed both the Skel and the storage, but this is impossible due to self-referencial.
2. To workaround the above we could try directly using the storage, without the Skel (after setting rodata). But internally it uses an OwnedRef which will drop the storage if we don't keep the Skel around. Or we need mem::forget(skel) but that leaks memory (and does not look nice).

Then I thought maybe the skeleton isn't meant for such usage and we could build an Object manually. This is honestly fine... until we need to access rodata. Then basically we go full circle and re-implement a skeleton (as we need the Rust rodata definition to be automatically generated to stay in-sync).

I'm not really sure how to proceed. Hopefully I'm missing something :)

Thanks!

(The actual code is there, https://github.com/retis-org/retis/blob/59948504aa2f0da99cefb53057f7995baf8eed0b/retis/src/core/probe/kernel/kprobe.rs, but it's a bit complex to grab everything w/o reading quite a lot of code so hopefully the above explanation describes the generic issue).

[javierhonduco]

Thanks for opening this issue! This is something I also bumped into. My program also relies on access to rodata as well as various maps during the lifetime of the BPF program. We store the Skeleton in a field of the main structure of this application.

This is still a WIP, and not sure it's totally safe. I need to make sure Rust semantics can't cause any issues, so please don't use this code in production. This is how it roughly looks so far. Can update it once I get a more refined version, if this is useful. I ran it with ASAN and it didn't report any issues, but it's not a 100% guarantee that the code is correct.

// Warning!! The code below is quite brittle. Do *not* use it as-is.

// [...]

impl for Profiler {
  fn new() -> Self {
    // [...]

    let mut profiler_open_object = ManuallyDrop::new(Box::new(MaybeUninit::uninit()));
    let mut open_skel = skel_builder
        .open(&mut native_profiler_object)
        .expect("open skel");
    let native_unwinder_skel = ManuallyDrop::new(open_skel.load().expect("load skel"));
    // Safety: this has to be dropped before `profiler_open_object`, as long as this is guaranteed it should
    // be fine.
    let native_unwinder_skel = unsafe {
        std::mem::transmute::<ManuallyDrop<NativeProfilerSkel<'_>>, 
        ManuallyDrop<NativeProfilerSkel<'static>>>(native_unwinder_skel)
    };
    Profiler {
        // Rust guarantees drop order, and custom `drop` impls run first, so make sure to free
        // `profiler_open_object` after `native_unwinder_skel`.
        profiler_open_object,
        native_unwinder_skel,
    }
  }
}

// TODO: drop implementation in `Profiler` to free the `ManuallyDrop`s. Make sure the order
// does not do anything unsafe,

[danielocfb]

I think a dance like what @javierhonduco suggests may be necessary at the moment. We should check whether we can improve, but it would be helpful to have a minimal example reproducing the issue.

[atenart]

it would be helpful to have a minimal example reproducing the issue.

I wrote one at https://github.com/atenart/libbpf-rs-skel-embed. The code is pretty minimal. There's 3 commits:

1. Using libbpf-rs 0.22 and showing what was possible to write (embedding an Object in an internal structure).
2. Upgrading to libbpf-rs 0.24, this does not compile.
3. Adding a workaround to make things to work. The idea would be not to need this, but might be interesting to investigate possible solutions (to add to libbpf-rs).

Speaking of reproducers, is there any example reproducing the issues fixed by the introduction of the storage (MaybeUninit<OpenObject>)? That would help finding a solution that fits both needs.

[javierhonduco]

@atenart I believe this was the repro for the UAF #852

[danielocfb]

Sorry for the delayed response, I was on vacation. As c01e9b9 states, the main driver for the rework is the desire to access multiple maps/programs mutably.

Here is a trivial refactoring that fails to build due to one such borrowing conflict on top of v0.23.3:

--- examples/tcp_ca/src/main.rs
+++ examples/tcp_ca/src/main.rs
@@ -137,13 +137,9 @@ fn test(name_to_register: Option<&CStr>, name_to_use: &CStr, verbose: bool) -> R
         let () = ca_update.name[len..].fill(0);
     }

-    let ca_update_cong_control2 = open_skel
-        .progs()
-        .ca_update_cong_control2()
-        .as_libbpf_object()
-        .as_ptr();
+    let ca_update_cong_control2 = open_skel.progs().ca_update_cong_control2();
     let ca_update = open_skel.struct_ops.ca_update_mut();
-    ca_update.cong_control = ca_update_cong_control2;
+    ca_update.cong_control = ca_update_cong_control2.as_libbpf_object().as_ptr();

     let mut skel = open_skel.load()?;
     let mut maps = skel.maps_mut();

The storage logic is necessary to get around all the self-referentiality issues that we otherwise run into when objects are no longer hidden behind method calls.

Real world example of ergonomic hit: the state before sched-ext/scx#445

I don't have every detail paged in anymore, but I believe the use-after-free may have been solvable without this rework -- it certainly wasn't the driving factor.

[d-e-s-o]

it would be helpful to have a minimal example reproducing the issue.

I wrote one at https://github.com/atenart/libbpf-rs-skel-embed.

Many thanks for that!

The code is pretty minimal. There's 3 commits:

1. Using libbpf-rs 0.22 and showing what was possible to write (embedding an Object in an internal structure).
2. Upgrading to libbpf-rs 0.24, this does not compile.

Can you pass in the object storage instead or where does the requirement of storing it alongside the skeleton come from? Given that KprobeManager already accounts for a lifetime, we may as well actually use it, would be my guess:

diff --git src/main.rs src/main.rs
index 15f048..e07a90 100644
--- src/main.rs
+++ src/main.rs
@@ -1,10 +1,13 @@
+use std::mem::MaybeUninit;
+
 use anyhow::Result;

 mod manager;
 use manager::*;

 fn main() -> Result<()> {
-    let mut mgr = KprobeManager::new()?;
+    let mut obj = MaybeUninit::uninit();
+    let mut mgr = KprobeManager::new(&mut obj)?;

     mgr.attach("pskb_expand_head")?;
     mgr.attach("kfree_skb_partial")?;
diff --git src/manager.rs src/manager.rs
index 3034b2..db8522 100644
--- src/manager.rs
+++ src/manager.rs
@@ -1,7 +1,7 @@
 use std::mem::MaybeUninit;

 use anyhow::{anyhow, Result};
-use libbpf_rs::skel::{OpenSkel, Skel, SkelBuilder};
+use libbpf_rs::{skel::{OpenSkel, Skel, SkelBuilder}, OpenObject};

 mod bpf {
     include!("bpf/kprobe.skel.rs");
@@ -9,21 +9,18 @@ mod bpf {
 use bpf::*;

 pub(crate) struct KprobeManager<'a> {
-    storage: Box<MaybeUninit<libbpf_rs::OpenObject>>,
     skel: KprobeSkel<'a>,
     links: Vec<libbpf_rs::Link>,
 }

-impl KprobeManager<'_> {
-    pub(crate) fn new() -> Result<Self> {
-        let mut storage = Box::new(MaybeUninit::uninit());
-        let skel = KprobeSkelBuilder::default().open(&mut storage)?;
+impl<'a> KprobeManager<'a> {
+    pub(crate) fn new(storage: &'a mut MaybeUninit<OpenObject>) -> Result<Self> {
+        let skel = KprobeSkelBuilder::default().open(storage)?;

         // Set rodata.
         skel.maps.rodata_data.log_level = 0;

         Ok(Self {
-            storage,
             skel: skel.load()?,
             links: Vec::new(),
         })

[atenart]

The code is pretty minimal. There's 3 commits:

1. Using libbpf-rs 0.22 and showing what was possible to write (embedding an Object in an internal structure).
2. Upgrading to libbpf-rs 0.24, this does not compile.

Can you pass in the object storage instead or where does the requirement of storing it alongside the skeleton come from? Given that KprobeManager already accounts for a lifetime, we may as well actually use it, would be my guess:

I thought about this but although it works and does not look to bad in the small reproducer, in our real world application things are more tricky. The main issue is those objects building probes (KprobeManager in the reproducer) can be added dynamically at runtime. Another issue would be only cosmetic but IMO wouldn't look very nice: those managers are built quite far away from main.

[danielocfb]

Okay, short of that you can just fake a 'static lifetime, as per my understanding:

--- src/manager.rs
+++ src/manager.rs
@@ -1,4 +1,4 @@
-use std::mem::MaybeUninit;
+use std::mem::{transmute, MaybeUninit};

 use anyhow::{anyhow, Result};
 use libbpf_rs::skel::{OpenSkel, Skel, SkelBuilder};
@@ -8,23 +8,34 @@ mod bpf {
 }
 use bpf::*;

-pub(crate) struct KprobeManager<'a> {
-    storage: Box<MaybeUninit<libbpf_rs::OpenObject>>,
-    skel: KprobeSkel<'a>,
+pub(crate) struct KprobeManager {
+    // SAFETY: We must not hand out references with a 'static lifetime to
+    //         this member. Rather, they should never outlive `self`.
+    //         Furthermore, this member has to be listed before
+    //         `_storage` to make sure we never end up with a dangling
+    //         reference.
+    skel: KprobeSkel<'static>,
+    _storage: Box<MaybeUninit<libbpf_rs::OpenObject>>,
     links: Vec<libbpf_rs::Link>,
 }

-impl KprobeManager<'_> {
+impl KprobeManager {
     pub(crate) fn new() -> Result<Self> {
         let mut storage = Box::new(MaybeUninit::uninit());
-        let skel = KprobeSkelBuilder::default().open(&mut storage)?;
+        // SAFETY: We keep a reference to the box around for the
+        //         lifetime of self and we never hand out 'static
+        //         references to this object.
+        let obj_ref = unsafe {
+            transmute::<&mut MaybeUninit<_>, &'static mut MaybeUninit<_>>(storage.as_mut())
+        };
+        let skel = KprobeSkelBuilder::default().open(obj_ref)?;

         // Set rodata.
         skel.maps.rodata_data.log_level = 0;

         Ok(Self {
-            storage,
             skel: skel.load()?,
+            _storage: storage,
             links: Vec::new(),
         })
     }

This seems to be standard for self-referential stuff.

[atenart]

Okay, short of that you can just fake a 'static lifetime, as per my understanding:

Right, and control the order the storage and object are dropped. I used a third party object in the reproducer to hide the trick from the user, but your example should work too.

This seems to be standard for self-referential stuff.

I don't know if it's standard as we're tricking the compiler, but that's the common way to make self-referential working atm :)

[d-e-s-o]

Yeah, so that would be my recommendation. How you package it is really up to you. I suppose it may be possible to put all of that into the generated skeletons themselves and I think that would be fine. Not sure if I'll have the cycles to work on that any time soon, though, so if you need it in the near future it would be best to open a pull request.

[atenart]

I suppose it may be possible to put all of that into the generated skeletons themselves and I think that would be fine.

Do you mean embedding the storage in the skeleton and hiding it from the user? Something like the workaround from the reproducer but integrated.

Not sure if I'll have the cycles to work on that any time soon, though, so if you need it in the near future it would be best to open a pull request.

Yeah, no worries, this was about raising awareness and discussing a potential way forward (and also making sure I wasn't missing something obvious). I might as well look into this at some point.

[d-e-s-o]

Do you mean embedding the storage in the skeleton and hiding it from the user? Something like the workaround from the reproducer but integrated.

Yep, that's what I meant. I don't know how involved it is, to be honest, but perhaps with a hidden intermediate type similar to what you have it won't be too bad.

[telbizov]

Thanks for this discussion @atenart and @d-e-s-o . I'm also in the process of trying to upgrade to the latest libbpf-rs and the bump broke my code. I tried the suggested transmute fix which allows me to "package" both the skeleton and its storage within my own structure. Cheers for that:

pub struct Bpf {
    pub skel: skel::MyOwnSkel<'static>,
    storage: Box<MaybeUninit<OpenObject>>,
}

Another breakage from this change is that I am no longer able to have a global static (behind LazyLock) of Bpf kind. This is because of std::ptr::NonNull not Sync, which is needed. I am still looking into this, but wanted to bring this up and get your input and suggestions on that matter.

[danielocfb]

Another breakage from this change is that I am no longer able to have a global static (behind LazyLock) of Bpf kind. This is because of std::ptr::NonNull not Sync, which is needed. I am still looking into this, but wanted to bring this up and get your input and suggestions on that matter.

If the skeleton (or some embedded type) was Sync beforehand and no longer is we should fix that by adding the implementation. Feel free to open a PR or spin off a separate issue if more discussion is necessary.

[atenart]

I had a chance to do a few experiments lately.

In order to embed the storage into the (Open)Skel we need the ability to:

• Get references from the inner object.
• Take the inner object (required for libbpf_rs::Object::from_ptr).
• Move the wrapper without actually moving its inner storage.

There's no need for self reference tricks as {name}Maps and {name}Progs already have a workaround. This is actually a requirement for allowing maps & progs as publicly accessible members.

This could look like (not even sure if MaybeUninit is really needed, I don't recall if I had a version without it; but that's another story anyway):

struct ObjStorage<O> {{
    storage: Box<std::mem::MaybeUninit<O>>,
    taken: bool,
}}

impl<O> ObjStorage<O> {{
    fn new(obj: O) -> Self {{
        Self {{
            storage: Box::new(std::mem::MaybeUninit::new(obj)),
            taken: false,
        }}
    }}

    fn as_ref(&self) -> &O {{
        unsafe {{ self.storage.assume_init_ref() }}
    }}

    fn as_mut(&mut self) -> &mut O {{
        unsafe {{ self.storage.assume_init_mut() }}
    }}

    fn take(mut self) -> O {{
        let storage = std::mem::replace(&mut *self.storage, std::mem::MaybeUninit::uninit());
        self.taken = true;
        // Safety: the inner storage is always initialized.
        unsafe {{ storage.assume_init() }}
    }}
}}

impl<O> Drop for ObjStorage<O> {{
    fn drop(&mut self) {{
        if self.taken {{
            return;
        }}

        // Safety: we just checked the inner storage wasn't taken.
        unsafe {{ self.storage.assume_init_drop() }}
    }}
}}

• (partial diff)

@@ -980,15 +981,13 @@ fn gen_skel_contents(raw_obj_name: &str, obj_file_path: &Path) -> Result<String>
                 // SAFETY: `obj_ptr` points to an opened object after
                 //         skeleton open.
                 let obj = unsafe {{ libbpf_rs::OpenObject::from_ptr(obj_ptr) }};
-                let _obj = object.write(obj);
-                // SAFETY: We just wrote initialized data to `object`.
-                let mut obj_ref = unsafe {{ OwnedRef::new(object) }};
+                let mut object = ObjStorage::new(obj);
 
                 #[allow(unused_mut)]
                 let mut skel = Open{name}Skel {{
-                    maps: unsafe {{ Open{name}Maps::new(&skel_config, obj_ref.as_mut())? }},
-                    progs: unsafe {{ Open{name}Progs::new(obj_ref.as_mut())? }},
-                    obj: obj_ref,
+                    maps: unsafe {{ Open{name}Maps::new(&skel_config, object.as_mut())? }},
+                    progs: unsafe {{ Open{name}Progs::new(object.as_mut())? }},
+                    obj: object,
                     // SAFETY: Our `struct_ops` type contains only pointers,
                     //         which are allowed to be NULL.
                     // TODO: Generate and use a `Default` representation
@@ -1089,30 +1086,16 @@ pub struct StructOps {{}}
                     return Err(libbpf_rs::Error::from_raw_os_error(-ret));
                 }}
 
-                let obj_ref = self.obj.take();
-                let open_obj = std::mem::replace(obj_ref, std::mem::MaybeUninit::uninit());
-                // SAFETY: `open_obj` is guaranteed to be properly
-                //         initialized as it came from an `OwnedRef`.
-                let obj_ptr = unsafe {{ open_obj.assume_init().take_ptr() }};
+                let obj_ptr = self.obj.take().take_ptr();
                 // SAFETY: `obj_ptr` points to a loaded object after
                 //         skeleton load.
                 let obj = unsafe {{ libbpf_rs::Object::from_ptr(obj_ptr) }};
-                // SAFETY: `OpenObject` and `Object` are guaranteed to
-                //         have the same memory layout.
-                let obj_ref = unsafe {{
-                    std::mem::transmute::<
-                        &'obj mut std::mem::MaybeUninit<libbpf_rs::OpenObject>,
-                        &'obj mut std::mem::MaybeUninit<libbpf_rs::Object>,
-                    >(obj_ref)
-                }};
-                let _obj = obj_ref.write(obj);
-                // SAFETY: We just wrote initialized data to `obj_ref`.
-                let mut obj_ref = unsafe {{ OwnedRef::new(obj_ref) }};
+                let mut object = ObjStorage::new(obj);
 
                 Ok({name}Skel {{
-                    maps: unsafe {{ {name}Maps::new(&self.skel_config, obj_ref.as_mut())? }},
+                    maps: unsafe {{ {name}Maps::new(&self.skel_config, object.as_mut())? }},
                     progs: {name}Progs::new(self.progs),
-                    obj: obj_ref,
+                    obj: object,
                     struct_ops: self.struct_ops,
                     skel_config: self.skel_config,
                     {links}

But... the above works well until we do the following:

let rodata = {
    let skel = FooSkelBuilder::default().open()?;
    skel.maps.rodata_data
};

println!("{:?}", rodata);	// Segfaults.

This is because there's no relationship between the inner members and the skel, thanks to the maps & progs self-reference tricks. If we add it back, the above does not compile, but then this defeats the purpose... The exact same issue arise if instead the (Open)Skel itself is wrapped with the storage into another object (like what was shared in previous comments): it works until access to the inner members is allowed.

Another idea would be to let the hidden storage live for as long as either skel members. Arc and friends come to mind but they hit another issue: the inner value cannot be taken if there are more than a single strong reference[1]. But that's one of the initial requirements. Also what happens for accesses made to the inner object while the OpenSkel -> Skel transition is happening...

Not a lot of progress overall overall :)

[1] There might be a way with an Arc<Mutex<> + deref_mut + swap (twice); but let's not do that for obvious reasons.