1 // SPDX-License-Identifier: GPL-2.0
2 /* Copyright (c) 2021 Facebook */
3 #include <linux/stddef.h>
4 #include <linux/bpf.h>
5 #include <bpf/bpf_helpers.h>
6 #include <bpf/bpf_tracing.h>
7 #include <../../../tools/include/linux/filter.h>
8 #include <linux/btf.h>
9 #include <string.h>
10 #include <errno.h>
11 
12 char _license[] SEC("license") = "GPL";
13 
14 struct bpf_map {
15 	int id;
16 }  __attribute__((preserve_access_index));
17 
18 struct args {
19 	__u64 log_buf;
20 	__u32 log_size;
21 	int max_entries;
22 	int map_fd;
23 	int prog_fd;
24 	int btf_fd;
25 };
26 
27 #define BTF_INFO_ENC(kind, kind_flag, vlen) \
28 	((!!(kind_flag) << 31) | ((kind) << 24) | ((vlen) & BTF_MAX_VLEN))
29 #define BTF_TYPE_ENC(name, info, size_or_type) (name), (info), (size_or_type)
30 #define BTF_INT_ENC(encoding, bits_offset, nr_bits) \
31 	((encoding) << 24 | (bits_offset) << 16 | (nr_bits))
32 #define BTF_TYPE_INT_ENC(name, encoding, bits_offset, bits, sz) \
33 	BTF_TYPE_ENC(name, BTF_INFO_ENC(BTF_KIND_INT, 0, 0), sz), \
34 	BTF_INT_ENC(encoding, bits_offset, bits)
35 
36 struct {
37 	__uint(type, BPF_MAP_TYPE_ARRAY);
38 	__type(key, int);
39 	__type(value, union bpf_attr);
40 	__uint(max_entries, 1);
41 } bpf_attr_array SEC(".maps");
42 
43 struct inner_map_type {
44 	__uint(type, BPF_MAP_TYPE_ARRAY);
45 	__uint(key_size, 4);
46 	__uint(value_size, 4);
47 	__uint(max_entries, 1);
48 } inner_map SEC(".maps");
49 
50 struct {
51 	__uint(type, BPF_MAP_TYPE_ARRAY_OF_MAPS);
52 	__type(key, int);
53 	__type(value, int);
54 	__uint(max_entries, 1);
55 	__array(values, struct inner_map_type);
56 } outer_array_map SEC(".maps") = {
57 	.values = {
58 		[0] = &inner_map,
59 	},
60 };
61 
62 static inline __u64 ptr_to_u64(const void *ptr)
63 {
64 	return (__u64) (unsigned long) ptr;
65 }
66 
67 static int btf_load(void)
68 {
69 	struct btf_blob {
70 		struct btf_header btf_hdr;
71 		__u32 types[8];
72 		__u32 str;
73 	} raw_btf = {
74 		.btf_hdr = {
75 			.magic = BTF_MAGIC,
76 			.version = BTF_VERSION,
77 			.hdr_len = sizeof(struct btf_header),
78 			.type_len = sizeof(__u32) * 8,
79 			.str_off = sizeof(__u32) * 8,
80 			.str_len = sizeof(__u32),
81 		},
82 		.types = {
83 			/* long */
84 			BTF_TYPE_INT_ENC(0, BTF_INT_SIGNED, 0, 64, 8),  /* [1] */
85 			/* unsigned long */
86 			BTF_TYPE_INT_ENC(0, 0, 0, 64, 8),  /* [2] */
87 		},
88 	};
89 	static union bpf_attr btf_load_attr = {
90 		.btf_size = sizeof(raw_btf),
91 	};
92 
93 	btf_load_attr.btf = (long)&raw_btf;
94 	return bpf_sys_bpf(BPF_BTF_LOAD, &btf_load_attr, sizeof(btf_load_attr));
95 }
96 
97 SEC("syscall")
98 int load_prog(struct args *ctx)
99 {
100 	static char license[] = "GPL";
101 	static struct bpf_insn insns[] = {
102 		BPF_ST_MEM(BPF_DW, BPF_REG_10, -8, 0),
103 		BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),
104 		BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -8),
105 		BPF_LD_MAP_FD(BPF_REG_1, 0),
106 		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_lookup_elem),
107 		BPF_MOV64_IMM(BPF_REG_0, 0),
108 		BPF_EXIT_INSN(),
109 	};
110 	static union bpf_attr map_create_attr = {
111 		.map_type = BPF_MAP_TYPE_HASH,
112 		.key_size = 8,
113 		.value_size = 8,
114 		.btf_key_type_id = 1,
115 		.btf_value_type_id = 2,
116 	};
117 	static union bpf_attr map_update_attr = { .map_fd = 1, };
118 	static __u64 key = 12;
119 	static __u64 value = 34;
120 	static union bpf_attr prog_load_attr = {
121 		.prog_type = BPF_PROG_TYPE_XDP,
122 		.insn_cnt = sizeof(insns) / sizeof(insns[0]),
123 	};
124 	int ret;
125 
126 	ret = btf_load();
127 	if (ret <= 0)
128 		return ret;
129 
130 	ctx->btf_fd = ret;
131 	map_create_attr.max_entries = ctx->max_entries;
132 	map_create_attr.btf_fd = ret;
133 
134 	prog_load_attr.license = ptr_to_u64(license);
135 	prog_load_attr.insns = ptr_to_u64(insns);
136 	prog_load_attr.log_buf = ctx->log_buf;
137 	prog_load_attr.log_size = ctx->log_size;
138 	prog_load_attr.log_level = 1;
139 
140 	ret = bpf_sys_bpf(BPF_MAP_CREATE, &map_create_attr, sizeof(map_create_attr));
141 	if (ret <= 0)
142 		return ret;
143 	ctx->map_fd = ret;
144 	insns[3].imm = ret;
145 
146 	map_update_attr.map_fd = ret;
147 	map_update_attr.key = ptr_to_u64(&key);
148 	map_update_attr.value = ptr_to_u64(&value);
149 	ret = bpf_sys_bpf(BPF_MAP_UPDATE_ELEM, &map_update_attr, sizeof(map_update_attr));
150 	if (ret < 0)
151 		return ret;
152 
153 	ret = bpf_sys_bpf(BPF_PROG_LOAD, &prog_load_attr, sizeof(prog_load_attr));
154 	if (ret <= 0)
155 		return ret;
156 	ctx->prog_fd = ret;
157 	return 1;
158 }
159 
160 SEC("syscall")
161 int update_outer_map(void *ctx)
162 {
163 	int zero = 0, ret = 0, outer_fd = -1, inner_fd = -1, err;
164 	const int attr_sz = sizeof(union bpf_attr);
165 	union bpf_attr *attr;
166 
167 	attr = bpf_map_lookup_elem((struct bpf_map *)&bpf_attr_array, &zero);
168 	if (!attr)
169 		goto out;
170 
171 	memset(attr, 0, attr_sz);
172 	attr->map_id = ((struct bpf_map *)&outer_array_map)->id;
173 	outer_fd = bpf_sys_bpf(BPF_MAP_GET_FD_BY_ID, attr, attr_sz);
174 	if (outer_fd < 0)
175 		goto out;
176 
177 	memset(attr, 0, attr_sz);
178 	attr->map_type = BPF_MAP_TYPE_ARRAY;
179 	attr->key_size = 4;
180 	attr->value_size = 4;
181 	attr->max_entries = 1;
182 	inner_fd = bpf_sys_bpf(BPF_MAP_CREATE, attr, attr_sz);
183 	if (inner_fd < 0)
184 		goto out;
185 
186 	memset(attr, 0, attr_sz);
187 	attr->map_fd = outer_fd;
188 	attr->key = ptr_to_u64(&zero);
189 	attr->value = ptr_to_u64(&inner_fd);
190 	err = bpf_sys_bpf(BPF_MAP_UPDATE_ELEM, attr, attr_sz);
191 	if (err)
192 		goto out;
193 
194 	memset(attr, 0, attr_sz);
195 	attr->map_fd = outer_fd;
196 	attr->key = ptr_to_u64(&zero);
197 	err = bpf_sys_bpf(BPF_MAP_DELETE_ELEM, attr, attr_sz);
198 	if (err)
199 		goto out;
200 	ret = 1;
201 out:
202 	if (inner_fd >= 0)
203 		bpf_sys_close(inner_fd);
204 	if (outer_fd >= 0)
205 		bpf_sys_close(outer_fd);
206 	return ret;
207 }
208