#!/usr/bin/env python # Copyright (c) PLUMgrid, Inc. # Licensed under the Apache License, Version 2.0 (the "License") from bcc import BPF from unittest import main, TestCase class TestClang(TestCase): def test_complex(self): b = BPF(src_file="test_clang_complex.c", debug=0) fn = b.load_func("handle_packet", BPF.SCHED_CLS) def test_printk(self): text = """ #include <bcc/proto.h> int handle_packet(void *ctx) { u8 *cursor = 0; struct ethernet_t *ethernet = cursor_advance(cursor, sizeof(*ethernet)); bpf_trace_printk("ethernet->dst = %llx, ethernet->src = %llx\\n", ethernet->dst, ethernet->src); return 0; } """ b = BPF(text=text, debug=0) fn = b.load_func("handle_packet", BPF.SCHED_CLS) def test_probe_read1(self): text = """ #include <linux/sched.h> #include <uapi/linux/ptrace.h> int count_sched(struct pt_regs *ctx, struct task_struct *prev) { pid_t p = prev->pid; return (p != -1); } """ b = BPF(text=text, debug=0) fn = b.load_func("count_sched", BPF.KPROBE) def test_probe_read2(self): text = """ #include <linux/sched.h> #include <uapi/linux/ptrace.h> int count_foo(struct pt_regs *ctx, unsigned long a, unsigned long b) { return (a != b); } """ b = BPF(text=text, debug=0) fn = b.load_func("count_foo", BPF.KPROBE) def test_probe_read_keys(self): text = """ #include <uapi/linux/ptrace.h> #include <linux/blkdev.h> BPF_HASH(start, struct request *); int do_request(struct pt_regs *ctx, struct request *req) { u64 ts = bpf_ktime_get_ns(); start.update(&req, &ts); return 0; } int do_completion(struct pt_regs *ctx, struct request *req) { u64 *tsp = start.lookup(&req); if (tsp != 0) { start.delete(&req); } return 0; } """ b = BPF(text=text, debug=0) fns = b.load_funcs(BPF.KPROBE) def test_sscanf(self): text = """ BPF_TABLE("hash", int, struct { u64 a; u64 b; u64 c:36; u64 d:28; struct { u32 a; u32 b; } s; }, stats, 10); int foo(void *ctx) { return 0; } """ b = BPF(text=text, debug=0) fn = b.load_func("foo", BPF.KPROBE) t = b.get_table("stats") s1 = t.key_sprintf(t.Key(2)) self.assertEqual(s1, b"0x2") s2 = t.leaf_sprintf(t.Leaf(2, 3, 4, 1, (5, 6))) l = t.leaf_scanf(s2) self.assertEqual(l.a, 2) self.assertEqual(l.b, 3) self.assertEqual(l.c, 4) self.assertEqual(l.d, 1) self.assertEqual(l.s.a, 5) self.assertEqual(l.s.b, 6) def test_iosnoop(self): text = """ #include <linux/blkdev.h> #include <uapi/linux/ptrace.h> struct key_t { struct request *req; }; BPF_TABLE("hash", struct key_t, u64, start, 1024); int do_request(struct pt_regs *ctx, struct request *req) { struct key_t key = {}; bpf_trace_printk("traced start %d\\n", req->__data_len); return 0; } """ b = BPF(text=text, debug=0) fn = b.load_func("do_request", BPF.KPROBE) def test_blk_start_request(self): text = """ #include <linux/blkdev.h> #include <uapi/linux/ptrace.h> int do_request(struct pt_regs *ctx, int req) { bpf_trace_printk("req ptr: 0x%x\\n", req); return 0; } """ b = BPF(text=text, debug=0) fn = b.load_func("do_request", BPF.KPROBE) def test_bpf_hash(self): text = """ BPF_HASH(table1); BPF_HASH(table2, u32); BPF_HASH(table3, u32, int); """ b = BPF(text=text, debug=0) def test_consecutive_probe_read(self): text = """ #include <linux/fs.h> #include <linux/mount.h> BPF_HASH(table1, struct super_block *); int trace_entry(struct pt_regs *ctx, struct file *file) { if (!file) return 0; struct vfsmount *mnt = file->f_path.mnt; if (mnt) { struct super_block *k = mnt->mnt_sb; u64 zero = 0; table1.update(&k, &zero); k = mnt->mnt_sb; table1.update(&k, &zero); } return 0; } """ b = BPF(text=text, debug=0) fn = b.load_func("trace_entry", BPF.KPROBE) def test_nested_probe_read(self): text = """ #include <linux/fs.h> int trace_entry(struct pt_regs *ctx, struct file *file) { if (!file) return 0; const char *name = file->f_path.dentry->d_name.name; bpf_trace_printk("%s\\n", name); return 0; } """ b = BPF(text=text, debug=0) fn = b.load_func("trace_entry", BPF.KPROBE) def test_char_array_probe(self): BPF(text="""#include <linux/blkdev.h> int kprobe__blk_update_request(struct pt_regs *ctx, struct request *req) { bpf_trace_printk("%s\\n", req->rq_disk->disk_name); return 0; }""") def test_probe_read_helper(self): b = BPF(text=""" #include <linux/fs.h> static void print_file_name(struct file *file) { if (!file) return; const char *name = file->f_path.dentry->d_name.name; bpf_trace_printk("%s\\n", name); } static void print_file_name2(int unused, struct file *file) { print_file_name(file); } int trace_entry1(struct pt_regs *ctx, struct file *file) { print_file_name(file); return 0; } int trace_entry2(struct pt_regs *ctx, int unused, struct file *file) { print_file_name2(unused, file); return 0; } """) fn = b.load_func("trace_entry1", BPF.KPROBE) fn = b.load_func("trace_entry2", BPF.KPROBE) def test_probe_struct_assign(self): b = BPF(text = """ #include <uapi/linux/ptrace.h> struct args_t { const char *filename; int flags; int mode; }; int kprobe__sys_open(struct pt_regs *ctx, const char *filename, int flags, int mode) { struct args_t args = {}; args.filename = filename; args.flags = flags; args.mode = mode; bpf_trace_printk("%s\\n", args.filename); return 0; }; """) def test_task_switch(self): b = BPF(text=""" #include <uapi/linux/ptrace.h> #include <linux/sched.h> struct key_t { u32 prev_pid; u32 curr_pid; }; BPF_TABLE("hash", struct key_t, u64, stats, 1024); int kprobe__finish_task_switch(struct pt_regs *ctx, struct task_struct *prev) { struct key_t key = {}; u64 zero = 0, *val; key.curr_pid = bpf_get_current_pid_tgid(); key.prev_pid = prev->pid; val = stats.lookup_or_init(&key, &zero); (*val)++; return 0; } """) def test_probe_simple_assign(self): b = BPF(text=""" #include <uapi/linux/ptrace.h> #include <linux/gfp.h> struct leaf { size_t size; }; BPF_HASH(simple_map, u32, struct leaf); int kprobe____kmalloc(struct pt_regs *ctx, size_t size) { u32 pid = bpf_get_current_pid_tgid(); struct leaf* leaf = simple_map.lookup(&pid); if (leaf) leaf->size += size; return 0; }""", debug=4) if __name__ == "__main__": main()