/*
   Copyright (c)  2014 IOnU Security Inc. All rights reserved
   Created April 2014 by Kendrick Webster

   K2Client/servers.c - implementation of servers.h
*/
#include <sys/types.h>
#ifndef WIN32
#include <netdb.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <strings.h>
#include <unistd.h>
#include <sys/time.h>
#else
#include <Winsock2.h>
#include <ws2tcpip.h>
#pragma comment (lib, "Ws2_32.lib")
#define strdup _strdup
#include "eyelog.h"
using sequencelogic::EyeLog;
#endif
#include <ctype.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include "constants.h"
#include "util.h"
#include "platform_binding.h"
#include "servers.h"

int is_using_proxy; /* non-zero if using K2Proxy (HTTP/TCP) instaed of direct UDP connection with K2Daemon */
static void (*callback_onServerSelected)(const void* address, const char* host);
static const char * current_server;
static char * servers;

/* pointers into <servers> for each UDP (non-proxy) server */
size_t n_udp_servers;
static char const ** udp_servers;

/* pointers into <servers> for each proxy server */
size_t n_proxy_servers;
static char const ** proxy_servers;

/* two concatenated shuffled lists ordered with non-proxy (UDP) servers first, NUL terminated */
static char const ** shuffled_servers;
static char const * const * next_server;

static int is_proxy_port(unsigned int n)
{
    switch (n)
    {
#define X(port) case port: return 1;
        K2_PROXY_PORTS
#undef X
    default:
        return 0;
    }
}

static void show_server(const char * s)
{
    const char * p = "connecting to ";
    int n = (int)(strlen(p) + strlen(s) + 1);
    char * buf = (char *)malloc(n);
    snprintf(buf, n, "%s%s", p, s);
    upcall_k2cli_message(buf);
    free(buf);
}

/*
   Display a struct addrinfo, show if the address was ignored (wrong type)
   or selected (the K2Daemon address being used).
*/
#define SA_TYPE_SELECTED     "(selected)"
#define SA_TYPE_IGNORED      "(ignored) "
#define SA_TYPE_NOT_SELECTED ".........."
static void show_address(const struct addrinfo * a, unsigned int index, const char * type)
{
    char buf[ADDRESS_STRING_LENGTH + 64];
    buf[0] = '\0';
    char addrstr[ADDRESS_STRING_LENGTH];
    struct sockaddr_in* sin;
    struct sockaddr_in6* sin6;
    switch (a->ai_family) {
    case AF_INET:
        sin = (struct sockaddr_in*)(a->ai_addr);
        if (inet_ntop (AF_INET, &sin->sin_addr, addrstr, sizeof (addrstr))) {
            if (ntohs (sin->sin_port) != 0) {
                snprintf(buf, sizeof(buf), "address %2d %s %s:%d", index, type, addrstr, ntohs (sin->sin_port));
            }
        }
        break;
    case AF_INET6:
        sin6 = (struct sockaddr_in6*)(a->ai_addr);
        if (inet_ntop (AF_INET6, &sin6->sin6_addr, addrstr, sizeof (addrstr))) {
            if (ntohs (sin6->sin6_port) != 0) {
                snprintf(buf, sizeof(buf), "address %2d %s [%s]:%d", index, type, addrstr, ntohs (sin6->sin6_port));
            }
        }
        break;
    }
    upcall_k2cli_message(buf);
}

/* return non-zero if the address is of a type we can use */
static int addr_is_usable(const struct addrinfo * a)
{
    return (((is_using_proxy ? SOCK_STREAM : SOCK_DGRAM) == a->ai_socktype)) ? 1 : 0;
}

static unsigned int count_usable_addresses(const struct addrinfo * a)
{
    unsigned int n = 0;
    while (NULL != a)
    {
        if (addr_is_usable(a))
        {
            ++n;
        }
        a = a->ai_next;
    }
    return n;
}

/*
   This function does two things:
   1). Selects the usable address with index <i> (zero-based index into usable addresses)
   2). Shows all addresses and their status (selected, usable vs ignored)
*/
static void select_usable_address_and_show_all(unsigned int i, const struct addrinfo * a)
{
    unsigned int j = 0, n = 0;
    while (NULL != a)
    {
        ++j;
        if (addr_is_usable(a))
        {
            if (i == n)
            {
                show_address(a, j, SA_TYPE_SELECTED);
                if (callback_onServerSelected)
                {
                    callback_onServerSelected(a->ai_addr, current_server);
                }
            }
            else
            {
                show_address(a, j, SA_TYPE_NOT_SELECTED);
            }
            ++n;
        }
        else
        {
            show_address(a, j, SA_TYPE_IGNORED);
        }
        a = a->ai_next;
    }
}

/*
   This function does two things:
   1). Selects an address at random that matches our needs (address type, etc.)
   2). Shows all addresses and their status (selected, usable vs ignored)
*/
static void choose_address_and_show_all(const struct addrinfo * a)
{
    unsigned int i = get_random(count_usable_addresses(a));
    select_usable_address_and_show_all(i, a);
}

/*
   Fetch and save the IP address/port of the specified K2Daemon server
   (or, if <port> is in K2_PROXY_PORTS, the HTTP/TCP (K2Proxy) server).

   Sets the <is_using_proxy> global flag.

   The <host> parameter may be a numerical IP address or a hostname.
   If it is a hostname, pick one addresses (i.e. from multiple DNS
   A records) randomly.
*/
static void use_server_(const char* host, const char* port)
{
    int r;
    struct addrinfo hints, *result;

    is_using_proxy = is_proxy_port(atoi(port));

    /* Obtain a list of addresses for host/port */
    memset(&hints, 0, sizeof(struct addrinfo));
    hints.ai_family = AF_UNSPEC;      // Could be IPv4 or IPv6, we don't care
    hints.ai_socktype = is_using_proxy ? SOCK_STREAM : SOCK_DGRAM;
    r = getaddrinfo(host, port, &hints, &result);
    if (0 != r)
    {
        char buf[512];
        snprintf (buf, 512, "getaddrinfo (%s, %s)", host, port);
        upcall_error_string(buf, gai_strerror(r));
        return;
    }

    choose_address_and_show_all(result);
    freeaddrinfo(result);
}

/* split string "host:port", call callback(host, port) */
#define xstr(s) str(s)
#define str(s) #s
static void split_host_port(const char* s, void (*callback)(const char* host, const char* port))
{
    const char *port = xstr(K2IPC_DEFAULT_UDP_PORT);
    int colons = 0;
    char *p, *host;
    p = host = (char *)malloc(strlen(s) + 1);

    // Skip leading whitespace and other characters not allowed for domain name or ip address
    while (*s && *s != '[' && *s != '-' && !isalnum(*s))
        s++;

    // Skip http://, https://, and set default port accordingly, but allow http.*
    if (strncmp (s, "http", 4) == 0) {
        s += 4;
        port = xstr(K2IPC_DEFAULT_HTTP_PORT);
        if (*s == 's') {
            s++;
            port = xstr(K2IPC_DEFAULT_HTTPS_PORT);
        }
        if (*s == ':') {
            s++;
            while (*s == '/')
                s++;
        }
        else {
            if (*(s-1) == 's')
                s -= 5;
            else
                s -= 4;
            port = xstr(K2IPC_DEFAULT_UDP_PORT);
        }
    }

    const char* tptr = s;
    while (*tptr++) {
        if (*tptr == ':')
            colons++;
    }

    //IPv6 [addr]:port
    if (*s == '[' && colons > 1) {
        while (*s != ']' && *s) {
           *p++ = *++s;
        }
        if (*s)
            s++;
        if (K2_SERVERS_PORT_DELIMITER == *s)
            port = s + 1;    // skip port delimiter
        *(p-1) = '\0';   // terminate host before ']'
        // Validate and trim string for port number
        p = (char*)port;
        while (*p++) {
            if (!isdigit (*p)) {
                *p = '\0';
                break;
            }
        }
    }
    // IPv4 or domain name based address with port specification
    else if (colons == 1) {
        while (*s)
        {
            if (K2_SERVERS_PORT_DELIMITER == *s)
            {
                port = s + 1;
                break;
            }
            else
            {
                *p = *s;
            }
            ++s;
            ++p;
        }
        *p = '\0';
        p = (char*)port;
        // Validate and trim string for port number
        while (*p++) {
            if (!isdigit (*p)) {
                *p = '\0';
                break;
            }
        }
    }
    // Domain name, IPv4 or IPv6 with no port specification
    else {
        strcpy (host, s);
    }

    if (strlen (port) < 1)
        port = xstr(K2IPC_DEFAULT_UDP_PORT);

    //printf ("%s %s\n", host, port);
    callback(host, port);
    free(host);
}

static void use_server(const char* s)
{
    split_host_port(s, use_server_);
}

static void count_server(const char* host, const char* port)
{
    if (is_proxy_port(atoi(port)))
    {
        ++n_proxy_servers;
    }
    else
    {
        ++n_udp_servers;
    }
}

/* counts server substrings in <servers>, also NUL-terminates them, sets <n_udp_servers> and <n_proxy_servers> */
static void count_servers(void)
{
    char *p = servers;
    const char *q = p;
    n_udp_servers = 0;
    n_proxy_servers = 0;
    do
    {
        if (',' == *p || '|' == *p)
        {
            *p++ = 0;
            split_host_port(q, count_server);
            q = p;
        }
    }
    while (*p++);
    split_host_port(q, count_server);
}

static int is_proxy;
static void check_is_proxy(const char * host, const char * port)
{
    is_proxy = is_proxy_port(atoi(port));
}

static void index_servers(void)
{
    size_t n;
    char const * p = servers;
    char const ** udp = NULL;
    char const ** proxy = NULL;
    count_servers();
    if (n_udp_servers)
    {
        udp = udp_servers = (char const **)malloc(sizeof(char*) * n_udp_servers);
    }
    if (n_proxy_servers)
    {
        proxy = proxy_servers = (char const **)malloc(sizeof(char*) * n_proxy_servers);
    }
    n = n_udp_servers + n_proxy_servers;
    shuffled_servers = (char const **)malloc(sizeof(char*) * (n + 1));
    while (n--)
    {
        split_host_port(p, check_is_proxy);
        if (is_proxy)
        {
            if (proxy) *proxy++ = p;
        }
        else
        {
            if (udp) *udp++ = p;
        }
        p += strlen(p) + 1;
    }
    next_server = NULL;
}

static void shuffle(char const ** p, size_t n)
{
    char const *t;
    size_t i, j;
    for (i = 0; i < n; i++)
    {
        j = get_random((unsigned int)n);
        if (i != j)
        {
            t = p[i];
            p[i] = p[j];
            p[j] = t;
        }
    }
}

static void shuffle_servers(void)
{
    if (shuffled_servers)
    {
        if (n_udp_servers)
        {
            shuffle(udp_servers, n_udp_servers);
            memcpy(shuffled_servers, udp_servers, n_udp_servers * sizeof(char*));
        }
        if (n_proxy_servers)
        {
            shuffle(proxy_servers, n_proxy_servers);
            memcpy(shuffled_servers + n_udp_servers, proxy_servers, n_proxy_servers * sizeof(char*));
        }
        shuffled_servers[n_udp_servers + n_proxy_servers] = NULL;
    }
    else
    {
        log_info("shuffle_servers() was called while shuffled_servers is NULL");
    }
    next_server = shuffled_servers;
}

const char* get_next_server(void)
{
    if (!(next_server && *next_server))
    {
        shuffle_servers();
    }
    if (next_server && *next_server)
    {
        return *next_server++;
    }
    return NULL;
}


/* ----------------------- Module interface functions ----------------------- */

void set_servers(const char* servers_list)
{
    clear_servers();
    log_info("servers::set_servers(%s)", servers_list);
    if (servers_list)
    {
        servers = strdup(servers_list);
        index_servers();
    }
    else
    {
        upcall_error_string("K2_Start", "NULL servers string");
    }
}

void clear_servers(void)
{
    log_info("servers::clear_servers");
    free(shuffled_servers);
    shuffled_servers = NULL;
    next_server = NULL;
    free(udp_servers);
    udp_servers = NULL;
    free(proxy_servers);
    proxy_servers = NULL;
    free(servers);
    servers = NULL;
}

/*
   Randomly select a server from the server list string, example string:
   "cg2.ionu.com:32000|cg2.ionu.com:80|cg3.ionu.com:32000|cg3.ionu.com:80"

   Servers on ports in K2_PROXY_PORTS are HTTP proxies which should be used only if a UDP connection
   cannot be established.  Thus the 'random' selection of a server first goes through a shuffled list
   of UDP hosts, then through a shuffled list of HTTP/TCP hosts.
*/
void select_server(void (*onServerSelected)(const void* address, const char* host))
{
    log_debug("servers::select_server");
    const char *s = get_next_server();
    callback_onServerSelected = onServerSelected;
    if (s)
    {
        current_server = s;
        show_server(s);
        use_server(s);
    }
}