Merge pull request #725 from dave-tucker/enum64

aya, aya-obj: Implement ENUM64 fixups

aya-obj/src/btf/btf.rs

@@ -17,7 +17,7 @@ use crate::{
         info::{FuncSecInfo, LineSecInfo},
         relocation::Relocation,
         Array, BtfEnum, BtfKind, BtfMember, BtfType, Const, Enum, FuncInfo, FuncLinkage, Int,
-        IntEncoding, LineInfo, Struct, Typedef, VarLinkage,
+        IntEncoding, LineInfo, Struct, Typedef, Union, VarLinkage,
     },
     generated::{btf_ext_header, btf_header},
     util::{bytes_of, HashMap},
@@ -171,6 +171,7 @@ pub struct BtfFeatures {
     btf_float: bool,
     btf_decl_tag: bool,
     btf_type_tag: bool,
+    btf_enum64: bool,
 }
 
 impl BtfFeatures {
@@ -182,6 +183,7 @@ impl BtfFeatures {
         btf_float: bool,
         btf_decl_tag: bool,
         btf_type_tag: bool,
+        btf_enum64: bool,
     ) -> Self {
         BtfFeatures {
             btf_func,
@@ -190,6 +192,7 @@ impl BtfFeatures {
             btf_float,
             btf_decl_tag,
             btf_type_tag,
+            btf_enum64,
         }
     }
 
@@ -227,6 +230,11 @@ impl BtfFeatures {
     pub fn btf_kind_func_proto(&self) -> bool {
         self.btf_func && self.btf_decl_tag
     }
+
+    /// Returns true if the BTF_KIND_ENUM64 is supported.
+    pub fn btf_enum64(&self) -> bool {
+        self.btf_enum64
+    }
 }
 
 /// Bpf Type Format metadata.
@@ -467,37 +475,59 @@ impl Btf {
         buf
     }
 
+    // This follows the same logic as libbpf's bpf_object__sanitize_btf() function.
+    // https://github.com/libbpf/libbpf/blob/05f94ddbb837f5f4b3161e341eed21be307eaa04/src/libbpf.c#L2701
+    //
+    // Fixup: The loader needs to adjust values in the BTF before it's loaded into the kernel.
+    // Sanitize: Replace an unsupported BTF type with a placeholder type.
+    //
+    // In addition to the libbpf logic, it performs some fixups to the BTF generated by bpf-linker
+    // for Aya programs. These fixups are gradually moving into bpf-linker itself.
     pub(crate) fn fixup_and_sanitize(
         &mut self,
         section_infos: &HashMap<String, (SectionIndex, u64)>,
         symbol_offsets: &HashMap<String, u64>,
         features: &BtfFeatures,
     ) -> Result<(), BtfError> {
+        // ENUM64 placeholder type needs to be added before we take ownership of
+        // self.types to ensure that the offsets in the BtfHeader are correct.
+        let placeholder_name = self.add_string("enum64_placeholder");
+        let enum64_placeholder_id = (!features.btf_enum64
+            && self.types().any(|t| t.kind() == BtfKind::Enum64))
+        .then(|| {
+            self.add_type(BtfType::Int(Int::new(
+                placeholder_name,
+                1,
+                IntEncoding::None,
+                0,
+            )))
+        });
         let mut types = mem::take(&mut self.types);
         for i in 0..types.types.len() {
             let t = &mut types.types[i];
             let kind = t.kind();
             match t {
-                // Fixup PTR for Rust
-                // LLVM emits names for Rust pointer types, which the kernel doesn't like
+                // Fixup PTR for Rust.
+                //
+                // LLVM emits names for Rust pointer types, which the kernel doesn't like.
                 // While I figure out if this needs fixing in the Kernel or LLVM, we'll
-                // do a fixup here
+                // do a fixup here.
                 BtfType::Ptr(ptr) => {
                     ptr.name_offset = 0;
                 }
-                // Sanitize VAR if they are not supported
+                // Sanitize VAR if they are not supported.
                 BtfType::Var(v) if !features.btf_datasec => {
                     types.types[i] = BtfType::Int(Int::new(v.name_offset, 1, IntEncoding::None, 0));
                 }
-                // Sanitize DATASEC if they are not supported
+                // Sanitize DATASEC if they are not supported.
                 BtfType::DataSec(d) if !features.btf_datasec => {
                     debug!("{}: not supported. replacing with STRUCT", kind);
 
                     // STRUCT aren't allowed to have "." in their name, fixup this if needed.
                     let mut name_offset = d.name_offset;
                     let name = self.string_at(name_offset)?;
 
-                    // Handle any "." characters in struct names
+                    // Handle any "." characters in struct names.
                     // Example: ".maps"
                     let fixed_name = name.replace('.', "_");
                     if fixed_name != name {
@@ -521,29 +551,29 @@ impl Btf {
                     types.types[i] =
                         BtfType::Struct(Struct::new(name_offset, members, entries.len() as u32));
                 }
-                // Fixup DATASEC
-                // DATASEC sizes aren't always set by LLVM
-                // we need to fix them here before loading the btf to the kernel
+                // Fixup DATASEC.
+                //
+                // DATASEC sizes aren't always set by LLVM so we need to fix them
+                // here before loading the btf to the kernel.
                 BtfType::DataSec(d) if features.btf_datasec => {
                     // Start DataSec Fixups
                     let name = self.string_at(d.name_offset)?;
                     let name = name.into_owned();
 
-                    // Handle any "/" characters in section names
+                    // Handle any "/" characters in section names.
                     // Example: "maps/hashmap"
                     let fixed_name = name.replace('/', ".");
                     if fixed_name != name {
                         d.name_offset = self.add_string(&fixed_name);
                     }
 
-                    // There are some cases when the compiler does indeed populate the
-                    // size
+                    // There are some cases when the compiler does indeed populate the size.
                     if d.size > 0 {
                         debug!("{} {}: size fixup not required", kind, name);
                     } else {
-                        // We need to get the size of the section from the ELF file
+                        // We need to get the size of the section from the ELF file.
                         // Fortunately, we cached these when parsing it initially
-                        // and we can this up by name in section_infos
+                        // and we can this up by name in section_infos.
                         let size = match section_infos.get(&name) {
                             Some((_, size)) => size,
                             None => {
@@ -557,7 +587,7 @@ impl Btf {
                         // that need to have their offsets adjusted. To do this,
                         // we need to get the offset from the ELF file.
                         // This was also cached during initial parsing and
-                        // we can query by name in symbol_offsets
+                        // we can query by name in symbol_offsets.
                         let mut entries = mem::take(&mut d.entries);
                         let mut fixed_section = d.clone();
 
@@ -593,15 +623,15 @@ impl Btf {
                         types.types[i] = BtfType::DataSec(fixed_section);
                     }
                 }
-                // Fixup FUNC_PROTO
+                // Fixup FUNC_PROTO.
                 BtfType::FuncProto(ty) if features.btf_func => {
                     for (i, param) in ty.params.iter_mut().enumerate() {
                         if param.name_offset == 0 && param.btf_type != 0 {
                             param.name_offset = self.add_string(&format!("param{i}"));
                         }
                     }
                 }
-                // Sanitize FUNC_PROTO
+                // Sanitize FUNC_PROTO.
                 BtfType::FuncProto(ty) if !features.btf_func => {
                     debug!("{}: not supported. replacing with ENUM", kind);
                     let members: Vec<BtfEnum> = ty
@@ -612,13 +642,13 @@ impl Btf {
                             value: p.btf_type,
                         })
                         .collect();
-                    let enum_type = BtfType::Enum(Enum::new(ty.name_offset, members));
+                    let enum_type = BtfType::Enum(Enum::new(ty.name_offset, false, members));
                     types.types[i] = enum_type;
                 }
-                // Sanitize FUNC
+                // Sanitize FUNC.
                 BtfType::Func(ty) => {
                     let name = self.string_at(ty.name_offset)?;
-                    // Sanitize FUNC
+                    // Sanitize FUNC.
                     if !features.btf_func {
                         debug!("{}: not supported. replacing with TYPEDEF", kind);
                         let typedef_type =
@@ -648,25 +678,48 @@ impl Btf {
                         }
                     }
                 }
-                // Sanitize FLOAT
+                // Sanitize FLOAT.
                 BtfType::Float(ty) if !features.btf_float => {
                     debug!("{}: not supported. replacing with STRUCT", kind);
                     let struct_ty = BtfType::Struct(Struct::new(0, vec![], ty.size));
                     types.types[i] = struct_ty;
                 }
-                // Sanitize DECL_TAG
+                // Sanitize DECL_TAG.
                 BtfType::DeclTag(ty) if !features.btf_decl_tag => {
                     debug!("{}: not supported. replacing with INT", kind);
                     let int_type = BtfType::Int(Int::new(ty.name_offset, 1, IntEncoding::None, 0));
                     types.types[i] = int_type;
                 }
-                // Sanitize TYPE_TAG
+                // Sanitize TYPE_TAG.
                 BtfType::TypeTag(ty) if !features.btf_type_tag => {
                     debug!("{}: not supported. replacing with CONST", kind);
                     let const_type = BtfType::Const(Const::new(ty.btf_type));
                     types.types[i] = const_type;
                 }
-                // The type does not need fixing up or sanitization
+                // Sanitize Signed ENUMs.
+                BtfType::Enum(ty) if !features.btf_enum64 && ty.is_signed() => {
+                    debug!("{}: signed ENUMs not supported. Marking as unsigned", kind);
+                    ty.set_signed(false);
+                }
+                // Sanitize ENUM64.
+                BtfType::Enum64(ty) if !features.btf_enum64 => {
+                    debug!("{}: not supported. replacing with UNION", kind);
+                    let placeholder_id =
+                        enum64_placeholder_id.expect("enum64_placeholder_id must be set");
+                    let members: Vec<BtfMember> = ty
+                        .variants
+                        .iter()
+                        .map(|v| BtfMember {
+                            name_offset: v.name_offset,
+                            btf_type: placeholder_id,
+                            offset: 0,
+                        })
+                        .collect();
+                    let union_type =
+                        BtfType::Union(Union::new(ty.name_offset, members.len() as u32, members));
+                    types.types[i] = union_type;
+                }
+                // The type does not need fixing up or sanitization.
                 _ => {}
             }
         }
@@ -1054,7 +1107,8 @@ pub(crate) struct SecInfo<'a> {
 mod tests {
     use super::*;
     use crate::btf::{
-        BtfParam, DataSec, DataSecEntry, DeclTag, Float, Func, FuncProto, Ptr, TypeTag, Var,
+        BtfEnum64, BtfParam, DataSec, DataSecEntry, DeclTag, Enum64, Float, Func, FuncProto, Ptr,
+        TypeTag, Var,
     };
     use assert_matches::assert_matches;
 
@@ -1676,4 +1730,96 @@ mod tests {
         let u32_ty = btf.type_by_id(u32_base).unwrap();
         assert_eq!(u32_ty.kind(), BtfKind::Int);
     }
+
+    #[test]
+    fn test_sanitize_signed_enum() {
+        let mut btf = Btf::new();
+        let name_offset = btf.add_string("signed_enum");
+        let name_a = btf.add_string("A");
+        let name_b = btf.add_string("B");
+        let name_c = btf.add_string("C");
+        let enum64_type = Enum::new(
+            name_offset,
+            true,
+            vec![
+                BtfEnum::new(name_a, -1i32 as u32),
+                BtfEnum::new(name_b, -2i32 as u32),
+                BtfEnum::new(name_c, -3i32 as u32),
+            ],
+        );
+        let enum_type_id = btf.add_type(BtfType::Enum(enum64_type));
+
+        let features = BtfFeatures {
+            btf_enum64: false,
+            ..Default::default()
+        };
+
+        btf.fixup_and_sanitize(&HashMap::new(), &HashMap::new(), &features)
+            .unwrap();
+
+        assert_matches!(btf.type_by_id(enum_type_id).unwrap(), BtfType::Enum(fixed) => {
+            assert!(!fixed.is_signed());
+            assert_matches!(fixed.variants[..], [
+                BtfEnum { name_offset: name1, value: 0xFFFF_FFFF },
+                BtfEnum { name_offset: name2, value: 0xFFFF_FFFE },
+                BtfEnum { name_offset: name3, value: 0xFFFF_FFFD },
+            ] => {
+                assert_eq!(name1, name_a);
+                assert_eq!(name2, name_b);
+                assert_eq!(name3, name_c);
+            });
+        });
+
+        // Ensure we can convert to bytes and back again.
+        let raw = btf.to_bytes();
+        Btf::parse(&raw, Endianness::default()).unwrap();
+    }
+
+    #[test]
+    fn test_sanitize_enum64() {
+        let mut btf = Btf::new();
+        let name_offset = btf.add_string("enum64");
+        let name_a = btf.add_string("A");
+        let name_b = btf.add_string("B");
+        let name_c = btf.add_string("C");
+        let enum64_type = Enum64::new(
+            name_offset,
+            false,
+            vec![
+                BtfEnum64::new(name_a, 1),
+                BtfEnum64::new(name_b, 2),
+                BtfEnum64::new(name_c, 3),
+            ],
+        );
+        let enum_type_id = btf.add_type(BtfType::Enum64(enum64_type));
+
+        let features = BtfFeatures {
+            btf_enum64: false,
+            ..Default::default()
+        };
+
+        btf.fixup_and_sanitize(&HashMap::new(), &HashMap::new(), &features)
+            .unwrap();
+
+        assert_matches!(btf.type_by_id(enum_type_id).unwrap(), BtfType::Union(fixed) => {
+            let placeholder = btf.id_by_type_name_kind("enum64_placeholder", BtfKind::Int)
+                .expect("enum64_placeholder type not found");
+            assert_matches!(fixed.members[..], [
+                BtfMember { name_offset: name_offset1, btf_type: btf_type1, offset: 0 },
+                BtfMember { name_offset: name_offset2, btf_type: btf_type2, offset: 0 },
+                BtfMember { name_offset: name_offset3, btf_type: btf_type3, offset: 0 },
+            ] => {
+                assert_eq!(name_offset1, name_a);
+                assert_eq!(btf_type1, placeholder);
+                assert_eq!(name_offset2, name_b);
+                assert_eq!(btf_type2, placeholder);
+                assert_eq!(name_offset3, name_c);
+                assert_eq!(btf_type3, placeholder);
+            });
+        });
+
+        // Ensure we can convert to bytes and back again.
+        let raw = btf.to_bytes();
+        Btf::parse(&raw, Endianness::default()).unwrap();
+    }
 }