xref: /openbmc/linux/samples/bpf/test_cgrp2_sock.c (revision 1af5e8b1)
1 /* eBPF example program:
2  *
3  * - Loads eBPF program
4  *
5  *   The eBPF program sets the sk_bound_dev_if index in new AF_INET{6}
6  *   sockets opened by processes in the cgroup.
7  *
8  * - Attaches the new program to a cgroup using BPF_PROG_ATTACH
9  */
10 
11 #define _GNU_SOURCE
12 
13 #include <stdio.h>
14 #include <stdlib.h>
15 #include <stddef.h>
16 #include <string.h>
17 #include <unistd.h>
18 #include <assert.h>
19 #include <errno.h>
20 #include <fcntl.h>
21 #include <net/if.h>
22 #include <inttypes.h>
23 #include <linux/bpf.h>
24 #include <bpf/bpf.h>
25 
26 #include "bpf_insn.h"
27 
28 char bpf_log_buf[BPF_LOG_BUF_SIZE];
29 
30 static int prog_load(__u32 idx, __u32 mark, __u32 prio)
31 {
32 	/* save pointer to context */
33 	struct bpf_insn prog_start[] = {
34 		BPF_MOV64_REG(BPF_REG_6, BPF_REG_1),
35 	};
36 	struct bpf_insn prog_end[] = {
37 		BPF_MOV64_IMM(BPF_REG_0, 1), /* r0 = verdict */
38 		BPF_EXIT_INSN(),
39 	};
40 
41 	/* set sk_bound_dev_if on socket */
42 	struct bpf_insn prog_dev[] = {
43 		BPF_MOV64_IMM(BPF_REG_3, idx),
44 		BPF_MOV64_IMM(BPF_REG_2, offsetof(struct bpf_sock, bound_dev_if)),
45 		BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3, offsetof(struct bpf_sock, bound_dev_if)),
46 	};
47 
48 	/* set mark on socket */
49 	struct bpf_insn prog_mark[] = {
50 		/* get uid of process */
51 		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
52 			     BPF_FUNC_get_current_uid_gid),
53 		BPF_ALU64_IMM(BPF_AND, BPF_REG_0, 0xffffffff),
54 
55 		/* if uid is 0, use given mark, else use the uid as the mark */
56 		BPF_MOV64_REG(BPF_REG_3, BPF_REG_0),
57 		BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 1),
58 		BPF_MOV64_IMM(BPF_REG_3, mark),
59 
60 		/* set the mark on the new socket */
61 		BPF_MOV64_REG(BPF_REG_1, BPF_REG_6),
62 		BPF_MOV64_IMM(BPF_REG_2, offsetof(struct bpf_sock, mark)),
63 		BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3, offsetof(struct bpf_sock, mark)),
64 	};
65 
66 	/* set priority on socket */
67 	struct bpf_insn prog_prio[] = {
68 		BPF_MOV64_REG(BPF_REG_1, BPF_REG_6),
69 		BPF_MOV64_IMM(BPF_REG_3, prio),
70 		BPF_MOV64_IMM(BPF_REG_2, offsetof(struct bpf_sock, priority)),
71 		BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3, offsetof(struct bpf_sock, priority)),
72 	};
73 	LIBBPF_OPTS(bpf_prog_load_opts, opts,
74 		.log_buf = bpf_log_buf,
75 		.log_size = BPF_LOG_BUF_SIZE,
76 	);
77 
78 	struct bpf_insn *prog;
79 	size_t insns_cnt;
80 	void *p;
81 	int ret;
82 
83 	insns_cnt = sizeof(prog_start) + sizeof(prog_end);
84 	if (idx)
85 		insns_cnt += sizeof(prog_dev);
86 
87 	if (mark)
88 		insns_cnt += sizeof(prog_mark);
89 
90 	if (prio)
91 		insns_cnt += sizeof(prog_prio);
92 
93 	p = prog = malloc(insns_cnt);
94 	if (!prog) {
95 		fprintf(stderr, "Failed to allocate memory for instructions\n");
96 		return EXIT_FAILURE;
97 	}
98 
99 	memcpy(p, prog_start, sizeof(prog_start));
100 	p += sizeof(prog_start);
101 
102 	if (idx) {
103 		memcpy(p, prog_dev, sizeof(prog_dev));
104 		p += sizeof(prog_dev);
105 	}
106 
107 	if (mark) {
108 		memcpy(p, prog_mark, sizeof(prog_mark));
109 		p += sizeof(prog_mark);
110 	}
111 
112 	if (prio) {
113 		memcpy(p, prog_prio, sizeof(prog_prio));
114 		p += sizeof(prog_prio);
115 	}
116 
117 	memcpy(p, prog_end, sizeof(prog_end));
118 	p += sizeof(prog_end);
119 
120 	insns_cnt /= sizeof(struct bpf_insn);
121 
122 	ret = bpf_prog_load(BPF_PROG_TYPE_CGROUP_SOCK, NULL, "GPL",
123 			    prog, insns_cnt, &opts);
124 
125 	free(prog);
126 
127 	return ret;
128 }
129 
130 static int get_bind_to_device(int sd, char *name, size_t len)
131 {
132 	socklen_t optlen = len;
133 	int rc;
134 
135 	name[0] = '\0';
136 	rc = getsockopt(sd, SOL_SOCKET, SO_BINDTODEVICE, name, &optlen);
137 	if (rc < 0)
138 		perror("setsockopt(SO_BINDTODEVICE)");
139 
140 	return rc;
141 }
142 
143 static unsigned int get_somark(int sd)
144 {
145 	unsigned int mark = 0;
146 	socklen_t optlen = sizeof(mark);
147 	int rc;
148 
149 	rc = getsockopt(sd, SOL_SOCKET, SO_MARK, &mark, &optlen);
150 	if (rc < 0)
151 		perror("getsockopt(SO_MARK)");
152 
153 	return mark;
154 }
155 
156 static unsigned int get_priority(int sd)
157 {
158 	unsigned int prio = 0;
159 	socklen_t optlen = sizeof(prio);
160 	int rc;
161 
162 	rc = getsockopt(sd, SOL_SOCKET, SO_PRIORITY, &prio, &optlen);
163 	if (rc < 0)
164 		perror("getsockopt(SO_PRIORITY)");
165 
166 	return prio;
167 }
168 
169 static int show_sockopts(int family)
170 {
171 	unsigned int mark, prio;
172 	char name[16];
173 	int sd;
174 
175 	sd = socket(family, SOCK_DGRAM, 17);
176 	if (sd < 0) {
177 		perror("socket");
178 		return 1;
179 	}
180 
181 	if (get_bind_to_device(sd, name, sizeof(name)) < 0) {
182 		close(sd);
183 		return 1;
184 	}
185 
186 	mark = get_somark(sd);
187 	prio = get_priority(sd);
188 
189 	close(sd);
190 
191 	printf("sd %d: dev %s, mark %u, priority %u\n", sd, name, mark, prio);
192 
193 	return 0;
194 }
195 
196 static int usage(const char *argv0)
197 {
198 	printf("Usage:\n");
199 	printf("  Attach a program\n");
200 	printf("  %s -b bind-to-dev -m mark -p prio cg-path\n", argv0);
201 	printf("\n");
202 	printf("  Detach a program\n");
203 	printf("  %s -d cg-path\n", argv0);
204 	printf("\n");
205 	printf("  Show inherited socket settings (mark, priority, and device)\n");
206 	printf("  %s [-6]\n", argv0);
207 	return EXIT_FAILURE;
208 }
209 
210 int main(int argc, char **argv)
211 {
212 	__u32 idx = 0, mark = 0, prio = 0;
213 	const char *cgrp_path = NULL;
214 	int cg_fd, prog_fd, ret;
215 	int family = PF_INET;
216 	int do_attach = 1;
217 	int rc;
218 
219 	while ((rc = getopt(argc, argv, "db:m:p:6")) != -1) {
220 		switch (rc) {
221 		case 'd':
222 			do_attach = 0;
223 			break;
224 		case 'b':
225 			idx = if_nametoindex(optarg);
226 			if (!idx) {
227 				idx = strtoumax(optarg, NULL, 0);
228 				if (!idx) {
229 					printf("Invalid device name\n");
230 					return EXIT_FAILURE;
231 				}
232 			}
233 			break;
234 		case 'm':
235 			mark = strtoumax(optarg, NULL, 0);
236 			break;
237 		case 'p':
238 			prio = strtoumax(optarg, NULL, 0);
239 			break;
240 		case '6':
241 			family = PF_INET6;
242 			break;
243 		default:
244 			return usage(argv[0]);
245 		}
246 	}
247 
248 	if (optind == argc)
249 		return show_sockopts(family);
250 
251 	cgrp_path = argv[optind];
252 	if (!cgrp_path) {
253 		fprintf(stderr, "cgroup path not given\n");
254 		return EXIT_FAILURE;
255 	}
256 
257 	if (do_attach && !idx && !mark && !prio) {
258 		fprintf(stderr,
259 			"One of device, mark or priority must be given\n");
260 		return EXIT_FAILURE;
261 	}
262 
263 	cg_fd = open(cgrp_path, O_DIRECTORY | O_RDONLY);
264 	if (cg_fd < 0) {
265 		printf("Failed to open cgroup path: '%s'\n", strerror(errno));
266 		return EXIT_FAILURE;
267 	}
268 
269 	if (do_attach) {
270 		prog_fd = prog_load(idx, mark, prio);
271 		if (prog_fd < 0) {
272 			printf("Failed to load prog: '%s'\n", strerror(errno));
273 			printf("Output from kernel verifier:\n%s\n-------\n",
274 			       bpf_log_buf);
275 			return EXIT_FAILURE;
276 		}
277 
278 		ret = bpf_prog_attach(prog_fd, cg_fd,
279 				      BPF_CGROUP_INET_SOCK_CREATE, 0);
280 		if (ret < 0) {
281 			printf("Failed to attach prog to cgroup: '%s'\n",
282 			       strerror(errno));
283 			return EXIT_FAILURE;
284 		}
285 	} else {
286 		ret = bpf_prog_detach(cg_fd, BPF_CGROUP_INET_SOCK_CREATE);
287 		if (ret < 0) {
288 			printf("Failed to detach prog from cgroup: '%s'\n",
289 			       strerror(errno));
290 			return EXIT_FAILURE;
291 		}
292 	}
293 
294 	close(cg_fd);
295 	return EXIT_SUCCESS;
296 }
297