diff --git a/src/rust/lqos_heimdall/src/flows.rs b/src/rust/lqos_heimdall/src/flows.rs index c88313e2..215e093d 100644 --- a/src/rust/lqos_heimdall/src/flows.rs +++ b/src/rust/lqos_heimdall/src/flows.rs @@ -1,43 +1,11 @@ use crate::{timeline::expire_timeline, FLOW_EXPIRE_SECS}; use dashmap::DashMap; use lqos_bus::{tos_parser, BusResponse, FlowTransport}; -use lqos_sys::bpf_per_cpu_map::BpfPerCpuMap; +use lqos_sys::heimdall_data::{HeimdallKey, HeimdallData}; use lqos_utils::{unix_time::time_since_boot, XdpIpAddress}; use once_cell::sync::Lazy; use std::{collections::HashSet, time::Duration}; -/// Representation of the eBPF `heimdall_key` type. -#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] -#[repr(C)] -pub struct HeimdallKey { - /// Mapped `XdpIpAddress` source for the flow. - pub src_ip: XdpIpAddress, - /// Mapped `XdpIpAddress` destination for the flow - pub dst_ip: XdpIpAddress, - /// IP protocol (see the Linux kernel!) - pub ip_protocol: u8, - /// Source port number, or ICMP type. - pub src_port: u16, - /// Destination port number. - pub dst_port: u16, -} - -/// Mapped representation of the eBPF `heimdall_data` type. -#[derive(Debug, Clone, Default)] -#[repr(C)] -pub struct HeimdallData { - /// Last seen, in nanoseconds (since boot time). - pub last_seen: u64, - /// Number of bytes since the flow started being tracked - pub bytes: u64, - /// Number of packets since the flow started being tracked - pub packets: u64, - /// IP header TOS value - pub tos: u8, - /// Reserved to pad the structure - pub reserved: [u8; 3], -} - #[derive(Clone, Debug, PartialEq, Eq, Hash)] struct FlowKey { src: XdpIpAddress, @@ -95,11 +63,12 @@ static FLOW_DATA: Lazy> = Lazy::new(DashMap::new); fn heimdall_for_each( callback: &mut dyn FnMut(&HeimdallKey, &[HeimdallData]), ) { - if let Ok(heimdall) = BpfPerCpuMap::::from_path( + /*if let Ok(heimdall) = BpfPerCpuMap::::from_path( "/sys/fs/bpf/heimdall", ) { heimdall.for_each(callback); - } + }*/ + lqos_sys::iterate_heimdall(callback); } diff --git a/src/rust/lqos_sys/src/bpf/common/heimdall.h b/src/rust/lqos_sys/src/bpf/common/heimdall.h index 9657c8fb..9d5b8342 100644 --- a/src/rust/lqos_sys/src/bpf/common/heimdall.h +++ b/src/rust/lqos_sys/src/bpf/common/heimdall.h @@ -67,6 +67,7 @@ struct heimdall_key __u8 ip_protocol; __u16 src_port; __u16 dst_port; + __u8 pad; }; struct heimdall_data { @@ -155,6 +156,7 @@ static __always_inline void update_heimdall(struct dissector_t *dissector, __u32 { bpf_debug("Failed to insert tracking"); } + //bpf_debug("Inserted tracking"); } } else if (mode == 2) { struct heimdall_event event = {0}; diff --git a/src/rust/lqos_sys/src/bpf/lqos_kern.c b/src/rust/lqos_sys/src/bpf/lqos_kern.c index 52e4353f..c4980a69 100644 --- a/src/rust/lqos_sys/src/bpf/lqos_kern.c +++ b/src/rust/lqos_sys/src/bpf/lqos_kern.c @@ -386,4 +386,32 @@ int rtt_reader(struct bpf_iter__bpf_map_elem *ctx) return 0; } +SEC("iter/bpf_map_elem") +int heimdall_reader(struct bpf_iter__bpf_map_elem *ctx) { + // The sequence file + struct seq_file *seq = ctx->meta->seq; + void *counter = ctx->value; + struct heimdall_key *ip = ctx->key; + __u32 num_cpus = NUM_CPUS; + + if (ctx->meta->seq_num == 0) { + bpf_seq_write(seq, &num_cpus, sizeof(__u32)); + bpf_seq_write(seq, &num_cpus, sizeof(__u32)); // Repeat for padding + } + + // Bail on end + if (counter == NULL || ip == NULL) { + return 0; + } + + bpf_seq_write(seq, ip, sizeof(struct heimdall_key)); + for (__u32 i=0; idownload_bytes, counter->upload_bytes); + return 0; +} + char _license[] SEC("license") = "GPL"; diff --git a/src/rust/lqos_sys/src/bpf_iterator.rs b/src/rust/lqos_sys/src/bpf_iterator.rs index 087709ce..25c8b044 100644 --- a/src/rust/lqos_sys/src/bpf_iterator.rs +++ b/src/rust/lqos_sys/src/bpf_iterator.rs @@ -1,6 +1,6 @@ use crate::{ kernel_wrapper::BPF_SKELETON, lqos_kernel::bpf, HostCounter, - RttTrackingEntry, + RttTrackingEntry, heimdall_data::{HeimdallKey, HeimdallData}, }; use lqos_utils::XdpIpAddress; use once_cell::sync::Lazy; @@ -90,7 +90,11 @@ where log::error!("{e:?}"); Err(BpfIteratorError::UnableToCreateIterator) } - Ok(_) => { + Ok(bytes) => { + if bytes == 0 { + // Not having any data is not an error + return Ok(()); + } let first_four_bytes: [u8; 4] = [buf[0], buf[1], buf[2], buf[3]]; let num_cpus = u32::from_ne_bytes(first_four_bytes) as usize; let mut index = 8; @@ -98,11 +102,13 @@ where let key_start = index; let key_end = key_start + Self::KEY_SIZE; let key_slice = &buf[key_start..key_end]; + //println!("{:?}", unsafe { &key_slice.align_to::() }); let (_head, key, _tail) = unsafe { &key_slice.align_to::() }; let value_start = key_end; let value_end = value_start + (num_cpus * Self::VALUE_SIZE); let value_slice = &buf[value_start..value_end]; + //println!("{:?}", unsafe { &value_slice.align_to::() }); let (_head, values, _tail) = unsafe { &value_slice.align_to::() }; debug_assert_eq!(values.len(), num_cpus); @@ -181,6 +187,10 @@ static mut RTT_TRACKER: Lazy< Option>, > = Lazy::new(|| None); +static mut HEIMDALL_TRACKER: Lazy< + Option>, +> = Lazy::new(|| None); + pub unsafe fn iterate_throughput( callback: &mut dyn FnMut(&XdpIpAddress, &[HostCounter]), ) { @@ -226,3 +236,28 @@ pub unsafe fn iterate_rtt( let _ = iter.for_each(callback); } } + +pub fn iterate_heimdall( + callback: &mut dyn FnMut(&HeimdallKey, &[HeimdallData]), +) { + unsafe { + if HEIMDALL_TRACKER.is_none() { + let lock = BPF_SKELETON.lock().unwrap(); + if let Some(skeleton) = lock.as_ref() { + let skeleton = skeleton.get_ptr(); + if let Ok(iter) = { + BpfMapIterator::new( + (*skeleton).progs.heimdall_reader, + (*skeleton).maps.heimdall, + ) + } { + *HEIMDALL_TRACKER = Some(iter); + } + } + } + + if let Some(iter) = HEIMDALL_TRACKER.as_mut() { + let _ = iter.for_each_per_cpu(callback); + } + } +} \ No newline at end of file diff --git a/src/rust/lqos_sys/src/bpf_per_cpu_map.rs b/src/rust/lqos_sys/src/bpf_per_cpu_map.rs deleted file mode 100644 index 5fdc31a2..00000000 --- a/src/rust/lqos_sys/src/bpf_per_cpu_map.rs +++ /dev/null @@ -1,77 +0,0 @@ -use anyhow::{Error, Result}; -use libbpf_sys::{ - bpf_map_get_next_key, bpf_map_lookup_elem, bpf_obj_get, -}; -use std::fmt::Debug; -use std::{ - ffi::{c_void, CString}, - marker::PhantomData, - ptr::null_mut, -}; - -use crate::num_possible_cpus; - -/// Represents an underlying BPF map, accessed via the filesystem. -/// `BpfMap` *only* talks to PER-CPU variants of maps. -/// -/// `K` is the *key* type, indexing the map. -/// `V` is the *value* type, and must exactly match the underlying C data type. -pub struct BpfPerCpuMap { - fd: i32, - _key_phantom: PhantomData, - _val_phantom: PhantomData, -} - -impl BpfPerCpuMap -where - K: Default + Clone, - V: Default + Clone + Debug, -{ - /// Connect to a PER-CPU BPF map via a filename. Connects the internal - /// file descriptor, which is held until the structure is - /// dropped. The index of the CPU is *not* specified. - pub fn from_path(filename: &str) -> Result { - let filename_c = CString::new(filename)?; - let fd = unsafe { bpf_obj_get(filename_c.as_ptr()) }; - if fd < 0 { - Err(Error::msg("Unable to open BPF map")) - } else { - Ok(Self { fd, _key_phantom: PhantomData, _val_phantom: PhantomData }) - } - } - - /// Instead of clonining into a vector - /// and allocating, calls `callback` for each key/value slice - /// with references to the data returned from C. - pub fn for_each(&self, callback: &mut dyn FnMut(&K, &[V])) { - let num_cpus = num_possible_cpus().unwrap(); - let mut prev_key: *mut K = null_mut(); - let mut key: K = K::default(); - let key_ptr: *mut K = &mut key; - let mut value = vec![V::default(); num_cpus as usize]; - let value_ptr = value.as_mut_ptr(); - - unsafe { - while bpf_map_get_next_key( - self.fd, - prev_key as *mut c_void, - key_ptr as *mut c_void, - ) == 0 - { - bpf_map_lookup_elem( - self.fd, - key_ptr as *mut c_void, - value_ptr as *mut c_void, - ); - callback(&key, &value); - prev_key = key_ptr; - } - } - } -} - -impl Drop for BpfPerCpuMap { - fn drop(&mut self) { - let _ = nix::unistd::close(self.fd); - } -} diff --git a/src/rust/lqos_sys/src/heimdall_data.rs b/src/rust/lqos_sys/src/heimdall_data.rs new file mode 100644 index 00000000..31be4bca --- /dev/null +++ b/src/rust/lqos_sys/src/heimdall_data.rs @@ -0,0 +1,33 @@ +use lqos_utils::XdpIpAddress; +use zerocopy::FromBytes; + +/// Representation of the eBPF `heimdall_key` type. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash, FromBytes)] +#[repr(C)] +pub struct HeimdallKey { + /// Mapped `XdpIpAddress` source for the flow. + pub src_ip: XdpIpAddress, + /// Mapped `XdpIpAddress` destination for the flow + pub dst_ip: XdpIpAddress, + /// IP protocol (see the Linux kernel!) + pub ip_protocol: u8, + /// Source port number, or ICMP type. + pub src_port: u16, + /// Destination port number. + pub dst_port: u16, + _padding: u8, +} + +/// Mapped representation of the eBPF `heimdall_data` type. +#[derive(Debug, Clone, Default, FromBytes)] +#[repr(C)] +pub struct HeimdallData { + /// Last seen, in nanoseconds (since boot time). + pub last_seen: u64, + /// Number of bytes since the flow started being tracked + pub bytes: u64, + /// Number of packets since the flow started being tracked + pub packets: u64, + /// IP header TOS value + pub tos: u8, +} \ No newline at end of file diff --git a/src/rust/lqos_sys/src/lib.rs b/src/rust/lqos_sys/src/lib.rs index c3c64aba..d8ac2c3e 100644 --- a/src/rust/lqos_sys/src/lib.rs +++ b/src/rust/lqos_sys/src/lib.rs @@ -11,10 +11,6 @@ mod bifrost_maps; /// built-in, compiled eBPF programs. This is very-low level and should /// be handled with caution. pub mod bpf_map; -/// Provides direct access to LibBPF functionality, as exposed by the -/// built-in, compiled eBPF programs. This is very-low level and should -/// be handled with caution. -pub mod bpf_per_cpu_map; mod cpu_map; mod ip_mapping; mod kernel_wrapper; @@ -23,6 +19,9 @@ mod tcp_rtt; mod throughput; mod linux; mod bpf_iterator; +/// Data shared between eBPF and Heimdall that needs local access +/// for map control. +pub mod heimdall_data; pub use ip_mapping::{ add_ip_to_tc, clear_ips_from_tc, del_ip_from_tc, list_mapped_ips, @@ -32,3 +31,4 @@ pub use linux::num_possible_cpus; pub use lqos_kernel::max_tracked_ips; pub use tcp_rtt::{rtt_for_each, RttTrackingEntry}; pub use throughput::{throughput_for_each, HostCounter}; +pub use bpf_iterator::iterate_heimdall; \ No newline at end of file