[linux-2.6-block.git] / samples / bpf / test_cgrp2_sock.c

/* eBPF example program:
 *
 * - Loads eBPF program
 *
 *   The eBPF program sets the sk_bound_dev_if index in new AF_INET{6}
 *   sockets opened by processes in the cgroup.
 *
 * - Attaches the new program to a cgroup using BPF_PROG_ATTACH
 */

#define _GNU_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <string.h>
#include <unistd.h>
#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <net/if.h>
#include <inttypes.h>
#include <linux/bpf.h>
#include <bpf/bpf.h>

#include "bpf_insn.h"

char bpf_log_buf[BPF_LOG_BUF_SIZE];

static int prog_load(__u32 idx, __u32 mark, __u32 prio)
{
	/* save pointer to context */
	struct bpf_insn prog_start[] = {
		BPF_MOV64_REG(BPF_REG_6, BPF_REG_1),
	};
	struct bpf_insn prog_end[] = {
		BPF_MOV64_IMM(BPF_REG_0, 1), /* r0 = verdict */
		BPF_EXIT_INSN(),
	};

	/* set sk_bound_dev_if on socket */
	struct bpf_insn prog_dev[] = {
		BPF_MOV64_IMM(BPF_REG_3, idx),
		BPF_MOV64_IMM(BPF_REG_2, offsetof(struct bpf_sock, bound_dev_if)),
		BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3, offsetof(struct bpf_sock, bound_dev_if)),
	};

	/* set mark on socket */
	struct bpf_insn prog_mark[] = {
		/* get uid of process */
		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
			     BPF_FUNC_get_current_uid_gid),
		BPF_ALU64_IMM(BPF_AND, BPF_REG_0, 0xffffffff),

		/* if uid is 0, use given mark, else use the uid as the mark */
		BPF_MOV64_REG(BPF_REG_3, BPF_REG_0),
		BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 1),
		BPF_MOV64_IMM(BPF_REG_3, mark),

		/* set the mark on the new socket */
		BPF_MOV64_REG(BPF_REG_1, BPF_REG_6),
		BPF_MOV64_IMM(BPF_REG_2, offsetof(struct bpf_sock, mark)),
		BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3, offsetof(struct bpf_sock, mark)),
	};

	/* set priority on socket */
	struct bpf_insn prog_prio[] = {
		BPF_MOV64_REG(BPF_REG_1, BPF_REG_6),
		BPF_MOV64_IMM(BPF_REG_3, prio),
		BPF_MOV64_IMM(BPF_REG_2, offsetof(struct bpf_sock, priority)),
		BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3, offsetof(struct bpf_sock, priority)),
	};
	LIBBPF_OPTS(bpf_prog_load_opts, opts,
		.log_buf = bpf_log_buf,
		.log_size = BPF_LOG_BUF_SIZE,
	);

	struct bpf_insn *prog;
	size_t insns_cnt;
	void *p;
	int ret;

	insns_cnt = sizeof(prog_start) + sizeof(prog_end);
	if (idx)
		insns_cnt += sizeof(prog_dev);

	if (mark)
		insns_cnt += sizeof(prog_mark);

	if (prio)
		insns_cnt += sizeof(prog_prio);

	p = prog = malloc(insns_cnt);
	if (!prog) {
		fprintf(stderr, "Failed to allocate memory for instructions\n");
		return EXIT_FAILURE;
	}

	memcpy(p, prog_start, sizeof(prog_start));
	p += sizeof(prog_start);

	if (idx) {
		memcpy(p, prog_dev, sizeof(prog_dev));
		p += sizeof(prog_dev);
	}

	if (mark) {
		memcpy(p, prog_mark, sizeof(prog_mark));
		p += sizeof(prog_mark);
	}

	if (prio) {
		memcpy(p, prog_prio, sizeof(prog_prio));
		p += sizeof(prog_prio);
	}

	memcpy(p, prog_end, sizeof(prog_end));
	p += sizeof(prog_end);

	insns_cnt /= sizeof(struct bpf_insn);

	ret = bpf_prog_load(BPF_PROG_TYPE_CGROUP_SOCK, NULL, "GPL",
			    prog, insns_cnt, &opts);

	free(prog);

	return ret;
}

static int get_bind_to_device(int sd, char *name, size_t len)
{
	socklen_t optlen = len;
	int rc;

	name[0] = '\0';
	rc = getsockopt(sd, SOL_SOCKET, SO_BINDTODEVICE, name, &optlen);
	if (rc < 0)
		perror("setsockopt(SO_BINDTODEVICE)");

	return rc;
}

static unsigned int get_somark(int sd)
{
	unsigned int mark = 0;
	socklen_t optlen = sizeof(mark);
	int rc;

	rc = getsockopt(sd, SOL_SOCKET, SO_MARK, &mark, &optlen);
	if (rc < 0)
		perror("getsockopt(SO_MARK)");

	return mark;
}

static unsigned int get_priority(int sd)
{
	unsigned int prio = 0;
	socklen_t optlen = sizeof(prio);
	int rc;

	rc = getsockopt(sd, SOL_SOCKET, SO_PRIORITY, &prio, &optlen);
	if (rc < 0)
		perror("getsockopt(SO_PRIORITY)");

	return prio;
}

static int show_sockopts(int family)
{
	unsigned int mark, prio;
	char name[16];
	int sd;

	sd = socket(family, SOCK_DGRAM, 17);
	if (sd < 0) {
		perror("socket");
		return 1;
	}

	if (get_bind_to_device(sd, name, sizeof(name)) < 0)
		return 1;

	mark = get_somark(sd);
	prio = get_priority(sd);

	close(sd);

	printf("sd %d: dev %s, mark %u, priority %u\n", sd, name, mark, prio);

	return 0;
}

static int usage(const char *argv0)
{
	printf("Usage:\n");
	printf("  Attach a program\n");
	printf("  %s -b bind-to-dev -m mark -p prio cg-path\n", argv0);
	printf("\n");
	printf("  Detach a program\n");
	printf("  %s -d cg-path\n", argv0);
	printf("\n");
	printf("  Show inherited socket settings (mark, priority, and device)\n");
	printf("  %s [-6]\n", argv0);
	return EXIT_FAILURE;
}

int main(int argc, char **argv)
{
	__u32 idx = 0, mark = 0, prio = 0;
	const char *cgrp_path = NULL;
	int cg_fd, prog_fd, ret;
	int family = PF_INET;
	int do_attach = 1;
	int rc;

	while ((rc = getopt(argc, argv, "db:m:p:6")) != -1) {
		switch (rc) {
		case 'd':
			do_attach = 0;
			break;
		case 'b':
			idx = if_nametoindex(optarg);
			if (!idx) {
				idx = strtoumax(optarg, NULL, 0);
				if (!idx) {
					printf("Invalid device name\n");
					return EXIT_FAILURE;
				}
			}
			break;
		case 'm':
			mark = strtoumax(optarg, NULL, 0);
			break;
		case 'p':
			prio = strtoumax(optarg, NULL, 0);
			break;
		case '6':
			family = PF_INET6;
			break;
		default:
			return usage(argv[0]);
		}
	}

	if (optind == argc)
		return show_sockopts(family);

	cgrp_path = argv[optind];
	if (!cgrp_path) {
		fprintf(stderr, "cgroup path not given\n");
		return EXIT_FAILURE;
	}

	if (do_attach && !idx && !mark && !prio) {
		fprintf(stderr,
			"One of device, mark or priority must be given\n");
		return EXIT_FAILURE;
	}

	cg_fd = open(cgrp_path, O_DIRECTORY | O_RDONLY);
	if (cg_fd < 0) {
		printf("Failed to open cgroup path: '%s'\n", strerror(errno));
		return EXIT_FAILURE;
	}

	if (do_attach) {
		prog_fd = prog_load(idx, mark, prio);
		if (prog_fd < 0) {
			printf("Failed to load prog: '%s'\n", strerror(errno));
			printf("Output from kernel verifier:\n%s\n-------\n",
			       bpf_log_buf);
			return EXIT_FAILURE;
		}

		ret = bpf_prog_attach(prog_fd, cg_fd,
				      BPF_CGROUP_INET_SOCK_CREATE, 0);
		if (ret < 0) {
			printf("Failed to attach prog to cgroup: '%s'\n",
			       strerror(errno));
			return EXIT_FAILURE;
		}
	} else {
		ret = bpf_prog_detach(cg_fd, BPF_CGROUP_INET_SOCK_CREATE);
		if (ret < 0) {
			printf("Failed to detach prog from cgroup: '%s'\n",
			       strerror(errno));
			return EXIT_FAILURE;
		}
	}

	close(cg_fd);
	return EXIT_SUCCESS;
}
Commit	Line	Data
ad2805dc DA	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>
fa38aa17	22	#include <inttypes.h>
ad2805dc	23	#include <linux/bpf.h>
8d930450	24	#include <bpf/bpf.h>
ad2805dc	25
8d930450	26	#include "bpf_insn.h"
ad2805dc	27
d40fc181 JS	28	char bpf_log_buf[BPF_LOG_BUF_SIZE];
d40fc181 JS	29
fa38aa17	30	static int prog_load(__u32 idx, __u32 mark, __u32 prio)
ad2805dc	31	{
fa38aa17 DA	32	/* save pointer to context */
fa38aa17 DA	33	struct bpf_insn prog_start[] = {
ad2805dc	34	BPF_MOV64_REG(BPF_REG_6, BPF_REG_1),
fa38aa17 DA	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[] = {
ad2805dc DA	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)),
ad2805dc DA	46	};
ad2805dc DA	47
fa38aa17 DA	48	/* set mark on socket */
fa38aa17 DA	49	struct bpf_insn prog_mark[] = {
0adc3dd9 DA	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),
fa38aa17	58	BPF_MOV64_IMM(BPF_REG_3, mark),
0adc3dd9 DA	59
	60	/* set the mark on the new socket */
	61	BPF_MOV64_REG(BPF_REG_1, BPF_REG_6),
fa38aa17 DA	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	};
c58f9815 AN	73	LIBBPF_OPTS(bpf_prog_load_opts, opts,
	74	.log_buf = bpf_log_buf,
	75	.log_size = BPF_LOG_BUF_SIZE,
	76	);
fa38aa17 DA	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
c58f9815 AN	122	ret = bpf_prog_load(BPF_PROG_TYPE_CGROUP_SOCK, NULL, "GPL",
c58f9815 AN	123	prog, insns_cnt, &opts);
fa38aa17 DA	124
	125	free(prog);
	126
	127	return ret;
ad2805dc DA	128	}
ad2805dc DA	129
f776d460 DA	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	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
ad2805dc DA	194	static int usage(const char *argv0)
ad2805dc DA	195	{
609b1c32 DA	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);
f776d460 DA	202	printf("\n");
	203	printf(" Show inherited socket settings (mark, priority, and device)\n");
	204	printf(" %s [-6]\n", argv0);
ad2805dc DA	205	return EXIT_FAILURE;
	206	}
	207
	208	int main(int argc, char **argv)
	209	{
fa38aa17 DA	210	__u32 idx = 0, mark = 0, prio = 0;
fa38aa17 DA	211	const char *cgrp_path = NULL;
ad2805dc	212	int cg_fd, prog_fd, ret;
f776d460	213	int family = PF_INET;
609b1c32	214	int do_attach = 1;
fa38aa17 DA	215	int rc;
fa38aa17 DA	216
f776d460	217	while ((rc = getopt(argc, argv, "db:m:p:6")) != -1) {
fa38aa17	218	switch (rc) {
609b1c32 DA	219	case 'd':
	220	do_attach = 0;
	221	break;
fa38aa17 DA	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;
f776d460 DA	238	case '6':
	239	family = PF_INET6;
	240	break;
fa38aa17 DA	241	default:
	242	return usage(argv[0]);
	243	}
	244	}
ad2805dc	245
fa38aa17	246	if (optind == argc)
f776d460	247	return show_sockopts(family);
ad2805dc	248
fa38aa17 DA	249	cgrp_path = argv[optind];
	250	if (!cgrp_path) {
	251	fprintf(stderr, "cgroup path not given\n");
ad2805dc DA	252	return EXIT_FAILURE;
	253	}
	254
609b1c32	255	if (do_attach && !idx && !mark && !prio) {
fa38aa17 DA	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);
ad2805dc DA	262	if (cg_fd < 0) {
	263	printf("Failed to open cgroup path: '%s'\n", strerror(errno));
	264	return EXIT_FAILURE;
	265	}
	266
609b1c32 DA	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	}
ad2805dc	275
609b1c32 DA	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	}
ad2805dc DA	290	}
ad2805dc DA	291
609b1c32	292	close(cg_fd);
ad2805dc DA	293	return EXIT_SUCCESS;
ad2805dc DA	294	}