aya-obj: apply enum64-to-union fixup in reloc

This code is just awful.

aya-obj/src/btf/btf.rs

@@ -18,8 +18,9 @@ use object::{Endianness, SectionIndex};
 use crate::{
     Object,
     btf::{
-        Array, BtfEnum, BtfKind, BtfMember, BtfType, Const, DataSec, DataSecEntry, Enum, Enum64,
-        FuncInfo, FuncLinkage, Int, IntEncoding, LineInfo, Struct, Typedef, Union, Var, VarLinkage,
+        Array, BtfEnum, BtfEnum64, BtfKind, BtfMember, BtfType, Const, DataSec, DataSecEntry, Enum,
+        Enum64, Enum64Fallback, Enum64VariantFallback, FuncInfo, FuncLinkage, Int, IntEncoding,
+        LineInfo, Struct, Typedef, Union, Var, VarLinkage,
         info::{FuncSecInfo, LineSecInfo},
         relocation::Relocation,
     },
@@ -752,6 +753,7 @@ impl Btf {
                         // `ty` is borrowed from `types` and we use that borrow
                         // below, so we must not borrow it again in the
                         // get_or_init closure.
+                        let is_signed = ty.is_signed();
                         let Enum64 {
                             name_offset,
                             size,
@@ -760,6 +762,27 @@ impl Btf {
                         } = ty;
                         let (name_offset, size, variants) =
                             (*name_offset, *size, mem::take(variants));
+
+                        let fallback = Enum64Fallback {
+                            signed: is_signed,
+                            variants: variants
+                                .iter()
+                                .copied()
+                                .map(
+                                    |BtfEnum64 {
+                                         name_offset,
+                                         value_high,
+                                         value_low,
+                                     }| Enum64VariantFallback {
+                                        name_offset,
+                                        value: (u64::from(value_high) << 32) | u64::from(value_low),
+                                    },
+                                )
+                                .collect(),
+                        };
+
+                        // The rewritten UNION still needs a concrete member type. Share a single
+                        // synthetic INT placeholder between every downgraded ENUM64.
                         let placeholder_id = enum64_placeholder_id.get_or_init(|| {
                             let placeholder_name = self.add_string("enum64_placeholder");
                             add_type(
@@ -779,7 +802,7 @@ impl Btf {
 
                         // Must reborrow here because we borrow `types` above.
                         let t = &mut types.types[i];
-                        *t = BtfType::Union(Union::new(name_offset, size, members));
+                        *t = BtfType::Union(Union::new(name_offset, size, members, Some(fallback)));
                     }
                 }
                 // The type does not need fixing up or sanitization.

aya-obj/src/btf/relocation.rs

@@ -13,8 +13,8 @@ use object::SectionIndex;
 use crate::{
     Function, Object,
     btf::{
-        Array, Btf, BtfError, BtfMember, BtfType, IntEncoding, MAX_SPEC_LEN, Struct, Union,
-        fields_are_compatible, types_are_compatible,
+        Array, Btf, BtfError, BtfKind, BtfMember, BtfType, IntEncoding, MAX_SPEC_LEN, Struct,
+        Union, fields_are_compatible, types_are_compatible,
     },
     generated::{
         BPF_ALU, BPF_ALU64, BPF_B, BPF_CALL, BPF_DW, BPF_H, BPF_JMP, BPF_K, BPF_LD, BPF_LDX,
@@ -409,10 +409,26 @@ fn find_candidates<'target>(
 ) -> Result<Vec<Candidate<'target>>, BtfError> {
     let mut candidates = Vec::new();
     let local_name = flavorless_name(local_name);
-    for (type_id, ty) in target_btf.types().enumerate() {
-        if local_ty.kind() != ty.kind() {
-            continue;
-        }
+    let local_kind = local_ty.kind();
+
+    // When we downgrade an ENUM64 to a UNION we still want to match the enum
+    // definition recorded in the target BTF. If the sanitized type has a
+    // fallback, allow ENUM64 candidates through the kind check.
+    //
+    // Note that we do not sanitize the target BTF so the kinds will not
+    // naturally match!
+    let allow_enum_match = matches!(
+        local_ty,
+        BtfType::Union(Union {
+            enum64_fallback: Some(_),
+            ..
+        })
+    );
+
+    for (type_id, ty) in target_btf.types().enumerate().filter(|(_, ty)| {
+        let candidate_kind = ty.kind();
+        candidate_kind == local_kind || (allow_enum_match && candidate_kind == BtfKind::Enum64)
+    }) {
         let name = &*target_btf.type_name(ty)?;
         if local_name != flavorless_name(name) {
             continue;
@@ -488,6 +504,14 @@ fn match_candidate<'target>(
                 BtfType::Enum64(en) => {
                     match_enum(&mut en.variants.iter().map(|member| member.name_offset))
                 }
+                BtfType::Union(Union {
+                    enum64_fallback: Some(fallback),
+                    ..
+                }) => {
+                    // Local ENUM64 types become UNIONs during sanitisation; the fallback retains
+                    // their original variant names so we can line them up with target enums.
+                    match_enum(&mut fallback.variants.iter().map(|variant| variant.name_offset))
+                }
                 _ => Ok(None),
             }
         }
@@ -708,6 +732,17 @@ impl<'a> AccessSpec<'a> {
                             index,
                         )
                     }
+                    BtfType::Union(Union {
+                        enum64_fallback: Some(fallback),
+                        ..
+                    }) => {
+                        let index = index()?;
+                        (
+                            fallback.variants.len(),
+                            fallback.variants.get(index).map(|v| v.name_offset),
+                            index,
+                        )
+                    }
                     _ => {
                         return Err(RelocationError::InvalidRelocationKindForType {
                             relocation_number: relocation.number,
@@ -861,7 +896,7 @@ struct ComputedRelocation {
     target: Option<ComputedRelocationValue>,
 }
 
-#[derive(Debug)]
+#[derive(Clone, Copy, Debug)]
 struct ComputedRelocationValue {
     value: u64,
     size: u32,
@@ -1059,6 +1094,17 @@ impl ComputedRelocation {
                         let variant = &en.variants[accessor.index];
                         (u64::from(variant.value_high) << 32) | u64::from(variant.value_low)
                     }
+                    BtfType::Union(Union {
+                        enum64_fallback: Some(fallback),
+                        ..
+                    }) => {
+                        let variant = &fallback.variants[accessor.index];
+                        if fallback.signed {
+                            (variant.value as i64) as u64
+                        } else {
+                            variant.value
+                        }
+                    }
                     // candidate selection ensures that rel_kind == local_kind == target_kind
                     _ => unreachable!(),
                 }

aya-obj/src/btf/types.rs

@@ -473,7 +473,7 @@ impl Enum {
 }
 
 #[repr(C)]
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub struct BtfEnum64 {
     pub(crate) name_offset: u32,
     pub(crate) value_low: u32,
@@ -651,6 +651,22 @@ impl Struct {
     }
 }
 
+/// Snapshot of a single `ENUM64` variant so we can recover its 64-bit constant
+/// after the type is rewritten into a UNION.
+#[derive(Clone, Debug)]
+pub(crate) struct Enum64VariantFallback {
+    pub(crate) name_offset: u32,
+    pub(crate) value: u64,
+}
+
+/// Aggregate of the metadata we need to faithfully reconstruct a downgraded
+/// `ENUM64` during CO-RE relocation.
+#[derive(Clone, Debug)]
+pub(crate) struct Enum64Fallback {
+    pub(crate) signed: bool,
+    pub(crate) variants: Vec<Enum64VariantFallback>,
+}
+
 #[repr(C)]
 #[derive(Clone, Debug)]
 pub struct Union {
@@ -658,10 +674,16 @@ pub struct Union {
     info: u32,
     pub(crate) size: u32,
     pub(crate) members: Vec<BtfMember>,
+    pub(crate) enum64_fallback: Option<Enum64Fallback>,
 }
 
 impl Union {
-    pub(crate) fn new(name_offset: u32, size: u32, members: Vec<BtfMember>) -> Self {
+    pub(crate) fn new(
+        name_offset: u32,
+        size: u32,
+        members: Vec<BtfMember>,
+        enum64_fallback: Option<Enum64Fallback>,
+    ) -> Self {
         let mut info = (BtfKind::Union as u32) << 24;
         info |= (members.len() as u32) & 0xFFFF;
         Self {
@@ -669,6 +691,7 @@ impl Union {
             info,
             size,
             members,
+            enum64_fallback,
         }
     }
 
@@ -678,6 +701,7 @@ impl Union {
             info,
             size,
             members,
+            enum64_fallback: _,
         } = self;
         [
             bytes_of::<u32>(name_offset),
@@ -1224,6 +1248,7 @@ impl BtfType {
                 info: ty[1],
                 size: ty[2],
                 members: unsafe { read_array::<BtfMember>(data, vlen)? },
+                enum64_fallback: None,
             }),
             BtfKind::FuncProto => Self::FuncProto(FuncProto {
                 name_offset: ty[0],
@@ -1566,8 +1591,12 @@ pub(crate) fn fields_are_compatible(
 
 fn bytes_of<T>(val: &T) -> &[u8] {
     // Safety: all btf types are POD
+    //
+    // TODO: This is a fragile assumption and we should stop doing this. We should also remove
+    // repr(C) from our types, it doesn't make sense to rely on this.
     unsafe { crate::util::bytes_of(val) }
 }
+
 #[cfg(test)]
 mod tests {
     use assert_matches::assert_matches;
@@ -1635,7 +1664,7 @@ mod tests {
             btf_type: 0x68,
             offset: 0,
         }];
-        let bpf_type = BtfType::Union(Union::new(0, 4, members));
+        let bpf_type = BtfType::Union(Union::new(0, 4, members, None));
         let data: &[u8] = &bpf_type.to_bytes();
         assert_matches!(unsafe { BtfType::read(data, endianness) }.unwrap(), BtfType::Union(got) => {
             assert_eq!(got.to_bytes(), data);

xtask/public-api/aya-obj.txt

@@ -442,6 +442,7 @@ 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
 pub fn aya_obj::btf::BtfEnum64::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+impl core::marker::Copy for aya_obj::btf::BtfEnum64
 impl core::marker::Freeze for aya_obj::btf::BtfEnum64
 impl core::marker::Send for aya_obj::btf::BtfEnum64
 impl core::marker::Sync for aya_obj::btf::BtfEnum64