samples/bpf/test_cgrp2_sock.c

   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
  28 char bpf_log_buf[BPF_LOG_BUF_SIZE];
  29
  30 static 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
  74         struct bpf_insn *prog;
  75         size_t insns_cnt;
  76         void *p;
  77         int ret;
  78
  79         insns_cnt = sizeof(prog_start) + sizeof(prog_end);
  80         if (idx)
  81                 insns_cnt += sizeof(prog_dev);
  82
  83         if (mark)
  84                 insns_cnt += sizeof(prog_mark);
  85
  86         if (prio)
  87                 insns_cnt += sizeof(prog_prio);
  88
  89         p = prog = malloc(insns_cnt);
  90         if (!prog) {
  91                 fprintf(stderr, "Failed to allocate memory for instructions\n");
  92                 return EXIT_FAILURE;
  93         }
  94
  95         memcpy(p, prog_start, sizeof(prog_start));
  96         p += sizeof(prog_start);
  97
  98         if (idx) {
  99                 memcpy(p, prog_dev, sizeof(prog_dev));
 100                 p += sizeof(prog_dev);
 101         }
 102
 103         if (mark) {
 104                 memcpy(p, prog_mark, sizeof(prog_mark));
 105                 p += sizeof(prog_mark);
 106         }
 107
 108         if (prio) {
 109                 memcpy(p, prog_prio, sizeof(prog_prio));
 110                 p += sizeof(prog_prio);
 111         }
 112
 113         memcpy(p, prog_end, sizeof(prog_end));
 114         p += sizeof(prog_end);
 115
 116         insns_cnt /= sizeof(struct bpf_insn);
 117
 118         ret = bpf_load_program(BPF_PROG_TYPE_CGROUP_SOCK, prog, insns_cnt,
 119                                 "GPL", 0, bpf_log_buf, BPF_LOG_BUF_SIZE);
 120
 121         free(prog);
 122
 123         return ret;
 124 }
 125
 126 static int get_bind_to_device(int sd, char *name, size_t len)
 127 {
 128         socklen_t optlen = len;
 129         int rc;
 130
 131         name[0] = '\0';
 132         rc = getsockopt(sd, SOL_SOCKET, SO_BINDTODEVICE, name, &optlen);
 133         if (rc < 0)
 134                 perror("setsockopt(SO_BINDTODEVICE)");
 135
 136         return rc;
 137 }
 138
 139 static unsigned int get_somark(int sd)
 140 {
 141         unsigned int mark = 0;
 142         socklen_t optlen = sizeof(mark);
 143         int rc;
 144
 145         rc = getsockopt(sd, SOL_SOCKET, SO_MARK, &mark, &optlen);
 146         if (rc < 0)
 147                 perror("getsockopt(SO_MARK)");
 148
 149         return mark;
 150 }
 151
 152 static unsigned int get_priority(int sd)
 153 {
 154         unsigned int prio = 0;
 155         socklen_t optlen = sizeof(prio);
 156         int rc;
 157
 158         rc = getsockopt(sd, SOL_SOCKET, SO_PRIORITY, &prio, &optlen);
 159         if (rc < 0)
 160                 perror("getsockopt(SO_PRIORITY)");
 161
 162         return prio;
 163 }
 164
 165 static int show_sockopts(int family)
 166 {
 167         unsigned int mark, prio;
 168         char name[16];
 169         int sd;
 170
 171         sd = socket(family, SOCK_DGRAM, 17);
 172         if (sd < 0) {
 173                 perror("socket");
 174                 return 1;
 175         }
 176
 177         if (get_bind_to_device(sd, name, sizeof(name)) < 0)
 178                 return 1;
 179
 180         mark = get_somark(sd);
 181         prio = get_priority(sd);
 182
 183         close(sd);
 184
 185         printf("sd %d: dev %s, mark %u, priority %u\n", sd, name, mark, prio);
 186
 187         return 0;
 188 }
 189
 190 static int usage(const char *argv0)
 191 {
 192         printf("Usage:\n");
 193         printf("  Attach a program\n");
 194         printf("  %s -b bind-to-dev -m mark -p prio cg-path\n", argv0);
 195         printf("\n");
 196         printf("  Detach a program\n");
 197         printf("  %s -d cg-path\n", argv0);
 198         printf("\n");
 199         printf("  Show inherited socket settings (mark, priority, and device)\n");
 200         printf("  %s [-6]\n", argv0);
 201         return EXIT_FAILURE;
 202 }
 203
 204 int main(int argc, char **argv)
 205 {
 206         __u32 idx = 0, mark = 0, prio = 0;
 207         const char *cgrp_path = NULL;
 208         int cg_fd, prog_fd, ret;
 209         int family = PF_INET;
 210         int do_attach = 1;
 211         int rc;
 212
 213         while ((rc = getopt(argc, argv, "db:m:p:6")) != -1) {
 214                 switch (rc) {
 215                 case 'd':
 216                         do_attach = 0;
 217                         break;
 218                 case 'b':
 219                         idx = if_nametoindex(optarg);
 220                         if (!idx) {
 221                                 idx = strtoumax(optarg, NULL, 0);
 222                                 if (!idx) {
 223                                         printf("Invalid device name\n");
 224                                         return EXIT_FAILURE;
 225                                 }
 226                         }
 227                         break;
 228                 case 'm':
 229                         mark = strtoumax(optarg, NULL, 0);
 230                         break;
 231                 case 'p':
 232                         prio = strtoumax(optarg, NULL, 0);
 233                         break;
 234                 case '6':
 235                         family = PF_INET6;
 236                         break;
 237                 default:
 238                         return usage(argv[0]);
 239                 }
 240         }
 241
 242         if (optind == argc)
 243                 return show_sockopts(family);
 244
 245         cgrp_path = argv[optind];
 246         if (!cgrp_path) {
 247                 fprintf(stderr, "cgroup path not given\n");
 248                 return EXIT_FAILURE;
 249         }
 250
 251         if (do_attach && !idx && !mark && !prio) {
 252                 fprintf(stderr,
 253                         "One of device, mark or priority must be given\n");
 254                 return EXIT_FAILURE;
 255         }
 256
 257         cg_fd = open(cgrp_path, O_DIRECTORY | O_RDONLY);
 258         if (cg_fd < 0) {
 259                 printf("Failed to open cgroup path: '%s'\n", strerror(errno));
 260                 return EXIT_FAILURE;
 261         }
 262
 263         if (do_attach) {
 264                 prog_fd = prog_load(idx, mark, prio);
 265                 if (prog_fd < 0) {
 266                         printf("Failed to load prog: '%s'\n", strerror(errno));
 267                         printf("Output from kernel verifier:\n%s\n-------\n",
 268                                bpf_log_buf);
 269                         return EXIT_FAILURE;
 270                 }
 271
 272                 ret = bpf_prog_attach(prog_fd, cg_fd,
 273                                       BPF_CGROUP_INET_SOCK_CREATE, 0);
 274                 if (ret < 0) {
 275                         printf("Failed to attach prog to cgroup: '%s'\n",
 276                                strerror(errno));
 277                         return EXIT_FAILURE;
 278                 }
 279         } else {
 280                 ret = bpf_prog_detach(cg_fd, BPF_CGROUP_INET_SOCK_CREATE);
 281                 if (ret < 0) {
 282                         printf("Failed to detach prog from cgroup: '%s'\n",
 283                                strerror(errno));
 284                         return EXIT_FAILURE;
 285                 }
 286         }
 287
 288         close(cg_fd);
 289         return EXIT_SUCCESS;
 290 }