1 // SPDX-License-Identifier: GPL-2.0
2 /* Copyright (c) 2025 Meta Platforms, Inc. and affiliates. */
3
4 #include <vmlinux.h>
5 #include <string.h>
6 #include <stdbool.h>
7 #include <bpf/bpf_helpers.h>
8 #include <bpf/bpf_tracing.h>
9 #include "bpf_misc.h"
10 #include "errno.h"
11
12 char _license[] SEC("license") = "GPL";
13
14 const void *user_ptr = NULL;
15
16 struct elem {
17 char data[128];
18 struct bpf_task_work tw;
19 };
20
21 struct {
22 __uint(type, BPF_MAP_TYPE_HASH);
23 __uint(map_flags, BPF_F_NO_PREALLOC);
24 __uint(max_entries, 1);
25 __type(key, int);
26 __type(value, struct elem);
27 } hmap SEC(".maps");
28
29 struct {
30 __uint(type, BPF_MAP_TYPE_ARRAY);
31 __uint(max_entries, 1);
32 __type(key, int);
33 __type(value, struct elem);
34 } arrmap SEC(".maps");
35
36 struct {
37 __uint(type, BPF_MAP_TYPE_LRU_HASH);
38 __uint(max_entries, 1);
39 __type(key, int);
40 __type(value, struct elem);
41 } lrumap SEC(".maps");
42
process_work(struct bpf_map * map,void * key,void * value)43 static int process_work(struct bpf_map *map, void *key, void *value)
44 {
45 struct elem *work = value;
46
47 bpf_copy_from_user_str(work->data, sizeof(work->data), (const void *)user_ptr, 0);
48 return 0;
49 }
50
51 int key = 0;
52
53 SEC("perf_event")
oncpu_hash_map(struct pt_regs * args)54 int oncpu_hash_map(struct pt_regs *args)
55 {
56 struct elem empty_work = { .data = { 0 } };
57 struct elem *work;
58 struct task_struct *task;
59 int err;
60
61 task = bpf_get_current_task_btf();
62 err = bpf_map_update_elem(&hmap, &key, &empty_work, BPF_NOEXIST);
63 if (err)
64 return 0;
65 work = bpf_map_lookup_elem(&hmap, &key);
66 if (!work)
67 return 0;
68
69 bpf_task_work_schedule_resume_impl(task, &work->tw, &hmap, process_work, NULL);
70 return 0;
71 }
72
73 SEC("perf_event")
oncpu_array_map(struct pt_regs * args)74 int oncpu_array_map(struct pt_regs *args)
75 {
76 struct elem *work;
77 struct task_struct *task;
78
79 task = bpf_get_current_task_btf();
80 work = bpf_map_lookup_elem(&arrmap, &key);
81 if (!work)
82 return 0;
83 bpf_task_work_schedule_signal_impl(task, &work->tw, &arrmap, process_work, NULL);
84 return 0;
85 }
86
87 SEC("perf_event")
oncpu_lru_map(struct pt_regs * args)88 int oncpu_lru_map(struct pt_regs *args)
89 {
90 struct elem empty_work = { .data = { 0 } };
91 struct elem *work;
92 struct task_struct *task;
93 int err;
94
95 task = bpf_get_current_task_btf();
96 work = bpf_map_lookup_elem(&lrumap, &key);
97 if (work)
98 return 0;
99 err = bpf_map_update_elem(&lrumap, &key, &empty_work, BPF_NOEXIST);
100 if (err)
101 return 0;
102 work = bpf_map_lookup_elem(&lrumap, &key);
103 if (!work || work->data[0])
104 return 0;
105 bpf_task_work_schedule_resume_impl(task, &work->tw, &lrumap, process_work, NULL);
106 return 0;
107 }
108