diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..336cb72 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,54 @@ +name: CI + +on: + push: + branches: + - main + + pull_request: + branches: + - main + +env: + CARGO_TERM_COLOR: always + +jobs: + build: + runs-on: ubuntu-20.04 + strategy: + fail-fast: false + matrix: + program: + - kprobe + - kretprobe + - uprobe + - uretprobe + - sock_ops + - sk_msg + - xdp + - classifier + - cgroup_skb + - tracepoint + + steps: + - uses: actions/checkout@v2 + + - name: Install latest stable + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + + - name: Install latest nightly + uses: actions-rs/toolchain@v1 + with: + toolchain: nightly + components: rust-src + + - name: Install bpf-linker + run: cargo +nightly install bpf-linker + + - name: Install Cargo Generate + run: cargo install --git https://github.com/cargo-generate/cargo-generate + + - name: Run tests + run: ./test.sh ${{ github.workspace }} ${{ matrix.program }} diff --git a/cargo-generate.toml b/cargo-generate.toml new file mode 100644 index 0000000..cd91e5d --- /dev/null +++ b/cargo-generate.toml @@ -0,0 +1,41 @@ +[template] +cargo_generate_version = ">=0.9.0" +ignore = [".github", "test.sh"] + +[placeholders.program_type] +type = "string" +prompt = "Which type of eBPF program?" +choices = ["kprobe", "kretprobe", "uprobe", "uretprobe", "sock_ops", "sk_msg", "xdp", "classifier", "cgroup_skb", "tracepoint"] +default = "xdp" + +[conditional.'program_type == "kprobe" || program_type == "kretprobe"'.placeholders.kprobe] +type = "string" +prompt = "Where to attach the (k|kret)probe? (e.g try_to_wake_up)" + +[conditional.'program_type == "uprobe" || program_type == "uretprobe"'.placeholders.uprobe_target] +type = "string" +prompt = "Target to attach the (u|uret)probe? (e.g libc)" + +[conditional.'program_type == "uprobe" || program_type == "uretprobe"'.placeholders.uprobe_fn_name] +type = "string" +prompt = "Function name to attach the (u|uret)probe? (e.g getaddrinfo)" + +[conditional.'program_type == "cgroup_skb" || program_type == "classifier"'.placeholders.direction] +type = "string" +prompt = "Attach direction?" +choices = [ "Ingress", "Egress" ] + +[conditional.'program_type == "sk_msg"'.placeholders.sock_map] +type = "string" +prompt = "Map Name (UPPER_CASE)?" +regex = "[A-Z_]+" + +[conditional.'program_type == "tracepoint"'.placeholders.tracepoint_category] +type = "string" +prompt = "Which tracepoint category? (e.g sched, net etc...)" +regex = "[a-z]+" + +[conditional.'program_type == "tracepoint"'.placeholders.tracepoint_name] +type = "string" +prompt = "Which tracepoint name? (e.g sched_switch, net_dev_queue)" +regex = "[a-z]+" \ No newline at end of file diff --git a/test.sh b/test.sh new file mode 100755 index 0000000..7d642d3 --- /dev/null +++ b/test.sh @@ -0,0 +1,42 @@ +#!/bin/bash + +set -ex + +TEMPLATE_DIR=$1 +if [ -z "$TEMPLATE_DIR" ]; then echo "template dir required"; exit 1; fi +PROG_TYPE=$2 +if [ -z "$PROG_TYPE" ]; then echo "program type required"; exit 1; fi + +TMP_DIR=$(mktemp -d) +clean_up() { + rm -rf "${TMP_DIR}" +} +trap clean_up EXIT + +pushd $TMP_DIR +case "$PROG_TYPE" in + "kprobe"|"kretprobe") + ADDITIONAL_ARGS="-d kprobe=test" + ;; + "uprobe"|"uretprobe") + ADDITIONAL_ARGS="-d uprobe_target=testlib -d uprobe_fn_name=testfn" + ;; + "tracepoint") + ADDITIONAL_ARGS="-d tracepoint_category=net -d tracepoint_name=net_dev_queue" + ;; + "classifier"|"cgroup_skb") + ADDITIONAL_ARGS="-d direction=Ingress" + ;; + "sk_msg") + ADDITIONAL_ARGS="-d sock_map=TEST" + ;; + *) + ADDITIONAL_ARGS='' +esac + +cargo generate -v --path "${TEMPLATE_DIR}" -n test -d program_type="${PROG_TYPE}" ${ADDITIONAL_ARGS} +pushd test +cargo build +cargo xtask build-ebpf +popd +exit 0 diff --git a/{{project-name}}-common/src/lib.rs b/{{project-name}}-common/src/lib.rs index 2e7f0d4..bca8a6f 100644 --- a/{{project-name}}-common/src/lib.rs +++ b/{{project-name}}-common/src/lib.rs @@ -1 +1,15 @@ -#![no_std] \ No newline at end of file +#![no_std] +{%- if program_type == "sk_msg" %} + +#[derive(Copy, Clone)] +#[repr(C)] +pub struct SockKey { + pub remote_ip4: u32, + pub local_ip4: u32, + pub remote_port: u32, + pub local_port: u32, +} + +#[cfg(feature = "userspace")] +unsafe impl aya::Pod for SockKey {} +{%- endif -%} \ No newline at end of file diff --git a/{{project-name}}-ebpf/src/main.rs b/{{project-name}}-ebpf/src/main.rs index 997377a..2c8cd30 100644 --- a/{{project-name}}-ebpf/src/main.rs +++ b/{{project-name}}-ebpf/src/main.rs @@ -1,5 +1,183 @@ #![no_std] #![no_main] +{% case program_type -%} +{%- when "kprobe" %} +use aya_bpf::{ + macros::kprobe, + programs::ProbeContext, +}; + +#[kprobe(name="{{crate_name}}")] +pub fn {{crate_name}}(ctx: ProbeContext) -> u32 { + match unsafe { try_{{crate_name}}(ctx) } { + Ok(ret) => ret, + Err(ret) => ret, + } +} + +unsafe fn try_{{crate_name}}(_ctx: ProbeContext) -> Result { + Ok(0) +} +{%- when "kretprobe" %} +use aya_bpf::{ + macros::kretprobe, + programs::ProbeContext, +}; + +#[kretprobe(name="{{crate_name}}")] +pub fn {{crate_name}}(ctx: ProbeContext) -> u32 { + match unsafe { try_{{crate_name}}(ctx) } { + Ok(ret) => ret, + Err(ret) => ret, + } +} + +unsafe fn try_{{crate_name}}(_ctx: ProbeContext) -> Result { + Ok(0) +} +{%- when "uprobe" %} +use aya_bpf::{ + macros::uprobe, + programs::ProbeContext, +}; + +#[uprobe(name="{{crate_name}}")] +pub fn {{crate_name}}(ctx: ProbeContext) -> u32 { + match unsafe { try_{{crate_name}}(ctx) } { + Ok(ret) => ret, + Err(ret) => ret, + } +} + +unsafe fn try_{{crate_name}}(_ctx: ProbeContext) -> Result { + Ok(0) +} +{%- when "uretprobe" %} +use aya_bpf::{ + macros::uretprobe, + programs::ProbeContext, +}; + +#[uretprobe(name="{{crate_name}}")] +pub fn {{crate_name}}(ctx: ProbeContext) -> u32 { + match unsafe { try_{{crate_name}}(ctx) } { + Ok(ret) => ret, + Err(ret) => ret, + } +} + +unsafe fn try_{{crate_name}}(_ctx: ProbeContext) -> Result { + Ok(0) +} +{%- when "sock_ops" %} +use aya_bpf::{ + macros::sock_ops, + programs::SockOpsContext, +}; + +#[sock_ops(name="{{crate_name}}")] +pub fn {{crate_name}}(ctx: SockOpsContext) -> u32 { + match unsafe { try_{{crate_name}}(ctx) } { + Ok(ret) => ret, + Err(ret) => ret, + } +} + +unsafe fn try_{{crate_name}}(_ctx: SockOpsContext) -> Result { + Ok(0) +} +{%- when "sk_msg" %} +use aya_bpf::{ + macros::{map, sk_msg}, + maps::SockHash, + programs::SkMsgContext, +}; +use {{crate_name}}_common::SockKey; + +#[map(name="{{sock_map}}")] +static mut {{sock_map}}: SockHash = SockHash::::with_max_entries(1024, 0); + +#[sk_msg(name="{{crate_name}}")] +pub fn {{crate_name}}(ctx: SkMsgContext) -> u32 { + match unsafe { try_{{crate_name}}(ctx) } { + Ok(ret) => ret, + Err(ret) => ret, + } +} + +unsafe fn try_{{crate_name}}(_ctx: SkMsgContext) -> Result { + Ok(0) +} +{%- when "xdp" %} +use aya_bpf::{ + bindings::xdp_action, + macros::xdp, + programs::XdpContext, +}; + +#[xdp(name="{{crate_name}}")] +pub fn {{crate_name}}(ctx: XdpContext) -> u32 { + match unsafe { try_{{crate_name}}(ctx) } { + Ok(ret) => ret, + Err(_) => xdp_action::XDP_ABORTED, + } +} + +unsafe fn try_{{crate_name}}(_ctx: XdpContext) -> Result { + Ok(xdp_action::XDP_PASS) +} +{%- when "classifier" %} +use aya_bpf::{ + macros::classifier, + programs::SkSkbContext, +}; + +#[classifier(name="{{crate_name}}")] +pub fn {{crate_name}}(ctx: SkSkbContext) -> i32 { + match unsafe { try_{{crate_name}}(ctx) } { + Ok(ret) => ret, + Err(ret) => ret, + } +} + +unsafe fn try_{{crate_name}}(_ctx: SkSkbContext) -> Result { + Ok(0) +} +{%- when "cgroup_skb" %} +use aya_bpf::{ + macros::cgroup_skb, + programs::SkSkbContext, +}; + +#[cgroup_skb(name="{{crate_name}}")] +pub fn {{crate_name}}(ctx: SkSkbContext) -> i32 { + match unsafe { try_{{crate_name}}(ctx) } { + Ok(ret) => ret, + Err(ret) => ret, + } +} + +unsafe fn try_{{crate_name}}(_ctx: SkSkbContext) -> Result { + Ok(0) +} +{%- when "tracepoint" %} +use aya_bpf::{ + macros::tracepoint, + programs::TracePointContext, +}; + +#[tracepoint(name="{{crate_name}}")] +pub fn {{crate_name}}(ctx: TracePointContext) -> u32 { + match unsafe { try_{{crate_name}}(ctx) } { + Ok(ret) => ret, + Err(ret) => ret, + } +} + +unsafe fn try_{{crate_name}}(_ctx: TracePointContext) -> Result { + Ok(0) +} +{%- endcase %} #[panic_handler] fn panic(_info: &core::panic::PanicInfo) -> ! { diff --git a/{{project-name}}/Cargo.toml b/{{project-name}}/Cargo.toml index 95015ba..b8275f4 100644 --- a/{{project-name}}/Cargo.toml +++ b/{{project-name}}/Cargo.toml @@ -8,6 +8,9 @@ publish = false aya = { git = "https://github.com/alessandrod/aya", branch="main" } {{project-name}}-common = { path = "../{{project-name}}-common", features=["userspace"] } anyhow = "1.0.42" +ctrlc = "3.2" +{% if program_type == "uprobe" %}libc = "0.2.102"{% endif %} +structopt = { version = "0.3"} [[bin]] name = "{{project-name}}" diff --git a/{{project-name}}/src/main.rs b/{{project-name}}/src/main.rs index c0e24fd..0b9f8c6 100644 --- a/{{project-name}}/src/main.rs +++ b/{{project-name}}/src/main.rs @@ -1,9 +1,106 @@ +use aya::Bpf; +{% case program_type -%} +{%- when "kprobe", "kretprobe" -%} +use aya::programs::KProbe; +{%- when "uprobe", "uretprobe" -%} +use aya::programs::UProbe; +{%- when "sock_ops" -%} +use aya::programs::SockOps; +{%- when "sk_msg" -%} +use aya::maps::{MapRefMut,SockHash}; +use aya::programs::SkMsg; +use {{crate_name}}_common::SockKey; +{%- when "xdp" -%} +use aya::programs::{Xdp, XdpFlags}; +{%- when "classifier" -%} +use aya::programs::{tc, SchedClassifier, TcAttachType}; +{%- when "cgroup_skb" -%} +use aya::programs::{CgroupSkb, CgroupSkbAttachType}; +{%- when "tracepoint" -%} +use aya::programs::TracePoint; +{%- endcase %} +use std::{ + convert::{TryFrom,TryInto}, + sync::Arc, + sync::atomic::{AtomicBool, Ordering}, +}; +use structopt::StructOpt; + fn main() { if let Err(e) = try_main() { eprintln!("error: {:#}", e); } } +#[derive(Debug, StructOpt)] +struct Opt { + #[structopt(short, long)] + path: String, + {% if program_type == "xdp" or program_type == "classifier" -%} + #[structopt(short, long, default_value = "eth0")] + iface: String, + {%- elsif program_type == "sock_ops" or program_type == "cgroup_skb" -%} + #[structopt(short, long, default_value = "/sys/fs/cgroup/unified")] + cgroup_path: String, + {%- elsif program_type == "uprobe" or program_type == "uretprobe" -%} + #[structopt(short, long)] + pid: Option + {%- endif %} +} + fn try_main() -> Result<(), anyhow::Error> { + let opt = Opt::from_args(); + let mut bpf = Bpf::load_file(&opt.path)?; + {% case program_type -%} + {%- when "kprobe", "kretprobe" -%} + let program: &mut KProbe = bpf.program_mut("{{crate_name}}")?.try_into()?; + program.load()?; + program.attach("{{probe}}", 0)?; + {%- when "uprobe", "uretprobe" -%} + let program: &mut UProbe = bpf.program_mut("{{crate_name}}")?.try_into()?; + program.load()?; + program.attach(Some("{{uprobe_fn_name}}"), 0, "{{uprobe_target}}", opt.pid.try_into()?)?; + {%- when "sock_ops" -%} + let program: &mut SockOps = bpf.program_mut("{{crate_name}}")?.try_into()?; + let cgroup = std::fs::File::open(opt.cgroup_path)?; + program.load()?; + program.attach(cgroup)?; + {%- when "sk_msg" -%} + let sock_map = SockHash::::try_from(bpf.map_mut("{{sock_map}}")?)?; + let prog: &mut SkMsg = bpf.program_mut("{{crate_name}}")?.try_into()?; + prog.load()?; + prog.attach(&sock_map)?; + // insert sockets to the map using sock_map.insert here, or from a sock_ops program + {%- when "xdp" -%} + let program: &mut Xdp = bpf.program_mut("{{crate_name}}")?.try_into()?; + program.load()?; + program.attach(&opt.iface, XdpFlags::default())?; + {%- when "classifier" -%} + tc::qdisc_add_clsact(&opt.iface)?; + let program: &mut SchedClassifier = bpf.program_mut("{{crate_name}}")?.try_into()?; + program.load()?; + program.attach(&opt.iface, TcAttachType::{{direction}})?; + {%- when "cgroup_skb" -%} + let program: &mut CgroupSkb = bpf.program_mut("{{crate_name}}")?.try_into()?; + let cgroup = std::fs::File::open(opt.cgroup_path)?; + program.load()?; + program.attach(cgroup, CgroupSkbAttachType::{{direction}})?; + {%- when "tracepoint" -%} + let program: &mut TracePoint = bpf.program_mut("{{crate_name}}")?.try_into()?; + program.load()?; + program.attach("{{tracepoint_category}}", "{{tracepoint_name}}")?; + {%- endcase %} + + let running = Arc::new(AtomicBool::new(true)); + let r = running.clone(); + + ctrlc::set_handler(move || { + r.store(false, Ordering::SeqCst); + }).expect("Error setting Ctrl-C handler"); + + println!("Waiting for Ctrl-C..."); + while running.load(Ordering::SeqCst) {} + println!("Exiting..."); + Ok(()) } \ No newline at end of file