Merge pull request #3 from aya-rs/skel

Use cargo-generate conditionals to provide skeletons for all supported program types

.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 }}

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]+"

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

{{project-name}}-common/src/lib.rs

@@ -1 +1,15 @@
 #![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 -%}

{{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<u32, u32> {
+    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<u32, u32> {
+    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<u32, u32> {
+    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<u32, u32> {
+    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<u32, u32> {
+    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<SockKey> = SockHash::<SockKey>::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<u32, u32> {
+    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<u32, u32> {
+    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<i32, i32> {
+    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<i32, i32> {
+    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<u32, u32> {
+    Ok(0)
+}
+{%- endcase %}
 
 #[panic_handler]
 fn panic(_info: &core::panic::PanicInfo) -> ! {

{{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}}"

{{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<i32>
+    {%- 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::<MapRefMut, SockKey>::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(())
 }