multiplex.cpp

   1 /*
   2   Copyright (C) 2015  Paul Fertser <fercerpav@gmail.com>
   3
   4   This program is free software: you can redistribute it and/or modify
   5   it under the terms of the GNU General Public License as published by
   6   the Free Software Foundation, either version 3 of the License, or
   7   (at your option) any later version.
   8
   9   This program is distributed in the hope that it will be useful,
  10   but WITHOUT ANY WARRANTY; without even the implied warranty of
  11   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12   GNU General Public License for more details.
  13
  14   You should have received a copy of the GNU General Public License
  15   along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16 */
  17
  18 #include <unistd.h>
  19 #include <fcntl.h>
  20 #include <signal.h>
  21 #include <sys/select.h>
  22 #include <sys/socket.h>
  23 #include <netinet/in.h>
  24 #include <arpa/inet.h>
  25
  26 #include <iostream>
  27 #include <errno.h>
  28 #include <cstring>
  29 #include <vector>
  30 #include <algorithm>
  31
  32 using namespace std;
  33
  34 #define BUF_SIZE        8192
  35
  36 static int usage(const string &cmdname)
  37 {
  38         cout << "Usage: " << cmdname << " <port>" << endl
  39              << "Read continuous stream of data from standard input and multiplex it" << endl
  40              << "to all the connected TCP clients." << endl << endl
  41              << "<port> is the TCP port number to listen on" << endl << endl
  42              << "Example: tail -F tracelog.bin | " << cmdname << " 7777" << endl << endl
  43              << "Copyright (C) 2015  Paul Fertser <fercerpav@gmail.com>, distributed under GPLv3+ conditions" << endl;
  44         return 0;
  45 }
  46
  47 static int report_error(const string &description)
  48 {
  49         cout << description << strerror(errno) << endl;
  50         return errno;
  51 }
  52
  53 static void close_socket(int fd)
  54 {
  55         shutdown(fd, SHUT_RDWR);
  56         close(fd);
  57 }
  58
  59 int main(int argc, char *argv[])
  60 {
  61         if (argc != 2)
  62                 return usage(argv[0]);
  63
  64         struct sigaction sa = { { SIG_IGN } };
  65         if (sigaction(SIGPIPE, &sa, NULL) == -1)
  66                 return report_error("sigaction() error: ");
  67
  68         auto sock = socket(AF_INET6, SOCK_STREAM, 0);
  69         if (sock == -1)
  70                 return report_error("Can not create socket: ");
  71
  72         int so_reuseaddr = 1;
  73         if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &so_reuseaddr, sizeof so_reuseaddr) == -1)
  74                 return report_error("Can not set SO_REUSEADDR: ");
  75
  76         struct sockaddr_in6 saddr = {};
  77         saddr.sin6_family = AF_INET6;
  78         saddr.sin6_port = htons(stoi(argv[1]));
  79         saddr.sin6_addr = in6addr_any;
  80
  81         if (bind(sock, reinterpret_cast<sockaddr*>(&saddr), sizeof(saddr)) == -1)
  82                 return report_error("Can not bind socket: ");
  83
  84         if (listen(sock, 128) == -1)
  85                 return report_error("Can not enable listening: ");
  86
  87         vector<int> clients;
  88         clients.reserve(32);
  89         fd_set rfds;
  90         FD_ZERO(&rfds);
  91
  92         while (true) {
  93                 FD_SET(STDIN_FILENO, &rfds);
  94                 FD_SET(sock, &rfds);
  95
  96                 if (select(sock + 1, &rfds, NULL, NULL, NULL) == -1)
  97                         return report_error("select() error: ");
  98
  99                 if (FD_ISSET(sock, &rfds)) {
 100                         auto newclient = accept(sock, NULL, NULL);
 101                         if (newclient == -1)
 102                                 return report_error("accept() error: ");
 103
 104                         fcntl(newclient, F_SETFL, O_NONBLOCK);
 105                         clients.push_back(newclient);
 106                 }
 107
 108                 if (FD_ISSET(STDIN_FILENO, &rfds)) {
 109                         char b[BUF_SIZE];
 110                         ssize_t n;
 111                         n = read(STDIN_FILENO, b, sizeof(b));
 112                         if (n == -1)
 113                                 return report_error("read() error: ");
 114                         else if (n == 0)
 115                                 break; /* EOF received */
 116
 117                         clients.erase(remove_if(clients.begin(), clients.end(), [&](int fd) {
 118                                                 if (write(fd, b, n) != n) {
 119                                                         close_socket(fd);
 120                                                         return true;
 121                                                 }
 122                                                 return false;
 123                                         }), clients.end());
 124                 }
 125         }
 126
 127         for_each(clients.begin(), clients.end(), close_socket);
 128         close(sock);
 129
 130         return 0;
 131 }