Initial commit
[multiplex:multiplex.git] / 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 }