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