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.
@@ -473,6 +481,22 @@ impl Btf {
         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 self
+            .types()
+            .any(|t| t.kind() == BtfKind::Enum64 && !features.btf_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 +636,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
@@ -662,12 +686,34 @@ impl Btf {
                     let int_type = BtfType::Int(Int::new(ty.name_offset, 1, IntEncoding::None, 0));
                     types.types[i] = int_type;
                 }
-                // 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;
                 }
+                BtfType::Enum(ty) if !features.btf_enum64 && ty.is_signed() => {
+                    debug!("{}: signed ENUMs not supported. Marking as unsigned", kind);
+                    let mut ty = ty.clone();
+                    ty.set_signed(false);
+                    types.types[i] = BtfType::Enum(ty);
+                }
+                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 +1102,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 +1722,72 @@ 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 enum64_type = Enum::new(
+            name_offset,
+            true,
+            vec![
+                BtfEnum::new(btf.add_string("A"), -1i32 as u32),
+                BtfEnum::new(btf.add_string("B"), -2i32 as u32),
+                BtfEnum::new(btf.add_string("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_eq!(fixed.variants.len(), 3);
+           assert_eq!(fixed.variants[0].value, 0xFFFF_FFFF);
+           assert_eq!(fixed.variants[1].value, 0xFFFF_FFFE);
+           assert_eq!(fixed.variants[2].value, 0xFFFF_FFFD);
+        });
+
+        // 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 enum64_type = Enum64::new(
+            name_offset,
+            false,
+            vec![
+                BtfEnum64::new(btf.add_string("A"), 1),
+                BtfEnum64::new(btf.add_string("B"), 2),
+                BtfEnum64::new(btf.add_string("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) => {
+            assert_eq!(fixed.members.len(), 3);
+        });
+
+        // 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

@@ -389,8 +389,21 @@ impl Int {
 #[repr(C)]
 #[derive(Debug, Clone)]
 pub struct BtfEnum {
-    pub name_offset: u32,
+    pub(crate) name_offset: u32,
-    pub value: u32,
+    pub(crate) value: u32,
+}
+
+impl BtfEnum {
+    pub fn new(name_offset: u32, value: u32) -> Self {
+        BtfEnum { name_offset, value }
+    }
+
+    pub(crate) fn to_bytes(&self) -> Vec<u8> {
+        let mut buf = vec![];
+        buf.extend(bytes_of::<u32>(&self.name_offset));
+        buf.extend(bytes_of::<u32>(&self.value));
+        buf
+    }
 }
 
 #[repr(C)]
@@ -409,8 +422,7 @@ impl Enum {
         buf.extend(bytes_of::<u32>(&self.info));
         buf.extend(bytes_of::<u32>(&self.size));
         for v in &self.variants {
-            buf.extend(bytes_of::<u32>(&v.name_offset));
+            buf.extend(v.to_bytes());
-            buf.extend(bytes_of::<u32>(&v.value));
         }
         buf
     }
@@ -423,9 +435,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;
+        }
         Enum {
             name_offset,
             info,
@@ -437,6 +452,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)]
@@ -447,6 +470,24 @@ pub struct BtfEnum64 {
     pub(crate) value_high: u32,
 }
 
+impl BtfEnum64 {
+    pub fn new(name_offset: u32, value: u64) -> Self {
+        BtfEnum64 {
+            name_offset,
+            value_low: value as u32,
+            value_high: (value >> 32) as u32,
+        }
+    }
+
+    pub(crate) fn to_bytes(&self) -> Vec<u8> {
+        let mut buf = vec![];
+        buf.extend(bytes_of::<u32>(&self.name_offset));
+        buf.extend(bytes_of::<u32>(&self.value_low));
+        buf.extend(bytes_of::<u32>(&self.value_high));
+        buf
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct Enum64 {
@@ -463,9 +504,7 @@ impl Enum64 {
         buf.extend(bytes_of::<u32>(&self.info));
         buf.extend(bytes_of::<u32>(&self.size));
         for v in &self.variants {
-            buf.extend(bytes_of::<u32>(&v.name_offset));
+            buf.extend(v.to_bytes())
-            buf.extend(bytes_of::<u32>(&v.value_high));
-            buf.extend(bytes_of::<u32>(&v.value_low));
         }
         buf
     }
@@ -481,6 +520,20 @@ 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,
+            size: 8, // docs say size is 1/2/4/8 but we'll assume 8
+            variants,
+        }
+    }
 }
 
 #[repr(C)]
@@ -491,6 +544,16 @@ pub(crate) struct BtfMember {
     pub(crate) offset: u32,
 }
 
+impl BtfMember {
+    pub(crate) fn to_bytes(&self) -> Vec<u8> {
+        let mut buf = vec![];
+        buf.extend(bytes_of::<u32>(&self.name_offset));
+        buf.extend(bytes_of::<u32>(&self.btf_type));
+        buf.extend(bytes_of::<u32>(&self.offset));
+        buf
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct Struct {
@@ -507,9 +570,7 @@ impl Struct {
         buf.extend(bytes_of::<u32>(&self.info));
         buf.extend(bytes_of::<u32>(&self.size));
         for v in &self.members {
-            buf.extend(bytes_of::<u32>(&v.name_offset));
+            buf.extend(v.to_bytes());
-            buf.extend(bytes_of::<u32>(&v.btf_type));
-            buf.extend(bytes_of::<u32>(&v.offset));
         }
         buf
     }
@@ -562,15 +623,23 @@ 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 mut buf = vec![];
         buf.extend(bytes_of::<u32>(&self.name_offset));
         buf.extend(bytes_of::<u32>(&self.info));
         buf.extend(bytes_of::<u32>(&self.size));
         for v in &self.members {
-            buf.extend(bytes_of::<u32>(&v.name_offset));
+            buf.extend(v.to_bytes());
-            buf.extend(bytes_of::<u32>(&v.btf_type));
-            buf.extend(bytes_of::<u32>(&v.offset));
         }
         buf
     }
@@ -1672,4 +1741,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_enum64_supported, is_btf_float_supported, is_btf_func_global_supported,
-        is_btf_supported, is_btf_type_tag_supported, is_perf_link_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

@@ -382,8 +382,8 @@ pub fn aya_obj::btf::Btf::borrow_mut(&mut self) -> &mut T
 impl<T> core::convert::From<T> for aya_obj::btf::Btf
 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
+impl aya_obj::btf::BtfEnum
-pub aya_obj::btf::BtfEnum::value: u32
+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 +414,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 +479,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 +704,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 +735,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