Test client for grnvsa3

gistfile1.c

/*
Streaming speed test.
WTFPL!
*/

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <signal.h>
#include <time.h>
#include <errno.h>
#include <sys/socket.h>
#include <sys/epoll.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <netdb.h>
#include <stdbool.h>

#include "global.h"

#define MAXEVENTS 5

char epollcodes [2048];
const char * strepollcode(int c) {
        epollcodes[0] = 0;
        char * ccp = epollcodes;
#define ADJSTPTR ccp += strlen(ccp)
        if(c & EPOLL_CLOEXEC) strcpy(ccp, "EPOLL_CLOEXEC ");
        ADJSTPTR;
        if(c & EPOLL_NONBLOCK) strcpy(ccp, "EPOLL_NONBLOCK ");
        ADJSTPTR;
        if(c & EPOLLIN) strcpy(ccp, "EPOLLIN ");
        ADJSTPTR;
        if(c & EPOLLPRI) strcpy(ccp, "EPOLLPRI ");
        ADJSTPTR;
        if(c & EPOLLOUT) strcpy(ccp, "EPOLLOUT ");
        ADJSTPTR;
        if(c & EPOLLRDNORM) strcpy(ccp, "EPOLLRDNORM ");
        ADJSTPTR;
        if(c & EPOLLRDBAND) strcpy(ccp, "EPOLLRDBAND ");
        ADJSTPTR;
        if(c & EPOLLWRNORM) strcpy(ccp, "EPOLLWRNORM ");
        ADJSTPTR;
        if(c & EPOLLWRBAND) strcpy(ccp, "EPOLLWRBAND ");
        ADJSTPTR;
        if(c & EPOLLMSG) strcpy(ccp, "EPOLLMSG ");
        ADJSTPTR;
        if(c & EPOLLERR) strcpy(ccp, "EPOLLERR ");
        ADJSTPTR;
        if(c & EPOLLHUP) strcpy(ccp, "EPOLLHUP ");
        ADJSTPTR;
        if(c & EPOLLRDHUP) strcpy(ccp, "EPOLLRDHUP ");
        ADJSTPTR;
        if(c & EPOLLONESHOT) strcpy(ccp, "EPOLLONESHOT ");
        ADJSTPTR;
        if(c & EPOLLET) strcpy(ccp, "EPOLLET ");
        *(ccp + strlen(ccp) - 1) = 0;
        return epollcodes;
#undef ADJSTPTR
}

/**
 * Print usage on command line
 */
static void printUsage(char *name)
{
        fprintf(stderr, "Usage: %s <address> <port> <activeclients> [<stallclients>]\n\n", name);
        exit(-EINVAL);
}

struct _rtpclient;
struct _rtpclient {
        int fd;
        enum clienttype { active, stall } type;
        uint64_t input;
        struct timespec start;
        bool connected;
        struct _rtpclient * next; // intrusive list
        struct _rtpclient * prev;
};
typedef struct _rtpclient rtpclient;
rtpclient * clients = NULL;
rtpclient * clients_last = NULL;

static void setnonblocking(int fd) {
        int opts;
        opts = fcntl(fd, F_GETFL);
        if (opts < 0) {
                perror("fcntl(F_GETFL)");
                exit(EXIT_FAILURE);
        }
        opts = (opts | O_NONBLOCK | O_ASYNC);
        if (fcntl(fd, F_SETFL, opts) < 0) {
                perror("fcntl(F_SETFL)");
                exit(EXIT_FAILURE);
        }
}

static int epfd; // epoll descriptor
struct sockaddr_in remote;
int remote_port;


static rtpclient * new_epoll_conn(void) {
        int client;
        if ( 0 > (client=socket(AF_INET,SOCK_STREAM,IPPROTO_TCP)) )
                DIE("socket() failed: %d - %s\n",errno,strerror(errno));
        setnonblocking(client);
        if ( 0 > connect(client,(struct sockaddr *)&remote,sizeof(remote)) && errno != EINPROGRESS )
                DIE("connect() failed: %d - %s\n",errno,strerror(errno));
        struct epoll_event ev; bzero(&ev,sizeof(struct epoll_event));
        ev.data.ptr = malloc(sizeof(rtpclient));
        ev.events = EPOLLIN | EPOLLPRI | EPOLLERR | EPOLLRDHUP | EPOLLHUP | EPOLLET | EPOLLOUT;
        if (epoll_ctl(epfd, EPOLL_CTL_ADD, client, &ev) < 0) {
                DIE("epoll set insertion error: fd=%d0", client);
        }
        rtpclient * conn = (rtpclient*)ev.data.ptr;
        bzero(conn, sizeof(rtpclient));
        conn->fd = client;
        conn->next = NULL;
        conn->input = 0;
        conn->connected = false;
        if(clock_gettime(CLOCK_MONOTONIC, &(conn->start)))
                DIE("Could not retrieve monotonic time.");
        if(clients) {
                conn->prev = clients_last;
                clients_last->next = conn;
                clients_last = conn;
        } else {
                conn->prev = NULL;
                clients = clients_last = conn;
        }
        return conn;
}

static void remove_conn(rtpclient * conn) {
        struct timespec now;
        if(clock_gettime(CLOCK_MONOTONIC, &now))
                DIE("Could not retrieve monotonic time.");
        printf("Removing client %d: %.3fkB/s connected: %s\n", conn->fd, (double)conn->input / 1000 / ((double)now.tv_sec - conn->start.tv_sec + ((double)now.tv_nsec - conn->start.tv_nsec)/1000000000), conn->connected?"yes":"no");
        close(conn->fd);
        epoll_ctl(epfd, EPOLL_CTL_DEL, conn->fd, NULL);
        if(conn->prev)
                conn->prev->next = conn->next;
        else
                clients = conn->next;
        if(conn->next)
                conn->next->prev = conn->prev;
        else
                clients_last = conn->prev;
        free(conn);
}
static void epoll_remove_conn(const struct epoll_event * e) { remove_conn((rtpclient*)(e->data.ptr)); } // cute little helper for lazyass me

static void epoll_conn_event(const struct epoll_event * e) {
        rtpclient * conn = (rtpclient*)(e->data.ptr);
        //printf("Event 0x%X - %s from %d\n", e->events, strepollcode(e->events), ((rtpclient*)(e->data.ptr))->fd);
        if(e->events & (EPOLLHUP | EPOLLRDHUP | EPOLLERR)) { // any error or incoming data event yields a disconnect
                epoll_remove_conn(e);
        } else {
                if(e->events & (EPOLLIN | EPOLLRDNORM | EPOLLPRI)) {
                        // do I look like I actually care what was received?
                        char buffer[4096];
                        int len=recv(conn->fd,buffer,4096,0);
                        if(len > 0)
                                conn->input += len;
                        else
                                printf("recv from %d returned %d: e%d - %s.\n", conn->fd, len, errno, strerror(errno));
                }
                if(!conn->connected && e->events & EPOLLOUT)
                        conn->connected = true;
        }
}

/**
 * Signal handler to tidy up on SIGTERM & friends.
 */
static void cleanup(int sig)
{
        signal(sig, SIG_DFL);
        printf("Exiting...\n");
        while(clients)
                remove_conn(clients);
        killpg(0, sig);
}

int main(int argc, char **argv)
{
        // Check arguments
        if (argc != 4 && argc != 5)
                printUsage(argv[0]);

        remote_port  = atoi(argv[2]);

        struct hostent * h;
        if ( NULL == (h=gethostbyname2(argv[1],AF_INET)) ) {
                fprintf(stderr,"unable to resolve hostname %s\n",argv[1]);
                exit(1);
        }
        memcpy(&remote.sin_addr.s_addr,h->h_addr_list[0],h->h_length);
    remote.sin_family      = AF_INET;
    remote.sin_port        = htons(remote_port);

        int activeclients = atoi(argv[3]),
            stallclients  = atoi((argc > 4)?argv[4]:"0");


        // Setup signal handlers
        signal(SIGINT,  cleanup);
        signal(SIGTERM, cleanup);
        signal(SIGHUP,  cleanup);
        signal(SIGCLD, SIG_IGN);
        signal(SIGPIPE, SIG_IGN);
        setpgrp();

        epfd = epoll_create1(0);
        /*struct epoll_event ev; bzero(&ev, sizeof(struct epoll_event));
        ev.events = EPOLLIN | EPOLLET | EPOLLERR | EPOLLHUP;
        ev.data.ptr = NULL;
        if(epoll_ctl(epfd, EPOLL_CTL_ADD, sockfd, &ev) < 0) perror("epoll_ctl");*/

        int i;
        for(i = 0; i < activeclients; ++i)
                new_epoll_conn();

        while (clients) {
                struct epoll_event events[MAXEVENTS];
                //printf("Waiting for events... "); fflush(stdout);
                int nfds = epoll_wait(epfd, events, MAXEVENTS, -1);
                //printf("got %d.\n", nfds);
                if (nfds < 0) {
                        if(errno == EINTR) continue;
                        DIE("read epoll: %d - %s\n", errno, strerror(errno));
                }
                int j = 0;
                for(; j < nfds; ++j)
                        epoll_conn_event(events + j);
        }

        printf("All clients dead.\n");
        exit(0);
}