samples/sockmap/sockmap_user.c

   1 /* Copyright (c) 2017 Covalent IO, Inc. http://covalent.io
   2  *
   3  * This program is free software; you can redistribute it and/or
   4  * modify it under the terms of version 2 of the GNU General Public
   5  * License as published by the Free Software Foundation.
   6  *
   7  * This program is distributed in the hope that it will be useful, but
   8  * WITHOUT ANY WARRANTY; without even the implied warranty of
   9  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  10  * General Public License for more details.
  11  */
  12 #include <stdio.h>
  13 #include <stdlib.h>
  14 #include <sys/socket.h>
  15 #include <sys/ioctl.h>
  16 #include <sys/select.h>
  17 #include <netinet/in.h>
  18 #include <arpa/inet.h>
  19 #include <unistd.h>
  20 #include <string.h>
  21 #include <errno.h>
  22 #include <sys/ioctl.h>
  23 #include <stdbool.h>
  24 #include <signal.h>
  25 #include <fcntl.h>
  26
  27 #include <sys/time.h>
  28 #include <sys/types.h>
  29
  30 #include <linux/netlink.h>
  31 #include <linux/socket.h>
  32 #include <linux/sock_diag.h>
  33 #include <linux/bpf.h>
  34 #include <linux/if_link.h>
  35 #include <assert.h>
  36 #include <libgen.h>
  37
  38 #include "../bpf/bpf_load.h"
  39 #include "../bpf/bpf_util.h"
  40 #include "../bpf/libbpf.h"
  41
  42 int running;
  43 void running_handler(int a);
  44
  45 /* randomly selected ports for testing on lo */
  46 #define S1_PORT 10000
  47 #define S2_PORT 10001
  48
  49 static int sockmap_test_sockets(int rate, int dot)
  50 {
  51         int i, sc, err, max_fd, one = 1;
  52         int s1, s2, c1, c2, p1, p2;
  53         struct sockaddr_in addr;
  54         struct timeval timeout;
  55         char buf[1024] = {0};
  56         int *fds[4] = {&s1, &s2, &c1, &c2};
  57         fd_set w;
  58
  59         s1 = s2 = p1 = p2 = c1 = c2 = 0;
  60
  61         /* Init sockets */
  62         for (i = 0; i < 4; i++) {
  63                 *fds[i] = socket(AF_INET, SOCK_STREAM, 0);
  64                 if (*fds[i] < 0) {
  65                         perror("socket s1 failed()");
  66                         err = *fds[i];
  67                         goto out;
  68                 }
  69         }
  70
  71         /* Allow reuse */
  72         for (i = 0; i < 2; i++) {
  73                 err = setsockopt(*fds[i], SOL_SOCKET, SO_REUSEADDR,
  74                                  (char *)&one, sizeof(one));
  75                 if (err) {
  76                         perror("setsockopt failed()");
  77                         goto out;
  78                 }
  79         }
  80
  81         /* Non-blocking sockets */
  82         for (i = 0; i < 4; i++) {
  83                 err = ioctl(*fds[i], FIONBIO, (char *)&one);
  84                 if (err < 0) {
  85                         perror("ioctl s1 failed()");
  86                         goto out;
  87                 }
  88         }
  89
  90         /* Bind server sockets */
  91         memset(&addr, 0, sizeof(struct sockaddr_in));
  92         addr.sin_family = AF_INET;
  93         addr.sin_addr.s_addr = inet_addr("127.0.0.1");
  94
  95         addr.sin_port = htons(S1_PORT);
  96         err = bind(s1, (struct sockaddr *)&addr, sizeof(addr));
  97         if (err < 0) {
  98                 perror("bind s1 failed()\n");
  99                 goto out;
 100         }
 101
 102         addr.sin_port = htons(S2_PORT);
 103         err = bind(s2, (struct sockaddr *)&addr, sizeof(addr));
 104         if (err < 0) {
 105                 perror("bind s2 failed()\n");
 106                 goto out;
 107         }
 108
 109         /* Listen server sockets */
 110         addr.sin_port = htons(S1_PORT);
 111         err = listen(s1, 32);
 112         if (err < 0) {
 113                 perror("listen s1 failed()\n");
 114                 goto out;
 115         }
 116
 117         addr.sin_port = htons(S2_PORT);
 118         err = listen(s2, 32);
 119         if (err < 0) {
 120                 perror("listen s1 failed()\n");
 121                 goto out;
 122         }
 123
 124         /* Initiate Connect */
 125         addr.sin_port = htons(S1_PORT);
 126         err = connect(c1, (struct sockaddr *)&addr, sizeof(addr));
 127         if (err < 0 && errno != EINPROGRESS) {
 128                 perror("connect c1 failed()\n");
 129                 goto out;
 130         }
 131
 132         addr.sin_port = htons(S2_PORT);
 133         err = connect(c2, (struct sockaddr *)&addr, sizeof(addr));
 134         if (err < 0 && errno != EINPROGRESS) {
 135                 perror("connect c2 failed()\n");
 136                 goto out;
 137         }
 138
 139         /* Accept Connecrtions */
 140         p1 = accept(s1, NULL, NULL);
 141         if (p1 < 0) {
 142                 perror("accept s1 failed()\n");
 143                 goto out;
 144         }
 145
 146         p2 = accept(s2, NULL, NULL);
 147         if (p2 < 0) {
 148                 perror("accept s1 failed()\n");
 149                 goto out;
 150         }
 151
 152         max_fd = p2;
 153         timeout.tv_sec = 10;
 154         timeout.tv_usec = 0;
 155
 156         printf("connected sockets: c1 <-> p1, c2 <-> p2\n");
 157         printf("cgroups binding: c1(%i) <-> s1(%i) - - - c2(%i) <-> s2(%i)\n",
 158                 c1, s1, c2, s2);
 159
 160         /* Ping/Pong data from client to server */
 161         sc = send(c1, buf, sizeof(buf), 0);
 162         if (sc < 0) {
 163                 perror("send failed()\n");
 164                 goto out;
 165         }
 166
 167         do {
 168                 int s, rc, i;
 169
 170                 /* FD sets */
 171                 FD_ZERO(&w);
 172                 FD_SET(c1, &w);
 173                 FD_SET(c2, &w);
 174                 FD_SET(p1, &w);
 175                 FD_SET(p2, &w);
 176
 177                 s = select(max_fd + 1, &w, NULL, NULL, &timeout);
 178                 if (s == -1) {
 179                         perror("select()");
 180                         break;
 181                 } else if (!s) {
 182                         fprintf(stderr, "unexpected timeout\n");
 183                         break;
 184                 }
 185
 186                 for (i = 0; i <= max_fd && s > 0; ++i) {
 187                         if (!FD_ISSET(i, &w))
 188                                 continue;
 189
 190                         s--;
 191
 192                         rc = recv(i, buf, sizeof(buf), 0);
 193                         if (rc < 0) {
 194                                 if (errno != EWOULDBLOCK) {
 195                                         perror("recv failed()\n");
 196                                         break;
 197                                 }
 198                         }
 199
 200                         if (rc == 0) {
 201                                 close(i);
 202                                 break;
 203                         }
 204
 205                         sc = send(i, buf, rc, 0);
 206                         if (sc < 0) {
 207                                 perror("send failed()\n");
 208                                 break;
 209                         }
 210                 }
 211                 sleep(rate);
 212                 if (dot) {
 213                         printf(".");
 214                         fflush(stdout);
 215
 216                 }
 217         } while (running);
 218
 219 out:
 220         close(s1);
 221         close(s2);
 222         close(p1);
 223         close(p2);
 224         close(c1);
 225         close(c2);
 226         return err;
 227 }
 228
 229 int main(int argc, char **argv)
 230 {
 231         int rate = 1, dot = 1;
 232         char filename[256];
 233         int err, cg_fd;
 234         char *cg_path;
 235
 236         cg_path = argv[argc - 1];
 237         snprintf(filename, sizeof(filename), "%s_kern.o", argv[0]);
 238
 239         running = 1;
 240
 241         /* catch SIGINT */
 242         signal(SIGINT, running_handler);
 243
 244         if (load_bpf_file(filename)) {
 245                 fprintf(stderr, "load_bpf_file: (%s) %s\n",
 246                         filename, strerror(errno));
 247                 return 1;
 248         }
 249
 250         /* Cgroup configuration */
 251         cg_fd = open(cg_path, O_DIRECTORY, O_RDONLY);
 252         if (cg_fd < 0) {
 253                 fprintf(stderr, "ERROR: (%i) open cg path failed: %s\n",
 254                         cg_fd, cg_path);
 255                 return cg_fd;
 256         }
 257
 258         /* Attach programs to sockmap */
 259         err = bpf_prog_attach(prog_fd[0], map_fd[0],
 260                                 BPF_SK_SKB_STREAM_PARSER, 0);
 261         if (err) {
 262                 fprintf(stderr, "ERROR: bpf_prog_attach (sockmap): %d (%s)\n",
 263                         err, strerror(errno));
 264                 return err;
 265         }
 266
 267         err = bpf_prog_attach(prog_fd[1], map_fd[0],
 268                                 BPF_SK_SKB_STREAM_VERDICT, 0);
 269         if (err) {
 270                 fprintf(stderr, "ERROR: bpf_prog_attach (sockmap): %d (%s)\n",
 271                         err, strerror(errno));
 272                 return err;
 273         }
 274
 275         /* Attach to cgroups */
 276         err = bpf_prog_attach(prog_fd[2], cg_fd, BPF_CGROUP_SOCK_OPS, 0);
 277         if (err) {
 278                 fprintf(stderr, "ERROR: bpf_prog_attach (groups): %d (%s)\n",
 279                         err, strerror(errno));
 280                 return err;
 281         }
 282
 283         err = sockmap_test_sockets(rate, dot);
 284         if (err) {
 285                 fprintf(stderr, "ERROR: test socket failed: %d\n", err);
 286                 return err;
 287         }
 288         return 0;
 289 }
 290
 291 void running_handler(int a)
 292 {
 293         running = 0;
 294 }