/*
 * Godmar's version of the netcat(nc) utility.
 *
 * Yet another version of netcat. See http://nc110.sourceforge.net/
 *
 * netcat connects to a server, or accepts a client, and subsequently
 * shovels data between the network and its own stdin/stdout.
 * This version differs in how it handles EOF.
 *
 * When it encounters EOF on stdin, it'll shutdown the send direction
 * of the network socket, then continue copying from the network socket
 * to stdout until EOF is seen on the network socket.
 *
 * When it encounters EOF on the network socket, it'll close stdout,
 * then continue copying from stdin to the network socket until
 * EOF is seen on stdin.
 *
 * This allows it to be used in connection with the dpipe program, which
 * connects its stdin/stdout to another program.  Closing stdout is necessary
 * to drive programs that do not output anything until they encounter EOF
 * on their stdin (e.g., fmt, indent, etc. etc.)
 *
 * This program uses an old-school select loop.
 *
 * This error checking in this program is deficient for production code;
 * it would lead to deduction were this program to be submitted for
 * grading.
 *
 * Godmar Back, CS 3214 Fall 2011
 *
 * Updated Spring 2022 to provide better diagnostics and to support IPv6 only.
 * added support for Unix domain sockets.
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <netdb.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stddef.h>
#include <arpa/inet.h>
#include <assert.h>
#include <errno.h>
#include <unistd.h>

static struct addrinfo * 
get_addrinfo_from_port(char *port_number_string)
{
    struct addrinfo *info;
    struct addrinfo hint;
    memset(&hint, 0, sizeof hint);

    hint.ai_flags = AI_PASSIVE | AI_NUMERICSERV | AI_ADDRCONFIG;
    hint.ai_protocol = IPPROTO_TCP; // only interested in TCP
    hint.ai_family = AF_INET6;   // IPv6 only
    int rc = getaddrinfo(NULL, port_number_string, &hint, &info);
    if (rc != 0) {
        gai_strerror(rc);
        exit(EXIT_FAILURE);
    }
    return info;
}

static int 
accept_client(struct addrinfo *info) 
{
    int sockets[10], nsockets = 0;
    struct addrinfo *pinfo;
    int clientfd;

    for (pinfo = info; pinfo; pinfo = pinfo->ai_next) {
        int s = socket(pinfo->ai_family, pinfo->ai_socktype, pinfo->ai_protocol);
        if (s == -1) {
            perror("socket");
            exit(EXIT_FAILURE);
        }

        int opt = 1;
        int rc = setsockopt (s, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof (opt));
        if (rc != 0) {
            perror("setsockopt");
            exit(EXIT_FAILURE);
        }

        rc = bind(s, pinfo->ai_addr, pinfo->ai_addrlen);
        if (rc == -1) {
            int err = errno; 
            perror("bind");
            if (err == EADDRINUSE) {
                fprintf(stderr, "I cannot bind to this port, likely because some other process is still using it\n");
                fprintf(stderr, "Run 'ps x | grep dpipe' and check that you do not have any other instances running on this machine.\n");
            }
            exit(EXIT_FAILURE);
        }

        rc = listen(s, 1024);
        if (rc == -1) {
            perror("listen");
            exit(EXIT_FAILURE);
        }

        assert(nsockets < sizeof(sockets)/sizeof(sockets[0]));
        sockets[nsockets++] = s;
    }
    freeaddrinfo(info);

    assert(nsockets == 1);

    struct sockaddr_storage peer;
    socklen_t peersize = sizeof(peer);

    clientfd = accept (sockets[0], (struct sockaddr *) &peer, &peersize);
    if (clientfd == -1) {
        perror ("accept");
        exit(EXIT_FAILURE);
    }

    char peer_addr[1024], peer_port[10];
    int rc = getnameinfo((struct sockaddr *) &peer, peersize,
                         peer_addr, sizeof peer_addr, peer_port, sizeof peer_port,
                         NI_NUMERICSERV);
    if (rc != 0) {
        fprintf(stderr, "getnameinfo error: %s\n", gai_strerror(rc));
        exit(EXIT_FAILURE);
    }
    fprintf(stderr, "Accepted client from %s:%s\n", peer_addr, peer_port);

    close(sockets[0]);
    return clientfd;
}

static int
connect_to(char *host, char *port)
{
    struct addrinfo *info, *pinfo;
    struct addrinfo hint;
    memset(&hint, 0, sizeof hint);
    hint.ai_flags = AI_CANONNAME | AI_NUMERICSERV | AI_ADDRCONFIG;
    hint.ai_protocol = IPPROTO_TCP;
    hint.ai_family = AF_INET6;       // IPv6 only
    int rc = getaddrinfo(host, port, &hint, &info);
    if (rc != 0) {
        fprintf(stderr, "%s: %s %d\n", host, gai_strerror(rc), rc);
        fprintf(stderr, "It appears that %s is not the name of an IPv6 host\n"
                        "Run 'host -t aaaa %s' to verify this.\n", host, host);
        exit(EXIT_FAILURE);
    }

    for (pinfo = info; pinfo; pinfo = pinfo->ai_next) {
        int s = socket(pinfo->ai_family, pinfo->ai_socktype, pinfo->ai_protocol);
        if (s < 0) {
            perror("socket");
            exit(EXIT_FAILURE);
        }
        
        if (connect(s, pinfo->ai_addr, pinfo->ai_addrlen) == 0)
            return s;

        // silently ignore connect errors, just keep trying the next protocol.
        // (right now only one IPv4 should be returned since we set ai_family
        // above).
    }
    int err = errno;
    perror("connect");
    if (err == ECONNREFUSED) {
        fprintf(stderr, 
            "This error indicates that the server you were trying to reach\n"
            "is not currently accepting connections, very likely because it\n"
            "is not yet or no longer running, or perhaps because the host \n"
            "and/or port numbers don't match.\n");
    }
    exit(EXIT_FAILURE);
}

static void usage(char *p)
{
    fprintf(stderr, "Usage: %s [-l srvport] [hostname] [dstport]\n", p);
    fprintf(stderr, "\n");
    fprintf(stderr, "    gnetcat can be used as a server or a client.\n");
    fprintf(stderr, "    To use in client mode, specify hostname and dstport\n");
    fprintf(stderr, "    To use in server mode, specify -l srvport\n");
    fprintf(stderr, "\n");
    exit(EXIT_FAILURE); 
}

static int
copy(int from, int to)
{
    char buf[8192];
    ssize_t bytesread = read (from, buf, sizeof (buf));
    if (bytesread > 0) {
        ssize_t byteswritten = write (to, buf, bytesread);
        if (bytesread != byteswritten) {
            perror("write");
            exit(EXIT_FAILURE);
        }
    }

    return bytesread;
}

int
main(int ac, char *av[])
{
    struct addrinfo *serverinfo = NULL;
    char *unixserver = NULL, *unixclient = NULL;

    int c;
    while ((c = getopt (ac, av, "l:S:r:")) != -1) {
        switch (c) {
        case 'l':
            serverinfo = get_addrinfo_from_port(optarg);
            break;
        case 'S':
            unixserver = optarg;
            break;
        case 'r':
            unixclient = optarg;
            break;
        default:
            usage(av[0]);
        }
    }

    int netfd;
    int unixfd;
    struct sockaddr_un addr;
    if (unixserver || unixclient) {
        unixfd = socket(AF_UNIX, SOCK_STREAM, 0);
        if (unixfd == -1) {
            perror("socket");
            exit(EXIT_FAILURE);
        }

        memset(&addr, 0, sizeof(addr));
        addr.sun_family = AF_UNIX;
        snprintf(addr.sun_path, sizeof addr.sun_path, "%s", unixserver ? unixserver : unixclient);
    }

    if (unixserver) {
        fprintf(stderr, "Running in server mode (Unix domain)\n");
        int rc = bind(unixfd, (struct sockaddr*)&addr, sizeof(addr));
        if (rc == -1) {
            perror("bind");
            exit(EXIT_FAILURE);
        }
        rc = listen(unixfd, 1024);
        if (rc == -1) {
            perror("listen");
            exit(EXIT_FAILURE);
        }
        struct sockaddr_storage peer;
        socklen_t peersize = sizeof(peer);

        netfd = accept (unixfd, (struct sockaddr *) &peer, &peersize);
        if (netfd == -1) {
            perror ("accept");
            exit(EXIT_FAILURE);
        }

        unlink(addr.sun_path);
    } else if (unixclient) {
        fprintf(stderr, "Running in client mode (Unix domain)\n");
        netfd = unixfd;
        int rc = connect(netfd, (struct sockaddr*)&addr, sizeof(addr));
        if (rc == -1) {
            perror("bind");
            exit(EXIT_FAILURE);
        }
    } else if (serverinfo) {
        fprintf(stderr, "Running in server mode\n");
        netfd = accept_client(serverinfo);
    } else {
        if (optind + 1 >= ac)
            usage(av[0]);
        fprintf(stderr, "Running in client mode\n");
        netfd = connect_to(av[optind], av[optind+1]);
    }

    while (1) {
        fd_set fds;
        FD_ZERO(&fds);
        FD_SET(0, &fds);
        FD_SET(netfd, &fds);
        int r = select(netfd+1, &fds, NULL, NULL, NULL);
        if (r < 0) {
            perror("select");
            exit(EXIT_FAILURE);
        }
        assert (r > 0);

        if (FD_ISSET(0, &fds)) {
            int r = copy(0, netfd);
            if (r < 0) {
                perror("read");
                exit(EXIT_FAILURE);
            }

            if (r == 0) {
                shutdown(netfd, SHUT_WR);
                while (copy(netfd, 1) > 0)
                    ; 
                break;
            }
        }

        if (FD_ISSET(netfd, &fds)) {
            int r = copy(netfd, 1);
            if (r < 0) {
                perror("read");
                exit(EXIT_FAILURE);
            }

            if (r == 0) {
                close(1);
                while (copy(0, netfd) > 0)
                    ; 
                break;
            }
        }
    }

    return EXIT_SUCCESS;
}