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
28char bpf_log_buf[BPF_LOG_BUF_SIZE];
29
30static 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
74	struct bpf_insn *prog;
75	size_t insns_cnt;
76	void *p;
77	int ret;
78
79	insns_cnt = sizeof(prog_start) + sizeof(prog_end);
80	if (idx)
81		insns_cnt += sizeof(prog_dev);
82
83	if (mark)
84		insns_cnt += sizeof(prog_mark);
85
86	if (prio)
87		insns_cnt += sizeof(prog_prio);
88
89	p = prog = malloc(insns_cnt);
90	if (!prog) {
91		fprintf(stderr, "Failed to allocate memory for instructions\n");
92		return EXIT_FAILURE;
93	}
94
95	memcpy(p, prog_start, sizeof(prog_start));
96	p += sizeof(prog_start);
97
98	if (idx) {
99		memcpy(p, prog_dev, sizeof(prog_dev));
100		p += sizeof(prog_dev);
101	}
102
103	if (mark) {
104		memcpy(p, prog_mark, sizeof(prog_mark));
105		p += sizeof(prog_mark);
106	}
107
108	if (prio) {
109		memcpy(p, prog_prio, sizeof(prog_prio));
110		p += sizeof(prog_prio);
111	}
112
113	memcpy(p, prog_end, sizeof(prog_end));
114	p += sizeof(prog_end);
115
116	insns_cnt /= sizeof(struct bpf_insn);
117
118	ret = bpf_load_program(BPF_PROG_TYPE_CGROUP_SOCK, prog, insns_cnt,
119				"GPL", 0, bpf_log_buf, BPF_LOG_BUF_SIZE);
120
121	free(prog);
122
123	return ret;
124}
125
126static int get_bind_to_device(int sd, char *name, size_t len)
127{
128	socklen_t optlen = len;
129	int rc;
130
131	name[0] = '\0';
132	rc = getsockopt(sd, SOL_SOCKET, SO_BINDTODEVICE, name, &optlen);
133	if (rc < 0)
134		perror("setsockopt(SO_BINDTODEVICE)");
135
136	return rc;
137}
138
139static unsigned int get_somark(int sd)
140{
141	unsigned int mark = 0;
142	socklen_t optlen = sizeof(mark);
143	int rc;
144
145	rc = getsockopt(sd, SOL_SOCKET, SO_MARK, &mark, &optlen);
146	if (rc < 0)
147		perror("getsockopt(SO_MARK)");
148
149	return mark;
150}
151
152static unsigned int get_priority(int sd)
153{
154	unsigned int prio = 0;
155	socklen_t optlen = sizeof(prio);
156	int rc;
157
158	rc = getsockopt(sd, SOL_SOCKET, SO_PRIORITY, &prio, &optlen);
159	if (rc < 0)
160		perror("getsockopt(SO_PRIORITY)");
161
162	return prio;
163}
164
165static int show_sockopts(int family)
166{
167	unsigned int mark, prio;
168	char name[16];
169	int sd;
170
171	sd = socket(family, SOCK_DGRAM, 17);
172	if (sd < 0) {
173		perror("socket");
174		return 1;
175	}
176
177	if (get_bind_to_device(sd, name, sizeof(name)) < 0)
178		return 1;
179
180	mark = get_somark(sd);
181	prio = get_priority(sd);
182
183	close(sd);
184
185	printf("sd %d: dev %s, mark %u, priority %u\n", sd, name, mark, prio);
186
187	return 0;
188}
189
190static int usage(const char *argv0)
191{
192	printf("Usage:\n");
193	printf("  Attach a program\n");
194	printf("  %s -b bind-to-dev -m mark -p prio cg-path\n", argv0);
195	printf("\n");
196	printf("  Detach a program\n");
197	printf("  %s -d cg-path\n", argv0);
198	printf("\n");
199	printf("  Show inherited socket settings (mark, priority, and device)\n");
200	printf("  %s [-6]\n", argv0);
201	return EXIT_FAILURE;
202}
203
204int main(int argc, char **argv)
205{
206	__u32 idx = 0, mark = 0, prio = 0;
207	const char *cgrp_path = NULL;
208	int cg_fd, prog_fd, ret;
209	int family = PF_INET;
210	int do_attach = 1;
211	int rc;
212
213	while ((rc = getopt(argc, argv, "db:m:p:6")) != -1) {
214		switch (rc) {
215		case 'd':
216			do_attach = 0;
217			break;
218		case 'b':
219			idx = if_nametoindex(optarg);
220			if (!idx) {
221				idx = strtoumax(optarg, NULL, 0);
222				if (!idx) {
223					printf("Invalid device name\n");
224					return EXIT_FAILURE;
225				}
226			}
227			break;
228		case 'm':
229			mark = strtoumax(optarg, NULL, 0);
230			break;
231		case 'p':
232			prio = strtoumax(optarg, NULL, 0);
233			break;
234		case '6':
235			family = PF_INET6;
236			break;
237		default:
238			return usage(argv[0]);
239		}
240	}
241
242	if (optind == argc)
243		return show_sockopts(family);
244
245	cgrp_path = argv[optind];
246	if (!cgrp_path) {
247		fprintf(stderr, "cgroup path not given\n");
248		return EXIT_FAILURE;
249	}
250
251	if (do_attach && !idx && !mark && !prio) {
252		fprintf(stderr,
253			"One of device, mark or priority must be given\n");
254		return EXIT_FAILURE;
255	}
256
257	cg_fd = open(cgrp_path, O_DIRECTORY | O_RDONLY);
258	if (cg_fd < 0) {
259		printf("Failed to open cgroup path: '%s'\n", strerror(errno));
260		return EXIT_FAILURE;
261	}
262
263	if (do_attach) {
264		prog_fd = prog_load(idx, mark, prio);
265		if (prog_fd < 0) {
266			printf("Failed to load prog: '%s'\n", strerror(errno));
267			printf("Output from kernel verifier:\n%s\n-------\n",
268			       bpf_log_buf);
269			return EXIT_FAILURE;
270		}
271
272		ret = bpf_prog_attach(prog_fd, cg_fd,
273				      BPF_CGROUP_INET_SOCK_CREATE, 0);
274		if (ret < 0) {
275			printf("Failed to attach prog to cgroup: '%s'\n",
276			       strerror(errno));
277			return EXIT_FAILURE;
278		}
279	} else {
280		ret = bpf_prog_detach(cg_fd, BPF_CGROUP_INET_SOCK_CREATE);
281		if (ret < 0) {
282			printf("Failed to detach prog from cgroup: '%s'\n",
283			       strerror(errno));
284			return EXIT_FAILURE;
285		}
286	}
287
288	close(cg_fd);
289	return EXIT_SUCCESS;
290}
291