Skip to content

Commit

Permalink
bpf: add k/kr for bpf_map_lookup/update/delete
Browse files Browse the repository at this point in the history
  • Loading branch information
jschwinger233 committed Dec 2, 2024
1 parent 141e298 commit e6de8f9
Showing 1 changed file with 150 additions and 12 deletions.
162 changes: 150 additions & 12 deletions bpf/kprobe_pwru.c
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ struct event_t {
u64 ts;
u64 print_skb_id;
u64 print_shinfo_id;
//u64 print_bpf_map_id;
u64 print_bpfmap_id;
struct skb_meta meta;
struct tuple tuple;
s64 print_stack_id;
Expand Down Expand Up @@ -144,8 +144,7 @@ struct config {
u8 output_shinfo: 1;
u8 output_stack: 1;
u8 output_caller: 1;
u8 output_bpf_map: 1;
u8 output_unused: 1;
u8 output_unused: 2;
u8 is_set: 1;
u8 track_skb: 1;
u8 track_skb_by_stackid: 1;
Expand All @@ -166,6 +165,13 @@ struct {
__uint(value_size, MAX_STACK_DEPTH * sizeof(u64));
} print_stack_map SEC(".maps");

struct {
__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
__uint(max_entries, 1);
__type(key, u32);
__type(value, struct event_t);
} event_stash SEC(".maps");

struct print_skb_value {
u32 len;
char str[PRINT_SKB_STR_SIZE];
Expand Down Expand Up @@ -199,6 +205,28 @@ struct {
__type(value, struct print_shinfo_value);
} print_shinfo_map SEC(".maps");

struct print_bpfmap_value {
u32 id;
char name[16];
u32 key_size;
u32 value_size;
u8 key[128];
u8 value[128];
} __attribute__((packed));

struct {
__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
__uint(max_entries, 1);
__type(key, u32);
__type(value, u32);
} print_bpfmap_id_map SEC(".maps");

struct {
__uint(type, BPF_MAP_TYPE_HASH);
__type(key, u64);
__type(value, struct print_bpfmap_value);
} print_bpfmap_map SEC(".maps");

static __always_inline u32
get_netns(struct sk_buff *skb) {
u32 netns = BPF_CORE_READ(skb, dev, nd_net.net, ns.inum);
Expand Down Expand Up @@ -517,6 +545,7 @@ handle_everything(struct sk_buff *skb, void *ctx, struct event_t *event, u64 *_s
bpf_map_update_elem(&skb_stackid, &skb, &stackid, BPF_ANY);
}

event->skb_addr = (u64) skb;
event->pid = bpf_get_current_pid_tgid() >> 32;
event->ts = bpf_ktime_get_ns();
event->cpu_id = bpf_get_smp_processor_id();
Expand All @@ -531,20 +560,12 @@ kprobe_skb(struct sk_buff *skb, struct pt_regs *ctx, const bool has_get_func_ip,
if (!handle_everything(skb, ctx, &event, _stackid, true))
return BPF_OK;

event.skb_addr = (u64) skb;
event.addr = has_get_func_ip ? bpf_get_func_ip(ctx) : PT_REGS_IP(ctx);
event.param_second = PT_REGS_PARM2(ctx);
event.param_third = PT_REGS_PARM3(ctx);
if (CFG.output_caller)
bpf_probe_read_kernel(&event.caller_addr, sizeof(event.caller_addr), (void *)PT_REGS_SP(ctx));

if (CFG.output_bpf_map) {
// TODO@gray: kernel>=5.15
__u64 cookie = bpf_get_attach_cookie(ctx);
if (cookie)
set_bpf_map(ctx, cookie, NULL);
}

bpf_map_push_elem(&events, &event, BPF_EXIST);

return BPF_OK;
Expand Down Expand Up @@ -631,7 +652,6 @@ int BPF_PROG(fentry_tc, struct sk_buff *skb) {
if (!handle_everything(skb, ctx, &event, NULL, false))
return BPF_OK;

event.skb_addr = (u64) skb;
event.addr = BPF_PROG_ADDR;
event.type = EVENT_TYPE_TC;
bpf_map_push_elem(&events, &event, BPF_EXIST);
Expand Down Expand Up @@ -781,4 +801,122 @@ int kretprobe_veth_convert_skb_to_xdp_buff(struct pt_regs *ctx) {
return BPF_OK;
}

static __always_inline void
set_common_bpfmap_info(struct pt_regs *ctx, u64 *event_id,
struct print_bpfmap_value *map_value) {
*event_id = sync_fetch_and_add(&print_bpfmap_id_map);

struct bpf_map *map = (struct bpf_map *)PT_REGS_PARM1(ctx);
BPF_CORE_READ_INTO(&map_value->id, map, id);
BPF_CORE_READ_STR_INTO(&map_value->name, map, name);
BPF_CORE_READ_INTO(&map_value->key_size, map, key_size);
BPF_CORE_READ_INTO(&map_value->value_size, map, value_size);
bpf_probe_read_kernel(&map_value->key, sizeof(map_value->key), (void *)PT_REGS_PARM2(ctx));
}

SEC("kprobe/bpf_map_lookup")
int kprobe_bpf_map_lookup(struct pt_regs *ctx) {
u64 stackid = get_stackid(ctx, true);

struct sk_buff **skb = bpf_map_lookup_elem(&stackid_skb, &stackid);
if (skb && *skb) {
struct event_t event = {};

event.addr = PT_REGS_IP(ctx);
if (!handle_everything(*skb, ctx, &event, &stackid, true))
return BPF_OK;

if (CFG.output_caller)
bpf_probe_read_kernel(&event.caller_addr,
sizeof(event.caller_addr),
(void *)PT_REGS_SP(ctx));

struct print_bpfmap_value map_value = {};
set_common_bpfmap_info(ctx, &event.print_bpfmap_id, &map_value);

u64 pid_tgid = bpf_get_current_pid_tgid();
bpf_map_update_elem(&print_bpfmap_map, &event.print_bpfmap_id, &map_value, BPF_ANY);
bpf_map_update_elem(&event_stash, &ZERO, &event, BPF_ANY);
}

return BPF_OK;
}

SEC("kretprobe/bpf_map_lookup")
int kretprobe_bpf_map_lookup(struct pt_regs *ctx) {
struct event_t *event = bpf_map_lookup_elem(&event_stash, &ZERO);
if (!event)
return BPF_OK;

struct print_bpfmap_value *map_value = bpf_map_lookup_elem(&print_bpfmap_map,
&event->print_bpfmap_id);
if (!map_value)
return BPF_OK;

bpf_probe_read_kernel(&map_value->value,
sizeof(map_value->value),
(void *)PT_REGS_RC(ctx));

bpf_map_push_elem(&events, &event, BPF_EXIST);
return BPF_OK;
}

SEC("kprobe/bpf_map_update")
int kprobe_bpf_map_update(struct pt_regs *ctx) {
u64 stackid = get_stackid(ctx, true);

struct sk_buff **skb = bpf_map_lookup_elem(&stackid_skb, &stackid);
if (skb && *skb) {
struct event_t event = {};

event.addr = PT_REGS_IP(ctx);
if (!handle_everything(*skb, ctx, &event, &stackid, true))
return BPF_OK;

if (CFG.output_caller)
bpf_probe_read_kernel(&event.caller_addr,
sizeof(event.caller_addr),
(void *)PT_REGS_SP(ctx));

// todo@gray: static?
static struct print_bpfmap_value map_value = {};
set_common_bpfmap_info(ctx, &event.print_bpfmap_id, &map_value);

bpf_probe_read_kernel(&map_value.value,
sizeof(map_value.value),
(void *)PT_REGS_PARM3(ctx));
bpf_map_update_elem(&print_bpfmap_map, &event.print_bpfmap_id, &map_value, BPF_ANY);
bpf_map_push_elem(&events, &event, BPF_EXIST);
}

return BPF_OK;
}

SEC("kprobe/bpf_map_delete")
int kprobe_bpf_map_delete(struct pt_regs *ctx) {
u64 stackid = get_stackid(ctx, true);

struct sk_buff **skb = bpf_map_lookup_elem(&stackid_skb, &stackid);
if (skb && *skb) {
struct event_t event = {};

event.addr = PT_REGS_IP(ctx);
if (!handle_everything(*skb, ctx, &event, &stackid, true))
return BPF_OK;

if (CFG.output_caller)
bpf_probe_read_kernel(&event.caller_addr,
sizeof(event.caller_addr),
(void *)PT_REGS_SP(ctx));

static struct print_bpfmap_value map_value = {};
set_common_bpfmap_info(ctx, &event.print_bpfmap_id, &map_value);

bpf_map_update_elem(&print_bpfmap_map, &event.print_bpfmap_id, &map_value, BPF_ANY);
bpf_map_push_elem(&events, &event, BPF_EXIST);
}

return BPF_OK;
}

char __license[] SEC("license") = "Dual BSD/GPL";

0 comments on commit e6de8f9

Please sign in to comment.