Merge pull request #721 from dave-tucker/fix-funcinfo

aya: Fix (func|line)_info multiple progs in section

aya-obj/src/obj.rs

@@ -614,19 +614,9 @@ impl Object {
         name: String,
         symbol: &Symbol,
     ) -> Result<(Program, Function), ParseError> {
+        let offset = symbol.address as usize - section.address as usize;
         let (func_info, line_info, func_info_rec_size, line_info_rec_size) =
-            if let Some(btf_ext) = &self.btf_ext {
+            get_func_and_line_info(self.btf_ext.as_ref(), symbol, section, offset, true);
-                let func_info = btf_ext.func_info.get(section.name);
-                let line_info = btf_ext.line_info.get(section.name);
-                (
-                    func_info,
-                    line_info,
-                    btf_ext.func_info_rec_size(),
-                    btf_ext.line_info_rec_size(),
-                )
-            } else {
-                (FuncSecInfo::default(), LineSecInfo::default(), 0, 0)
-            };
 
         let start = symbol.address as usize;
         let end = (symbol.address + symbol.size) as usize;
@@ -690,28 +680,7 @@ impl Object {
             }
 
             let (func_info, line_info, func_info_rec_size, line_info_rec_size) =
-                if let Some(btf_ext) = &self.btf_ext {
+                get_func_and_line_info(self.btf_ext.as_ref(), sym, &section, offset, false);
-                    let bytes_offset = offset as u32 / INS_SIZE as u32;
-                    let section_size_bytes = sym.size as u32 / INS_SIZE as u32;
-
-                    let mut func_info = btf_ext.func_info.get(section.name);
-                    func_info.func_info.retain(|f| f.insn_off == bytes_offset);
-
-                    let mut line_info = btf_ext.line_info.get(section.name);
-                    line_info.line_info.retain(|l| {
-                        l.insn_off >= bytes_offset
-                            && l.insn_off < (bytes_offset + section_size_bytes)
-                    });
-
-                    (
-                        func_info,
-                        line_info,
-                        btf_ext.func_info_rec_size(),
-                        btf_ext.line_info_rec_size(),
-                    )
-                } else {
-                    (FuncSecInfo::default(), LineSecInfo::default(), 0, 0)
-                };
 
             self.functions.insert(
                 (section.index.0, sym.address),
@@ -1369,6 +1338,50 @@ pub fn copy_instructions(data: &[u8]) -> Result<Vec<bpf_insn>, ParseError> {
     Ok(instructions)
 }
 
+fn get_func_and_line_info(
+    btf_ext: Option<&BtfExt>,
+    symbol: &Symbol,
+    section: &Section,
+    offset: usize,
+    rewrite_insn_off: bool,
+) -> (FuncSecInfo, LineSecInfo, usize, usize) {
+    btf_ext
+        .map(|btf_ext| {
+            let instruction_offset = (offset / INS_SIZE) as u32;
+            let symbol_size_instructions = (symbol.size as usize / INS_SIZE) as u32;
+
+            let mut func_info = btf_ext.func_info.get(section.name);
+            func_info.func_info.retain_mut(|f| {
+                let retain = f.insn_off == instruction_offset;
+                if retain && rewrite_insn_off {
+                    f.insn_off = 0;
+                }
+                retain
+            });
+
+            let mut line_info = btf_ext.line_info.get(section.name);
+            line_info
+                .line_info
+                .retain_mut(|l| match l.insn_off.checked_sub(instruction_offset) {
+                    None => false,
+                    Some(insn_off) => {
+                        let retain = insn_off < symbol_size_instructions;
+                        if retain && rewrite_insn_off {
+                            l.insn_off = insn_off
+                        }
+                        retain
+                    }
+                });
+            (
+                func_info,
+                line_info,
+                btf_ext.func_info_rec_size(),
+                btf_ext.line_info_rec_size(),
+            )
+        })
+        .unwrap_or_default()
+}
+
 #[cfg(test)]
 mod tests {
     use alloc::vec;

test/integration-test/bpf/reloc.bpf.c

@@ -29,9 +29,7 @@ __attribute__((noinline)) int field_global() {
   return set_output(__builtin_preserve_access_index(s.b));
 }
 
-SEC("uprobe/field") int field(void *ctx) {
+SEC("uprobe") int field(void *ctx) { return field_global(); }
-  return field_global();
-}
 
 struct relocated_struct_with_pointer {
   struct relocated_struct_with_pointer *first;
@@ -46,9 +44,7 @@ __attribute__((noinline)) int pointer_global() {
   return set_output((__u64)__builtin_preserve_access_index(s.first));
 }
 
-SEC("uprobe/pointer") int pointer(void *ctx) {
+SEC("uprobe") int pointer(void *ctx) { return pointer_global(); }
-  return pointer_global();
-}
 
 __attribute__((noinline)) int struct_flavors_global() {
   struct relocated_struct_with_scalars s = {1, 2, 3};
@@ -59,9 +55,7 @@ __attribute__((noinline)) int struct_flavors_global() {
   }
 }
 
-SEC("uprobe/struct_flavors") int struct_flavors(void *ctx) {
+SEC("uprobe") int struct_flavors(void *ctx) { return struct_flavors_global(); }
-  return struct_flavors_global();
-}
 
 enum relocated_enum_unsigned_32 { U32 = 0xAAAAAAAA };
 
@@ -69,8 +63,7 @@ __attribute__((noinline)) int enum_unsigned_32_global() {
   return set_output(bpf_core_enum_value(enum relocated_enum_unsigned_32, U32));
 }
 
-SEC("uprobe/enum_unsigned_32")
+SEC("uprobe") int enum_unsigned_32(void *ctx) {
-int enum_unsigned_32(void *ctx) {
   return enum_unsigned_32_global();
 }
 
@@ -80,9 +73,7 @@ __attribute__((noinline)) int enum_signed_32_global() {
   return set_output(bpf_core_enum_value(enum relocated_enum_signed_32, S32));
 }
 
-SEC("uprobe/enum_signed_32") int enum_signed_32(void *ctx) {
+SEC("uprobe") int enum_signed_32(void *ctx) { return enum_signed_32_global(); }
-  return enum_signed_32_global();
-}
 
 enum relocated_enum_unsigned_64 { U64 = 0xAAAAAAAABBBBBBBB };
 
@@ -90,8 +81,7 @@ __attribute__((noinline)) int enum_unsigned_64_global() {
   return set_output(bpf_core_enum_value(enum relocated_enum_unsigned_64, U64));
 }
 
-SEC("uprobe/enum_unsigned_64")
+SEC("uprobe") int enum_unsigned_64(void *ctx) {
-int enum_unsigned_64(void *ctx) {
   return enum_unsigned_64_global();
 }
 
@@ -101,6 +91,4 @@ __attribute__((noinline)) int enum_signed_64_global() {
   return set_output(bpf_core_enum_value(enum relocated_enum_signed_64, u64));
 }
 
-SEC("uprobe/enum_signed_64") int enum_signed_64(void *ctx) {
+SEC("uprobe") int enum_signed_64(void *ctx) { return enum_signed_64_global(); }
-  return enum_signed_64_global();
-}