Skip to main content

ghostscope_compiler/ebpf/
maps.rs

1use aya_ebpf_bindings::bindings::bpf_map_type;
2use inkwell::context::Context;
3use inkwell::debug_info::{AsDIScope, DebugInfoBuilder};
4use inkwell::module::Linkage;
5use inkwell::module::Module;
6use inkwell::values::PointerValue;
7use inkwell::AddressSpace;
8// AddressSpace was used for pointer-typed BTF encodings; no longer needed after int-field BTF
9use std::collections::{HashMap, HashSet};
10use tracing::{error, info};
11
12#[derive(Debug, Clone, Copy)]
13pub enum BpfMapType {
14    Ringbuf,
15    Array,
16    PerCpuArray,
17    Hash,
18    PerfEventArray,
19    ProgramArray,
20}
21
22impl BpfMapType {
23    fn to_aya_map_type(self) -> u32 {
24        match self {
25            BpfMapType::Ringbuf => bpf_map_type::BPF_MAP_TYPE_RINGBUF,
26            BpfMapType::PerCpuArray => bpf_map_type::BPF_MAP_TYPE_PERCPU_ARRAY,
27            BpfMapType::Array => bpf_map_type::BPF_MAP_TYPE_ARRAY,
28            BpfMapType::Hash => bpf_map_type::BPF_MAP_TYPE_HASH,
29            BpfMapType::PerfEventArray => bpf_map_type::BPF_MAP_TYPE_PERF_EVENT_ARRAY,
30            BpfMapType::ProgramArray => bpf_map_type::BPF_MAP_TYPE_PROG_ARRAY,
31        }
32    }
33}
34
35#[derive(Debug, Clone)]
36pub struct SizedType {
37    pub size: u64, // size in bits
38    pub is_none: bool,
39}
40
41impl SizedType {
42    pub fn none() -> Self {
43        SizedType {
44            size: 0,
45            is_none: true,
46        }
47    }
48
49    pub fn integer(size: u64) -> Self {
50        SizedType {
51            size,
52            is_none: false,
53        }
54    }
55}
56
57pub struct MapManager<'ctx> {
58    context: &'ctx Context,
59    map_types: HashMap<String, BpfMapType>,
60    pinned_maps: HashSet<String>,
61}
62
63#[derive(Debug, thiserror::Error)]
64pub enum MapError {
65    #[error("Map not found: {0}")]
66    MapNotFound(String),
67
68    #[error("Builder error: {0}")]
69    Builder(String),
70
71    #[error("Debug info error: {0}")]
72    DebugInfo(String),
73}
74
75impl From<&str> for MapError {
76    fn from(err: &str) -> Self {
77        MapError::DebugInfo(err.to_string())
78    }
79}
80
81pub type Result<T> = std::result::Result<T, MapError>;
82
83impl<'ctx> MapManager<'ctx> {
84    pub fn new(context: &'ctx Context) -> Self {
85        MapManager {
86            context,
87            map_types: HashMap::new(),
88            pinned_maps: HashSet::new(),
89        }
90    }
91
92    fn map_is_pinned_by_name(name: &str) -> bool {
93        matches!(
94            name,
95            "proc_module_offsets" | "pid_aliases" | "proc_module_range_meta" | "proc_module_ranges"
96        )
97    }
98
99    pub fn mark_pinned_map(&mut self, name: &str) {
100        self.pinned_maps.insert(name.to_string());
101    }
102
103    fn map_is_pinned(&self, name: &str) -> bool {
104        Self::map_is_pinned_by_name(name) || self.pinned_maps.contains(name)
105    }
106
107    fn map_definition_field_count_for(&self, name: &str, map_type: BpfMapType) -> usize {
108        match map_type {
109            BpfMapType::Ringbuf => 2,
110            _ if self.map_is_pinned(name) => 5,
111            _ => 4,
112        }
113    }
114
115    #[cfg(test)]
116    fn map_definition_field_count(name: &str, map_type: BpfMapType) -> usize {
117        match map_type {
118            BpfMapType::Ringbuf => 2,
119            _ if Self::map_is_pinned_by_name(name) => 5,
120            _ => 4,
121        }
122    }
123
124    #[allow(clippy::too_many_arguments)]
125    pub fn create_map_definition(
126        &mut self,
127        module: &Module<'ctx>,
128        di_builder: &DebugInfoBuilder<'ctx>,
129        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
130        name: &str,
131        map_type: BpfMapType,
132        max_entries: u64,
133        key_type: SizedType,
134        value_type: SizedType,
135    ) -> Result<()> {
136        info!(
137            "Creating map definition: {} (type: {:?}, max_entries: {}, key_type: {:?}, value_type: {:?})",
138            name, map_type, max_entries, key_type, value_type
139        );
140
141        // Store map type information
142        self.map_types.insert(name.to_string(), map_type);
143
144        // Use the original map name directly (like "ringbuf")
145        let var_name = name.to_string();
146        info!("Map variable name: {}", var_name);
147
148        // Create BPF map definition structure that aya expects
149        // Match clang-style: fields are pointers (64-bit); actual values are
150        // encoded via BTF pointer-to-array lengths, not via initializers.
151        let ptr_ty = self.context.ptr_type(inkwell::AddressSpace::default());
152
153        // Values are conveyed by BTF; initializers can be null pointers.
154
155        // Keep the concrete map variable layout in sync with the BTF map
156        // definition. Pinned maps include the optional `pinning` field.
157        let field_count = self.map_definition_field_count_for(&var_name, map_type);
158        let elements: Vec<_> = (0..field_count).map(|_| ptr_ty.into()).collect();
159        let initializer_values: Vec<_> = (0..field_count)
160            .map(|_| ptr_ty.const_null().into())
161            .collect();
162        let struct_type = self.context.struct_type(&elements, false);
163        let initializer = struct_type.const_named_struct(&initializer_values);
164
165        // Create BTF type information for the map
166        // This is critical for aya to understand the map structure
167        let map_di_type = self.create_map_btf_info(
168            di_builder,
169            compile_unit,
170            &var_name,
171            map_type,
172            max_entries,
173            key_type,
174            value_type,
175        )?;
176
177        // Create the global variable
178        let map_var = module.add_global(struct_type, None, &var_name);
179
180        // Set the proper initializer
181        map_var.set_initializer(&initializer);
182
183        // Set section to .maps
184        map_var.set_section(Some(".maps"));
185
186        // Set linkage to External so aya can find and relocate the map symbols
187        // Private linkage hides symbols from relocations, which breaks aya
188        map_var.set_linkage(Linkage::External);
189
190        // Associate the global variable with its debug type
191        // This ensures the BTF type information is properly linked
192        let file = compile_unit.get_file();
193        let di_global_variable = di_builder.create_global_variable_expression(
194            compile_unit.as_debug_info_scope(), // scope
195            &var_name,                          // name
196            &var_name,                          // linkage_name
197            file,                               // file
198            1,                                  // line_no
199            map_di_type,                        // ty
200            false,                              // is_local_to_unit
201            None,                               // expr
202            None,                               // decl
203            map_var.get_alignment(),            // align_in_bits
204        );
205
206        // Attach the debug info to the global variable using proper metadata API
207        // The kind_id for "dbg" in LLVM is typically 0
208        map_var.set_metadata(di_global_variable.as_metadata_value(self.context), 0);
209
210        info!(
211            "Successfully created map: {} with {} fields",
212            var_name, field_count
213        );
214        Ok(())
215    }
216
217    pub fn get_map(&self, module: &Module<'ctx>, name: &str) -> Result<PointerValue<'ctx>> {
218        let var_name = name.to_string(); // Use direct name like "ringbuf"
219        info!("Looking up map: {}", var_name);
220
221        if let Some(map_var) = module.get_global(&var_name) {
222            info!("Found map: {}", var_name);
223            Ok(map_var.as_pointer_value())
224        } else {
225            error!("Map not found: {}", var_name);
226            Err(MapError::MapNotFound(var_name))
227        }
228    }
229
230    pub fn create_ringbuf_map(
231        &mut self,
232        module: &Module<'ctx>,
233        di_builder: &DebugInfoBuilder<'ctx>,
234        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
235        name: &str,
236        ringbuf_size: u64,
237    ) -> Result<()> {
238        // For ringbuf, max_entries is the buffer size in bytes (must be power of 2)
239        // The parameter name is kept as perf_rb_pages for backward compatibility,
240        // but we now interpret it directly as the ringbuf size in bytes
241        let max_entries = ringbuf_size;
242        info!("Creating ringbuf map: {} with {} bytes", name, max_entries);
243        self.create_map_definition(
244            module,
245            di_builder,
246            compile_unit,
247            name,
248            BpfMapType::Ringbuf,
249            max_entries,
250            // Ringbuf map: key_size = 0, value_size = 0 for ringbuf
251            SizedType::none(),
252            SizedType::none(),
253        )
254    }
255
256    /// Create PerfEventArray map for event output (fallback when RingBuf not supported)
257    pub fn create_perf_event_array_map(
258        &mut self,
259        module: &Module<'ctx>,
260        di_builder: &DebugInfoBuilder<'ctx>,
261        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
262        name: &str,
263    ) -> Result<()> {
264        info!("Creating PerfEventArray map: {}", name);
265        self.create_map_definition(
266            module,
267            di_builder,
268            compile_unit,
269            name,
270            BpfMapType::PerfEventArray,
271            0, // max_entries = 0 means auto-detect number of CPUs
272            // PerfEventArray: key = u32 (CPU index), value = u32 (FD)
273            SizedType::integer(32),
274            SizedType::integer(32),
275        )
276    }
277
278    /// Create the per-(pid,module) section offsets map used for ASLR address calculation
279    pub fn create_proc_module_offsets_map(
280        &mut self,
281        module: &Module<'ctx>,
282        di_builder: &DebugInfoBuilder<'ctx>,
283        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
284        name: &str,
285        max_entries: u64,
286    ) -> Result<()> {
287        // Key: {pid:u32, pad:u32, cookie:u64} => 16 bytes => 128 bits
288        // Value: {text, rodata, data, bss, base, size: u64} => 48 bytes => 384 bits
289        self.create_map_definition(
290            module,
291            di_builder,
292            compile_unit,
293            name,
294            BpfMapType::Hash,
295            max_entries,
296            SizedType::integer(128),
297            SizedType::integer(384),
298        )
299    }
300
301    pub fn create_pid_aliases_map(
302        &mut self,
303        module: &Module<'ctx>,
304        di_builder: &DebugInfoBuilder<'ctx>,
305        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
306        name: &str,
307        max_entries: u64,
308    ) -> Result<()> {
309        self.create_map_definition(
310            module,
311            di_builder,
312            compile_unit,
313            name,
314            BpfMapType::Hash,
315            max_entries,
316            SizedType::integer(32),
317            SizedType::integer(32),
318        )
319    }
320
321    pub fn create_proc_module_range_meta_map(
322        &mut self,
323        module: &Module<'ctx>,
324        di_builder: &DebugInfoBuilder<'ctx>,
325        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
326        name: &str,
327        max_entries: u64,
328    ) -> Result<()> {
329        self.create_map_definition(
330            module,
331            di_builder,
332            compile_unit,
333            name,
334            BpfMapType::Hash,
335            max_entries,
336            SizedType::integer(32),
337            SizedType::integer(ghostscope_protocol::PROC_MODULE_RANGE_META_SIZE as u64 * 8),
338        )
339    }
340
341    pub fn create_proc_module_ranges_map(
342        &mut self,
343        module: &Module<'ctx>,
344        di_builder: &DebugInfoBuilder<'ctx>,
345        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
346        name: &str,
347        max_entries: u64,
348    ) -> Result<()> {
349        self.create_map_definition(
350            module,
351            di_builder,
352            compile_unit,
353            name,
354            BpfMapType::Hash,
355            max_entries,
356            SizedType::integer(ghostscope_protocol::PROC_MODULE_RANGE_KEY_SIZE as u64 * 8),
357            SizedType::integer(ghostscope_protocol::PROC_MODULE_RANGE_VALUE_SIZE as u64 * 8),
358        )
359    }
360
361    pub fn create_event_loss_counter_map(
362        &mut self,
363        module: &Module<'ctx>,
364        di_builder: &DebugInfoBuilder<'ctx>,
365        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
366        name: &str,
367        max_entries: u64,
368    ) -> Result<()> {
369        info!(
370            "Creating event loss counter map: {} with {} max entries",
371            name, max_entries
372        );
373        self.create_map_definition(
374            module,
375            di_builder,
376            compile_unit,
377            name,
378            BpfMapType::PerCpuArray,
379            max_entries,
380            SizedType::integer(32),
381            SizedType::integer(64),
382        )
383    }
384
385    /// Create BTF type information for a BPF map matching clang's output format
386    /// This allows aya to understand the map's key and value types
387    #[allow(clippy::too_many_arguments)]
388    fn create_map_btf_info(
389        &self,
390        di_builder: &DebugInfoBuilder<'ctx>,
391        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
392        map_name: &str,
393        map_type: BpfMapType,
394        max_entries: u64,
395        key_type: SizedType,
396        value_type: SizedType,
397    ) -> Result<inkwell::debug_info::DIType<'ctx>> {
398        info!(
399            "Creating BTF info for map: {} (type: {:?})",
400            map_name, map_type
401        );
402
403        // Create basic types needed for the map structure
404        let i32_type = di_builder.create_basic_type("int", 32, 0x05, 0)?; // DW_ATE_signed = 0x05
405
406        let file = compile_unit.get_file();
407        let scope = compile_unit.as_debug_info_scope();
408
409        // Create the map structure based on map type, matching clang/aya BTF format:
410        // fields are pointers to arrays whose nr_elems encode values.
411        let map_type_id = map_type.to_aya_map_type();
412
413        // Helper: pointer to array with given element count (encoded in range)
414        let mk_ptr_to_array = |name: &str, nr_elems: i64| {
415            let range = 0..nr_elems;
416            let arr = di_builder.create_array_type(
417                i32_type.as_type(),
418                64,
419                32,
420                std::slice::from_ref(&range),
421            );
422            di_builder.create_pointer_type(name, arr.as_type(), 64, 64, AddressSpace::default())
423        };
424
425        let type_ptr = mk_ptr_to_array("type", map_type_id as i64);
426
427        let members = match map_type {
428            BpfMapType::Ringbuf => {
429                info!("Creating ringbuf BTF with 2 fields (type, max_entries) as pointer-to-array");
430                let max_entries_ptr = mk_ptr_to_array("max_entries", max_entries as i64);
431                vec![
432                    di_builder.create_member_type(
433                        scope,
434                        "type",
435                        file,
436                        0,
437                        64,
438                        64,
439                        0,
440                        0,
441                        type_ptr.as_type(),
442                    ),
443                    di_builder.create_member_type(
444                        scope,
445                        "max_entries",
446                        file,
447                        0,
448                        64,
449                        64,
450                        64,
451                        0,
452                        max_entries_ptr.as_type(),
453                    ),
454                ]
455            }
456            _ => {
457                info!("Creating array/hash BTF with pointer-to-array fields for aya compatibility");
458                let key_size_val = if key_type.is_none {
459                    0
460                } else {
461                    (key_type.size / 8) as i64
462                };
463                let value_size_val = if value_type.is_none {
464                    0
465                } else {
466                    (value_type.size / 8) as i64
467                };
468                let key_size_ptr = mk_ptr_to_array("key_size", key_size_val);
469                let value_size_ptr = mk_ptr_to_array("value_size", value_size_val);
470                let max_entries_ptr = mk_ptr_to_array("max_entries", max_entries as i64);
471                let mut v = vec![
472                    di_builder.create_member_type(
473                        scope,
474                        "type",
475                        file,
476                        0,
477                        64,
478                        64,
479                        0,
480                        0,
481                        type_ptr.as_type(),
482                    ),
483                    di_builder.create_member_type(
484                        scope,
485                        "key_size",
486                        file,
487                        0,
488                        64,
489                        64,
490                        64,
491                        0,
492                        key_size_ptr.as_type(),
493                    ),
494                    di_builder.create_member_type(
495                        scope,
496                        "value_size",
497                        file,
498                        0,
499                        64,
500                        64,
501                        128,
502                        0,
503                        value_size_ptr.as_type(),
504                    ),
505                    di_builder.create_member_type(
506                        scope,
507                        "max_entries",
508                        file,
509                        0,
510                        64,
511                        64,
512                        192,
513                        0,
514                        max_entries_ptr.as_type(),
515                    ),
516                ];
517                // For pinned maps, include optional 'pinning' to signal Aya ByName pinning.
518                if self.map_is_pinned(map_name) {
519                    // ByName is typically encoded as 1 in aya_obj::maps::PinningType
520                    let pinning_ptr = mk_ptr_to_array("pinning", 1);
521                    v.push(di_builder.create_member_type(
522                        scope,
523                        "pinning",
524                        file,
525                        0,
526                        64,
527                        64,
528                        256,
529                        0,
530                        pinning_ptr.as_type(),
531                    ));
532                }
533                v
534            }
535        };
536
537        // Convert members to DIType vector
538        let member_types: Vec<_> = members.iter().map(|m| m.as_type()).collect();
539
540        // Total structure size: pointers (64-bit) per field.
541        let field_count = self.map_definition_field_count_for(map_name, map_type);
542        let total_size_bits = (field_count as u64) * 64;
543
544        // Create the map structure type (anonymous like reference)
545        let map_struct_type = di_builder.create_struct_type(
546            scope,           // scope
547            "",              // name - empty for anonymous struct
548            file,            // file
549            0,               // line_number
550            total_size_bits, // size_in_bits
551            32,              // align_in_bits
552            0,               // flags
553            None,            // derived_from
554            &member_types,   // elements
555            0,               // runtime_lang
556            None,            // vtable_holder
557            "",              // unique_id
558        );
559
560        info!(
561            "Created BTF struct type for map: {} with {} fields, {} total bits",
562            map_name, field_count, total_size_bits
563        );
564        Ok(map_struct_type.as_type())
565    }
566
567    /// Get ringbuf map by name
568    pub fn get_ringbuf_map(&self, module: &Module<'ctx>, name: &str) -> Result<PointerValue<'ctx>> {
569        self.get_map(module, name)
570    }
571
572    /// Get a perf event array map by name
573    pub fn get_perf_map(&self, module: &Module<'ctx>, name: &str) -> Result<PointerValue<'ctx>> {
574        self.get_map(module, name)
575    }
576
577    /// Create a Per-CPU Array map (key=u32, value arbitrary size)
578    pub fn create_percpu_array_map(
579        &mut self,
580        module: &Module<'ctx>,
581        di_builder: &DebugInfoBuilder<'ctx>,
582        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
583        name: &str,
584        max_entries: u64,
585        value_size_bytes: u64,
586    ) -> Result<()> {
587        self.create_map_definition(
588            module,
589            di_builder,
590            compile_unit,
591            name,
592            BpfMapType::PerCpuArray,
593            max_entries,
594            SizedType::integer(32),
595            SizedType::integer(value_size_bytes * 8),
596        )
597    }
598
599    /// Create a regular Array map (key=u32, value arbitrary size).
600    pub fn create_array_map(
601        &mut self,
602        module: &Module<'ctx>,
603        di_builder: &DebugInfoBuilder<'ctx>,
604        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
605        name: &str,
606        max_entries: u64,
607        value_size_bytes: u64,
608    ) -> Result<()> {
609        self.create_map_definition(
610            module,
611            di_builder,
612            compile_unit,
613            name,
614            BpfMapType::Array,
615            max_entries,
616            SizedType::integer(32),
617            SizedType::integer(value_size_bytes * 8),
618        )
619    }
620
621    /// Create a regular Hash map with caller-specified key/value sizes.
622    pub fn create_hash_map(
623        &mut self,
624        module: &Module<'ctx>,
625        di_builder: &DebugInfoBuilder<'ctx>,
626        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
627        name: &str,
628        max_entries: u64,
629        key_value_size_bytes: (u64, u64),
630    ) -> Result<()> {
631        let (key_size_bytes, value_size_bytes) = key_value_size_bytes;
632        self.create_map_definition(
633            module,
634            di_builder,
635            compile_unit,
636            name,
637            BpfMapType::Hash,
638            max_entries,
639            SizedType::integer(key_size_bytes * 8),
640            SizedType::integer(value_size_bytes * 8),
641        )
642    }
643
644    /// Create a ProgramArray map for eBPF tail calls.
645    pub fn create_program_array_map(
646        &mut self,
647        module: &Module<'ctx>,
648        di_builder: &DebugInfoBuilder<'ctx>,
649        compile_unit: &inkwell::debug_info::DICompileUnit<'ctx>,
650        name: &str,
651        max_entries: u64,
652    ) -> Result<()> {
653        self.create_map_definition(
654            module,
655            di_builder,
656            compile_unit,
657            name,
658            BpfMapType::ProgramArray,
659            max_entries,
660            SizedType::integer(32),
661            SizedType::integer(32),
662        )
663    }
664}
665
666#[cfg(test)]
667mod tests {
668    use super::{BpfMapType, MapManager};
669
670    #[test]
671    fn pinned_maps_include_pinning_field_in_concrete_layout() {
672        assert_eq!(
673            MapManager::map_definition_field_count("proc_module_offsets", BpfMapType::Hash),
674            5
675        );
676        assert_eq!(
677            MapManager::map_definition_field_count("pid_aliases", BpfMapType::Hash),
678            5
679        );
680        assert_eq!(
681            MapManager::map_definition_field_count("event_accum_buffer", BpfMapType::PerCpuArray),
682            4
683        );
684        assert_eq!(
685            MapManager::map_definition_field_count("event_loss_counters", BpfMapType::PerCpuArray),
686            4
687        );
688        assert_eq!(
689            MapManager::map_definition_field_count("ringbuf", BpfMapType::Ringbuf),
690            2
691        );
692        assert_eq!(
693            MapManager::map_definition_field_count("bt_prog_array", BpfMapType::ProgramArray),
694            4
695        );
696    }
697}