1/* eBPF example program:
  2 *
  3 * - Creates arraymap in kernel with 4 bytes keys and 8 byte values
  4 *
  5 * - Loads eBPF program
  6 *
  7 *   The eBPF program accesses the map passed in to store two pieces of
  8 *   information. The number of invocations of the program, which maps
  9 *   to the number of packets received, is stored to key 0. Key 1 is
 10 *   incremented on each iteration by the number of bytes stored in
 11 *   the skb.
 12 *
 13 * - Attaches the new program to a cgroup using BPF_PROG_ATTACH
 14 *
 15 * - Every second, reads map[0] and map[1] to see how many bytes and
 16 *   packets were seen on any socket of tasks in the given cgroup.
 17 */
 18
 19#define _GNU_SOURCE
 20
 21#include <stdio.h>
 22#include <stdlib.h>
 23#include <stddef.h>
 24#include <string.h>
 25#include <unistd.h>
 26#include <assert.h>
 27#include <errno.h>
 28#include <fcntl.h>
 29
 30#include <linux/bpf.h>
 31#include <bpf/bpf.h>
 32
 33#include "bpf_insn.h"
 34
 35enum {
 36	MAP_KEY_PACKETS,
 37	MAP_KEY_BYTES,
 38};
 39
 40char bpf_log_buf[BPF_LOG_BUF_SIZE];
 41
 42static int prog_load(int map_fd, int verdict)
 43{
 44	struct bpf_insn prog[] = {
 45		BPF_MOV64_REG(BPF_REG_6, BPF_REG_1), /* save r6 so it's not clobbered by BPF_CALL */
 46
 47		/* Count packets */
 48		BPF_MOV64_IMM(BPF_REG_0, MAP_KEY_PACKETS), /* r0 = 0 */
 49		BPF_STX_MEM(BPF_W, BPF_REG_10, BPF_REG_0, -4), /* *(u32 *)(fp - 4) = r0 */
 50		BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),
 51		BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -4), /* r2 = fp - 4 */
 52		BPF_LD_MAP_FD(BPF_REG_1, map_fd), /* load map fd to r1 */
 53		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_lookup_elem),
 54		BPF_JMP_IMM(BPF_JEQ, BPF_REG_0, 0, 2),
 55		BPF_MOV64_IMM(BPF_REG_1, 1), /* r1 = 1 */
 56		BPF_ATOMIC_OP(BPF_DW, BPF_ADD, BPF_REG_0, BPF_REG_1, 0),
 57
 58		/* Count bytes */
 59		BPF_MOV64_IMM(BPF_REG_0, MAP_KEY_BYTES), /* r0 = 1 */
 60		BPF_STX_MEM(BPF_W, BPF_REG_10, BPF_REG_0, -4), /* *(u32 *)(fp - 4) = r0 */
 61		BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),
 62		BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -4), /* r2 = fp - 4 */
 63		BPF_LD_MAP_FD(BPF_REG_1, map_fd),
 64		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_lookup_elem),
 65		BPF_JMP_IMM(BPF_JEQ, BPF_REG_0, 0, 2),
 66		BPF_LDX_MEM(BPF_W, BPF_REG_1, BPF_REG_6, offsetof(struct __sk_buff, len)), /* r1 = skb->len */
 67
 68		BPF_ATOMIC_OP(BPF_DW, BPF_ADD, BPF_REG_0, BPF_REG_1, 0),
 69
 70		BPF_MOV64_IMM(BPF_REG_0, verdict), /* r0 = verdict */
 71		BPF_EXIT_INSN(),
 72	};
 73	size_t insns_cnt = sizeof(prog) / sizeof(struct bpf_insn);
 74
 75	return bpf_load_program(BPF_PROG_TYPE_CGROUP_SKB,
 76				prog, insns_cnt, "GPL", 0,
 77				bpf_log_buf, BPF_LOG_BUF_SIZE);
 78}
 79
 80static int usage(const char *argv0)
 81{
 82	printf("Usage: %s [-d] [-D] <cg-path> <egress|ingress>\n", argv0);
 83	printf("	-d	Drop Traffic\n");
 84	printf("	-D	Detach filter, and exit\n");
 85	return EXIT_FAILURE;
 86}
 87
 88static int attach_filter(int cg_fd, int type, int verdict)
 89{
 90	int prog_fd, map_fd, ret, key;
 91	long long pkt_cnt, byte_cnt;
 92
 93	map_fd = bpf_create_map(BPF_MAP_TYPE_ARRAY,
 94				sizeof(key), sizeof(byte_cnt),
 95				256, 0);
 96	if (map_fd < 0) {
 97		printf("Failed to create map: '%s'\n", strerror(errno));
 98		return EXIT_FAILURE;
 99	}
100
101	prog_fd = prog_load(map_fd, verdict);
102	printf("Output from kernel verifier:\n%s\n-------\n", bpf_log_buf);
103
104	if (prog_fd < 0) {
105		printf("Failed to load prog: '%s'\n", strerror(errno));
106		return EXIT_FAILURE;
107	}
108
109	ret = bpf_prog_attach(prog_fd, cg_fd, type, 0);
110	if (ret < 0) {
111		printf("Failed to attach prog to cgroup: '%s'\n",
112		       strerror(errno));
113		return EXIT_FAILURE;
114	}
115	while (1) {
116		key = MAP_KEY_PACKETS;
117		assert(bpf_map_lookup_elem(map_fd, &key, &pkt_cnt) == 0);
118
119		key = MAP_KEY_BYTES;
120		assert(bpf_map_lookup_elem(map_fd, &key, &byte_cnt) == 0);
121
122		printf("cgroup received %lld packets, %lld bytes\n",
123		       pkt_cnt, byte_cnt);
124		sleep(1);
125	}
126
127	return EXIT_SUCCESS;
128}
129
130int main(int argc, char **argv)
131{
132	int detach_only = 0, verdict = 1;
133	enum bpf_attach_type type;
134	int opt, cg_fd, ret;
135
136	while ((opt = getopt(argc, argv, "Dd")) != -1) {
137		switch (opt) {
138		case 'd':
139			verdict = 0;
140			break;
141		case 'D':
142			detach_only = 1;
143			break;
144		default:
145			return usage(argv[0]);
146		}
147	}
148
149	if (argc - optind < 2)
150		return usage(argv[0]);
151
152	if (strcmp(argv[optind + 1], "ingress") == 0)
153		type = BPF_CGROUP_INET_INGRESS;
154	else if (strcmp(argv[optind + 1], "egress") == 0)
155		type = BPF_CGROUP_INET_EGRESS;
156	else
157		return usage(argv[0]);
158
159	cg_fd = open(argv[optind], O_DIRECTORY | O_RDONLY);
160	if (cg_fd < 0) {
161		printf("Failed to open cgroup path: '%s'\n", strerror(errno));
162		return EXIT_FAILURE;
163	}
164
165	if (detach_only) {
166		ret = bpf_prog_detach(cg_fd, type);
167		printf("bpf_prog_detach() returned '%s' (%d)\n",
168		       strerror(errno), errno);
169	} else
170		ret = attach_filter(cg_fd, type, verdict);
171
172	return ret;
173}