samples/bpf/test_cgrp2_attach.c

0001 /* eBPF example program:
0002  *
0003  * - Creates arraymap in kernel with 4 bytes keys and 8 byte values
0004  *
0005  * - Loads eBPF program
0006  *
0007  *   The eBPF program accesses the map passed in to store two pieces of
0008  *   information. The number of invocations of the program, which maps
0009  *   to the number of packets received, is stored to key 0. Key 1 is
0010  *   incremented on each iteration by the number of bytes stored in
0011  *   the skb.
0012  *
0013  * - Attaches the new program to a cgroup using BPF_PROG_ATTACH
0014  *
0015  * - Every second, reads map[0] and map[1] to see how many bytes and
0016  *   packets were seen on any socket of tasks in the given cgroup.
0017  */
0018
0019 #define _GNU_SOURCE
0020
0021 #include <stdio.h>
0022 #include <stdlib.h>
0023 #include <stddef.h>
0024 #include <string.h>
0025 #include <unistd.h>
0026 #include <assert.h>
0027 #include <errno.h>
0028 #include <fcntl.h>
0029
0030 #include <linux/bpf.h>
0031 #include <bpf/bpf.h>
0032
0033 #include "bpf_insn.h"
0034 #include "bpf_util.h"
0035
0036 enum {
0037     MAP_KEY_PACKETS,
0038     MAP_KEY_BYTES,
0039 };
0040
0041 char bpf_log_buf[BPF_LOG_BUF_SIZE];
0042
0043 static int prog_load(int map_fd, int verdict)
0044 {
0045     struct bpf_insn prog[] = {
0046         BPF_MOV64_REG(BPF_REG_6, BPF_REG_1), /* save r6 so it's not clobbered by BPF_CALL */
0047
0048         /* Count packets */
0049         BPF_MOV64_IMM(BPF_REG_0, MAP_KEY_PACKETS), /* r0 = 0 */
0050         BPF_STX_MEM(BPF_W, BPF_REG_10, BPF_REG_0, -4), /* *(u32 *)(fp - 4) = r0 */
0051         BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),
0052         BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -4), /* r2 = fp - 4 */
0053         BPF_LD_MAP_FD(BPF_REG_1, map_fd), /* load map fd to r1 */
0054         BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_lookup_elem),
0055         BPF_JMP_IMM(BPF_JEQ, BPF_REG_0, 0, 2),
0056         BPF_MOV64_IMM(BPF_REG_1, 1), /* r1 = 1 */
0057         BPF_ATOMIC_OP(BPF_DW, BPF_ADD, BPF_REG_0, BPF_REG_1, 0),
0058
0059         /* Count bytes */
0060         BPF_MOV64_IMM(BPF_REG_0, MAP_KEY_BYTES), /* r0 = 1 */
0061         BPF_STX_MEM(BPF_W, BPF_REG_10, BPF_REG_0, -4), /* *(u32 *)(fp - 4) = r0 */
0062         BPF_MOV64_REG(BPF_REG_2, BPF_REG_10),
0063         BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, -4), /* r2 = fp - 4 */
0064         BPF_LD_MAP_FD(BPF_REG_1, map_fd),
0065         BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_map_lookup_elem),
0066         BPF_JMP_IMM(BPF_JEQ, BPF_REG_0, 0, 2),
0067         BPF_LDX_MEM(BPF_W, BPF_REG_1, BPF_REG_6, offsetof(struct __sk_buff, len)), /* r1 = skb->len */
0068
0069         BPF_ATOMIC_OP(BPF_DW, BPF_ADD, BPF_REG_0, BPF_REG_1, 0),
0070
0071         BPF_MOV64_IMM(BPF_REG_0, verdict), /* r0 = verdict */
0072         BPF_EXIT_INSN(),
0073     };
0074     size_t insns_cnt = ARRAY_SIZE(prog);
0075     LIBBPF_OPTS(bpf_prog_load_opts, opts,
0076         .log_buf = bpf_log_buf,
0077         .log_size = BPF_LOG_BUF_SIZE,
0078     );
0079
0080     return bpf_prog_load(BPF_PROG_TYPE_CGROUP_SKB, NULL, "GPL",
0081                  prog, insns_cnt, &opts);
0082 }
0083
0084 static int usage(const char *argv0)
0085 {
0086     printf("Usage: %s [-d] [-D] <cg-path> <egress|ingress>\n", argv0);
0087     printf("    -d  Drop Traffic\n");
0088     printf("    -D  Detach filter, and exit\n");
0089     return EXIT_FAILURE;
0090 }
0091
0092 static int attach_filter(int cg_fd, int type, int verdict)
0093 {
0094     int prog_fd, map_fd, ret, key;
0095     long long pkt_cnt, byte_cnt;
0096
0097     map_fd = bpf_map_create(BPF_MAP_TYPE_ARRAY, NULL,
0098                 sizeof(key), sizeof(byte_cnt),
0099                 256, NULL);
0100     if (map_fd < 0) {
0101         printf("Failed to create map: '%s'\n", strerror(errno));
0102         return EXIT_FAILURE;
0103     }
0104
0105     prog_fd = prog_load(map_fd, verdict);
0106     printf("Output from kernel verifier:\n%s\n-------\n", bpf_log_buf);
0107
0108     if (prog_fd < 0) {
0109         printf("Failed to load prog: '%s'\n", strerror(errno));
0110         return EXIT_FAILURE;
0111     }
0112
0113     ret = bpf_prog_attach(prog_fd, cg_fd, type, 0);
0114     if (ret < 0) {
0115         printf("Failed to attach prog to cgroup: '%s'\n",
0116                strerror(errno));
0117         return EXIT_FAILURE;
0118     }
0119     while (1) {
0120         key = MAP_KEY_PACKETS;
0121         assert(bpf_map_lookup_elem(map_fd, &key, &pkt_cnt) == 0);
0122
0123         key = MAP_KEY_BYTES;
0124         assert(bpf_map_lookup_elem(map_fd, &key, &byte_cnt) == 0);
0125
0126         printf("cgroup received %lld packets, %lld bytes\n",
0127                pkt_cnt, byte_cnt);
0128         sleep(1);
0129     }
0130
0131     return EXIT_SUCCESS;
0132 }
0133
0134 int main(int argc, char **argv)
0135 {
0136     int detach_only = 0, verdict = 1;
0137     enum bpf_attach_type type;
0138     int opt, cg_fd, ret;
0139
0140     while ((opt = getopt(argc, argv, "Dd")) != -1) {
0141         switch (opt) {
0142         case 'd':
0143             verdict = 0;
0144             break;
0145         case 'D':
0146             detach_only = 1;
0147             break;
0148         default:
0149             return usage(argv[0]);
0150         }
0151     }
0152
0153     if (argc - optind < 2)
0154         return usage(argv[0]);
0155
0156     if (strcmp(argv[optind + 1], "ingress") == 0)
0157         type = BPF_CGROUP_INET_INGRESS;
0158     else if (strcmp(argv[optind + 1], "egress") == 0)
0159         type = BPF_CGROUP_INET_EGRESS;
0160     else
0161         return usage(argv[0]);
0162
0163     cg_fd = open(argv[optind], O_DIRECTORY | O_RDONLY);
0164     if (cg_fd < 0) {
0165         printf("Failed to open cgroup path: '%s'\n", strerror(errno));
0166         return EXIT_FAILURE;
0167     }
0168
0169     if (detach_only) {
0170         ret = bpf_prog_detach(cg_fd, type);
0171         printf("bpf_prog_detach() returned '%s' (%d)\n",
0172                strerror(errno), errno);
0173     } else
0174         ret = attach_filter(cg_fd, type, verdict);
0175
0176     return ret;
0177 }