| /* Copyright 2015-2023 Rivoreo |
| |
| Permission is hereby granted, free of charge, to any person obtaining |
| a copy of this software and associated documentation files (the |
| "Software"), to deal in the Software without restriction, including |
| without limitation the rights to use, copy, modify, merge, publish, |
| distribute, sublicense, and/or sell copies of the Software, and to |
| permit persons to whom the Software is furnished to do so, subject to |
| the following conditions: |
| |
| The above copyright notice and this permission notice shall be |
| included in all copies or substantial portions of the Software. |
| |
| THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, |
| EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF |
| MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND |
| NONINFRINGEMENT. IN NO EVENT SHALL THE COPYRIGHT HOLDERS BE LIABLE |
| FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF |
| CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION |
| WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. |
| */ |
| |
| #include <sys/types.h> |
| #include <sys/mman.h> |
| #include <string.h> |
| #include <stdio.h> |
| #include <stdlib.h> |
| #include <unistd.h> |
| #include <ctype.h> |
| #include <sys/socket.h> |
| #include <sys/select.h> |
| #include <netinet/in.h> |
| #include <arpa/inet.h> |
| #include <errno.h> |
| #include <assert.h> |
| #include <fcntl.h> |
| #include <signal.h> |
| |
| static void *map_write_only(size_t length) { |
| #if 0 |
| int fd = open("/dev/zero", O_RDONLY); |
| if(fd == -1) { |
| perror("/dev/zero"); |
| return MAP_FAILED; |
| } |
| #endif |
| char *p = mmap(NULL, length, PROT_WRITE, |
| #ifdef MAP_NOSYNC |
| MAP_NOSYNC | |
| #endif |
| #ifdef MAP_NOCORE |
| MAP_NOCORE | |
| #endif |
| #ifdef MAP_NORESERVE |
| MAP_NORESERVE | |
| #endif |
| #ifdef MAP_UNINITIALIZED |
| MAP_UNINITIALIZED | |
| #endif |
| #if 1 |
| MAP_ANON | |
| #endif |
| #if 0 |
| MAP_PRIVATE, 0, 0); |
| close(fd); |
| #else |
| MAP_PRIVATE, -1, 0); |
| #endif |
| return p; |
| } |
| |
| static int open_connection_background(const struct sockaddr_in *addr) { |
| #ifdef SOCK_NONBLOCK |
| int fd = socket(PF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0); |
| #else |
| int fd = socket(PF_INET, SOCK_STREAM, 0); |
| #endif |
| if(fd == -1) { |
| perror("socket"); |
| return -1; |
| } |
| if(fd >= FD_SETSIZE) { |
| fprintf(stderr, "Warning: fd %d too large, ignoring\n", fd); |
| close(fd); |
| return -1; |
| } |
| #ifndef SOCK_NONBLOCK |
| int flags = fcntl(fd, F_GETFL, 0); |
| if(flags == -1) { |
| perror("fcntl get flag error"); |
| close(fd); |
| return -1; |
| } |
| if(fcntl(fd, F_SETFL, flags | O_NONBLOCK) < 0) { |
| perror("fcntl F_SETFL"); |
| close(fd); |
| return -1; |
| } |
| #endif |
| int value = 1; |
| setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &value, sizeof value); |
| #ifdef SO_REUSEPORT |
| setsockopt(fd, SOL_SOCKET, SO_REUSEPORT, &value, sizeof value); |
| #endif |
| setsockopt(fd, SOL_SOCKET, SO_KEEPALIVE, &value, sizeof value); |
| #if 0 |
| struct timeval timeout = { .tv_sec = 30 }; |
| setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof timeout); |
| #endif |
| while(connect(fd, (struct sockaddr *)addr, sizeof(struct sockaddr_in)) == -1) { |
| if(errno == EINTR || errno == EINPROGRESS) return fd; |
| perror("connect"); |
| close(fd); |
| return -1; |
| } |
| return fd; |
| } |
| |
| int main(int argc, char **argv) { |
| if(argc != 4) { |
| printf("Usage: %s <server> <port> 1-%d\n Open ports to receive data\n", argv[0], FD_SETSIZE); |
| exit(1); |
| } |
| |
| struct sockaddr_in addr = { |
| .sin_family = AF_INET, .sin_port = htons(atoi(argv[2])) |
| }; |
| if(inet_pton(AF_INET, argv[1], &addr.sin_addr) < 1) { |
| fprintf(stderr, "%s: Failed to resolve server address %s\n", argv[0], argv[1]); |
| return 1; |
| } |
| static const char payload[] = "binary 32Ki\nyes\n"; |
| |
| int nports = atoi(argv[3]); |
| if(nports < 1) { |
| fprintf(stderr, "Number of ports should be bigger than 0\n"); |
| exit(1); |
| } |
| if(nports > FD_SETSIZE) { |
| fprintf(stderr, "Number of ports can't be larger than %d\n", FD_SETSIZE); |
| return 1; |
| } |
| |
| struct sigaction act = { .sa_handler = SIG_IGN }; |
| if(sigaction(SIGPIPE, &act, NULL) < 0) perror("sigaction"); |
| |
| size_t buffer_size = sysconf(_SC_PAGESIZE); |
| if(buffer_size < 8192) buffer_size = 8192; |
| void *buffer = map_write_only(buffer_size); |
| if(buffer == MAP_FAILED) { |
| buffer = malloc(buffer_size); |
| if(!buffer) { |
| perror("malloc"); |
| return 1; |
| } |
| } |
| |
| int max_fd = -1; |
| int *sockets = malloc(nports * sizeof(int)); |
| if(!sockets) { |
| perror("malloc"); |
| return 1; |
| } |
| int i = 0; |
| while(i < nports) { |
| int fd = open_connection_background(&addr); |
| if(fd == -1) break; |
| //fprintf(stderr, "Opened fd %d\n", fd); |
| if(fd > max_fd) max_fd = fd; |
| sockets[i++] = fd; |
| } |
| if(i < 1) { |
| fprintf(stderr, "%s: Failed to create any connection\n", argv[0]); |
| exit(1); |
| } |
| if(i < nports) { |
| nports = i; |
| sockets = realloc(sockets, nports * sizeof(int)); |
| assert(sockets); |
| } |
| |
| int *nbytes_written = calloc(nports, sizeof(int)); |
| if(!nbytes_written) { |
| perror("calloc"); |
| return 1; |
| } |
| |
| fd_set orig_rfdset, orig_wfdset; |
| FD_ZERO(&orig_rfdset); |
| FD_ZERO(&orig_wfdset); |
| i = 0; |
| while(i < nports) { |
| if(sockets[i] != -1) FD_SET(sockets[i], &orig_wfdset); |
| i++; |
| } |
| int orig_nports = nports; |
| while(nports > 0) { |
| fd_set rfdset = orig_rfdset; |
| fd_set wfdset = orig_wfdset; |
| struct timeval timeout = { .tv_sec = 600 }; |
| int select_return = select(max_fd + 1, &rfdset, &wfdset, NULL, &timeout); |
| if(select_return == -1) { |
| if(errno == EINTR) continue; |
| perror("select error"); |
| exit(1); |
| } |
| if(!select_return) { |
| fprintf(stderr, "%s: Timed out waiting for %d remaining connection(s), reconnecting\n", |
| argv[0], nports); |
| for(i = 0; i < orig_nports; i++) { |
| int fd = sockets[i]; |
| if(fd != -1) { |
| close(fd); |
| FD_CLR(fd, &orig_rfdset); |
| FD_CLR(fd, &orig_wfdset); |
| fd = open_connection_background(&addr); |
| if(fd == -1) { |
| fprintf(stderr, "Failed to reopen socket for index %d\n", i); |
| nports--; |
| } |
| sockets[i] = fd; |
| FD_SET(fd, &orig_wfdset); |
| if(fd > max_fd) max_fd = fd; |
| nbytes_written[i] = 0; |
| } |
| } |
| continue; |
| } |
| for(i = 0; i < orig_nports && select_return > 0; i++) { |
| int fd = sockets[i]; |
| if(fd == -1) continue; |
| if(FD_ISSET(fd, &wfdset)) { |
| int ss = nbytes_written[i]; |
| assert(ss < sizeof payload); |
| int s; |
| do { |
| s = write(fd, payload + ss, sizeof payload - 1 - ss); |
| } while(s < 0 && errno == EINTR); |
| if(s < 0) { |
| fprintf(stderr, "index %d fd %d write error: %s\n", i, fd, strerror(errno)); |
| close(fd); |
| FD_CLR(fd, &orig_wfdset); |
| fd = open_connection_background(&addr); |
| if(fd == -1) { |
| fprintf(stderr, "Failed to reopen socket for index %d\n", i); |
| nports--; |
| fprintf(stderr, "%d connection(s) left\n", nports); |
| } |
| sockets[i] = fd; |
| FD_SET(fd, &orig_wfdset); |
| if(fd > max_fd) max_fd = fd; |
| nbytes_written[i] = 0; |
| } else { |
| nbytes_written[i] += s; |
| assert(nbytes_written[i] < sizeof payload); |
| if(nbytes_written[i] == sizeof payload - 1) { |
| FD_CLR(fd, &orig_wfdset); |
| FD_SET(fd, &orig_rfdset); |
| } |
| } |
| select_return--; |
| } else if(FD_ISSET(fd, &rfdset)) { |
| int s; |
| do { |
| s = read(fd, buffer, buffer_size); |
| } while(s < 0 && errno == EINTR); |
| if(s < 0) { |
| fprintf(stderr, "index %d fd %d read error: %s\n", i, fd, strerror(errno)); |
| close(fd); |
| FD_CLR(fd, &orig_rfdset); |
| fd = open_connection_background(&addr); |
| if(fd == -1) { |
| fprintf(stderr, "Failed to reopen socket for index %d\n", i); |
| nports--; |
| fprintf(stderr, "%d connection(s) left\n", nports); |
| } |
| sockets[i] = fd; |
| FD_SET(fd, &orig_wfdset); |
| if(fd > max_fd) max_fd = fd; |
| nbytes_written[i] = 0; |
| } else if(!s) { |
| fprintf(stderr, "index %d fd %d EOF\n", i, fd); |
| close(fd); |
| sockets[i] = -1; |
| nports--; |
| FD_CLR(fd, &orig_rfdset); |
| fprintf(stderr, "%d connection(s) left\n", nports); |
| } else { |
| fprintf(stderr, "index %d fd %d read %d byte\n", |
| i, fd, s); |
| } |
| select_return--; |
| } |
| } |
| } |
| |
| return 0; |
| } |