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
 27char bpf_log_buf[BPF_LOG_BUF_SIZE];
 28
 29static 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
125static 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
138static 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
151static 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
164static 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
189static 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
203int 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}