Skip to main content

kbpf_basic/helper/
mod.rs

1//! Basic eBPF helper functions module.
2//!
3pub(crate) mod ringbuf;
4use alloc::{collections::btree_map::BTreeMap, string::String, vec::Vec};
5use core::{
6    ffi::{c_char, c_int, c_void},
7    fmt::Write,
8};
9
10use consts::BPF_F_CURRENT_CPU;
11use lock_api::RawMutex;
12
13use crate::{
14    BpfError, BpfResult as Result, KernelAuxiliaryOps,
15    map::{BpfCallBackFn, UnifiedMap},
16};
17
18pub mod consts;
19
20/// Type alias for a raw BPF helper function.
21pub type RawBPFHelperFn = fn(u64, u64, u64, u64, u64) -> u64;
22
23/// Transmute a function pointer to a RawBPFHelperFn.
24macro_rules! helper_func {
25    ($name:ident::<$($generic:ident),*>) => {
26        unsafe {
27            core::mem::transmute::<usize, RawBPFHelperFn>($name::<$($generic),*> as *const () as usize)
28        }
29    };
30    ($name:ident) => {
31        unsafe {
32            core::mem::transmute::<usize, RawBPFHelperFn>($name as *const () as usize)
33        }
34    };
35}
36
37use printf_compat::{format, output};
38
39/// Printf according to the format string, function will return the number of bytes written(including '\0')
40///
41/// # Safety
42/// The caller must ensure that the format string and arguments are valid.
43pub unsafe extern "C" fn printf(w: &mut impl Write, str: *const c_char, args: ...) -> c_int {
44    let bytes_written = unsafe { format(str as _, args, output::fmt_write(w)) };
45    bytes_written + 1
46}
47
48fn extract_format_specifiers(format_str: &str) -> usize {
49    // let mut result = Vec::new();
50    let mut fmt_arg_count = 0;
51    let chars: Vec<char> = format_str.chars().collect();
52    let mut i = 0;
53
54    while i < chars.len() {
55        if chars[i] == '%' {
56            if i + 1 < chars.len() && chars[i + 1] == '%' {
57                // Skip literal %%
58                i += 2;
59            } else {
60                let start = i;
61                i += 1;
62
63                // Parse optional flags
64                while i < chars.len() && "-+#0 .0123456789lhL*".contains(chars[i]) {
65                    i += 1;
66                }
67
68                // Parse type specifier (a single letter)
69                if i < chars.len() && "cdieEfFgGosuxXpn".contains(chars[i]) {
70                    i += 1;
71                    let _spec: String = chars[start..i].iter().collect();
72                    // result.push(spec);
73                    fmt_arg_count += 1; // Count this format specifier
74                }
75            }
76        } else {
77            i += 1;
78        }
79    }
80
81    fmt_arg_count
82}
83
84/// See <https://ebpf-docs.dylanreimerink.nl/linux/helper-function/bpf_trace_printk/>
85///
86/// ## Warning
87/// The `arg3`, `arg4`, and `arg5` parameters are pointers to the stack, so we first read the value
88/// they point to before passing them to the `printf` function.
89pub fn trace_printf<F: KernelAuxiliaryOps>(
90    fmt_ptr: u64,
91    fmt_len: u64,
92    arg3: u64,
93    arg4: u64,
94    arg5: u64,
95) -> i64 {
96    struct FakeWriter<F: KernelAuxiliaryOps> {
97        _phantom: core::marker::PhantomData<F>,
98    }
99    impl<F: KernelAuxiliaryOps> FakeWriter<F> {
100        fn default() -> Self {
101            FakeWriter {
102                _phantom: core::marker::PhantomData,
103            }
104        }
105    }
106    impl<F: KernelAuxiliaryOps> Write for FakeWriter<F> {
107        fn write_str(&mut self, s: &str) -> core::fmt::Result {
108            F::ebpf_write_str(s).map_err(|_| core::fmt::Error)?;
109            Ok(())
110        }
111    }
112
113    let fmt_str = unsafe {
114        core::str::from_utf8_unchecked(core::slice::from_raw_parts(
115            fmt_ptr as *const u8,
116            fmt_len as usize,
117        ))
118    };
119    let fmt_arg_count = extract_format_specifiers(fmt_str);
120
121    let (arg3, arg4, arg5) = match fmt_arg_count {
122        0 => (0, 0, 0),
123        1 => (unsafe { (arg3 as *const u64).read() }, 0, 0),
124        2 => (
125            unsafe { (arg3 as *const u64).read() },
126            unsafe { (arg4 as *const u64).read() },
127            0,
128        ),
129        3 => (
130            unsafe { (arg3 as *const u64).read() },
131            unsafe { (arg4 as *const u64).read() },
132            unsafe { (arg5 as *const u64).read() },
133        ),
134        _ => {
135            log::error!("trace_printf: too many arguments, only 3 are supported");
136            return -1;
137        }
138    };
139
140    let mut fmt = FakeWriter::<F>::default();
141    unsafe { printf(&mut fmt, fmt_ptr as _, arg3, arg4, arg5) as _ }
142}
143
144/// See <https://ebpf-docs.dylanreimerink.nl/linux/helper-function/bpf_map_lookup_elem/>
145pub fn raw_map_lookup_elem<F: KernelAuxiliaryOps>(
146    map: *mut c_void,
147    key: *const c_void,
148) -> *const c_void {
149    let res = F::get_unified_map_from_ptr(map as *const u8, |unified_map| {
150        let meta = unified_map.map_meta();
151        let key_size = meta.key_size as usize;
152        let key = unsafe { core::slice::from_raw_parts(key as *const u8, key_size) };
153        let value = map_lookup_elem::<F>(unified_map, key)?;
154        Ok(value)
155    });
156    match res {
157        Ok(Some(value)) => value as _,
158        _ => core::ptr::null(),
159    }
160}
161
162/// Lookup an element in map.
163pub fn map_lookup_elem<F: KernelAuxiliaryOps>(
164    unified_map: &UnifiedMap<F::MapLock>,
165    key: &[u8],
166) -> Result<Option<*const u8>> {
167    let mut map = unified_map.map_mut();
168    let value = map.lookup_elem(key);
169    match value {
170        Ok(Some(value)) => Ok(Some(value.as_ptr())),
171        _ => Ok(None),
172    }
173}
174
175/// See https://ebpf-docs.dylanreimerink.nl/linux/helper-function/bpf_perf_event_output/
176///
177/// See https://man7.org/linux/man-pages/man7/bpf-helpers.7.html
178pub fn raw_perf_event_output<F: KernelAuxiliaryOps>(
179    ctx: *mut c_void,
180    map: *mut c_void,
181    flags: u64,
182    data: *mut c_void,
183    size: u64,
184) -> i64 {
185    let res = F::get_unified_map_from_ptr(map as *const u8, |unified_map| {
186        let data = unsafe { core::slice::from_raw_parts(data as *const u8, size as usize) };
187        perf_event_output::<F>(ctx, unified_map, flags, data)
188    });
189
190    match res {
191        Ok(_) => 0,
192        Err(e) => e as _,
193    }
194}
195
196/// Output data to a perf event.
197pub fn perf_event_output<F: KernelAuxiliaryOps>(
198    ctx: *mut c_void,
199    unified_map: &UnifiedMap<F::MapLock>,
200    flags: u64,
201    data: &[u8],
202) -> Result<()> {
203    let index = flags as u32;
204    let flags = (flags >> 32) as u32;
205    let key = if index == BPF_F_CURRENT_CPU as u32 {
206        F::current_cpu_id()
207    } else {
208        index
209    };
210    let mut map = unified_map.map_mut();
211    let fd = map
212        .lookup_elem(&key.to_ne_bytes())?
213        .ok_or(BpfError::ENOENT)?;
214    let fd = u32::from_ne_bytes(fd.try_into().map_err(|_| BpfError::EINVAL)?);
215    F::perf_event_output(ctx, fd, flags, data)?;
216    Ok(())
217}
218
219/// See https://ebpf-docs.dylanreimerink.nl/linux/helper-function/bpf_probe_read/
220fn raw_bpf_probe_read(dst: *mut c_void, size: u32, unsafe_ptr: *const c_void) -> i64 {
221    let (dst, src) = unsafe {
222        let dst = core::slice::from_raw_parts_mut(dst as *mut u8, size as usize);
223        let src = core::slice::from_raw_parts(unsafe_ptr as *const u8, size as usize);
224        (dst, src)
225    };
226    let res = bpf_probe_read(dst, src);
227    match res {
228        Ok(_) => 0,
229        Err(e) => e as _,
230    }
231}
232
233/// For tracing programs, safely attempt to read size
234/// bytes from kernel space address unsafe_ptr and
235/// store the data in dst.
236pub fn bpf_probe_read(dst: &mut [u8], src: &[u8]) -> Result<()> {
237    dst.copy_from_slice(src);
238    Ok(())
239}
240
241/// Update entry with key in map.
242///
243/// See <https://docs.ebpf.io/linux/helper-function/bpf_map_update_elem/>
244pub fn raw_map_update_elem<F: KernelAuxiliaryOps>(
245    map: *mut c_void,
246    key: *const c_void,
247    value: *const c_void,
248    flags: u64,
249) -> i64 {
250    let res = F::get_unified_map_from_ptr(map as *const u8, |unified_map| {
251        let meta = unified_map.map_meta();
252        let key_size = meta.key_size as usize;
253        let value_size = meta.value_size as usize;
254        let key = unsafe { core::slice::from_raw_parts(key as *const u8, key_size) };
255        let value = unsafe { core::slice::from_raw_parts(value as *const u8, value_size) };
256        map_update_elem::<F>(unified_map, key, value, flags)
257    });
258    match res {
259        Ok(_) => 0,
260        Err(e) => e as _,
261    }
262}
263
264/// Update entry with key in map.
265pub fn map_update_elem<F: KernelAuxiliaryOps>(
266    unified_map: &UnifiedMap<F::MapLock>,
267    key: &[u8],
268    value: &[u8],
269    flags: u64,
270) -> Result<()> {
271    let mut map = unified_map.map_mut();
272
273    map.update_elem(key, value, flags)
274}
275
276/// Delete entry with key from map.
277///
278/// The delete map element helper call is used to delete values from maps.
279pub fn raw_map_delete_elem<F: KernelAuxiliaryOps>(map: *mut c_void, key: *const c_void) -> i64 {
280    let res = F::get_unified_map_from_ptr(map as *const u8, |unified_map| {
281        let meta = unified_map.map_meta();
282        let key_size = meta.key_size as usize;
283        let key = unsafe { core::slice::from_raw_parts(key as *const u8, key_size) };
284        map_delete_elem::<F>(unified_map, key)
285    });
286    match res {
287        Ok(_) => 0,
288        Err(e) => e as _,
289    }
290}
291
292/// Delete entry with key from map.
293pub fn map_delete_elem<F: KernelAuxiliaryOps>(
294    unified_map: &UnifiedMap<F::MapLock>,
295    key: &[u8],
296) -> Result<()> {
297    let mut map = unified_map.map_mut();
298
299    map.delete_elem(key)
300}
301
302/// For each element in map, call callback_fn function with map, callback_ctx and other map-specific
303/// parameters. The callback_fn should be a static function and the callback_ctx should be a pointer
304/// to the stack. The flags is used to control certain aspects of the helper.  Currently, the flags must
305/// be 0.
306///
307/// The following are a list of supported map types and their respective expected callback signatures:
308/// - BPF_MAP_TYPE_HASH
309/// - BPF_MAP_TYPE_PERCPU_HASH
310/// - BPF_MAP_TYPE_LRU_HASH
311/// - BPF_MAP_TYPE_LRU_PERCPU_HASH
312/// - BPF_MAP_TYPE_ARRAY
313/// - BPF_MAP_TYPE_PERCPU_ARRAY
314///
315/// `long (*callback_fn)(struct bpf_map *map, const void key, void *value, void *ctx);`
316///
317/// For per_cpu maps, the map_value is the value on the cpu where the bpf_prog is running.
318pub fn raw_map_for_each_elem<F: KernelAuxiliaryOps>(
319    map: *mut c_void,
320    cb: *const c_void,
321    ctx: *const c_void,
322    flags: u64,
323) -> i64 {
324    if cb.is_null() {
325        return BpfError::EINVAL as _;
326    }
327    let res = F::get_unified_map_from_ptr(map as *const u8, |unified_map| {
328        let cb = unsafe { *(cb as *const BpfCallBackFn) };
329        map_for_each_elem::<F>(unified_map, cb, ctx as _, flags)
330    });
331    match res {
332        Ok(v) => v as i64,
333        Err(e) => e as _,
334    }
335}
336
337/// Do some action for each element in map.
338pub fn map_for_each_elem<F: KernelAuxiliaryOps>(
339    unified_map: &UnifiedMap<F::MapLock>,
340    cb: BpfCallBackFn,
341    ctx: *const u8,
342    flags: u64,
343) -> Result<u32> {
344    let mut map = unified_map.map_mut();
345
346    map.for_each_elem(cb, ctx, flags)
347}
348
349/// Perform a lookup in percpu map for an entry associated to key on cpu.
350///
351/// See https://ebpf-docs.dylanreimerink.nl/linux/helper-function/bpf_map_lookup_percpu_elem/
352pub fn raw_map_lookup_percpu_elem<F: KernelAuxiliaryOps>(
353    map: *mut c_void,
354    key: *const c_void,
355    cpu: u32,
356) -> *const c_void {
357    let res = F::get_unified_map_from_ptr(map as *const u8, |unified_map| {
358        let meta = unified_map.map_meta();
359        let key_size = meta.key_size as usize;
360        let key = unsafe { core::slice::from_raw_parts(key as *const u8, key_size) };
361        map_lookup_percpu_elem::<F>(unified_map, key, cpu)
362    });
363    match res {
364        Ok(Some(value)) => value as *const c_void,
365        _ => core::ptr::null_mut(),
366    }
367}
368
369/// Lookup an element in percpu map.
370pub fn map_lookup_percpu_elem<F: KernelAuxiliaryOps>(
371    unified_map: &UnifiedMap<F::MapLock>,
372    key: &[u8],
373    cpu: u32,
374) -> Result<Option<*const u8>> {
375    let mut map = unified_map.map_mut();
376    let value = map.lookup_percpu_elem(key, cpu);
377    match value {
378        Ok(Some(value)) => Ok(Some(value.as_ptr())),
379        _ => Ok(None),
380    }
381}
382/// Push an element value in map.
383///
384/// See https://ebpf-docs.dylanreimerink.nl/linux/helper-function/bpf_map_push_elem/
385pub fn raw_map_push_elem<F: KernelAuxiliaryOps>(
386    map: *mut c_void,
387    value: *const c_void,
388    flags: u64,
389) -> i64 {
390    let res = F::get_unified_map_from_ptr(map as *const u8, |unified_map| {
391        let meta = unified_map.map_meta();
392        let value_size = meta.value_size as usize;
393        let value = unsafe { core::slice::from_raw_parts(value as *const u8, value_size) };
394        map_push_elem::<F>(unified_map, value, flags)
395    });
396    match res {
397        Ok(_) => 0,
398        Err(e) => e as _,
399    }
400}
401
402/// Push an element value in map.
403pub fn map_push_elem<F: KernelAuxiliaryOps>(
404    unified_map: &UnifiedMap<F::MapLock>,
405    value: &[u8],
406    flags: u64,
407) -> Result<()> {
408    let mut map = unified_map.map_mut();
409
410    map.push_elem(value, flags)
411}
412
413/// Pop an element from map.
414///
415/// See https://ebpf-docs.dylanreimerink.nl/linux/helper-function/bpf_map_pop_elem/
416pub fn raw_map_pop_elem<F: KernelAuxiliaryOps>(map: *mut c_void, value: *mut c_void) -> i64 {
417    let res = F::get_unified_map_from_ptr(map as *const u8, |unified_map| {
418        let meta = unified_map.map_meta();
419        let value_size = meta.value_size as usize;
420        let value = unsafe { core::slice::from_raw_parts_mut(value as *mut u8, value_size) };
421        map_pop_elem::<F>(unified_map, value)
422    });
423    match res {
424        Ok(_) => 0,
425        Err(e) => e as _,
426    }
427}
428
429/// Pop an element from map.
430pub fn map_pop_elem<F: KernelAuxiliaryOps>(
431    unified_map: &UnifiedMap<F::MapLock>,
432    value: &mut [u8],
433) -> Result<()> {
434    let mut map = unified_map.map_mut();
435
436    map.pop_elem(value)
437}
438
439/// Get an element from map without removing it.
440///
441/// See https://ebpf-docs.dylanreimerink.nl/linux/helper-function/bpf_map_peek_elem/
442pub fn raw_map_peek_elem<F: KernelAuxiliaryOps>(map: *mut c_void, value: *mut c_void) -> i64 {
443    let res = F::get_unified_map_from_ptr(map as *const u8, |unified_map| {
444        let meta = unified_map.map_meta();
445        let value_size = meta.value_size as usize;
446        let value = unsafe { core::slice::from_raw_parts_mut(value as *mut u8, value_size) };
447        map_peek_elem::<F>(unified_map, value)
448    });
449    match res {
450        Ok(_) => 0,
451        Err(e) => e as _,
452    }
453}
454
455/// Get an element from map without removing it.
456pub fn map_peek_elem<F: KernelAuxiliaryOps>(
457    unified_map: &UnifiedMap<F::MapLock>,
458    value: &mut [u8],
459) -> Result<()> {
460    let mut map = unified_map.map_mut();
461
462    map.peek_elem(value)
463}
464
465/// Get the current kernel time in nanoseconds.
466pub fn bpf_ktime_get_ns<F: KernelAuxiliaryOps>() -> u64 {
467    F::ebpf_time_ns().unwrap_or_default()
468}
469
470/// Copy a NULL terminated string from an unsafe user address unsafe_ptr to dst.
471/// The size should include the terminating NULL byte. In case the string length is smaller than size,
472/// the target is not padded with further NULL bytes. If the string length is larger than size,
473/// just size-1 bytes are copied and the last byte is set to NULL.
474///
475/// On success, the strictly positive length of the output string, including the trailing NULL character. On error, a negative value
476///
477/// See https://docs.ebpf.io/linux/helper-function/bpf_probe_read_user_str/
478fn raw_probe_read_user_str<F: KernelAuxiliaryOps>(
479    dst: *mut c_void,
480    size: u32,
481    unsafe_ptr: *const c_void,
482) -> i64 {
483    let dst = unsafe { core::slice::from_raw_parts_mut(dst as *mut u8, size as usize) };
484    let res = probe_read_user_str::<F>(dst, unsafe_ptr as *const u8);
485    match res {
486        Ok(len) => len as i64,
487        Err(e) => e as _,
488    }
489}
490
491/// Copy a NULL terminated string from an unsafe user address unsafe_ptr to dst.
492pub fn probe_read_user_str<F: KernelAuxiliaryOps>(dst: &mut [u8], src: *const u8) -> Result<usize> {
493    if dst.is_empty() {
494        return Err(BpfError::EINVAL);
495    }
496    let str = F::string_from_user_cstr(src)?;
497    let len = str.len();
498    let copy_len = len.min(dst.len() - 1); // Leave space for NULL terminator
499    dst[..copy_len].copy_from_slice(&str.as_bytes()[..copy_len]);
500    dst[copy_len] = 0; // Null-terminate the string
501    Ok(copy_len + 1) // Return length including NULL terminator
502}
503
504/// Initialize the helper functions map.
505pub fn init_helper_functions<F: KernelAuxiliaryOps>() -> BTreeMap<u32, RawBPFHelperFn> {
506    use consts::*;
507    let mut map = BTreeMap::new();
508
509    // Map helpers::Generic map helpers
510    map.insert(
511        HELPER_MAP_LOOKUP_ELEM,
512        helper_func!(raw_map_lookup_elem::<F>),
513    );
514    map.insert(
515        HELPER_MAP_UPDATE_ELEM,
516        helper_func!(raw_map_update_elem::<F>),
517    );
518    map.insert(
519        HELPER_MAP_DELETE_ELEM,
520        helper_func!(raw_map_delete_elem::<F>),
521    );
522    map.insert(HELPER_KTIME_GET_NS, helper_func!(bpf_ktime_get_ns::<F>));
523    map.insert(
524        HELPER_MAP_FOR_EACH_ELEM,
525        helper_func!(raw_map_for_each_elem::<F>),
526    );
527    map.insert(
528        HELPER_MAP_LOOKUP_PERCPU_ELEM,
529        helper_func!(raw_map_lookup_percpu_elem::<F>),
530    );
531    // map.insert(93,define_func!(raw_bpf_spin_lock);
532    // map.insert(94,define_func!(raw_bpf_spin_unlock);
533    // Map helpers::Perf event array helpers
534    map.insert(
535        HELPER_PERF_EVENT_OUTPUT,
536        helper_func!(raw_perf_event_output::<F>),
537    );
538    // Probe and trace helpers::Memory helpers
539    map.insert(HELPER_BPF_PROBE_READ, helper_func!(raw_bpf_probe_read));
540    // Print helpers
541    map.insert(HELPER_TRACE_PRINTF, helper_func!(trace_printf::<F>));
542
543    // Map helpers::Queue and stack helpers
544    map.insert(HELPER_MAP_PUSH_ELEM, helper_func!(raw_map_push_elem::<F>));
545    map.insert(HELPER_MAP_POP_ELEM, helper_func!(raw_map_pop_elem::<F>));
546    map.insert(HELPER_MAP_PEEK_ELEM, helper_func!(raw_map_peek_elem::<F>));
547
548    // Map helpers::User space helpers
549    map.insert(
550        HELPER_PROBE_READ_USER_STR,
551        helper_func!(raw_probe_read_user_str::<F>),
552    );
553
554    use ringbuf::*;
555    // Ring Buffer helpers
556    map.insert(
557        HELPER_BPF_RINGBUF_OUTPUT,
558        helper_func!(raw_bpf_ringbuf_output::<F>),
559    );
560    map.insert(
561        HELPER_BPF_RINGBUF_RESERVE,
562        helper_func!(raw_bpf_ringbuf_reserve::<F>),
563    );
564    map.insert(
565        HELPER_BPF_RINGBUF_SUBMIT,
566        helper_func!(raw_bpf_ringbuf_submit::<F>),
567    );
568    map.insert(
569        HELPER_BPF_RINGBUF_DISCARD,
570        helper_func!(raw_bpf_ringbuf_discard::<F>),
571    );
572    map.insert(
573        HELPER_BPF_RINGBUF_QUERY,
574        helper_func!(raw_bpf_ringbuf_query::<F>),
575    );
576    map.insert(
577        HELPER_BPF_RINGBUF_RESERVE_DYNPTR,
578        helper_func!(raw_bpf_ringbuf_reserve_dynptr::<F>),
579    );
580    map.insert(
581        HELPER_BPF_RINGBUF_SUBMIT_DYNPTR,
582        helper_func!(raw_bpf_ringbuf_submit_dynptr::<F>),
583    );
584    map.insert(
585        HELPER_BPF_RINGBUF_DISCARD_DYNPTR,
586        helper_func!(raw_bpf_ringbuf_discard_dynptr::<F>),
587    );
588
589    map
590}