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,12 +475,32 @@ 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 mut enum64_placeholder_id = None;
+        if !features.btf_enum64 && self.types().any(|t| t.kind() == BtfKind::Enum64) {
+            let placeholder = BtfType::Int(Int::new(
+                self.add_string("enum64_placeholder"),
+                1,
+                IntEncoding::None,
+                0,
+            ));
+            enum64_placeholder_id = Some(self.add_type(placeholder));
+        }
         let mut types = mem::take(&mut self.types);
         for i in 0..types.types.len() {
             let t = &types.types[i];
@@ -612,7 +640,7 @@ 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
@@ -668,6 +696,31 @@ impl Btf {
                     let const_type = BtfType::Const(Const::new(ty.btf_type));
                     types.types[i] = const_type;
                 }
+                // Sanitize Signed ENUMs
+                BtfType::Enum(ty) if !features.btf_enum64 && ty.is_signed() => {
+                    debug!("{}: signed ENUMs not supported. Marking as unsigned", kind);
+                    if let BtfType::Enum(t) = &mut types.types[i] {
+                        t.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
                 _ => {}
             }
@@ -1056,7 +1109,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;
 
@@ -1675,4 +1729,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,24 @@ 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 +659,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;
+        Union {
+            name_offset,
+            info,
+            size,
+            members,
+        }
+    }
+
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
         let Self {
             name_offset,
@@ -1794,4 +1849,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");