{%- case program_type -%} {%- when "kprobe", "kretprobe" -%} use aya::programs::KProbe; {%- when "fentry" -%} use anyhow::Context as _; use aya::{programs::FEntry, Btf}; {%- when "fexit" -%} use anyhow::Context as _; use aya::{programs::FExit, Btf}; {%- when "uprobe", "uretprobe" -%} use aya::programs::UProbe; {%- when "sock_ops" -%} use anyhow::Context as _; use aya::programs::{links::CgroupAttachMode, SockOps}; {%- when "sk_msg" -%} use aya::{maps::SockHash, programs::SkMsg}; use {{crate_name}}_common::SockKey; {%- when "xdp" -%} use anyhow::Context as _; use aya::programs::{Xdp, XdpFlags}; {%- when "classifier" -%} use aya::programs::{tc, SchedClassifier, TcAttachType}; {%- when "cgroup_skb" -%} use anyhow::Context as _; use aya::programs::{links::CgroupAttachMode, CgroupSkb, CgroupSkbAttachType}; {%- when "cgroup_sysctl" -%} use anyhow::Context as _; use aya::programs::{links::CgroupAttachMode, CgroupSysctl}; {%- when "cgroup_sockopt" -%} use anyhow::Context as _; use aya::programs::{links::CgroupAttachMode, CgroupSockopt}; {%- when "tracepoint" -%} use aya::programs::TracePoint; {%- when "lsm" -%} use aya::{programs::Lsm, Btf}; {%- when "perf_event" -%} use aya::{ programs::{perf_event, PerfEvent}, util::online_cpus, }; {%- when "tp_btf" -%} use aya::{programs::BtfTracePoint, Btf}; {%- when "socket_filter" -%} use aya::programs::SocketFilter; {%- when "raw_tracepoint" -%} use aya::programs::RawTracePoint; {%- endcase %} {% if program_types_with_opts contains program_type -%} use clap::Parser; {% endif -%} #[rustfmt::skip] use log::{debug, warn}; use tokio::signal; {% if program_types_with_opts contains program_type -%} #[derive(Debug, Parser)] struct Opt { {%- case program_type -%} {%- when "xdp", "classifier" %} #[clap(short, long, default_value = "eth0")] iface: String, {%- when "sock_ops", "cgroup_skb", "cgroup_sysctl", "cgroup_sockopt" %} #[clap(short, long, default_value = "/sys/fs/cgroup")] cgroup_path: std::path::PathBuf, {%- when "uprobe", "uretprobe" %} #[clap(short, long)] pid: Option, {%- endcase %} } {% endif -%} #[tokio::main] async fn main() -> anyhow::Result<()> { {%- if program_types_with_opts contains program_type %} let opt = Opt::parse(); {% endif %} env_logger::init(); // Bump the memlock rlimit. This is needed for older kernels that don't use the // new memcg based accounting, see https://lwn.net/Articles/837122/ let rlim = libc::rlimit { rlim_cur: libc::RLIM_INFINITY, rlim_max: libc::RLIM_INFINITY, }; let ret = unsafe { libc::setrlimit(libc::RLIMIT_MEMLOCK, &rlim) }; if ret != 0 { debug!("remove limit on locked memory failed, ret is: {ret}"); } // This will include your eBPF object file as raw bytes at compile-time and load it at // runtime. This approach is recommended for most real-world use cases. If you would // like to specify the eBPF program at runtime rather than at compile-time, you can // reach for `Bpf::load_file` instead. let mut ebpf = aya::Ebpf::load(aya::include_bytes_aligned!(concat!( env!("OUT_DIR"), "/{{project-name}}_ebpf" )))?; if let Err(e) = aya_log::EbpfLogger::init(&mut ebpf) { // This can happen if you remove all log statements from your eBPF program. warn!("failed to initialize eBPF logger: {e}"); } {%- case program_type -%} {%- when "kprobe", "kretprobe" %} let program: &mut KProbe = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load()?; program.attach("{{kprobe}}", 0)?; {%- when "fentry" %} let btf = Btf::from_sys_fs().context("BTF from sysfs")?; let program: &mut FEntry = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load("{{fn_name}}", &btf)?; program.attach()?; {%- when "fexit" %} let btf = Btf::from_sys_fs().context("BTF from sysfs")?; let program: &mut FExit = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load("{{fn_name}}", &btf)?; program.attach()?; {%- when "uprobe", "uretprobe" %} let Opt { pid } = opt; let program: &mut UProbe = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load()?; program.attach(Some("{{uprobe_fn_name}}"), 0, "{{uprobe_target}}", pid)?; {%- when "sock_ops", "cgroup_skb", "cgroup_sysctl", "cgroup_sockopt" %} let Opt { cgroup_path } = opt; let cgroup = std::fs::File::open(&cgroup_path).with_context(|| format!("{}", cgroup_path.display()))?; {%- if program_type == "sock_ops" %} let program: &mut SockOps = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load()?; program.attach(cgroup, CgroupAttachMode::default())?; {%- elsif program_type == "cgroup_skb" %} let program: &mut CgroupSkb = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load()?; program.attach( cgroup, CgroupSkbAttachType::{{direction}}, CgroupAttachMode::default(), )?; {%- elsif program_type == "cgroup_sysctl" %} let program: &mut CgroupSysctl = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load()?; program.attach(cgroup, CgroupAttachMode::default())?; {%- elsif program_type == "cgroup_sockopt" %} let program: &mut CgroupSockopt = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load()?; program.attach(cgroup, CgroupAttachMode::default())?; {%- endif -%} {%- when "sk_msg" %} let sock_map: SockHash<_, SockKey> = ebpf.map("{{sock_map}}").unwrap().try_into()?; let map_fd = sock_map.fd().try_clone()?; let prog: &mut SkMsg = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; prog.load()?; prog.attach(&map_fd)?; // insert sockets to the map using sock_map.insert here, or from a sock_ops program {%- when "xdp" %} let Opt { iface } = opt; let program: &mut Xdp = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load()?; program.attach(&iface, XdpFlags::default()) .context("failed to attach the XDP program with default flags - try changing XdpFlags::default() to XdpFlags::SKB_MODE")?; {%- when "classifier" %} let Opt { iface } = opt; // error adding clsact to the interface if it is already added is harmless // the full cleanup can be done with 'sudo tc qdisc del dev eth0 clsact'. let _ = tc::qdisc_add_clsact(&iface); let program: &mut SchedClassifier = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load()?; program.attach(&iface, TcAttachType::{{direction}})?; {%- when "tracepoint" %} let program: &mut TracePoint = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load()?; program.attach("{{tracepoint_category}}", "{{tracepoint_name}}")?; {%- when "lsm" %} let btf = Btf::from_sys_fs()?; let program: &mut Lsm = ebpf.program_mut("{{lsm_hook}}").unwrap().try_into()?; program.load("{{lsm_hook}}", &btf)?; program.attach()?; {%- when "tp_btf" %} let btf = Btf::from_sys_fs()?; let program: &mut BtfTracePoint = ebpf.program_mut("{{tracepoint_name}}").unwrap().try_into()?; program.load("{{tracepoint_name}}", &btf)?; program.attach()?; {%- when "socket_filter" %} let listener = std::net::TcpListener::bind("localhost:0")?; let prog: &mut SocketFilter = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; prog.load()?; prog.attach(&listener)?; {%- when "perf_event" %} // This will raise scheduled events on each CPU at 1 HZ, triggered by the kernel based // on clock ticks. let program: &mut PerfEvent = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load()?; for cpu in online_cpus().map_err(|(_, error)| error)? { program.attach( perf_event::PerfTypeId::Software, perf_event::perf_sw_ids::PERF_COUNT_SW_CPU_CLOCK as u64, perf_event::PerfEventScope::AllProcessesOneCpu { cpu }, perf_event::SamplePolicy::Frequency(1), true, )?; } {%- when "raw_tracepoint" %} let program: &mut RawTracePoint = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?; program.load()?; program.attach("{{tracepoint_name}}")?; {%- endcase %} let ctrl_c = signal::ctrl_c(); println!("Waiting for Ctrl-C..."); ctrl_c.await?; println!("Exiting..."); Ok(()) }