/* Pure Load Balancer - (C)opyleft 2003 Jedi/Sector One <j@pureftpd.org> */

#include <config.h>
#define DEFINE_GLOBALS
#include "plb.h"
#include "parser.h"
#include "plb_globals.h"
#include "plb_p.h"
#ifdef WITH_DMALLOC
# include <dmalloc.h>
#endif

static int init_pool_add(Server **serverpool_last,
                         Server **serverpool_previous,
                         const char * const host, const char * const port)
{
    struct addrinfo hints, *res;
    int on;

    plb_log(LL_NOTIFY, "Adding [%s]:[%s] to the server pool", host, port);
    *serverpool_previous = *serverpool_last;
    if ((*serverpool_last = malloc(sizeof **serverpool_last)) == NULL) {
        plb_log(LL_ERROR, "Out of memory to fill the server pool : [%s]",
                strerror(errno));
        return -1;
    }
    (*serverpool_last)->next = NULL;
    if (*serverpool_previous == NULL) {
        serverpool_head = *serverpool_last;
    } else {
        (*serverpool_previous)->next = *serverpool_last;
    }
    memset(&hints, 0, sizeof hints);
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_addr = NULL;
    if ((on = getaddrinfo(host, port, &hints, &res)) != 0 ||
        (res->ai_family != AF_INET && res->ai_family != AF_INET6)) {
        plb_log(LL_ERROR, "Unable to add [%s]:[%s] to the server pool : [%s]",
            host, port, gai_strerror(on));
        return -1;
    }    
    (*serverpool_last)->ai_family = res->ai_family;
    if (res->ai_addrlen > sizeof (*serverpool_last)->ai_addr) {
        plb_log(LL_FATAL, "ai_addr overflow");
        return -1;
    }    
    memcpy(&(*serverpool_last)->ai_addr, res->ai_addr, res->ai_addrlen);
    (*serverpool_last)->ai_addrlen = res->ai_addrlen;
    (*serverpool_last)->status = server_retry;
    (*serverpool_last)->cleanup_fd = -1;
    if (((*serverpool_last)->name = strdup(host)) == NULL) {
        plb_log(LL_ERROR, "Out of memory when adding [%s]:[%s]", host, port);
        return -1;
    }
    
    return 0;
}

static int init_pool(const char * const poolstr_,
                     const char * const port)
{
    char *poolstr;
    Server *serverpool_last = NULL;
    Server *serverpool_previous;
    char *poolstrbegtok;
    char *poolstrtokpnt;
    char savechar;    

    if ((poolstr = strdup(poolstr_)) == NULL) {
        plb_log(LL_FATAL, "Out of memory to duplicate the server pool list");
        return -1;
    }
    poolstrbegtok = poolstrtokpnt = poolstr;
    for (;;) {
        while (*poolstrtokpnt != 0 &&
               !isspace((unsigned int) *poolstrtokpnt)) {
            poolstrtokpnt++;
        }
        savechar = *poolstrtokpnt;
        *poolstrtokpnt = 0;        
        if (init_pool_add(&serverpool_last, &serverpool_previous,
                          poolstrbegtok, port) != 0) {
            free(poolstr);
            return -1;
        }        
        if (savechar == 0) {
            break;
        }
        poolstrtokpnt++;
        while (*poolstrtokpnt != 0 &&
               isspace((unsigned int) *poolstrtokpnt)) {
            poolstrtokpnt++;
        }
        if (*poolstrtokpnt == 0) {
            break;
        }
        poolstrbegtok = poolstrtokpnt;
    }
    if ((serverpool_current = serverpool_head) == NULL) {
        free(poolstr);        
        return -1;
    }   
    free(poolstr);
    
    return 0;
}

int parse(const char * const file)
{
    if (generic_parser(file, plb_config_keywords) != 0) {
        plb_log(LL_FATAL, "Invalid configuration file : [%s]", file);
        return -1;
    }
    if (cfg_bind_ipv6 != NULL) {
        if (atoi(cfg_bind_ipv6) > 0) {
            bind_ipv6 = 1;
        }
    }
    if (cfg_protocol != NULL) {
        if (strcasecmp(cfg_protocol, "HTTP") == 0) {
            protocol = PROTOCOL_HTTP;
        } else if (strcasecmp(cfg_protocol, "SMTP") == 0) {
            protocol = PROTOCOL_SMTP;
        } else {
            plb_log(LL_FATAL, "Unknown protocol : [%s]", cfg_protocol);
            return -1;
        }
    }
    if (cfg_timeout_header_client_read != NULL) {
        timeout_header_client_read.tv_sec =
            (time_t) strtoul(cfg_timeout_header_client_read, NULL, 10);
    }
    if (cfg_timeout_header_client_write != NULL) {
        timeout_header_client_write.tv_sec =
            (time_t) strtoul(cfg_timeout_header_client_write, NULL, 10);
    }
    if (cfg_timeout_header_server_read != NULL) {
        timeout_header_server_read.tv_sec =
            (time_t) strtoul(cfg_timeout_header_server_read, NULL, 10);
    }
    if (cfg_timeout_header_server_write != NULL) {
        timeout_header_server_write.tv_sec =
            (time_t) strtoul(cfg_timeout_header_server_write, NULL, 10);
    }
    if (cfg_timeout_forward_client_read != NULL) {
        timeout_forward_client_read.tv_sec =
            (time_t) strtoul(cfg_timeout_forward_client_read, NULL, 10);
    }
    if (cfg_timeout_forward_client_write != NULL) {
        timeout_forward_client_write.tv_sec =
            (time_t) strtoul(cfg_timeout_forward_client_write, NULL, 10);
    }
    if (cfg_timeout_forward_server_read != NULL) {
        timeout_forward_server_read.tv_sec =
            (time_t) strtoul(cfg_timeout_forward_server_read, NULL, 10);
    }
    if (cfg_timeout_forward_server_write != NULL) {
        timeout_forward_server_write.tv_sec =
            (time_t) strtoul(cfg_timeout_forward_server_write, NULL, 10);
    }
    if (cfg_timeout_cleanup != NULL) {
        timeout_cleanup.tv_sec =
            (time_t) strtoul(cfg_timeout_cleanup, NULL, 10);
    }
    if (cfg_max_clients != NULL) {
        max_clients = (int) strtoul(cfg_max_clients, NULL, 10);
    }
    if (cfg_server_retry != NULL) {
        server_retry = (unsigned int) strtoul(cfg_server_retry, NULL, 10);
    }
    if (cfg_backlog != NULL) {
        backlog = (int) strtoul(cfg_backlog, NULL, 10);
    }
    if (cfg_log_level != NULL) {
        log_level = (LogLevel) (int) strtoul(cfg_log_level, NULL, 10);
    }
    /* TODO: check ranges */
    
    return 0;
}

static void usage(void)
{
    puts("Usage: plb [-B|--daemonize] [-c|--config <configuration file>]\n"
         "           [-d|--loglevel <verbosity level>] [-h --help]\n"
         "           [-l|--logfile <log file>] [-v|--version]");
}

static void version(void)
{
    puts("Pure Load Balancer v" VERSION ", compiled on " __DATE__);
}

static int sigterm(void)
{
    return -1;
}

static RETSIGTYPE sigterm_(int sig)
{    
    event_sigcb = sigterm;    
    event_gotsig = sig;
}

static void set_signals(void)
{
    struct sigaction sa;
    
    sigemptyset(&sa.sa_mask);    
    sa.sa_flags = SA_RESTART;    
    sa.sa_handler = SIG_IGN;
    (void) sigaction(SIGPIPE, &sa, NULL);
#ifdef SIGURG
    (void) sigaction(SIGURG, &sa, NULL);
#endif       
    sa.sa_flags = 0;
    sa.sa_handler = sigterm_;
    (void) sigaction(SIGTERM, &sa, NULL);
    (void) sigaction(SIGHUP, &sa, NULL);
    (void) sigaction(SIGQUIT, &sa, NULL);
    (void) sigaction(SIGINT, &sa, NULL);
#ifdef SIGXCPU
    (void) sigaction(SIGXCPU, &sa, NULL);
#endif    
}

int main(int argc, char *argv[])
{
    struct addrinfo hints, *res;
    struct event ev;
    struct event cleanup_ev;
    int option_index = 0;
    int on;    
    int listenfd;
    int fodder;
    int want_daemonization = 0;

    while ((fodder = getopt_long(argc, argv, GETOPT_OPTIONS, long_options,
                                 &option_index)) != -1) {
        switch (fodder) {
        case 'B':
            want_daemonization = 1;
            break;
        case 'c':
            if ((config_file = strdup(optarg)) == NULL) {
                plb_log(LL_FATAL,
                        "Out of memory to store the config file name");
                return 1;
            }
            break;
        case 'd':
            log_level = (LogLevel) (int) strtoul(optarg, NULL, 10);
            break;
        case 'g':
            if ((cfg_pid_file = strdup(optarg)) == NULL) {
                plb_log(LL_FATAL, "Out of memory to store the pid file name");
                return 1;
            }
            break;
        case 'h':
            usage();
            return 0;
        case 'l':
            if (strcmp(optarg, "-") == 0) {
                cfg_log_file = NULL;
            } else if ((cfg_log_file = strdup(optarg)) == NULL) {
                plb_log(LL_FATAL,
                        "Out of memory to store the log file name");
                return 1;
            }
            break;
        case 'v':
            version();
            return 0;
        default:
            plb_log(LL_ERROR, "Unrecognized command-line switch");
            return 1;
        }
    }
    init_timeouts();
    if (parse(config_file) != 0) {
        return 1;
    }
    if (protocol == PROTOCOL_HTTP) {
        protocol_header_end = HTTP_HEADER_END;
        protocol_header_end_len = sizeof HTTP_HEADER_END - (size_t) 1U;
    } else if (protocol == PROTOCOL_SMTP) {
        protocol_header_end = SMTP_HEADER_END;
        protocol_header_end_len = sizeof SMTP_HEADER_END - (size_t) 1U;        
    } else {
        plb_log(LL_FATAL, "Unknown protocol");
        return 1;
    }
    if (cfg_log_file != NULL && plb_open_log(cfg_log_file) != 0) {
        plb_log(LL_FATAL, "Unable to access the log file, exiting");
        return 1;
    }
    plb_log(LL_NOTIFY, "PLB " VERSION " has been started");    
    if (init_pool(cfg_servers_ip, cfg_servers_port) != 0) {
        plb_log(LL_FATAL, "No server pool defined");
        return 1;
    }
    if ((clients = malloc((size_t) max_clients * sizeof *clients)) == NULL) {
        plb_log(LL_FATAL, "Out of memory : max_clients too high?");
        return 1;
    }
    memset(&hints, 0, sizeof hints);
    hints.ai_flags = AI_PASSIVE;
    hints.ai_family = bind_ipv6 > 0 ? AF_INET6 : AF_INET;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_addr = NULL;
    if ((on = getaddrinfo(cfg_listen_ip, cfg_listen_port,
                          &hints, &res)) != 0 ||
        (res->ai_family != AF_INET && res->ai_family != AF_INET6)) {
        plb_log(LL_FATAL, "Unable to get the local address [%s]:[%s] : [%s]",
            cfg_listen_ip, cfg_listen_port, gai_strerror(on));
        return 1;
    }
    on = 1;
    if ((listenfd = socket(res->ai_family, SOCK_STREAM, IPPROTO_TCP)) == -1 ||
        setsockopt(listenfd, SOL_SOCKET, SO_REUSEADDR,
                   (char *) &on, sizeof on) != 0 ||
        bind(listenfd, (struct sockaddr *) res->ai_addr,
             (socklen_t) res->ai_addrlen) != 0 ||
        listen(listenfd, backlog) != 0) {
        plb_log(LL_FATAL, "Unable to listen to [%s]:[%s] : [%s]",
            cfg_listen_ip, cfg_listen_port, strerror(errno));            
        freeaddrinfo(res);
        return 1;
    }
    freeaddrinfo(res);
    
    if (want_daemonization != 0 && daemonize() != 0) {
        plb_log(LL_FATAL, "Unable to switch to background");
        return 1;
    }
    if (update_pid_file(cfg_pid_file) != 0) {
        plb_log(LL_FATAL, "Unable to create the pid file");
        return 1;        
    }
    if (plb_drop_caps(cfg_user, cfg_group, cfg_chroot_dir) != 0) {
        plb_log(LL_FATAL, "Unable to drop privileges");
        _exit(1);
    }
    set_signals();
    
    event_init();
    
    event_set(&ev, listenfd, EV_READ | EV_PERSIST, new_client, &ev);
    event_add(&ev, NULL);

    evtimer_set(&cleanup_ev, periodic_cleanup, &cleanup_ev);
    evtimer_add(&cleanup_ev, &timeout_cleanup);

    plb_log(LL_NOTIFY, "Server ready, now accepting connections on [%s]:[%s]",
            cfg_listen_ip, cfg_listen_port);
    
    event_dispatch();
    
    free(clients);
    plb_close_log();
    if (delete_pid_file() != 0) {
        plb_log(LL_NOTIFY, "Failed to remove the pid file");
    }
    
    return 0;
}


syntax highlighted by Code2HTML, v. 0.9.1