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	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
130static 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
143static 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
156static 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
169static 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		return 1;
183
184	mark = get_somark(sd);
185	prio = get_priority(sd);
186
187	close(sd);
188
189	printf("sd %d: dev %s, mark %u, priority %u\n", sd, name, mark, prio);
190
191	return 0;
192}
193
194static int usage(const char *argv0)
195{
196	printf("Usage:\n");
197	printf("  Attach a program\n");
198	printf("  %s -b bind-to-dev -m mark -p prio cg-path\n", argv0);
199	printf("\n");
200	printf("  Detach a program\n");
201	printf("  %s -d cg-path\n", argv0);
202	printf("\n");
203	printf("  Show inherited socket settings (mark, priority, and device)\n");
204	printf("  %s [-6]\n", argv0);
205	return EXIT_FAILURE;
206}
207
208int main(int argc, char **argv)
209{
210	__u32 idx = 0, mark = 0, prio = 0;
211	const char *cgrp_path = NULL;
212	int cg_fd, prog_fd, ret;
213	int family = PF_INET;
214	int do_attach = 1;
215	int rc;
216
217	while ((rc = getopt(argc, argv, "db:m:p:6")) != -1) {
218		switch (rc) {
219		case 'd':
220			do_attach = 0;
221			break;
222		case 'b':
223			idx = if_nametoindex(optarg);
224			if (!idx) {
225				idx = strtoumax(optarg, NULL, 0);
226				if (!idx) {
227					printf("Invalid device name\n");
228					return EXIT_FAILURE;
229				}
230			}
231			break;
232		case 'm':
233			mark = strtoumax(optarg, NULL, 0);
234			break;
235		case 'p':
236			prio = strtoumax(optarg, NULL, 0);
237			break;
238		case '6':
239			family = PF_INET6;
240			break;
241		default:
242			return usage(argv[0]);
243		}
244	}
245
246	if (optind == argc)
247		return show_sockopts(family);
248
249	cgrp_path = argv[optind];
250	if (!cgrp_path) {
251		fprintf(stderr, "cgroup path not given\n");
252		return EXIT_FAILURE;
253	}
254
255	if (do_attach && !idx && !mark && !prio) {
256		fprintf(stderr,
257			"One of device, mark or priority must be given\n");
258		return EXIT_FAILURE;
259	}
260
261	cg_fd = open(cgrp_path, O_DIRECTORY | O_RDONLY);
262	if (cg_fd < 0) {
263		printf("Failed to open cgroup path: '%s'\n", strerror(errno));
264		return EXIT_FAILURE;
265	}
266
267	if (do_attach) {
268		prog_fd = prog_load(idx, mark, prio);
269		if (prog_fd < 0) {
270			printf("Failed to load prog: '%s'\n", strerror(errno));
271			printf("Output from kernel verifier:\n%s\n-------\n",
272			       bpf_log_buf);
273			return EXIT_FAILURE;
274		}
275
276		ret = bpf_prog_attach(prog_fd, cg_fd,
277				      BPF_CGROUP_INET_SOCK_CREATE, 0);
278		if (ret < 0) {
279			printf("Failed to attach prog to cgroup: '%s'\n",
280			       strerror(errno));
281			return EXIT_FAILURE;
282		}
283	} else {
284		ret = bpf_prog_detach(cg_fd, BPF_CGROUP_INET_SOCK_CREATE);
285		if (ret < 0) {
286			printf("Failed to detach prog from cgroup: '%s'\n",
287			       strerror(errno));
288			return EXIT_FAILURE;
289		}
290	}
291
292	close(cg_fd);
293	return EXIT_SUCCESS;
294}