/* $Id: ncat_broker.c 12720 2009-03-27 19:48:00Z david $ */

#include "ncat.h"

#include <assert.h>
#ifndef WIN32
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#endif

#ifdef HAVE_OPENSSL
#include <openssl/ssl.h>
#include <openssl/err.h>
#endif

static fd_set master, read_fds;
static fd_list_t fdlist;
static int listen_socket;
static int conn_count = 0;
/* Has stdin seen EOF? */
static int stdin_eof = 0;

static void handle_connection(void);
static void read_and_broadcast(int recv_socket);
static int chat_announce_connect(int fd, const struct sockaddr_storage *ss);
static int chat_announce_disconnect(int fd);
static char *chat_filter(char *buf, size_t size, int fd, int *nwritten);

int ncat_broker(void)
{
#ifdef HAVE_OPENSSL
    SSL_CTX *ctx;
#endif

    /* clear out structs */
    FD_ZERO(&master);
    FD_ZERO(&read_fds);
    zmem(&fdlist, sizeof(fdlist));

#ifdef HAVE_OPENSSL
    if (o.ssl)
        ctx = setup_ssl_listen();
#endif

    /* setup the main listening socket */
    listen_socket = do_listen(SOCK_STREAM);

    /* Make our listening socket non-blocking because there are timing issues
     * which could cause us to block on accept() even though select() says it's
     * readable.  See UNPv1 2nd ed, p422 for more.
     */
    unblock_socket(listen_socket);

    /* setup select sets and max fd */
    FD_SET(listen_socket, &master);

    /* we need a list of fds to keep current fdmax */
    init_fdlist(&fdlist, sadd(o.conn_limit, 2));
    add_fd(&fdlist, listen_socket, NULL);
    add_fd(&fdlist, STDIN_FILENO, NULL);

    while (1) {
        int i;

        if (o.debug > 1)
            logdebug("Broker connection count is %d\n", conn_count);

        read_fds = master;
        fselect(fdlist.fdmax + 1, &read_fds, NULL, NULL, NULL);

        /* Loop through descriptors */
        for (i = 0; i <= fdlist.fdmax; i++) {
            /* Loop through descriptors until there's something to read */
            if (!FD_ISSET(i, &read_fds))
                continue;

            if (i == listen_socket) {
                /* we have a new connection request */
                handle_connection();
            } else {
                /* Handle incoming client data and distribute it. */
                read_and_broadcast(i);
            }
        }
    }

    return 0;
}

/* Accept a connection on a listening socket. Allow or deny the connection.
   If allowed, add the new socket to the watch set. */
static void handle_connection(void)
{
    struct sockaddr_storage remoteaddr;
    socklen_t ss_len;
    int fd;
#ifdef HAVE_OPENSSL
    SSL *tmpssl = NULL;
#else
    void *tmpssl = NULL;
#endif

    ss_len = sizeof(remoteaddr);
    errno = 0;
    fd = accept(listen_socket, (struct sockaddr *) &remoteaddr, &ss_len);

    if (fd < 0) {
        if (o.debug)
            logdebug("Error in accept: %s\n", strerror(errno));

        close(fd);
        return;
    }

    if (o.verbose) {
        if (o.chat)
            loguser("Connection from %s on file descriptor %d.\n", inet_socktop(&remoteaddr), fd);
        else
            loguser("Connection from %s.\n", inet_socktop(&remoteaddr));
    }

    /* Check conditions that might cause us to deny the connection. */
    if (conn_count >= o.conn_limit) {
        if (o.verbose)
            loguser("New connection denied: connection limit reached (%d)\n", conn_count);
        Close(fd);
        return;
    }
    if (!allow_access(&remoteaddr)) {
        if (o.verbose)
            loguser("New connection denied: not allowed\n");
        Close(fd);
        return;
    }

#ifdef HAVE_OPENSSL
    if (o.ssl) {
        tmpssl = new_ssl(fd);
        if (SSL_accept(tmpssl) != 1) {
            if (o.verbose) {
                loguser("Failed SSL connection from %s: %s\n",
                        inet_socktop(&remoteaddr), ERR_error_string(ERR_get_error(), NULL));
            }
            SSL_free(tmpssl);
            Close(fd);
            return;
        }
    }
#endif

    conn_count++;

    /* do this after we test for the limit */
    FD_SET(fd, &master);
    /* Now that a client is connected, pay attention to stdin. */
    if (!stdin_eof)
        FD_SET(STDIN_FILENO, &master);

    /* add it to our list of fds for maintaining maxfd */
    if(add_fd(&fdlist, fd, tmpssl) < 0)
         bye("add_fd() failed.");

    if (o.chat)
        chat_announce_connect(fd, &remoteaddr);
}

/* Read from recv_fd and broadcast whatever is read to all other descriptors in
   master, with the exception of stdin, listen_socket, and recv_fd itself.
   Handles CRLF translation and chat mode. On read error or end of stream,
   closes the socket and removes it from the master list. */
static void read_and_broadcast(int recv_fd)
{
    char buf[DEFAULT_TCP_BUF_LEN];
    char *chatbuf, *outbuf;
    struct fdinfo *fdn = get_fdinfo(&fdlist, recv_fd);
    ssize_t nbytes;
    fd_set fds;

    assert(fdn);

#ifdef HAVE_OPENSSL
readagain:
    if (o.ssl && fdn->ssl)
        nbytes = SSL_read(fdn->ssl, buf, sizeof(buf));
    else
#endif
    if (recv_fd == STDIN_FILENO) {
        /* Behavior differs depending on whether this is stdin or a socket. */
        nbytes = read(recv_fd, buf, sizeof(buf) - o.crlf);
        if (nbytes <= 0) {
            if (nbytes < 0 && o.verbose)
                logdebug("Error reading from stdin: %s\n", strerror(errno));
            if (nbytes == 0 && o.debug)
                logdebug("EOF on stdin\n");

            /* Don't close the file because that allows a socket to be fd 0. */
            FD_CLR(recv_fd, &master);
            /* Buf mark that we've seen EOF so it doesn't get re-added to the
               select list. */
            stdin_eof = 1;

            return;
        }

        /* We gave ourselves extra room in our special-case stdin read above */
        if (o.crlf && buf[nbytes - 1] == '\n' && buf[nbytes - 2] != '\r') {
            memcpy(&buf[nbytes - 1], "\r\n", 2);
            nbytes++;
        }
    } else {
        /* From a connected socket, not stdin. */
        nbytes = recv(recv_fd, buf, sizeof(buf), 0);

        if (nbytes <= 0) {
            if (o.debug)
                logdebug("Closing connection.\n");

#ifdef HAVE_OPENSSL
            if (o.ssl && fdn->ssl) {
                if (nbytes == 0)
                    SSL_shutdown(fdn->ssl);
                SSL_free(fdn->ssl);
            }
#endif

            Close(recv_fd);
            FD_CLR(recv_fd, &master);
            rm_fd(&fdlist, recv_fd);

            conn_count--;
            if (conn_count == 0)
                FD_CLR(STDIN_FILENO, &master);

            if (o.chat)
                chat_announce_disconnect(recv_fd);

            return;
        }

        /* If we received from something other than stdin, and --send-only was
           given, do no further processing. */
        if (o.sendonly)
            return;
    }

    if (o.debug > 1)
        logdebug("Handling data from client %d.\n", recv_fd);

    chatbuf = NULL;
    outbuf = buf;

    if (o.chat) {
        chatbuf = chat_filter(buf, nbytes, recv_fd, &nbytes);
        if (chatbuf == NULL) {
            if (o.verbose)
                logdebug("Error formatting chat message from fd %d\n", recv_fd);
        } else {
            outbuf = chatbuf;
        }
    }

    /* Write to everything in the master set, except the listener, sender, and
       stdin. */
    fds = master;
    FD_CLR(STDIN_FILENO, &fds);
    FD_CLR(listen_socket, &fds);
    FD_CLR(recv_fd, &fds);
    broadcast(&fds, &fdlist, outbuf, nbytes);

    free(chatbuf);

    zmem(buf, sizeof(buf));

#ifdef HAVE_OPENSSL
    /* SSL can buffer our input, so doing another select()
     * won't necessarily work for us.  We jump back up to
     * read any more data we can grab now
     */
    if (o.ssl && fdn->ssl && SSL_pending(fdn->ssl))
        goto readagain;
#endif
}

/* Do a sprintf at the given offset into a malloc-allocated buffer. Reallocates
   the buffer and updates the variables to make room if necessary. */
static int strbuf_sprintf(char **buf, size_t *size, size_t *offset,
    const char *fmt, ...) {
    va_list va;
    int n;

    assert(*offset <= *size);

    if (*buf == NULL) {
        *size = 1;
        *buf = (char*) safe_malloc(*size);
    }

    for (;;) {
        va_start(va, fmt);
        n = Vsnprintf(*buf + *offset, *size - *offset, fmt, va);
        va_end(va);
        if (n < 0)
            *size = MAX(*size, 1) * 2;
        else if (n >= *size - *offset)
            *size += n + 1;
        else
            break;
        *buf = (char*) safe_realloc(*buf, *size);
    }
    *offset += n;

    return n;
}

/* Announce the new connection and who is already connected. */
static int chat_announce_connect(int fd, const struct sockaddr_storage *ss)
{
    char *buf = NULL;
    size_t size = 0, offset = 0;
    fd_set fds;
    int i, count, ret;

    strbuf_sprintf(&buf, &size, &offset,
        "<announce> %s is connected as <user%d>.\n", inet_socktop(ss), fd);

    fds = master;
    FD_CLR(STDIN_FILENO, &fds);
    FD_CLR(listen_socket, &fds);

    strbuf_sprintf(&buf, &size, &offset, "<announce> already connected: ");
    count = 0;
    for (i = 0; i < fdlist.fdmax; i++) {
        struct sockaddr_storage ss;
        socklen_t len = sizeof(ss);

        if (i == fd || !FD_ISSET(i, &fds))
            continue;

        if (getpeername(i, (struct sockaddr *) &ss, &len) == -1)
            bye("getpeername for sd %d failed: %s.", strerror(errno));

        if (count > 0)
            strbuf_sprintf(&buf, &size, &offset, ", ");

        strbuf_sprintf(&buf, &size, &offset, "%s as <user%d>", inet_socktop(&ss), i);

        count++;
    }
    if (count == 0)
        strbuf_sprintf(&buf, &size, &offset, "nobody");
    strbuf_sprintf(&buf, &size, &offset, ".\n");

    ret = broadcast(&fds, &fdlist, buf, offset);

    free(buf);

    return ret;
}

static int chat_announce_disconnect(int fd)
{
    fd_set fds;
    char buf[128];
    int n;

    n = Snprintf(buf, sizeof(buf),
        "<announce> <user%d> is disconnected.\n", fd);
    if (n >= sizeof(buf) || n < 0)
        return -1;

    fds = master;
    FD_CLR(STDIN_FILENO, &fds);
    FD_CLR(listen_socket, &fds);

    return broadcast(&fds, &fdlist, buf, n);
}

/*
 * This is stupid. But it's just a bit of fun.
 *
 * The file descriptor of the sender is prepended to the
 * message sent to clients, so you can distinguish
 * each other with a degree of sanity. This gives a
 * similar effect to an IRC session. But stupider.
 */
static char *chat_filter(char *buf, size_t size, int fd, int *nwritten)
{
    char *result = NULL;
    size_t n = 0;
    const char *p;
    int i;

    n = 32;
    result = (char *) safe_malloc(n);
    i = Snprintf(result, n, "<user%d> ", fd);

    /* Escape control characters. */
    for (p = buf; p - buf < size; p++) {
        char repl[32];
        int repl_len;

        if (isprint(*p) || *p == '\r' || *p == '\n' || *p == '\t') {
            repl[0] = *p;
            repl_len = 1;
        } else {
            repl_len = Snprintf(repl, sizeof(repl), "\\%03o", (unsigned char) *p);
        }

        if (i + repl_len > n) {
            n = (i + repl_len) * 2;
            result = (char *) safe_realloc(result, n + 1);
        }
        memcpy(result + i, repl, repl_len);
        i += repl_len;
    }
    /* Trim to length. (Also does initial allocation when str is empty.) */
    result = (char *) safe_realloc(result, i + 1);
    result[i] = '\0';

    *nwritten = i;

    return result;
}
