#include "fdproxy.hpp"
#include "sock.hpp"
#include "socks.hpp"
#include "session.hpp"
#include "status.hpp"
#include "sys.hpp"
#include "common.hpp"
#include <signal.h>
#include <assert.h>


volatile int* SHUTDOWN = NULL; // set by signal handler


static int handle_conn(int fd, const sockaddr_in& cli, const sockaddr_in* original_dst, config_t& config) {
    int mark = 0;
    sockaddr_in dst;
    if (original_dst) {
        dst = *original_dst;
        mark = 42; // set socket mark if we got a transparent connection to prevent loops
    } else if (!socks_req(fd, &dst)) {
        close(fd);
        return 1;
    }

    int dfd = get_conn(config, &dst, mark);
    if (dfd == -1) {
        if (!original_dst) {
            (void)socks_rep(fd, false);
        }
        close(fd);
        return 1;
    } else if (!original_dst && !socks_rep(fd, true)) {
        close(fd);
        close(dfd);
        return 1;
    }

    if (!sock_nonblock(fd, true) || !sock_nonblock(dfd, true)) {
        close(fd);
        close(dfd);
        return 1;
    }

    LOG("...connected.");
    return fdproxy(fd, dfd)? 0: 1;
}


static bool detach_conn(int fd, const sockaddr_in& cli, const sockaddr_in* dst, const config_t& config) {
    pid_t pid = fork();
    if (pid == -1) {
        LOG_ERRNO("cannot fork() for new connection");
        close(fd);
        return false;
    } else if (pid > 0) { // parent
        close(fd);
        return true;
    } else { // child
        // no post-fork cleanup needed atm (keeping signals & fds intact)
        _exit(handle_conn(fd, cli, dst, (config_t&)config)); // const-cast after fork
        return false; // not reached
    }
}


int main(int argc, char** argv) {
    // parse args
    config_t config;
    in_port_t port = 0;
    if (!parse_config(config, port, argc, argv)) {
        STDERR("Usage: %s --port=N [--min=N|--pool=N] [--proxies=<filename>] [proxy1 [proxy2 ...]]", argv[0]);
        STDERR("       Accepts connections on localhost (SOCKS4 or transparent) and");
        STDERR("         connects to the destination by using the given intermediate proxies,");
        STDERR("         in the form of proto://ip:port (currently supported: http and socks4).");
        STDERR("       If provided, <min> allows to skip apparently defunct proxies up to this");
        STDERR("         minimum amount of proxies traversed (so 0 allows direct, default: all).");
        STDERR("       In <pool> mode, the proxy list is not processed as a chain with a lower");
        STDERR("         limit on its length, but the given number of proxies are randomly selected");
        STDERR("         instead, and apparently defunct entries are replaced by spare ones.");
        return 1;
    }

    // setup signals
    SHUTDOWN = setup_signals();
    assert(SHUTDOWN != NULL);
    close(STDIN_FILENO); // not needed and not inherited
    close(STDOUT_FILENO);

    // open listening socket & drop root
    const bool trans_mode = (getuid() == 0); // implicitly enabled when root
    const int fd = sock_listen(port, false, trans_mode); // blocking
    if (fd == -1) return 1;
    if (!su(NULL)) {
        close(fd);
        return 1;
    }
    (void)status_init(config.proxies.back().id);

    // main accept loop
    LOG("listening on %d...", port);
    int rv = 0;

    while (!*SHUTDOWN) {
        sockaddr_in cli;
        socklen_t clilen = sizeof(cli);
        int cfd = accept4(fd, (sockaddr*)&cli, &clilen, 0); // no SOCK_NONBLOCK
        if (cfd == -1) {
            if (errno == EINTR) continue;
            LOG_ERRNO("accept() failed");
            rv = 1;
            break;
        }

        bool trans = trans_mode;
        sockaddr_in dst;
        if (trans) {
            if (!sock_dst(cfd, &dst)) {
                LOG("invalid transparent destination, continuing anyways");
                trans = false;
            } else if (sock_mark(cfd) != 0) { // pretty useless as it does not work for our own connections, though
                LOG("received marked connection - transparent interception loop? closing.");
                close(cfd);
                continue;
            }
        }

        if (!detach_conn(cfd, cli, trans? &dst: NULL, config)) {
            rv = 1;
            break;
        }
    }

    // done
    LOG("stopping service.");
    close(fd);
    kill(0, SIGTERM); // to every process in the process group
    //while (waitpid(-1, &status, 0) != -1 || errno != ECHILD) {}

    for (std::vector<proxy_t>::const_iterator it = config.proxies.begin(); it != config.proxies.end(); ++it) {
        int rating = status_rating(it->id);
        if (rating >= 0) {
            LOG("%-21s %3d%% ok", it->addrstr, status_rating(it->id));
        }
    }
    status_deinit();

    return rv;
}