Writing eBPF Kprobe Programs with Rust Aya: Complete Developer Guide#

1
arch: "x86_64"
2
cpus: 4
3
memory: "8GiB"
4
disk: "50GiB"
5

6
images:
7
  - location: "https://cloud-images.ubuntu.com/releases/22.04/release-20240821/ubuntu-22.04-server-cloudimg-amd64.img"
8
    arch: "x86_64"
9

10
provision:
11
  - mode: system
12
    script: |
13
      #!/bin/bash
14
      apt-get update
15
      apt-get install -y curl build-essential
16

17
      # Install Rust
18
      curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
19
      source /root/.cargo/env
20
      rustup install nightly
21
      rustup default nightly
22

23
      # Install eBPF tools
24
      cargo install bpf-linker bindgen-cli
25

26
      # Install system dependencies
27
      apt-get install -y clang llvm libelf-dev libz-dev libbpf-dev linux-headers-generic bpftool
28

29
  - mode: user
30
    script: |
31
      #!/bin/bash
32
      curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y

1
[package]
2
name = "kprobe-observer-ebpf"
3
version = "0.1.0"
4
edition = "2021"
5

6
[dependencies]
7
aya-ebpf = "0.1"
8
aya-log-ebpf = "0.1"
9

10
[[bin]]
11
name = "kprobe-observer"
12
path = "src/main.rs"
13

14
[profile.dev]
15
opt-level = 3
16
debug = false
17
debug-assertions = false
18
overflow-checks = false
19
lto = true
20
panic = "abort"
21
incremental = false
22
codegen-units = 1
23
rpath = false
24

25
[profile.release]
26
lto = true
27
panic = "abort"
28
codegen-units = 1

1
[package]
2
name = "xtask"
3
version = "0.1.0"
4
edition = "2021"
5

6
[dependencies]
7
anyhow = "1.0"
8
clap = { version = "4.0", features = ["derive"] }
9
aya-tool = "0.1"

1
/*
2
 * wake_up_new_task - wake up a newly created task for the first time.
3
 *
4
 * This function will do some initial scheduler statistics housekeeping
5
 * that must be done for every newly created context, then puts the task
6
 * on the runqueue and wakes it.
7
 */
8
void wake_up_new_task(struct task_struct *p)
9
{
10
    struct rq_flags rf;
11
    struct rq *rq;
12

13
    // ... function implementation
14
}

1
use anyhow::Result;
2
use aya_tool::generate::InputFile;
3
use std::fs::File;
4
use std::io::Write;
5
use std::path::PathBuf;
6

7
pub fn generate_bindings() -> Result<()> {
8
    println!("Generating kernel bindings...");
9

10
    let dir = PathBuf::from("ebpf/src");
11
    let names: Vec<&str> = vec![
12
        "task_struct",
13
        "pid_t",
14
        "cred",
15
        "mm_struct",
16
        "files_struct",
17
        "fs_struct",
18
        "signal_struct",
19
        "sighand_struct",
20
        "thread_info",
21
        "cpu_context_save",
22
        "thread_struct",
23
    ];
24

25
    let bindings = aya_tool::generate(
26
        InputFile::Btf(PathBuf::from("/sys/kernel/btf/vmlinux")),
27
        &names,
28
        &[],
29
    )?;
30

31
    let mut output = File::create(dir.join("bindings.rs"))?;
32
    write!(output, "{}", bindings)?;
33

34
    println!("Kernel bindings generated successfully!");
35
    Ok(())
36
}

1
use anyhow::Result;
2
use std::process::Command;
3

4
pub fn build_ebpf() -> Result<()> {
5
    println!("Building eBPF program...");
6

7
    let output = Command::new("cargo")
8
        .args(&[
9
            "build",
10
            "--target=bpfel-unknown-none",
11
            "--release",
12
        ])
13
        .current_dir("ebpf")
14
        .output()?;
15

16
    if !output.status.success() {
17
        anyhow::bail!(
18
            "Failed to build eBPF program:\n{}",
19
            String::from_utf8_lossy(&output.stderr)
20
        );
21
    }
22

23
    println!("eBPF program built successfully!");
24
    Ok(())
25
}

1
mod build;
2
mod codegen;
3

4
use anyhow::Result;
5
use clap::{Parser, Subcommand};
6

7
#[derive(Parser)]
8
#[command(name = "xtask")]
9
#[command(about = "Build and development tasks")]
10
struct Cli {
11
    #[command(subcommand)]
12
    command: Commands,
13
}
14

15
#[derive(Subcommand)]
16
enum Commands {
17
    /// Generate kernel type bindings
18
    Codegen,
19
    /// Build eBPF program
20
    Build,
21
    /// Build and run the observer
22
    Run,
23
}
24

25
fn main() -> Result<()> {
26
    let cli = Cli::parse();
27

28
    match cli.command {
29
        Commands::Codegen => codegen::generate_bindings(),
30
        Commands::Build => build::build_ebpf(),
31
        Commands::Run => {
32
            build::build_ebpf()?;
33
            println!("Running observer...");
34
            let output = std::process::Command::new("sudo")
35
                .args(&["./target/release/kprobe-observer"])
36
                .output()?;
37

38
            println!("{}", String::from_utf8_lossy(&output.stdout));
39
            if !output.stderr.is_empty() {
40
                eprintln!("{}", String::from_utf8_lossy(&output.stderr));
41
            }
42
            Ok(())
43
        }
44
    }
45
}

1
#![no_std]
2
#![no_main]
3

4
mod bindings;
5

6
use aya_ebpf::{
7
    helpers::bpf_get_current_pid_tgid,
8
    macros::{kprobe, map},
9
    maps::RingBuf,
10
    programs::ProbeContext,
11
    EbpfContext,
12
};
13
use aya_log_ebpf::info;
14
use bindings::{task_struct, pid_t};
15

16
// Event structure for user space communication
17
#[repr(C)]
18
#[derive(Clone, Copy)]
19
pub struct TaskEvent {
20
    pub caller_pid: u32,
21
    pub caller_tgid: u32,
22
    pub new_task_pid: u32,
23
    pub new_task_tgid: u32,
24
    pub timestamp: u64,
25
    pub comm: [u8; 16],
26
}
27

28
// Ring buffer for efficient event streaming
29
#[map]
30
static TASK_EVENTS: RingBuf = RingBuf::with_byte_size(1024 * 1024, 0);
31

32
// Statistics tracking
33
#[map]
34
static mut STATS: aya_ebpf::maps::Array<u64> = aya_ebpf::maps::Array::with_max_entries(4, 0);
35

36
// Statistics indices
37
const STAT_TOTAL_EVENTS: u32 = 0;
38
const STAT_SUCCESSFUL_EVENTS: u32 = 1;
39
const STAT_ERROR_EVENTS: u32 = 2;
40
const STAT_LAST_TIMESTAMP: u32 = 3;
41

42
#[kprobe]
43
pub fn wake_up_new_task(ctx: ProbeContext) -> u32 {
44
    match try_wake_up_new_task(ctx) {
45
        Ok(ret) => ret,
46
        Err(ret) => {
47
            // Update error statistics
48
            unsafe {
49
                if let Ok(mut stat) = STATS.get_ptr_mut(STAT_ERROR_EVENTS) {
50
                    *stat += 1;
51
                }
52
            }
53
            ret
54
        }
55
    }
56
}
57

58
fn try_wake_up_new_task(ctx: ProbeContext) -> Result<u32, u32> {
59
    // Get the task_struct pointer from the first argument
60
    let task: *const task_struct = ctx.arg(0).ok_or(1)?;
61

62
    // Read task information safely
63
    let new_task_pid = unsafe {
64
        core::ptr::read_volatile(&(*task).pid as *const pid_t)
65
    };
66
    let new_task_tgid = unsafe {
67
        core::ptr::read_volatile(&(*task).tgid as *const pid_t)
68
    };
69

70
    // Get caller information
71
    let caller_pid_tgid = bpf_get_current_pid_tgid();
72
    let caller_pid = (caller_pid_tgid & 0xFFFFFFFF) as u32;
73
    let caller_tgid = (caller_pid_tgid >> 32) as u32;
74

75
    // Get current timestamp
76
    let timestamp = unsafe {
77
        aya_ebpf::helpers::bpf_ktime_get_ns()
78
    };
79

80
    // Read process name from task_struct
81
    let mut comm = [0u8; 16];
82
    unsafe {
83
        let comm_ptr = &(*task).comm as *const [::aya_ebpf::cty::c_char; 16];
84
        for i in 0..16 {
85
            let c = core::ptr::read_volatile(&(*comm_ptr)[i]);
86
            comm[i] = c as u8;
87
            if c == 0 { break; }
88
        }
89
    }
90

91
    // Create event for user space
92
    let event = TaskEvent {
93
        caller_pid,
94
        caller_tgid,
95
        new_task_pid: new_task_pid as u32,
96
        new_task_tgid: new_task_tgid as u32,
97
        timestamp,
98
        comm,
99
    };
100

101
    // Send event to user space via ring buffer
102
    if let Some(mut entry) = TASK_EVENTS.reserve::<TaskEvent>(0) {
103
        entry.write(event);
104
        entry.submit(0);
105

106
        // Update statistics
107
        unsafe {
108
            if let Ok(mut stat) = STATS.get_ptr_mut(STAT_SUCCESSFUL_EVENTS) {
109
                *stat += 1;
110
            }
111
            if let Ok(mut stat) = STATS.get_ptr_mut(STAT_LAST_TIMESTAMP) {
112
                *stat = timestamp;
113
            }
114
        }
115
    }
116

117
    // Update total events counter
118
    unsafe {
119
        if let Ok(mut stat) = STATS.get_ptr_mut(STAT_TOTAL_EVENTS) {
120
            *stat += 1;
121
        }
122
    }
123

124
    // Log the event (visible in /sys/kernel/debug/tracing/trace_pipe)
125
    info!(
126
        &ctx,
127
        "wake_up_new_task: caller PID {}, new task PID {}, TGID {}",
128
        caller_pid,
129
        new_task_pid,
130
        new_task_tgid
131
    );
132

133
    Ok(0)
134
}
135

136
#[panic_handler]
137
fn panic(_info: &core::panic::PanicInfo) -> ! {
138
    unsafe { core::hint::unreachable_unchecked() }
139
}

1
use anyhow::Result;
2
use aya::{
3
    include_bytes_aligned,
4
    maps::RingBuf,
5
    programs::KProbe,
6
    Bpf,
7
};
8
use aya_log::BpfLogger;
9
use bytes::BytesMut;
10
use clap::Parser;
11
use log::{info, warn, error};
12
use std::{
13
    convert::TryInto,
14
    sync::{
15
        atomic::{AtomicBool, Ordering},
16
        Arc,
17
    },
18
    time::{Duration, SystemTime, UNIX_EPOCH},
19
};
20
use tokio::{signal, time::sleep};
21

22
#[derive(Parser, Debug)]
23
#[command(name = "kprobe-observer")]
24
#[command(about = "eBPF Kprobe observer for wake_up_new_task")]
25
struct Args {
26
    /// Enable verbose logging
27
    #[arg(short, long)]
28
    verbose: bool,
29

30
    /// Statistics reporting interval (seconds)
31
    #[arg(short, long, default_value = "30")]
32
    stats_interval: u64,
33

34
    /// Maximum events to process (0 = unlimited)
35
    #[arg(short, long, default_value = "0")]
36
    max_events: u64,
37
}
38

39
// Event structure matching eBPF program
40
#[repr(C)]
41
#[derive(Clone, Copy, Debug)]
42
struct TaskEvent {
43
    pub caller_pid: u32,
44
    pub caller_tgid: u32,
45
    pub new_task_pid: u32,
46
    pub new_task_tgid: u32,
47
    pub timestamp: u64,
48
    pub comm: [u8; 16],
49
}
50

51
// Application statistics
52
#[derive(Default, Debug)]
53
struct Statistics {
54
    total_events: u64,
55
    events_per_second: f64,
56
    unique_callers: std::collections::HashSet<u32>,
57
    unique_processes: std::collections::HashSet<String>,
58
    start_time: Option<SystemTime>,
59
    last_event_time: Option<SystemTime>,
60
}
61

62
impl Statistics {
63
    fn new() -> Self {
64
        Self {
65
            start_time: Some(SystemTime::now()),
66
            ..Default::default()
67
        }
68
    }
69

70
    fn update(&mut self, event: &TaskEvent) {
71
        self.total_events += 1;
72
        self.unique_callers.insert(event.caller_pid);
73

74
        // Convert comm to string
75
        let comm = String::from_utf8_lossy(&event.comm)
76
            .trim_end_matches('\0')
77
            .to_string();
78
        if !comm.is_empty() {
79
            self.unique_processes.insert(comm);
80
        }
81

82
        self.last_event_time = Some(SystemTime::now());
83

84
        // Calculate events per second
85
        if let Some(start) = self.start_time {
86
            if let Ok(duration) = SystemTime::now().duration_since(start) {
87
                if duration.as_secs() > 0 {
88
                    self.events_per_second = self.total_events as f64 / duration.as_secs() as f64;
89
                }
90
            }
91
        }
92
    }
93

94
    fn print_summary(&self) {
95
        println!("\n=== eBPF Kprobe Observer Statistics ===");
96
        println!("Total events processed: {}", self.total_events);
97
        println!("Unique callers: {}", self.unique_callers.len());
98
        println!("Unique processes: {}", self.unique_processes.len());
99
        println!("Events per second: {:.2}", self.events_per_second);
100

101
        if let Some(start) = self.start_time {
102
            if let Ok(duration) = SystemTime::now().duration_since(start) {
103
                println!("Runtime: {:.2} seconds", duration.as_secs_f64());
104
            }
105
        }
106

107
        println!("======================================\n");
108
    }
109
}
110

111
async fn load_and_attach_ebpf() -> Result<(Bpf, RingBuf<&'static mut [u8]>)> {
112
    // Load the eBPF program
113
    let mut bpf = Bpf::load(include_bytes_aligned!(
114
        "../target/bpfel-unknown-none/release/kprobe-observer"
115
    ))?;
116

117
    // Initialize BPF logger for eBPF program logs
118
    if let Err(e) = BpfLogger::init(&mut bpf) {
119
        warn!("Failed to initialize BPF logger: {}", e);
120
    }
121

122
    // Load and attach the kprobe program
123
    let program: &mut KProbe = bpf.program_mut("wake_up_new_task").unwrap().try_into()?;
124
    program.load()?;
125
    program.attach("wake_up_new_task", 0)?;
126

127
    info!("eBPF program loaded and attached successfully");
128

129
    // Get the ring buffer for event communication
130
    let ring_buf = RingBuf::try_from(bpf.map_mut("TASK_EVENTS").unwrap())?;
131

132
    Ok((bpf, ring_buf))
133
}
134

135
fn parse_task_event(data: &[u8]) -> Result<TaskEvent> {
136
    if data.len() < std::mem::size_of::<TaskEvent>() {
137
        anyhow::bail!("Invalid event data size: {} bytes", data.len());
138
    }
139

140
    let event = unsafe {
141
        std::ptr::read_unaligned(data.as_ptr() as *const TaskEvent)
142
    };
143

144
    Ok(event)
145
}
146

147
async fn process_events(
148
    mut ring_buf: RingBuf<&'static mut [u8]>,
149
    running: Arc<AtomicBool>,
150
    args: &Args,
151
) -> Result<()> {
152
    let mut stats = Statistics::new();
153
    let mut events_processed = 0u64;
154

155
    info!("Starting event processing...");
156

157
    // Spawn statistics reporting task
158
    let stats_running = running.clone();
159
    let stats_interval = args.stats_interval;
160
    let stats_handle = tokio::spawn(async move {
161
        let mut interval = tokio::time::interval(Duration::from_secs(stats_interval));
162

163
        while stats_running.load(Ordering::Relaxed) {
164
            interval.tick().await;
165
            // Statistics will be printed by main loop
166
        }
167
    });
168

169
    while running.load(Ordering::Relaxed) {
170
        // Poll for events with timeout
171
        match ring_buf.next() {
172
            Some(item) => {
173
                match parse_task_event(&item) {
174
                    Ok(event) => {
175
                        events_processed += 1;
176
                        stats.update(&event);
177

178
                        if args.verbose {
179
                            print_event_details(&event);
180
                        } else {
181
                            print_event_summary(&event);
182
                        }
183

184
                        // Check if we should stop based on max_events
185
                        if args.max_events > 0 && events_processed >= args.max_events {
186
                            info!("Reached maximum events limit ({}), stopping...", args.max_events);
187
                            break;
188
                        }
189
                    }
190
                    Err(e) => {
191
                        warn!("Failed to parse event: {}", e);
192
                    }
193
                }
194
            }
195
            None => {
196
                // No events available, sleep briefly
197
                sleep(Duration::from_millis(10)).await;
198
            }
199
        }
200

201
        // Print periodic statistics
202
        if events_processed % 1000 == 0 && events_processed > 0 {
203
            stats.print_summary();
204
        }
205
    }
206

207
    // Final statistics
208
    stats.print_summary();
209

210
    // Clean up
211
    stats_handle.abort();
212

213
    Ok(())
214
}
215

216
fn print_event_summary(event: &TaskEvent) {
217
    let comm = String::from_utf8_lossy(&event.comm).trim_end_matches('\0');
218
    let timestamp_secs = event.timestamp / 1_000_000_000;
219
    let timestamp_us = (event.timestamp % 1_000_000_000) / 1_000;
220

221
    println!(
222
        "[{}.{:06}] wake_up_new_task: caller PID {}, new task PID {} ({}), TGID {}",
223
        timestamp_secs,
224
        timestamp_us,
225
        event.caller_pid,
226
        event.new_task_pid,
227
        comm,
228
        event.new_task_tgid
229
    );
230
}
231

232
fn print_event_details(event: &TaskEvent) {
233
    let comm = String::from_utf8_lossy(&event.comm).trim_end_matches('\0');
234
    let timestamp_secs = event.timestamp / 1_000_000_000;
235
    let timestamp_us = (event.timestamp % 1_000_000_000) / 1_000;
236

237
    println!("=== Task Wake-up Event ===");
238
    println!("Timestamp: {}.{:06}", timestamp_secs, timestamp_us);
239
    println!("Caller PID: {}", event.caller_pid);
240
    println!("Caller TGID: {}", event.caller_tgid);
241
    println!("New Task PID: {}", event.new_task_pid);
242
    println!("New Task TGID: {}", event.new_task_tgid);
243
    println!("Process Name: {}", comm);
244
    println!("========================\n");
245
}
246

247
#[tokio::main]
248
async fn main() -> Result<()> {
249
    let args = Args::parse();
250

251
    // Initialize logging
252
    env_logger::Builder::from_default_env()
253
        .filter_level(if args.verbose {
254
            log::LevelFilter::Debug
255
        } else {
256
            log::LevelFilter::Info
257
        })
258
        .init();
259

260
    info!("Starting eBPF Kprobe Observer");
261

262
    // Check if we're running as root
263
    if !nix::unistd::Uid::effective().is_root() {
264
        error!("This program requires root privileges to load eBPF programs");
265
        std::process::exit(1);
266
    }
267

268
    // Load and attach eBPF program
269
    let (_bpf, ring_buf) = load_and_attach_ebpf().await?;
270

271
    // Set up signal handling
272
    let running = Arc::new(AtomicBool::new(true));
273
    let r = running.clone();
274

275
    tokio::spawn(async move {
276
        signal::ctrl_c().await.expect("Failed to listen for Ctrl+C");
277
        info!("Received Ctrl+C, shutting down...");
278
        r.store(false, Ordering::Relaxed);
279
    });
280

281
    info!("Monitoring wake_up_new_task events... Press Ctrl+C to exit");
282

283
    // Process events
284
    process_events(ring_buf, running, &args).await?;
285

286
    info!("Observer shutting down");
287

288
    Ok(())
289
}