aya, aya-obj: Implement ENUM64 fixups

This commit adds:

- A probe to see if the ENUM64 feature is supported
- Fixups for the use of signed enums, or enum64 types
  on systems where enum64 is not supported

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>

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,29 +475,51 @@ 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);
 
@@ -497,7 +527,7 @@ impl Btf {
                     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,7 +623,7 @@ 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 {
@@ -601,7 +631,7 @@ impl Btf {
                         }
                     }
                 }
-                // 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();
+    }
 }

aya-obj/src/btf/types.rs

@@ -400,6 +400,12 @@ pub struct BtfEnum {
     pub value: u32,
 }
 
+impl BtfEnum {
+    pub fn new(name_offset: u32, value: u32) -> Self {
+        Self { name_offset, value }
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct Enum {
@@ -439,9 +445,12 @@ impl Enum {
         mem::size_of::<Fwd>() + mem::size_of::<BtfEnum>() * self.variants.len()
     }
 
-    pub fn new(name_offset: u32, variants: Vec<BtfEnum>) -> Self {
+    pub fn new(name_offset: u32, signed: bool, variants: Vec<BtfEnum>) -> Self {
         let mut info = (BtfKind::Enum as u32) << 24;
         info |= (variants.len() as u32) & 0xFFFF;
+        if signed {
+            info |= 1 << 31;
+        }
         Self {
             name_offset,
             info,
@@ -453,6 +462,14 @@ impl Enum {
     pub(crate) fn is_signed(&self) -> bool {
         self.info >> 31 == 1
     }
+
+    pub(crate) fn set_signed(&mut self, signed: bool) {
+        if signed {
+            self.info |= 1 << 31;
+        } else {
+            self.info &= !(1 << 31);
+        }
+    }
 }
 
 #[repr(C)]
@@ -463,6 +480,16 @@ pub struct BtfEnum64 {
     pub(crate) value_high: u32,
 }
 
+impl BtfEnum64 {
+    pub fn new(name_offset: u32, value: u64) -> Self {
+        Self {
+            name_offset,
+            value_low: value as u32,
+            value_high: (value >> 32) as u32,
+        }
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct Enum64 {
@@ -515,6 +542,26 @@ impl Enum64 {
     pub(crate) fn is_signed(&self) -> bool {
         self.info >> 31 == 1
     }
+
+    pub fn new(name_offset: u32, signed: bool, variants: Vec<BtfEnum64>) -> Self {
+        let mut info = (BtfKind::Enum64 as u32) << 24;
+        if signed {
+            info |= 1 << 31
+        };
+        info |= (variants.len() as u32) & 0xFFFF;
+        Enum64 {
+            name_offset,
+            info,
+            // According to the documentation:
+            //
+            // https://www.kernel.org/doc/html/next/bpf/btf.html
+            //
+            // The size may be 1/2/4/8. Since BtfEnum64::new() takes a u64, we
+            // can assume that the size is 8.
+            size: 8,
+            variants,
+        }
+    }
 }
 
 #[repr(C)]
@@ -614,6 +661,16 @@ pub struct Union {
 }
 
 impl Union {
+    pub(crate) fn new(name_offset: u32, size: u32, members: Vec<BtfMember>) -> Self {
+        let info = (BtfKind::Union as u32) << 24;
+        Self {
+            name_offset,
+            info,
+            size,
+            members,
+        }
+    }
+
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
         let Self {
             name_offset,
@@ -1794,4 +1851,21 @@ mod tests {
         assert!(types_are_compatible(&btf, u32t, &btf, u64t).unwrap());
         assert!(types_are_compatible(&btf, array_type, &btf, array_type).unwrap());
     }
+
+    #[test]
+    pub fn test_read_btf_type_enum64() {
+        let endianness = Endianness::default();
+        let data: &[u8] = &[
+            0x00, 0x00, 0x00, 0x00, // name offset
+            0x01, 0x00, 0x00, 0x13, // info: vlen, type_kind
+            0x08, 0x00, 0x00, 0x00, // size
+            0xd7, 0x06, 0x00, 0x00, // enum variant name offset
+            0xbb, 0xbb, 0xbb, 0xbb, // enum variant low
+            0xaa, 0xaa, 0xaa, 0xaa, // enum variant high
+        ];
+
+        assert_matches!(unsafe { BtfType::read(data, endianness) }.unwrap(), BtfType::Enum64(got) => {
+            assert_eq!(got.to_bytes(), data);
+        });
+    }
 }

aya/src/bpf.rs

@@ -39,8 +39,8 @@ use crate::{
     sys::{
         bpf_load_btf, bpf_map_freeze, bpf_map_update_elem_ptr, is_bpf_cookie_supported,
         is_bpf_global_data_supported, is_btf_datasec_supported, is_btf_decl_tag_supported,
-        is_btf_float_supported, is_btf_func_global_supported, is_btf_func_supported,
-        is_btf_supported, is_btf_type_tag_supported, is_perf_link_supported,
+        is_btf_enum64_supported, is_btf_float_supported, is_btf_func_global_supported,
+        is_btf_func_supported, is_btf_supported, is_btf_type_tag_supported, is_perf_link_supported,
         is_probe_read_kernel_supported, is_prog_name_supported, retry_with_verifier_logs,
         SyscallError,
     },
@@ -84,6 +84,7 @@ fn detect_features() -> Features {
             is_btf_float_supported(),
             is_btf_decl_tag_supported(),
             is_btf_type_tag_supported(),
+            is_btf_enum64_supported(),
         ))
     } else {
         None

aya/src/sys/bpf.rs

@@ -10,6 +10,7 @@ use std::{
 use crate::util::KernelVersion;
 use libc::{c_char, c_long, close, ENOENT, ENOSPC};
 use obj::{
+    btf::{BtfEnum64, Enum64},
     maps::{bpf_map_def, LegacyMap},
     BpfSectionKind, VerifierLog,
 };
@@ -873,6 +874,22 @@ pub(crate) fn is_btf_datasec_supported() -> bool {
     bpf_load_btf(btf_bytes.as_slice(), &mut [], Default::default()).is_ok()
 }
 
+pub(crate) fn is_btf_enum64_supported() -> bool {
+    let mut btf = Btf::new();
+    let name_offset = btf.add_string("enum64");
+
+    let enum_64_type = BtfType::Enum64(Enum64::new(
+        name_offset,
+        true,
+        vec![BtfEnum64::new(btf.add_string("a"), 1)],
+    ));
+    btf.add_type(enum_64_type);
+
+    let btf_bytes = btf.to_bytes();
+
+    bpf_load_btf(btf_bytes.as_slice(), &mut [], Default::default()).is_ok()
+}
+
 pub(crate) fn is_btf_float_supported() -> bool {
     let mut btf = Btf::new();
     let name_offset = btf.add_string("float");

xtask/public-api/aya-obj.txt

@@ -384,6 +384,8 @@ pub fn aya_obj::btf::Btf::from(t: T) -> T
 #[repr(C)] pub struct aya_obj::btf::BtfEnum
 pub aya_obj::btf::BtfEnum::name_offset: u32
 pub aya_obj::btf::BtfEnum::value: u32
+impl aya_obj::btf::BtfEnum
+pub fn aya_obj::btf::BtfEnum::new(name_offset: u32, value: u32) -> Self
 impl core::clone::Clone for aya_obj::btf::BtfEnum
 pub fn aya_obj::btf::BtfEnum::clone(&self) -> aya_obj::btf::BtfEnum
 impl core::fmt::Debug for aya_obj::btf::BtfEnum
@@ -414,6 +416,8 @@ pub fn aya_obj::btf::BtfEnum::borrow_mut(&mut self) -> &mut T
 impl<T> core::convert::From<T> for aya_obj::btf::BtfEnum
 pub fn aya_obj::btf::BtfEnum::from(t: T) -> T
 #[repr(C)] pub struct aya_obj::btf::BtfEnum64
+impl aya_obj::btf::BtfEnum64
+pub fn aya_obj::btf::BtfEnum64::new(name_offset: u32, value: u64) -> Self
 impl core::clone::Clone for aya_obj::btf::BtfEnum64
 pub fn aya_obj::btf::BtfEnum64::clone(&self) -> aya_obj::btf::BtfEnum64
 impl core::fmt::Debug for aya_obj::btf::BtfEnum64
@@ -477,6 +481,7 @@ pub struct aya_obj::btf::BtfFeatures
 impl aya_obj::btf::BtfFeatures
 pub fn aya_obj::btf::BtfFeatures::btf_datasec(&self) -> bool
 pub fn aya_obj::btf::BtfFeatures::btf_decl_tag(&self) -> bool
+pub fn aya_obj::btf::BtfFeatures::btf_enum64(&self) -> bool
 pub fn aya_obj::btf::BtfFeatures::btf_float(&self) -> bool
 pub fn aya_obj::btf::BtfFeatures::btf_func(&self) -> bool
 pub fn aya_obj::btf::BtfFeatures::btf_func_global(&self) -> bool
@@ -701,7 +706,7 @@ impl<T> core::convert::From<T> for aya_obj::btf::DeclTag
 pub fn aya_obj::btf::DeclTag::from(t: T) -> T
 #[repr(C)] pub struct aya_obj::btf::Enum
 impl aya_obj::btf::Enum
-pub fn aya_obj::btf::Enum::new(name_offset: u32, variants: alloc::vec::Vec<aya_obj::btf::BtfEnum>) -> Self
+pub fn aya_obj::btf::Enum::new(name_offset: u32, signed: bool, variants: alloc::vec::Vec<aya_obj::btf::BtfEnum>) -> Self
 impl core::clone::Clone for aya_obj::btf::Enum
 pub fn aya_obj::btf::Enum::clone(&self) -> aya_obj::btf::Enum
 impl core::fmt::Debug for aya_obj::btf::Enum
@@ -732,6 +737,8 @@ pub fn aya_obj::btf::Enum::borrow_mut(&mut self) -> &mut T
 impl<T> core::convert::From<T> for aya_obj::btf::Enum
 pub fn aya_obj::btf::Enum::from(t: T) -> T
 #[repr(C)] pub struct aya_obj::btf::Enum64
+impl aya_obj::btf::Enum64
+pub fn aya_obj::btf::Enum64::new(name_offset: u32, signed: bool, variants: alloc::vec::Vec<aya_obj::btf::BtfEnum64>) -> Self
 impl core::clone::Clone for aya_obj::btf::Enum64
 pub fn aya_obj::btf::Enum64::clone(&self) -> aya_obj::btf::Enum64
 impl core::fmt::Debug for aya_obj::btf::Enum64