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
  25 #include "libbpf.h"
  26
  27 char bpf_log_buf[BPF_LOG_BUF_SIZE];
  28
  29 static 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
 125 static 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
 138 static 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
 151 static 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
 164 static 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
 189 static 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
 203 int 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 }