/*
 * Copyright (c) 2020 Roger Seguin <rgrsgn@gmail.com>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

static char rcsid[] = "$Id: cnxbenchmark.C,v 1.0 2020/06/29 04:13:04 RogerSeguin Exp $";

#include <stdlib.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <memory.h>
#include <time.h>
#include <errno.h>
#include <netdb.h>
#include <poll.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <arpa/inet.h>

#include <iostream>
#include <sstream>
#include <string>
#include <iomanip>

/*
cnxbenchmark - TCP client/server measuring connection speed
Usage
- as a server: cnxbenchmark [-p port]
- as a client: cnxbenchmark [-4|-6] [-p port] [-G msgsizeGiB] server_host
Examples
- Server listening to any connection request on port 4996
        cnxbenchmark -p 4996
- IPv4 Client sending 8 GiB of data to server on same host through port 4996
        cnxbenchmark -4 -p 4996 -G 8 localhost
*/

using namespace std;

/**************************************************************************/
static int usage(const char *prg)
/**************************************************************************/
{
    cout << rcsid + 5 << endl;
    cout << prg << " - TCP client/server measuring connection speed" << endl;
    cout << "Usage" << endl;
    cout << "- as a server: " << prg  << " [-p port]" << endl;
    cout << "- as a client: " << prg
         <<  " [-4|-6] [-p port] [-G msgsizeGiB] server_host" << endl;
    cout << "Examples" << endl;
    cout << "- "  << "Server listening to any connection request on port 4996" << endl;
    cout << "\t" << prg << " -p 4996" << endl;
    cout << "- "  << "IPv4 Client sending 8 GiB of data to server on same host through port 4996" << endl;
    cout << "\t" << prg << " -4 -p 4996 -G 8 localhost" << endl;
    return (0);
}                               // usage()


/**************************************************************************/
namespace {
    /* Default values */
    const int CNXPORT = 4996;
    int m_dataGiB = 8;	// Data trasnfer of 8 GiB

    /* Segmented message sent through the socket */
    char m_segmentBuffer[1024 * 1024];
}
/**************************************************************************/

#define Min(x,y)	((x)<(y)?(x):(y))

/**************************************************************************/
/**** Server ****/
static int startServer(const char* prgname, int cnxPort);
static int processClientRequests(struct pollfd fds[], int nfds, int timeout_ms);
static int readMsg(int sockid);

/**** Client ****/
static int cnx2server(const char* prgname,
                      const char * server, int port, bool isIPv6);
// Return socket id of connection with the server

static int getHostAddr(const char *hostname, bool isIPv6, int port,
                       struct sockaddr **addr);
static int cnx(int domain, struct sockaddr *addr, size_t addrsz);
static int sendMsg(int sockid);
/**************************************************************************/

/**************************************************************************/
int main(int argc, char**argv)
/**************************************************************************/
{
    const char* prgname = basename(argv[0]);  // assuming _GNU_SOURCE is defined
    bool isIPv6 = false;
    int cnxPort = CNXPORT;
    bool fUsage = false;
    int status;

    for (char c; !fUsage && ((c = getopt(argc, argv, "46hG:p:")) != EOF);)
        switch(c) {
            case '4':
                isIPv6 = false;
                break;
            case '6':
                isIPv6 = true;
                break;
            case 'p':
                cnxPort = cnxPort = atoi(optarg);
                break;
            case 'G':
                /* How many gibabytes to transfer through the TCP socket */
                m_dataGiB  = atoi(optarg);
                break;
            case 'h':          // Help
            case '?':
            default:
                fUsage = true;
        }

    /* If server, no parameters expected,
     otherwise (i.e. client) only one parameter expected */
    fUsage = fUsage || ((argc - optind) > 1);
    if (fUsage) {
        usage(prgname);
        return (1);
    };

    if (optind == argc)
        /* No parameter ==> Server */
        /* Will loop forever waiting for and processing clients requests */
        status = startServer(prgname, cnxPort);

    else {
        /* Client: send data to server and terminate */
        int sockid = cnx2server(prgname, argv[optind], cnxPort, isIPv6);
        if (sockid > 0) {
            sendMsg(sockid);
            close (sockid);
        }
        else
            status = -1;
    }

    return (status);
}	// main()

/**************************************************************************
	SERVER
**************************************************************************/

/**************************************************************************/
int startServer(const char* prgname, int cnxPort)
/**************************************************************************/
{
    cout << "Running (server) " << prgname << " -p " << cnxPort << endl;

    const int boolTrue = 1;
    int sockid4;
    int sockid6;
    struct sockaddr_in addr4;
    struct sockaddr_in6 addr6;
    int status;
    bool error4 = false,
         error6 = false;

    /* Prepare accepting IPv4 connections */
    sockid4 = socket(AF_INET, SOCK_STREAM, 0);
    if (sockid4 <= 0) {
        perror("socket(AF_INET)");
        error4 = true;
    }
    if (!error4) {
        status = setsockopt(sockid4, SOL_SOCKET,
                            SO_REUSEADDR, &boolTrue, sizeof(boolTrue));
        if (status == -1) {
            perror("setsockopt(AF_INET)");
            error4 = true;
        }
    }
    if (!error4) {
        addr4.sin_family = AF_INET;
        addr4.sin_port = htons(cnxPort);
        addr4.sin_addr.s_addr = htonl(INADDR_ANY);
        status = bind(sockid4, (struct sockaddr *) &addr4, sizeof(addr4));
        if (status == -1) {
            perror("bind(AF_INET)");
            error4 = true;
        }
    }
    if (!error4) {
        status = listen(sockid4, SOMAXCONN);
        if (status == -1) {
            perror("listen(AF_INET)");
            error4 = true;
        }
    }
    if (error4) {
        if (sockid4 != -1)
            close(sockid4);
        sockid4 = -1;
        cerr << "Won't accept IPv4 connections" << endl;
    }

    /* Prepare accepting IPv6 connections */
    sockid6 = socket(AF_INET6, SOCK_STREAM, 0);
    if (sockid6 <= 0) {
        perror("socket(AF_INET6)");
        error6 = true;
    }
    if (!error6) {
        status = setsockopt(sockid6, SOL_SOCKET, SO_REUSEADDR,
                            &boolTrue, sizeof(boolTrue));
        if (status == -1) {
            perror("setsockopt(AF_INET6)");
            error6 = true;
        }
    }
    if (!error6) {
        /* Allow both IPv4 and IPv6 sockets to bind to same port */
        status = setsockopt(sockid6, IPPROTO_IPV6, IPV6_V6ONLY,
                            &boolTrue, sizeof(boolTrue));
        if (status == -1) {
            perror("setsockopt(IPV6_V6ONLY)");
            error6 = true;
        }
    }
    if (!error6) {
        addr6.sin6_family = AF_INET6;
        addr6.sin6_port = htons(cnxPort);
        addr6.sin6_addr = in6addr_any;
        status = bind(sockid6, (struct sockaddr *) &addr6, sizeof(addr6));
        if (status == -1) {
            perror("bind(AF_INET6)");
            error6 = true;
        }
    }
    if (!error6) {
        status = listen(sockid6, SOMAXCONN);
        if (status == -1) {
            perror("listen(AF_INET6)");
            error6 = true;
        }
    }
    if (error6) {
        if (sockid6 != -1)
            close(sockid6);
        sockid6 = -1;
        cerr << "Won't accept IPv6 connections" << endl;
    }

    if (error4 && error6)
        /* No working socket for listening to clients requests */
        return (-1);

    /* Client can request connection using either IPv4 or IPv6 */
    struct pollfd fds[2];
    const int timeout_ms = 5 * 60 * 1000;	// 5 mn
    int    nfds = 0;

    memset(fds, 0, sizeof(fds));
    if (sockid4 != -1) {
        fds[nfds].fd = sockid4;
        fds[nfds].events = POLLIN;
        nfds++;
    }
    if (sockid6 != -1) {
        fds[nfds].fd = sockid6;
        fds[nfds].events = POLLIN;
        nfds++;
    }

    /* Loop waiting for clients to connect and process their requests */
    while ((status = processClientRequests(fds, nfds, timeout_ms)) != -1);

    if (sockid4 != -1)
        close (sockid4);
    if (sockid6 != -1)
        close (sockid6);

    return (status);
}	// startServer()

/**************************************************************************/
int processClientRequests(struct pollfd fds[], int nfds, int timeout_ms)
/**************************************************************************/
{
    int status;

    status = poll(fds, nfds, timeout_ms);
    switch(status) {
        case -1:
            perror("poll()");
            return (-1);
        case 0:
            cout << "poll() timed out" << endl;
            return (1);
        default:
            // Socket to be read
            break;
    }
    for (int i = 0; i < nfds; i++)
        if (fds[i].revents & POLLIN) {
            cout << (i == 0 ? "IPv4" : "IPv6") << " client connected" << endl;
            int fd = accept(fds[i].fd, 0, 0);
            if (fd == -1) {
                perror("accept()");
                status = -1;
            }
            else {
                readMsg(fd);
                close (fd);
            }
        }
    return(status == -1 ? -1 : 0);
}	// processClientRequests()

/**************************************************************************/
int readMsg(int sockid)
/**************************************************************************/
{
    const uint64_t NANOS = 1000 * 1000 * 1000; // How many nanoseconds in a sec
    uint64_t fullMsgSize;
    uint64_t alreadyread = 0;
    ssize_t nbytes;
    struct timespec t1, t2, dt;
    uint64_t dt_ms;

    /* The client first sends the size of the data it is going to transfer,
    then expects that size sent back before starting the actual transfer */
    nbytes = recv(sockid, &fullMsgSize, sizeof(fullMsgSize), 0);
    if (nbytes != sizeof(fullMsgSize)) {
        if (nbytes == -1)
            perror("recv()");
        else
            cerr << "Couldn't get message size from client" << endl;
        return (-1);
    }
    nbytes = send(sockid, &fullMsgSize, sizeof(fullMsgSize), 0);
    if (nbytes != sizeof(fullMsgSize)) {
        if (nbytes == -1)
            perror("send()");
        else
            cerr << "Couldn't send acknowledgment to client" << endl;
        return (-1);
    }

    clock_gettime (CLOCK_REALTIME, &t1);
    for (ssize_t bytes2read = sizeof(m_segmentBuffer); bytes2read > 0;) {
        bytes2read = Min (fullMsgSize - alreadyread, sizeof(m_segmentBuffer));
        if (bytes2read <= 0)
            break;
        nbytes = recv(sockid, m_segmentBuffer, bytes2read, 0);
        switch (nbytes) {
            case -1:
                perror("recv()");
                return (-1);
            case 0:
                /* The client process may have been interrupted */
                bytes2read = 0;
                break;
            default:	// Keep reading...
                alreadyread += nbytes;
        }
    }
    clock_gettime (CLOCK_REALTIME, &t2);

    if (t2.tv_nsec < t1.tv_nsec)
        dt.tv_sec = t2.tv_sec - t1.tv_sec - 1,
        dt.tv_nsec = t2.tv_nsec + (NANOS - t1.tv_nsec);
    else
        dt.tv_sec = t2.tv_sec - t1.tv_sec,
        dt.tv_nsec = t2.tv_nsec - t1.tv_nsec;
    dt_ms = dt.tv_sec * 1000 + dt.tv_nsec / (1000 * 1000);
    double speedGbit = (double) alreadyread * 8 / (1000 * 1000 * dt_ms);
    cout << "  " << alreadyread / (1024 * 1024) << " MiB transferred in "
         << dt_ms << " ms (" << setprecision(3) << speedGbit << " Gbit/s)" << endl;
    return (0);
}	// readMsg()

/**************************************************************************
	CLIENT
**************************************************************************/

/**************************************************************************/
int cnx2server(const char* prgname, const char * server, int port, bool isIPv6)
/**************************************************************************/
{
    /* Connectwith the server, return the connection socket */
    const char ipvers = (isIPv6 ? '6' : '4');
    cout << "Running (client) " << prgname << " -" << ipvers << " -p " << port
         << " -G" << m_dataGiB << " " << server << endl;
    struct sockaddr *addr;
    int sockid = -1;
    int status;

    if (isIPv6) {
        status = getHostAddr(server, true, port, &addr);
        if (status == -1)
            return (-1);
        struct sockaddr_in6 *addr_in = (struct sockaddr_in6 *) addr;
        char buffer [INET6_ADDRSTRLEN];
        const char *saddr = inet_ntop(AF_INET6, & (addr_in->sin6_addr),
                                      buffer, sizeof(buffer));
        if (!saddr) {
            perror("inet_ntop");
            return (-1);
        }
        cout << "Server @ " << saddr << endl;
        sockid = cnx(AF_INET6, addr, sizeof(struct sockaddr_in6));
        if (sockid == -1)
            return (-1);
    }

    else {
        /* IPv4 */
        status = getHostAddr(server, false, port, &addr);
        if (status == -1)
            return (-1);
        struct sockaddr_in *addr_in = (struct sockaddr_in *) addr;
        const char *saddr =  inet_ntoa(addr_in->sin_addr);
        if (!saddr) {
            perror("inet_ntoa");
            return (-1);
        }
        cout << "Server @ " << saddr << endl;
        sockid = cnx(AF_INET, addr, sizeof(struct sockaddr_in));
        if (sockid == -1)
            return (-1);
    }

    return (sockid);
}	// cnx2server()

/**************************************************************************/
int getHostAddr(const char *hostname, bool isIPv6, int port, struct sockaddr **addr)
/**************************************************************************/
{
    /* Translate host name to host address, either IPv4 or IPv6 */
    struct addrinfo *psAddrInfo = 0,    //
                         *psInfo,                //
                         hints;
    ostringstream ssport;
    bool found = false;
    int status;
    errno = 0;

    memset(&hints, 0, sizeof(hints));
    hints.ai_flags = AI_CANONNAME;
    hints.ai_family = (isIPv6 ? AF_INET6 : AF_INET);
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = 0;
    ssport << port;

    status =
        getaddrinfo(hostname, ssport.str().c_str(), &hints,
                    &psAddrInfo);
    if (status || !psAddrInfo) {
        if (errno)
            perror(hostname);
        return (-1);
    }
    for (psInfo = psAddrInfo; !found && psInfo;) {
        switch(psInfo->ai_family) {
            case AF_INET:
                *addr = psInfo->ai_addr;
                found = !isIPv6;
                break;
            case AF_INET6:
                *addr = psInfo->ai_addr;
                found = isIPv6;
                break;
        }
        if (!found)
            psInfo = psInfo->ai_next;
    }
    if (!found) {
        cerr << hostname << ": Addr Info not found" << endl;
        freeaddrinfo(psAddrInfo), psAddrInfo = 0;
        return (-1);
    }

    return (0);
}	// getHostAddr()

/**************************************************************************/
int cnx(int domain, struct sockaddr *addr, size_t addrsz)
/**************************************************************************/
{
    int sockid = socket(domain, SOCK_STREAM, 0);
    if (sockid <= 0) {
        perror("socket");
        return (-1);
    }
    int status = connect(sockid, addr, addrsz);
    if (status == -1) {
        perror("connect");
        close(sockid);
        return (-1);
    }
    return (sockid);
}	// cnx()

/**************************************************************************/
int sendMsg(int sockid)
/**************************************************************************/
{
    /* Send data message to the server */
    const uint64_t GiB = (uint64_t) 1024 * 1024 * 1024;
    uint64_t fullMsgSize = m_dataGiB * GiB;
    ssize_t nbytes;

    /* First tell the server the size of the message that will be transmitted
     and wait for the server to acknowledge before starting the actual transfer */
    nbytes = send(sockid, &fullMsgSize, sizeof(fullMsgSize), 0);
    if (nbytes != sizeof(fullMsgSize)) {
        if (nbytes == -1)
            perror("send()");
        else
            cerr << "Couldn't send size of message" << endl;
        return (-1);
    }
    nbytes = recv(sockid, &fullMsgSize, sizeof(fullMsgSize), 0);
    if (nbytes != sizeof(fullMsgSize)) {
        if (nbytes == -1)
            perror("recv()");
        else
            cerr << "Couldn't get acknowledgment from server" << endl;
        return (-1);
    }

    const uint64_t N = fullMsgSize / sizeof(m_segmentBuffer);
    for (uint64_t i = 0; i < N; i++) {
        nbytes = send(sockid, m_segmentBuffer, sizeof(m_segmentBuffer), 0);
        if (nbytes == -1) {
            perror("send()");
            return (-1);
        }
        else if (nbytes != sizeof(m_segmentBuffer))
            cerr << "Missed sending " << nbytes << " bytes" << endl;
    }

    return (0);
}	// sendMsg()

/*
$Log: cnxbenchmark.C,v $
Revision 1.0  2020/06/29 04:13:04  RogerSeguin
Initial revision

*/