/* $Id: ncat_broker.c 11795 2009-01-21 22:16:20Z 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

int ncat_broker(void)
{
    char buf[DEFAULT_TCP_BUF_LEN] = {0};
    int nbytes = 0, i = 0, j = 0, conn_count = 0;
    int readsock = 0,   listen_sock = 0, new_fd = 0;
    socklen_t   ss_len = 0;
    pid_t   pid;
    fd_set master,  read_fds;
    fd_list_t   fdlist;
    struct sockaddr_storage remoteaddr;
#ifdef HAVE_OPENSSL
    SSL_CTX *ctx;
    SSL *tmpssl;
#else
    void *tmpssl = NULL;
#endif

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

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

    /* setup the main listening socket */
    listen_sock = 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_sock);

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

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

    while (1) {
	    read_fds = master;

	    if (verbose_flag > 0)
	        Fprintf(stderr, "DEBUG: broker connection count is %d\n", conn_count);

	    readsock = 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;

            /* we have a new connection request */
	        if (i == listen_sock) {
		        ss_len = sizeof(remoteaddr);

		        new_fd = accept(listen_sock, (struct sockaddr *) &remoteaddr, &ss_len);

		        if (new_fd < 0) {
		            if (errno != EAGAIN && errno != EWOULDBLOCK)
                                die("accept");

		            close(new_fd);
		            continue;
		        }
		    
		        if (verbose_flag > 0 && o.talk)
			        Fprintf(stderr,
				            "DEBUG: Connection from %s on file descriptor %d\n",
				            inet_socktop(&remoteaddr), new_fd);
		        else if (verbose_flag > 0)
			        Fprintf(stderr, "DEBUG: Connection from %s\n",
				            inet_socktop(&remoteaddr));

		        /* Check number of connections and deny list */
		        if (++conn_count > o.conn_limit || !allow_access(&remoteaddr)) {
		            if (verbose_flag > 1)
		                Fprintf(stderr, "DEBUG: New connection denied: %s\n",
		                         (conn_count >= o.conn_limit) ? 
		                          "Max connections reached" : "ACL denial");
		            Close(new_fd);
		            conn_count--;

		            continue;
		        }
 
#ifdef HAVE_OPENSSL
                        if (o.ssl) {
                            tmpssl = new_ssl(new_fd);
                            if (SSL_accept(tmpssl) != 1)
                                bye("SSL_accept(): %s", ERR_error_string(ERR_get_error(), NULL));
                        }
#endif

		        /* If there is a command to execute. */
		        if (o.cmdexec) {
                            pid = Fork();

		            if (pid == 0)
                                if ((netexec(new_fd, o.cmdexec)) == -1)
                                    die("execve");

                           /* parent */
                           Close(new_fd);
                           conn_count--;
                           continue;
                        }

                        /* do this after we test for the limit and cmdexec */
                        FD_SET(new_fd, &master);

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

	        } else {
	                /* Handle incoming client data and distribute it. */
		        char chatbuf[DEFAULT_TCP_BUF_LEN + 10];
                        struct fdinfo *fdn = get_fdinfo(&fdlist, i), *fdn2;

                        assert(fdn);

#ifdef HAVE_OPENSSL
                readagain:
                        if (o.ssl && fdn->ssl)
                            nbytes = SSL_read(fdn->ssl, buf, sizeof(buf));
                        else
#endif
                            if (i == STDIN_FILENO)
                                nbytes = read(i, buf, sizeof(buf) - o.crlf);
                            else
                                nbytes = recv(i, buf, sizeof(buf), 0);

                        if (nbytes <= 0) {
		            if (verbose_flag > 1)
		                Fprintf(stderr, "DEBUG: Closing connection.\n");

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

                            Close(i);
                            FD_CLR(i, &master);

                            conn_count--;

                            rm_fd(&fdlist, i);
                            continue;
                        }

		        if (verbose_flag > 1)
		            Fprintf(stderr, "DEBUG: Handling data from client %d.\n", i);

                        /* We gave ourselves extra room in our special-case
                         * stdin read() above
                         */
                        if (i == STDIN_FILENO && o.crlf && buf[nbytes - 1] == '\n' && buf[nbytes - 2] != '\r') {
                            memcpy(&buf[nbytes - 1], "\r\n", 2);
                            nbytes++;
                        }
		
		        if (o.normlogfd != -1)
                            Write(o.normlogfd, buf, nbytes);

		        if (o.hexlogfd != -1)
		            ncat_hexdump(o.hexlogfd, buf, nbytes);

		        for (j = 0; j <= fdlist.fdmax; j++) {
		            /* 
                             * write to everything in the master set, except
                             * listener, sender and stdin
                             */
		            if (FD_ISSET(j, &master) && j != listen_sock && j != i && j != STDIN_FILENO) {
                                    fdn2 = get_fdinfo(&fdlist, j);

			            if (o.talk) {
                                        /* 
                                         * 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.
                                         */
			                nbytes = Snprintf(chatbuf, sizeof(chatbuf), "<user%d> %s", i, buf);
                                        if(nbytes < 0 || nbytes > sizeof(chatbuf))
                                            nbytes = sizeof(chatbuf);

#ifdef HAVE_OPENSSL
                                        if (o.ssl && fdn2->ssl)
                                            SSL_write(fdn2->ssl, chatbuf, nbytes);
                                        else
#endif
                                            Send(j, chatbuf, nbytes, 0);
                                        zmem(chatbuf, sizeof(chatbuf));
			            } else
#ifdef HAVE_OPENSSL
                                        if (o.ssl && fdn2->ssl)
                                            SSL_write(fdn2->ssl, buf, nbytes);
                                        else
#endif
                                            Send(j, buf, nbytes, 0);
                            }
                        }

                        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
            }
        }
    }

    return 0;
}
