rsyslog/tests/minitcpsrvr.c

/* a very simplistic tcp receiver for the rsyslog testbench.
 *
 * Copyright 2016,2024 Rainer Gerhards and Adiscon GmbH.
 *
 * This file is part of the rsyslog project.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *       -or-
 *       see COPYING.ASL20 in the source distribution
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "config.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <poll.h>
#include <fcntl.h>
#include <unistd.h>
#include <arpa/inet.h>
#if defined(__FreeBSD__)
    #include <netinet/in.h>
#endif

#define MAX_CONNECTIONS 10
#define BUFFER_SIZE 1024

/* OK, we use a lot of global. But after all, this is "just" a small
 * testing ... that has evolved a bit ;-)
 */
char wrkBuf[4096];
ssize_t nRead;
size_t nRcv = 0;
int nConnDrop; /* how often has the connection already been dropped? */
int abortListener = 0; /* act like listener was totally aborted */
int sleepAfterConnDrop = 0; /* number of seconds to sleep() when a connection was dropped */
int dropConnection_NbrRcv = 0;
int dropConnection_MaxTimes = 0;
int opt;
int sleepStartup = 0;
char *targetIP = NULL;
int targetPort = -1;
char *portFileName = NULL;
size_t totalWritten = 0;
int listen_fd, conn_fd, fd, file_fd, nfds, port = 8080;
struct sockaddr_in server_addr;
struct pollfd fds[MAX_CONNECTIONS + 1];  // +1 for the listen socket
char buffer[MAX_CONNECTIONS][BUFFER_SIZE];
int buffer_offs[MAX_CONNECTIONS];

static void errout(char *reason) {
    perror(reason);
    fprintf(stderr, "minitcpsrv ABORTS!!!\n\n\n");
    exit(1);
}

static void usage(void) {
    fprintf(stderr, "usage: minitcpsrv -t ip-addr -p port -P portFile -f outfile\n");
    exit(1);
}

static void createListenSocket(void) {
    struct sockaddr_in srvAddr;
    unsigned int srvAddrLen;
    static int portFileWritten = 0;
    int r;

    listen_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (listen_fd < 0) {
        errout("Failed to create listen socket");
    }
    /* Set SO_REUSEADDR and SO_REUSEPORT options - these are vital for some
     * Tests. If not both are supported by the OS (e.g. Solaris 10), some tests
     * will fail. Those need to be excluded.
     */
    int opt = 1;
    if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        errout("setsockopt failed for SO_REUSEADDR");
    }
#ifdef SO_REUSEPORT
    if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEPORT, &opt, sizeof(opt)) < 0) {
        errout("setsockopt failed for SO_REUSEPORT");
    }
#endif

    fprintf(stderr, "listen on target port %d\n", targetPort);
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = inet_addr(targetIP);  // htonl(INADDR_ANY);
    server_addr.sin_port = htons(targetPort);
    srvAddrLen = sizeof(server_addr);


    int sockBound = 0;
    int try = 0;
    do {
        r = bind(listen_fd, (struct sockaddr *)&server_addr, sizeof(server_addr));
        if (r < 0) {
            if (errno == EADDRINUSE) {
                perror("minitcpsrv bind listen socket");
                if (try++ < 10) {
                    sleep(1);
                } else {
                    fprintf(stderr,
                            "minitcpsrv could not bind socket "
                            "after trying hard - terminating\n\n");
                    exit(2);
                }
            } else {
                errout("bind");
            }
        } else {
            sockBound = 1;
        }
    } while (!sockBound);

    if (listen(listen_fd, MAX_CONNECTIONS) < 0) {
        errout("Listen failed");
    }

    if (getsockname(listen_fd, (struct sockaddr *)&srvAddr, &srvAddrLen) == -1) {
        errout("getsockname");
    }
    targetPort = ntohs(srvAddr.sin_port);

    if (portFileName != NULL && !portFileWritten) {
        FILE *fp;
        if (getsockname(listen_fd, (struct sockaddr *)&srvAddr, &srvAddrLen) == -1) {
            errout("getsockname");
        }
        if ((fp = fopen(portFileName, "w+")) == NULL) {
            errout(portFileName);
        }
        fprintf(fp, "%d", ntohs(srvAddr.sin_port));
        fclose(fp);
        portFileWritten = 1;
    }

    fds[0].fd = listen_fd;
    fds[0].events = POLLIN;
}


int main(int argc, char *argv[]) {
    int fdc;
    int fdf = -1;
    struct sockaddr_in cliAddr;
    unsigned int cliAddrLen;
    memset(fds, 0, sizeof(fds));
    memset(buffer_offs, 0, sizeof(buffer_offs));

    while ((opt = getopt(argc, argv, "aB:D:t:p:P:f:s:S:")) != -1) {
        switch (opt) {
            case 'a':  // abort listener: act like the server has died (shutdown and re-open listen socket)
                abortListener = 1;
                break;
            case 'S':  // sleep time after connection drop
                sleepAfterConnDrop = atoi(optarg);
                break;
            case 's':  // sleep time immediately after startup
                sleepStartup = atoi(optarg);
                break;
            case 'B':  // max number of time the connection shall be dropped
                dropConnection_MaxTimes = atoi(optarg);
                break;
            case 'D':  // drop connection after this number of recv() operations (not messages)
                dropConnection_NbrRcv = atoi(optarg);
                break;
            case 't':
                targetIP = optarg;
                break;
            case 'p':
                targetPort = atoi(optarg);
                break;
            case 'P':
                portFileName = optarg;
                break;
            case 'f':
                if (!strcmp(optarg, "-")) {
                    fdf = 1;
                } else {
                    fprintf(stderr, "writing to file %s\n", optarg);
                    fdf = open(optarg, O_WRONLY | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR);
                    if (fdf == -1) errout(argv[3]);
                }
                break;
            default:
                fprintf(stderr, "invalid option '%c' or value missing - terminating...\n", opt);
                usage();
                break;
        }
    }

    if (targetIP == NULL) {
        fprintf(stderr, "-t parameter missing -- terminating\n");
        usage();
    }
    if (targetPort == -1) {
        fprintf(stderr, "-p parameter missing -- terminating\n");
        usage();
    }
    if (fdf == -1) {
        fprintf(stderr, "-f parameter missing -- terminating\n");
        usage();
    }

    if (sleepStartup) {
        printf("minitcpsrv: deliberate sleep of %d seconds\n", sleepStartup);
        sleep(sleepStartup);
        printf("minitcpsrv: end sleep\n");
    }


    createListenSocket();
    nfds = 1;

    int bKeepRunning;
    while (1) {
        int poll_count = poll(fds, nfds, -1);  // -1 means no timeout
        if (poll_count < 0) {
            errout("poll");
        }

        bKeepRunning = 0; /* terminate on last connection close */

        for (int i = 0; i < nfds; i++) {
            if (fds[i].revents == 0) continue;

            if (fds[i].revents != POLLIN) {
                fprintf(stderr, "Error! revents = %d\n", fds[i].revents);
                exit(EXIT_FAILURE);
            }

            if (fds[i].fd == listen_fd) {
                // Accept new connection
                fprintf(stderr, "minitcpsrv: NEW CONNECT\n");
                conn_fd = accept(listen_fd, (struct sockaddr *)NULL, NULL);
                if (conn_fd < 0) {
                    perror("Accept failed");
                    continue;
                }

                fds[nfds].fd = conn_fd;
                fds[nfds].events = POLLIN;
                nfds++;
            } else {
                // Handle data from a client
                fd = fds[i].fd;
                const size_t bytes_to_read = sizeof(buffer[i]) - buffer_offs[i] - 1;
                int read_bytes = read(fd, &(buffer[i][buffer_offs[i]]), bytes_to_read);
                nRcv += read_bytes;
                if (read_bytes < 0) {
                    perror("Read error");
                    close(fd);
                    fds[i].fd = -1;  // Remove from poll set
                } else if (read_bytes == 0) {
                    // Connection closed
                    close(fd);
                    fds[i].fd = -1;  // Remove from poll set
                } else {
                    // last valid char in rcv buffer
                    const int last_byte = read_bytes + buffer_offs[i] - 1;
                    int last_lf = last_byte;
                    while (last_lf > -1 && buffer[i][last_lf] != '\n') {
                        --last_lf;
                    }
                    if (last_lf == -1) { /* no LF found at all */
                        buffer_offs[i] = last_byte;
                    } else {
                        const int bytes_to_write = last_lf + 1;

                        const int written_bytes = write(fdf, buffer[i], bytes_to_write);
                        if (written_bytes != bytes_to_write) errout("write");
                        totalWritten += bytes_to_write;
                        if (bytes_to_write == last_byte + 1) {
                            buffer_offs[i] = 0;
                        } else {
                            /* keep data in buffer, move it to start
                             * and adjust offset.
                             */
                            int unfinished_bytes = last_byte - bytes_to_write + 1;
                            memmove(buffer[i], &(buffer[i][bytes_to_write]), unfinished_bytes);
                            buffer_offs[i] = unfinished_bytes;
                        }
                    }
                }

                /* simulate connection abort, if requested */
                if ((buffer_offs[i] == 0) && (dropConnection_NbrRcv > 0) && (nRcv >= dropConnection_NbrRcv) &&
                    (dropConnection_MaxTimes > 0) && (nConnDrop < dropConnection_MaxTimes)) {
                    nConnDrop++;
                    nRcv = 0;
                    if (abortListener) {
                        fprintf(stderr, "minitcpsrv: simulating died client\n");
                        shutdown(listen_fd, SHUT_RDWR);
                    }
                    close(fd);
                    fds[i].fd = -1;  // Remove from poll set

                    if (sleepAfterConnDrop > 0) {
                        sleep(sleepAfterConnDrop);
                    }
                    if (abortListener) {
                        /* we hope that when we close and immediately re-open, we will
                         * can use the some port again.
                         */
                        close(listen_fd);
                        createListenSocket();
                        fprintf(stderr, "minitcpsrv: reactivated\n");
                    }
                    bKeepRunning = 1; /* do not abort if sole connection! */
                }
            }
        }

        // Compact the array of monitored file descriptors
        for (int i = 0; i < nfds; i++) {
            if (fds[i].fd == -1) {
                for (int j = i; j < nfds - 1; j++) {
                    fds[j] = fds[j + 1];
                }
                i--;
                nfds--;
            }
        }
        // terminate if all connections have been closed
        if (nfds == 1 && !bKeepRunning) break;
    }
    /* -------------------------------------------------- */
    /* let the OS do the cleanup */
    fprintf(stderr, "minitcpsrv on port %d terminates itself, %zu bytes written\n", targetPort, totalWritten);
    return 0;
}