/*
 *  Copyright (C) 2013-2019 Cisco Systems, Inc. and/or its affiliates. All rights reserved.
 *  Copyright (C) 2008-2013 Sourcefire, Inc.
 *
 *  Author: aCaB <acab@clamav.net>
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License version 2 as
 *  published by the Free Software Foundation.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 *  MA 02110-1301, USA.
 */

#if HAVE_CONFIG_H
#include "clamav-config.h"
#endif

#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/un.h>
#include <netdb.h>
#include <unistd.h>
#include <time.h>
#include <netdb.h>

#include "shared/optparser.h"
#include "shared/output.h"
#include "shared/misc.h"

#include "connpool.h"
#include "netcode.h"

#define SETGAI(k, v)            \
    {                           \
        (k)->gai = (void *)(v); \
    }                           \
    while (0)
#define FREESRV(k)                 \
    {                              \
        if ((k).gai)               \
            freeaddrinfo((k).gai); \
        else if ((k).server)       \
            free((k).server);      \
    }                              \
    while (0)

#if __GNUC__ >= 3 || (__GNUC__ == 2 && __GNUC_MINOR__ >= 7)
#define _UNUSED_ __attribute__((__unused__))
#else
#define _UNUSED_
#endif

struct CPOOL *cp               = NULL;
static pthread_cond_t mon_cond = PTHREAD_COND_INITIALIZER;
static int quitting            = 1;
static pthread_t probe_th;

static int cpool_addunix(char *path)
{
    struct sockaddr_un *srv;
    struct CP_ENTRY *cpe = &cp->pool[cp->entries - 1];

    if (!cli_is_abspath(path)) {
        logg("!Unix clamd socket must be an absolute path\n");
        return 1;
    }
    if (!(srv = (struct sockaddr_un *)malloc(sizeof(*srv)))) {
        logg("!Out of memory allocating unix socket space\n");
        return 1;
    }

    srv->sun_family = AF_UNIX;
    strncpy(srv->sun_path, path, sizeof(srv->sun_path));
    srv->sun_path[sizeof(srv->sun_path) - 1] = '\0';
    cpe->type                                = 0;
    cpe->dead                                = 1;
    cpe->local                               = 1;
    cpe->last_poll                           = 0;
    cpe->server                              = (struct sockaddr *)srv;
    cpe->socklen                             = sizeof(*srv);
    SETGAI(cpe, NULL);
    if (!cp->local_cpe) cp->local_cpe = cpe;
    logg("*Local socket unix:%s added to the pool (slot %d)\n", srv->sun_path, cp->entries);
    return 0;
}

static int islocal(struct sockaddr *sa, socklen_t addrlen)
{
    int s = socket(sa->sa_family, SOCK_STREAM, 0);
    int ret;
    if (s < 0) return 0;
    ret = (bind(s, sa, addrlen) == 0);
    close(s);
    return ret;
}

static int cpool_addtcp(char *addr, char *port)
{
    struct addrinfo hints, *res, *res2;
    struct CP_ENTRY *cpe = (struct CP_ENTRY *)&cp->pool[cp->entries - 1];

    memset(&hints, 0, sizeof(hints));
    hints.ai_family   = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;

    if (getaddrinfo(addr, port ? port : "3310", &hints, &res)) {
        logg("^Can't resolve hostname %s\n", addr ? addr : "");
        return 1;
    }
    cpe->type = 1;
    cpe->dead = 1;

    memset(&hints, 0, sizeof(hints));
    hints.ai_flags    = AI_PASSIVE;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_family   = AF_UNSPEC;
    if (!getaddrinfo(addr, NULL, &hints, &res2)) {
        cpe->local = islocal(res2->ai_addr, res2->ai_addrlen);
        freeaddrinfo(res2);
    } else
        cpe->local = 0;
    cpe->last_poll = 0;
    cpe->server    = res->ai_addr;
    cpe->socklen   = res->ai_addrlen;
    SETGAI(cpe, res);
    logg("*%s socket tcp:%s:%s added to the pool (slot %d)\n", cpe->local ? "Local" : "Remote", addr ? addr : "localhost", port ? port : "3310", cp->entries);
    return 0;
}

static int addslot(void)
{
    struct CP_ENTRY *cpe;

    if (!(cpe = realloc(cp->pool, (cp->entries + 1) * sizeof(struct CP_ENTRY)))) {
        logg("!Out of memory while initializing the connection pool\n");
        cpool_free();
        return 1;
    }
    if (cp->local_cpe)
        cp->local_cpe = (struct CP_ENTRY *)((char *)cp->local_cpe + ((char *)cpe - (char *)cp->pool));
    memset(&cpe[cp->entries], 0, sizeof(*cpe));
    cp->pool = cpe;
    cp->entries++;
    return 0;
}

/* Probe strategy:
- wake up every minute
- probe alive if last check > 15 min
- probe dead if (last check > 2 min || no clamd available)
*/

static void cpool_probe(void)
{
    unsigned int i, dead = 0;
    struct CP_ENTRY *cpe = cp->pool;
    time_t now           = time(NULL);

    for (i = 1; i <= cp->entries; i++) {
        if ((cpe->dead && (cpe->last_poll < now - 120 || !cp->alive)) || cpe->last_poll < now - 15 * 60 * 60) {
            cpe->last_poll = time(NULL);
            nc_ping_entry(cpe);
            logg("*Probe for slot %u returned: %s\n", i, cpe->dead ? "failed" : "success");
        }
        dead += cpe->dead;
        cpe++;
    }
    cp->alive = cp->entries - dead;

    if (!cp->alive)
        logg("^No clamd server appears to be available\n");
}

static void *cpool_mon(_UNUSED_ void *v)
{
    pthread_mutex_t conv;

    pthread_mutex_init(&conv, NULL);
    pthread_mutex_lock(&conv);

    while (!quitting) {
        struct timespec t;

        cpool_probe();
        t.tv_sec  = time(NULL) + 60;
        t.tv_nsec = 0;
        pthread_cond_timedwait(&mon_cond, &conv, &t);
    }
    pthread_mutex_unlock(&conv);
    pthread_mutex_destroy(&conv);
    return NULL;
}

void cpool_init(struct optstruct *opts)
{
    const struct optstruct *opt;
    int failed = 0;

    if (!(cp = calloc(sizeof(*cp), 1))) {
        logg("!Out of memory while initializing the connection pool");
        return;
    }

    cp->local_cpe = NULL;

    if ((opt = optget(opts, "ClamdSocket"))->enabled) {
        while (opt) {
            char *socktype = opt->strarg;

            if (addslot()) return;
            if (!strncasecmp(socktype, "unix:", 5)) {
                failed = cpool_addunix(socktype + 5);
            } else if (!strncasecmp(socktype, "tcp:", 4)) {
                char *port = strrchr(socktype + 4, ':');
                if (port) {
                    *port = '\0';
                    port++;
                }
                failed = cpool_addtcp(socktype + 4, port);
            } else {
                logg("!Failed to parse ClamdSocket directive '%s'\n", socktype);
                failed = 1;
            }
            if (failed) break;
            opt = opt->nextarg;
        }
        if (failed) {
            cpool_free();
            return;
        }
    }

    if (!cp->entries) {
        logg("!No ClamdSocket specified\n");
        cpool_free();
        return;
    }
    quitting = 0;
    pthread_create(&probe_th, NULL, cpool_mon, NULL);
    srand(time(NULL));
}

void cpool_free(void)
{
    unsigned int i;

    if (!quitting) {
        logg("*Killing the monitor and stopping\n");
        quitting = 1;
        pthread_cond_signal(&mon_cond);
        pthread_join(probe_th, NULL);
    }

    if (cp) {
        if (cp->pool) {
            for (i = 0; i < cp->entries; i++)
                FREESRV(cp->pool[i]);
            free(cp->pool);
        }
        free(cp);
        cp = NULL;
    }
}

struct CP_ENTRY *cpool_get_rand(int *s)
{
    unsigned int start, i;
    struct CP_ENTRY *cpe;

    if (cp->alive) {
        start = rand() % cp->entries;
        for (i = 0; i < cp->entries; i++) {
            cpe = &cp->pool[(i + start) % cp->entries];
            if (cpe->dead) continue;
            if (cpe->local && cp->local_cpe && !cp->local_cpe->dead)
                cpe = cp->local_cpe;
            if ((*s = nc_connect_entry(cpe)) == -1) {
                cpe->dead = 1;
                continue;
            }
            return cpe;
        }
    }
    pthread_cond_signal(&mon_cond);
    return NULL;
}

/*
 * Local Variables:
 * mode: c
 * c-basic-offset: 4
 * tab-width: 8
 * End: 
 * vim: set cindent smartindent autoindent softtabstop=4 shiftwidth=4 tabstop=8: 
 */