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 #include <sys/wait.h>
  27 #include <time.h>
  28
  29 #include <sys/time.h>
  30 #include <sys/resource.h>
  31 #include <sys/types.h>
  32
  33 #include <linux/netlink.h>
  34 #include <linux/socket.h>
  35 #include <linux/sock_diag.h>
  36 #include <linux/bpf.h>
  37 #include <linux/if_link.h>
  38 #include <assert.h>
  39 #include <libgen.h>
  40
  41 #include <getopt.h>
  42
  43 #include "../bpf/bpf_load.h"
  44 #include "../bpf/bpf_util.h"
  45 #include "../bpf/libbpf.h"
  46
  47 int running;
  48 void running_handler(int a);
  49
  50 /* randomly selected ports for testing on lo */
  51 #define S1_PORT 10000
  52 #define S2_PORT 10001
  53
  54 /* global sockets */
  55 int s1, s2, c1, c2, p1, p2;
  56
  57 static const struct option long_options[] = {
  58         {"help",        no_argument,            NULL, 'h' },
  59         {"cgroup",      required_argument,      NULL, 'c' },
  60         {"rate",        required_argument,      NULL, 'r' },
  61         {"verbose",     no_argument,            NULL, 'v' },
  62         {"iov_count",   required_argument,      NULL, 'i' },
  63         {"length",      required_argument,      NULL, 'l' },
  64         {"test",        required_argument,      NULL, 't' },
  65         {0, 0, NULL, 0 }
  66 };
  67
  68 static void usage(char *argv[])
  69 {
  70         int i;
  71
  72         printf(" Usage: %s --cgroup <cgroup_path>\n", argv[0]);
  73         printf(" options:\n");
  74         for (i = 0; long_options[i].name != 0; i++) {
  75                 printf(" --%-12s", long_options[i].name);
  76                 if (long_options[i].flag != NULL)
  77                         printf(" flag (internal value:%d)\n",
  78                                 *long_options[i].flag);
  79                 else
  80                         printf(" -%c\n", long_options[i].val);
  81         }
  82         printf("\n");
  83 }
  84
  85 static int sockmap_init_sockets(void)
  86 {
  87         int i, err, one = 1;
  88         struct sockaddr_in addr;
  89         int *fds[4] = {&s1, &s2, &c1, &c2};
  90
  91         s1 = s2 = p1 = p2 = c1 = c2 = 0;
  92
  93         /* Init sockets */
  94         for (i = 0; i < 4; i++) {
  95                 *fds[i] = socket(AF_INET, SOCK_STREAM, 0);
  96                 if (*fds[i] < 0) {
  97                         perror("socket s1 failed()");
  98                         return errno;
  99                 }
 100         }
 101
 102         /* Allow reuse */
 103         for (i = 0; i < 2; i++) {
 104                 err = setsockopt(*fds[i], SOL_SOCKET, SO_REUSEADDR,
 105                                  (char *)&one, sizeof(one));
 106                 if (err) {
 107                         perror("setsockopt failed()");
 108                         return errno;
 109                 }
 110         }
 111
 112         /* Non-blocking sockets */
 113         for (i = 0; i < 2; i++) {
 114                 err = ioctl(*fds[i], FIONBIO, (char *)&one);
 115                 if (err < 0) {
 116                         perror("ioctl s1 failed()");
 117                         return errno;
 118                 }
 119         }
 120
 121         /* Bind server sockets */
 122         memset(&addr, 0, sizeof(struct sockaddr_in));
 123         addr.sin_family = AF_INET;
 124         addr.sin_addr.s_addr = inet_addr("127.0.0.1");
 125
 126         addr.sin_port = htons(S1_PORT);
 127         err = bind(s1, (struct sockaddr *)&addr, sizeof(addr));
 128         if (err < 0) {
 129                 perror("bind s1 failed()\n");
 130                 return errno;
 131         }
 132
 133         addr.sin_port = htons(S2_PORT);
 134         err = bind(s2, (struct sockaddr *)&addr, sizeof(addr));
 135         if (err < 0) {
 136                 perror("bind s2 failed()\n");
 137                 return errno;
 138         }
 139
 140         /* Listen server sockets */
 141         addr.sin_port = htons(S1_PORT);
 142         err = listen(s1, 32);
 143         if (err < 0) {
 144                 perror("listen s1 failed()\n");
 145                 return errno;
 146         }
 147
 148         addr.sin_port = htons(S2_PORT);
 149         err = listen(s2, 32);
 150         if (err < 0) {
 151                 perror("listen s1 failed()\n");
 152                 return errno;
 153         }
 154
 155         /* Initiate Connect */
 156         addr.sin_port = htons(S1_PORT);
 157         err = connect(c1, (struct sockaddr *)&addr, sizeof(addr));
 158         if (err < 0 && errno != EINPROGRESS) {
 159                 perror("connect c1 failed()\n");
 160                 return errno;
 161         }
 162
 163         addr.sin_port = htons(S2_PORT);
 164         err = connect(c2, (struct sockaddr *)&addr, sizeof(addr));
 165         if (err < 0 && errno != EINPROGRESS) {
 166                 perror("connect c2 failed()\n");
 167                 return errno;
 168         } else if (err < 0) {
 169                 err = 0;
 170         }
 171
 172         /* Accept Connecrtions */
 173         p1 = accept(s1, NULL, NULL);
 174         if (p1 < 0) {
 175                 perror("accept s1 failed()\n");
 176                 return errno;
 177         }
 178
 179         p2 = accept(s2, NULL, NULL);
 180         if (p2 < 0) {
 181                 perror("accept s1 failed()\n");
 182                 return errno;
 183         }
 184
 185         printf("connected sockets: c1 <-> p1, c2 <-> p2\n");
 186         printf("cgroups binding: c1(%i) <-> s1(%i) - - - c2(%i) <-> s2(%i)\n",
 187                 c1, s1, c2, s2);
 188         return 0;
 189 }
 190
 191 struct msg_stats {
 192         size_t bytes_sent;
 193         size_t bytes_recvd;
 194         struct timespec start;
 195         struct timespec end;
 196 };
 197
 198 static int msg_loop(int fd, int iov_count, int iov_length, int cnt,
 199                     struct msg_stats *s, bool tx)
 200 {
 201         struct msghdr msg = {0};
 202         int err, i, flags = MSG_NOSIGNAL;
 203         struct iovec *iov;
 204
 205         iov = calloc(iov_count, sizeof(struct iovec));
 206         if (!iov)
 207                 return errno;
 208
 209         for (i = 0; i < iov_count; i++) {
 210                 char *d = calloc(iov_length, sizeof(char));
 211
 212                 if (!d) {
 213                         fprintf(stderr, "iov_count %i/%i OOM\n", i, iov_count);
 214                         goto out_errno;
 215                 }
 216                 iov[i].iov_base = d;
 217                 iov[i].iov_len = iov_length;
 218         }
 219
 220         msg.msg_iov = iov;
 221         msg.msg_iovlen = iov_count;
 222
 223         if (tx) {
 224                 clock_gettime(CLOCK_MONOTONIC, &s->start);
 225                 for (i = 0; i < cnt; i++) {
 226                         int sent = sendmsg(fd, &msg, flags);
 227
 228                         if (sent < 0) {
 229                                 perror("send loop error:");
 230                                 goto out_errno;
 231                         }
 232                         s->bytes_sent += sent;
 233                 }
 234                 clock_gettime(CLOCK_MONOTONIC, &s->end);
 235         } else {
 236                 int slct, recv, max_fd = fd;
 237                 struct timeval timeout;
 238                 float total_bytes;
 239                 fd_set w;
 240
 241                 total_bytes = (float)iov_count * (float)iov_length * (float)cnt;
 242                 err = clock_gettime(CLOCK_MONOTONIC, &s->start);
 243                 if (err < 0)
 244                         perror("recv start time: ");
 245                 while (s->bytes_recvd < total_bytes) {
 246                         timeout.tv_sec = 1;
 247                         timeout.tv_usec = 0;
 248
 249                         /* FD sets */
 250                         FD_ZERO(&w);
 251                         FD_SET(fd, &w);
 252
 253                         slct = select(max_fd + 1, &w, NULL, NULL, &timeout);
 254                         if (slct == -1) {
 255                                 perror("select()");
 256                                 clock_gettime(CLOCK_MONOTONIC, &s->end);
 257                                 goto out_errno;
 258                         } else if (!slct) {
 259                                 fprintf(stderr, "unexpected timeout\n");
 260                                 errno = -EIO;
 261                                 clock_gettime(CLOCK_MONOTONIC, &s->end);
 262                                 goto out_errno;
 263                         }
 264
 265                         recv = recvmsg(fd, &msg, flags);
 266                         if (recv < 0) {
 267                                 if (errno != EWOULDBLOCK) {
 268                                         clock_gettime(CLOCK_MONOTONIC, &s->end);
 269                                         perror("recv failed()\n");
 270                                         goto out_errno;
 271                                 }
 272                         }
 273
 274                         s->bytes_recvd += recv;
 275                 }
 276                 clock_gettime(CLOCK_MONOTONIC, &s->end);
 277         }
 278
 279         for (i = 0; i < iov_count; i++)
 280                 free(iov[i].iov_base);
 281         free(iov);
 282         return 0;
 283 out_errno:
 284         for (i = 0; i < iov_count; i++)
 285                 free(iov[i].iov_base);
 286         free(iov);
 287         return errno;
 288 }
 289
 290 static float giga = 1000000000;
 291
 292 static inline float sentBps(struct msg_stats s)
 293 {
 294         return s.bytes_sent / (s.end.tv_sec - s.start.tv_sec);
 295 }
 296
 297 static inline float recvdBps(struct msg_stats s)
 298 {
 299         return s.bytes_recvd / (s.end.tv_sec - s.start.tv_sec);
 300 }
 301
 302 static int sendmsg_test(int iov_count, int iov_buf, int cnt,
 303                         int verbose, bool base)
 304 {
 305         float sent_Bps = 0, recvd_Bps = 0;
 306         int rx_fd, txpid, rxpid, err = 0;
 307         struct msg_stats s = {0};
 308         int status;
 309
 310         errno = 0;
 311
 312         if (base)
 313                 rx_fd = p1;
 314         else
 315                 rx_fd = p2;
 316
 317         rxpid = fork();
 318         if (rxpid == 0) {
 319                 err = msg_loop(rx_fd, iov_count, iov_buf, cnt, &s, false);
 320                 if (err)
 321                         fprintf(stderr,
 322                                 "msg_loop_rx: iov_count %i iov_buf %i cnt %i err %i\n",
 323                                 iov_count, iov_buf, cnt, err);
 324                 shutdown(p2, SHUT_RDWR);
 325                 shutdown(p1, SHUT_RDWR);
 326                 if (s.end.tv_sec - s.start.tv_sec) {
 327                         sent_Bps = sentBps(s);
 328                         recvd_Bps = recvdBps(s);
 329                 }
 330                 fprintf(stdout,
 331                         "rx_sendmsg: TX: %zuB %fB/s %fGB/s RX: %zuB %fB/s %fGB/s\n",
 332                         s.bytes_sent, sent_Bps, sent_Bps/giga,
 333                         s.bytes_recvd, recvd_Bps, recvd_Bps/giga);
 334                 exit(1);
 335         } else if (rxpid == -1) {
 336                 perror("msg_loop_rx: ");
 337                 return errno;
 338         }
 339
 340         txpid = fork();
 341         if (txpid == 0) {
 342                 err = msg_loop(c1, iov_count, iov_buf, cnt, &s, true);
 343                 if (err)
 344                         fprintf(stderr,
 345                                 "msg_loop_tx: iov_count %i iov_buf %i cnt %i err %i\n",
 346                                 iov_count, iov_buf, cnt, err);
 347                 shutdown(c1, SHUT_RDWR);
 348                 if (s.end.tv_sec - s.start.tv_sec) {
 349                         sent_Bps = sentBps(s);
 350                         recvd_Bps = recvdBps(s);
 351                 }
 352                 fprintf(stdout,
 353                         "tx_sendmsg: TX: %zuB %fB/s %f GB/s RX: %zuB %fB/s %fGB/s\n",
 354                         s.bytes_sent, sent_Bps, sent_Bps/giga,
 355                         s.bytes_recvd, recvd_Bps, recvd_Bps/giga);
 356                 exit(1);
 357         } else if (txpid == -1) {
 358                 perror("msg_loop_tx: ");
 359                 return errno;
 360         }
 361
 362         assert(waitpid(rxpid, &status, 0) == rxpid);
 363         assert(waitpid(txpid, &status, 0) == txpid);
 364         return err;
 365 }
 366
 367 static int forever_ping_pong(int rate, int verbose)
 368 {
 369         struct timeval timeout;
 370         char buf[1024] = {0};
 371         int sc;
 372
 373         timeout.tv_sec = 10;
 374         timeout.tv_usec = 0;
 375
 376         /* Ping/Pong data from client to server */
 377         sc = send(c1, buf, sizeof(buf), 0);
 378         if (sc < 0) {
 379                 perror("send failed()\n");
 380                 return sc;
 381         }
 382
 383         do {
 384                 int s, rc, i, max_fd = p2;
 385                 fd_set w;
 386
 387                 /* FD sets */
 388                 FD_ZERO(&w);
 389                 FD_SET(c1, &w);
 390                 FD_SET(c2, &w);
 391                 FD_SET(p1, &w);
 392                 FD_SET(p2, &w);
 393
 394                 s = select(max_fd + 1, &w, NULL, NULL, &timeout);
 395                 if (s == -1) {
 396                         perror("select()");
 397                         break;
 398                 } else if (!s) {
 399                         fprintf(stderr, "unexpected timeout\n");
 400                         break;
 401                 }
 402
 403                 for (i = 0; i <= max_fd && s > 0; ++i) {
 404                         if (!FD_ISSET(i, &w))
 405                                 continue;
 406
 407                         s--;
 408
 409                         rc = recv(i, buf, sizeof(buf), 0);
 410                         if (rc < 0) {
 411                                 if (errno != EWOULDBLOCK) {
 412                                         perror("recv failed()\n");
 413                                         return rc;
 414                                 }
 415                         }
 416
 417                         if (rc == 0) {
 418                                 close(i);
 419                                 break;
 420                         }
 421
 422                         sc = send(i, buf, rc, 0);
 423                         if (sc < 0) {
 424                                 perror("send failed()\n");
 425                                 return sc;
 426                         }
 427                 }
 428
 429                 if (rate)
 430                         sleep(rate);
 431
 432                 if (verbose) {
 433                         printf(".");
 434                         fflush(stdout);
 435
 436                 }
 437         } while (running);
 438
 439         return 0;
 440 }
 441
 442 enum {
 443         PING_PONG,
 444         SENDMSG,
 445         BASE,
 446 };
 447
 448 int main(int argc, char **argv)
 449 {
 450         int iov_count = 1, length = 1024, rate = 1, verbose = 0;
 451         struct rlimit r = {10 * 1024 * 1024, RLIM_INFINITY};
 452         int opt, longindex, err, cg_fd = 0;
 453         int test = PING_PONG;
 454         char filename[256];
 455
 456         while ((opt = getopt_long(argc, argv, "hvc:r:i:l:t:",
 457                                   long_options, &longindex)) != -1) {
 458                 switch (opt) {
 459                 /* Cgroup configuration */
 460                 case 'c':
 461                         cg_fd = open(optarg, O_DIRECTORY, O_RDONLY);
 462                         if (cg_fd < 0) {
 463                                 fprintf(stderr,
 464                                         "ERROR: (%i) open cg path failed: %s\n",
 465                                         cg_fd, optarg);
 466                                 return cg_fd;
 467                         }
 468                         break;
 469                 case 'r':
 470                         rate = atoi(optarg);
 471                         break;
 472                 case 'v':
 473                         verbose = 1;
 474                         break;
 475                 case 'i':
 476                         iov_count = atoi(optarg);
 477                         break;
 478                 case 'l':
 479                         length = atoi(optarg);
 480                         break;
 481                 case 't':
 482                         if (strcmp(optarg, "ping") == 0) {
 483                                 test = PING_PONG;
 484                         } else if (strcmp(optarg, "sendmsg") == 0) {
 485                                 test = SENDMSG;
 486                         } else if (strcmp(optarg, "base") == 0) {
 487                                 test = BASE;
 488                         } else {
 489                                 usage(argv);
 490                                 return -1;
 491                         }
 492                         break;
 493                 case 'h':
 494                 default:
 495                         usage(argv);
 496                         return -1;
 497                 }
 498         }
 499
 500         if (!cg_fd) {
 501                 fprintf(stderr, "%s requires cgroup option: --cgroup <path>\n",
 502                         argv[0]);
 503                 return -1;
 504         }
 505
 506         if (setrlimit(RLIMIT_MEMLOCK, &r)) {
 507                 perror("setrlimit(RLIMIT_MEMLOCK)");
 508                 return 1;
 509         }
 510
 511         snprintf(filename, sizeof(filename), "%s_kern.o", argv[0]);
 512
 513         running = 1;
 514
 515         /* catch SIGINT */
 516         signal(SIGINT, running_handler);
 517
 518         /* If base test skip BPF setup */
 519         if (test == BASE)
 520                 goto run;
 521
 522         if (load_bpf_file(filename)) {
 523                 fprintf(stderr, "load_bpf_file: (%s) %s\n",
 524                         filename, strerror(errno));
 525                 return 1;
 526         }
 527
 528         /* Attach programs to sockmap */
 529         err = bpf_prog_attach(prog_fd[0], map_fd[0],
 530                                 BPF_SK_SKB_STREAM_PARSER, 0);
 531         if (err) {
 532                 fprintf(stderr, "ERROR: bpf_prog_attach (sockmap): %d (%s)\n",
 533                         err, strerror(errno));
 534                 return err;
 535         }
 536
 537         err = bpf_prog_attach(prog_fd[1], map_fd[0],
 538                                 BPF_SK_SKB_STREAM_VERDICT, 0);
 539         if (err) {
 540                 fprintf(stderr, "ERROR: bpf_prog_attach (sockmap): %d (%s)\n",
 541                         err, strerror(errno));
 542                 return err;
 543         }
 544
 545         /* Attach to cgroups */
 546         err = bpf_prog_attach(prog_fd[2], cg_fd, BPF_CGROUP_SOCK_OPS, 0);
 547         if (err) {
 548                 fprintf(stderr, "ERROR: bpf_prog_attach (groups): %d (%s)\n",
 549                         err, strerror(errno));
 550                 return err;
 551         }
 552
 553 run:
 554         err = sockmap_init_sockets();
 555         if (err) {
 556                 fprintf(stderr, "ERROR: test socket failed: %d\n", err);
 557                 goto out;
 558         }
 559
 560         if (test == PING_PONG)
 561                 err = forever_ping_pong(rate, verbose);
 562         else if (test == SENDMSG)
 563                 err = sendmsg_test(iov_count, length, rate, verbose, false);
 564         else if (test == BASE)
 565                 err = sendmsg_test(iov_count, length, rate, verbose, true);
 566         else
 567                 fprintf(stderr, "unknown test\n");
 568 out:
 569         close(s1);
 570         close(s2);
 571         close(p1);
 572         close(p2);
 573         close(c1);
 574         close(c2);
 575         close(cg_fd);
 576         return err;
 577 }
 578
 579 void running_handler(int a)
 580 {
 581         running = 0;
 582 }