/* Copyright (c) 2019 Dan Wilcox.
* For information on usage and redistribution, and for a DISCLAIMER OF ALL
* WARRANTIES, see the file, "LICENSE.txt," in this distribution.  */

#include "s_net.h"

#include <stdio.h>
#include <string.h>
#include <errno.h>
#ifndef _WIN32
#include <unistd.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/time.h>
#endif

#ifdef HAVE_SYS_FILIO_H
/* IllumOS defines FIONREAD in sys/filio.h */
#include <sys/filio.h>
#endif

    /* Windows XP winsock doesn't provide inet_ntop */
#ifdef _WIN32
const char* INET_NTOP(int af, const void *src, char *dst, socklen_t size) {
    struct sockaddr_storage addr;
    socklen_t addrlen;
    memset(&addr, 0, sizeof(struct sockaddr_storage));
    addr.ss_family = af;
    if (af == AF_INET6)
    {
        struct sockaddr_in6 *sa6 = (struct sockaddr_in6 *)&addr;
        memcpy(&(sa6->sin6_addr.s6_addr), src, sizeof(sa6->sin6_addr.s6_addr));
        addrlen = sizeof(struct sockaddr_in6);
    }
    else if (af == AF_INET)
    {
        struct sockaddr_in *sa4 = (struct sockaddr_in *)&addr;
        memcpy(&(sa4->sin_addr.s_addr), src, sizeof(sa4->sin_addr.s_addr));
        addrlen = sizeof(struct sockaddr_in);
    }
    else
        return NULL;
    if (WSAAddressToStringA((struct sockaddr *)&addr, addrlen, 0, dst,
        (LPDWORD)&size) != 0)
        return NULL;
    return dst;
}
#else /* _WIN32 */
#define INET_NTOP inet_ntop
#endif

int addrinfo_get_list(struct addrinfo **ailist, const char *hostname,
                             int port, int protocol) {
    struct addrinfo hints;
    int result;
    char portstr[10]; /* largest port is 65535 */
    memset(&hints, 0, sizeof hints);
    hints.ai_family = AF_UNSPEC; /* IPv4 or IPv6 */
    hints.ai_socktype = protocol;
    hints.ai_protocol = (protocol == SOCK_STREAM ? IPPROTO_TCP : IPPROTO_UDP);
    hints.ai_flags =
#ifdef AI_ALL
                     AI_ALL |       /* both IPv4 and IPv6 addrs */
#endif
#ifdef AI_V4MAPPED
                     AI_V4MAPPED |  /* fallback to IPv4-mapped IPv6 addrs */
#endif
                     AI_PASSIVE;    /* listen to any addr if hostname is NULL */

#if DISABLE_IPV6
    hints.ai_family = AF_INET;
    hints.ai_flags = AI_PASSIVE;
#endif

    portstr[0] = '\0';
    sprintf(portstr, "%d", port);
    result = getaddrinfo(hostname, portstr, &hints, ailist);
        /* There's currently a bug in the BSD libc where getaddrinfo()
         * will return EAI_BADFLAGS for the AI_ALL and AI_V4MAPPED flags.
         * NOTE: this also seems to affect Android!
         * In practice, this means we can't use dual stack sockets,
         * so we fall back to IPv4 networking... */
    if (result == EAI_BADFLAGS)
    {
        static int warned = 0;
        if (!warned)
        {
            fprintf(stderr, "Warning: can't create IPv6 dual-stack socket - falling "
                "back to IPv4. (This is a known bug in the BSD libc, which doesn't "
                "implement the AI_ALL and AI_V4MAPPED flags for getaddrinfo().)\n");
            warned = 1;
        }
        hints.ai_family = AF_INET;
        hints.ai_flags = AI_PASSIVE;
        result = getaddrinfo(hostname, portstr, &hints, ailist);
    }
    return result;
}

int addrinfo_ipv4_first(const struct addrinfo* ai1, const struct addrinfo* ai2)
{
    return (ai1->ai_family == AF_INET) ?
        ((ai2->ai_family == AF_INET) ? 0 : -1) : 1;
}

int addrinfo_ipv6_first(const struct addrinfo* ai1, const struct addrinfo* ai2){
    return (ai1->ai_family == AF_INET6) ?
        ((ai2->ai_family == AF_INET6) ? 0 : -1) : 1;
}

void addrinfo_sort_list(struct addrinfo **ailist,
    int (*compare)(const struct addrinfo*, const struct addrinfo*))
{
    struct addrinfo *result = NULL, *ai = (*ailist);
    while (ai)
    {
        struct addrinfo *temp = ai;
        ai = ai->ai_next;
        if (result)
        {
            struct addrinfo *ai2 = result, *last = NULL;
            while (ai2 && (compare(temp, ai2) >= 0))
            {
                last = ai2;
                ai2 = ai2->ai_next;
            }
            if (last)
            {
                last->ai_next = temp;
                temp->ai_next = ai2;
            }
            else
            {
                result = temp;
                temp->ai_next = ai2;
            }
        }
        else
        {
            result = temp;
            temp->ai_next = NULL;
        }
    }
    *ailist = result;
}

void addrinfo_print_list(const struct addrinfo *ailist)
{
    const struct addrinfo *ai;
    char addrstr[INET6_ADDRSTRLEN];
    for (ai = ailist; ai != NULL; ai = ai->ai_next)
    {
        char *ipver;
        void *addr;
        unsigned int port = 0;
        if (ai->ai_family == AF_INET6)
        {
            struct sockaddr_in6 *sa6 = (struct sockaddr_in6 *)ai->ai_addr;
            addr = &(sa6->sin6_addr);
            ipver = "IPv6";
            port = ntohs(sa6->sin6_port);
        }
        else if (ai->ai_family == AF_INET)
        {
            struct sockaddr_in *sa4 = (struct sockaddr_in *)ai->ai_addr;
            addr = &(sa4->sin_addr);
            ipver = "IPv4";
            port = ntohs(sa4->sin_port);
        }
        else continue;
        INET_NTOP(ai->ai_family, addr, addrstr, INET6_ADDRSTRLEN);
        printf("%s %s %d\n", ipver, addrstr, port);
    }
}

const char* sockaddr_get_addrstr(const struct sockaddr *sa, char *addrstr,
    int addrstrlen)
{
    const void *addr;
    addrstr[0] = '\0';
    if (sa->sa_family == AF_INET6)
    {
        struct sockaddr_in6 *sa6 = (struct sockaddr_in6 *)sa;
        addr = (const void *)&sa6->sin6_addr.s6_addr;
    }
    else if (sa->sa_family == AF_INET)
    {
        struct sockaddr_in *sa4 = (struct sockaddr_in *)sa;
        addr = (const void *)&sa4->sin_addr.s_addr;
    }
    else return NULL;
    return INET_NTOP(sa->sa_family, addr, addrstr, addrstrlen);
}

unsigned int sockaddr_get_port(const struct sockaddr *sa)
{
    if (sa->sa_family == AF_INET6)
    {
        struct sockaddr_in6 *sa6 = (struct sockaddr_in6 *)sa;
        return ntohs(sa6->sin6_port);
    }
    else if (sa->sa_family == AF_INET)
    {
        struct sockaddr_in *sa4 = (struct sockaddr_in *)sa;
        return ntohs(sa4->sin_port);
    }
    else
        return 0;
}

void sockaddr_set_port(const struct sockaddr *sa, unsigned int port) {
    if (sa->sa_family == AF_INET6)
    {
        struct sockaddr_in6 *sa6 = (struct sockaddr_in6 *)sa;
        sa6->sin6_port = htons(port);
    }
    else if (sa->sa_family == AF_INET)
    {
        struct sockaddr_in *sa4 = (struct sockaddr_in *)sa;
        sa4->sin_port = htons(port);
    }
}

int sockaddr_is_multicast(const struct sockaddr *sa)
{
    if (sa->sa_family == AF_INET6)
    {
        struct sockaddr_in6 *sa6 = (struct sockaddr_in6 *)sa;
        return (sa6->sin6_addr.s6_addr[0] == 0xFF);
    }
    else if (sa->sa_family == AF_INET)
    {
        struct sockaddr_in *sa4 = (struct sockaddr_in *)sa;
        return ((ntohl(sa4->sin_addr.s_addr) & 0xF0000000) == 0xE0000000);
    }
    return 0;
}

int socket_init(void)
{
#ifdef _WIN32
    static int initialized = 0;
    if (!initialized)
    {
        short version = MAKEWORD(2, 0);
        WSADATA nobby;
        if (WSAStartup(version, &nobby))
            return -1;
        initialized = 1;
    }
#endif
    return 0;
}

    /* kudos to https://stackoverflow.com/a/46062474/6063908 */
int socket_connect(int socket, const struct sockaddr *addr,
                   socklen_t addrlen, float timeout)
{
    /* set nonblocking and connect */
    socket_set_nonblocking(socket, 1);

    if (connect(socket, addr, addrlen) < 0)
    {
    #ifdef _WIN32
        if (socket_errno() != WSAEWOULDBLOCK)
    #else
        if (socket_errno() != EINPROGRESS)
    #endif
            return -1; /* break on "real" error */

        /* block with select() using timeout.
        NB: the watchdog may cause select() to be interrupted by a signal
        and thus fail with EINTR! In this case we just decrease the timeout
        value by the elapsed time and try again. */
        while (1)
        {
            int status, err;
            socklen_t len;
            struct timeval timeoutval;
        #ifndef _WIN32
            struct timeval t1;
        #endif
            fd_set writefds, errfds;
            if (timeout < 0) timeout = 0;
            timeoutval.tv_sec = (int)timeout;
            timeoutval.tv_usec = (timeout - timeoutval.tv_sec) * 1000000;
            /* the operation has been completed once the socket is writeable. */
            FD_ZERO(&writefds);
            FD_SET(socket, &writefds);
            /* also check for errors */
            FD_ZERO(&errfds);
            FD_SET(socket, &errfds);

        #ifndef _WIN32
            gettimeofday(&t1, NULL);
        #endif

            status = select(socket+1, NULL, &writefds, &errfds, &timeoutval);
            if (status < 0) /* select failed */
            {
            #ifndef _WIN32
                if (errno == EINTR)
                {
                    double elapsed;
                    struct timeval t2;
                    gettimeofday(&t2, NULL);
                    elapsed = ((t2.tv_sec - t1.tv_sec) +
                        (1./1000000.) * (t2.tv_usec - t1.tv_usec));
                    timeout -= elapsed;
                #if 0
                    fprintf(stderr, "socket_connect: interrupted after %f ms\n",
                            elapsed * 1000);
                    fprintf(stderr, "remaining timeout: %f ms\n", timeout * 1000);
                #endif
                    continue;
                }
            #endif
                fprintf(stderr, "socket_connect: select failed (%d)\n", errno);
                return -1;
            }
            else if (status == 0) /* connection timed out */
            {
            #ifdef _WIN32
                WSASetLastError(WSAETIMEDOUT);
            #else
                errno = ETIMEDOUT;
            #endif
                return -1;
            }

            /* Once the socket is writeable we must check the SO_ERROR socket option
            to determine whether the connection has been successful. We also wait
            for exceptions because it can give better (and earlier) error messages,
            but we must not rely on it! (See the documentation for EINPROGRESS in
            https://man7.org/linux/man-pages/man2/connect.2.html.) */
            len = sizeof(err);
            getsockopt(socket, SOL_SOCKET, SO_ERROR, (void *)&err, &len);
            if (err == 0)
                break; /* success! */
            else
            {
            #ifdef _WIN32
                WSASetLastError(err);
            #else
                errno = err;
            #endif
                return -1;
            }
        }
    }
    /* done, set blocking again */
    socket_set_nonblocking(socket, 0);
    return 0;
}

void socket_close(int socket)
{
#ifdef _WIN32
    closesocket(socket);
#else
    if (socket < 0) return;
    close(socket);
#endif
}

unsigned int socket_get_port(int socket)
{
    struct sockaddr_storage sa;
    socklen_t len = sizeof(sa);
    if (getsockname(socket, (struct sockaddr *)&sa, &len) < 0)
        return 0;

    if (sa.ss_family == AF_INET6)
    {
        struct sockaddr_in6 *sa6 = (struct sockaddr_in6 *)&sa;
        return ntohs(sa6->sin6_port);
    }
    else if (sa.ss_family == AF_INET)
    {
        struct sockaddr_in *sa4 = (struct sockaddr_in *)&sa;
        return ntohs(sa4->sin_port);
    }
    else
        return 0;
}

int socket_set_boolopt(int socket, int level, int option_name, int bool_value)
{
    return setsockopt(socket, level, option_name,
        (void *)&bool_value, sizeof(int));
}

int socket_set_nonblocking(int socket, int nonblocking)
{
#ifdef _WIN32
    u_long modearg = nonblocking;
    if (ioctlsocket(socket, FIONBIO, &modearg) != NO_ERROR)
        return -1;
#else
    int sockflags = fcntl(socket, F_GETFL, 0);
    if (nonblocking)
        sockflags |= O_NONBLOCK;
    else
        sockflags &= ~O_NONBLOCK;
    if (fcntl(socket, F_SETFL, sockflags) < 0)
        return -1;
#endif
    return 0;
}

int socket_bytes_available(int socket)
{
#ifdef _WIN32
    u_long n = 0;
    if (ioctlsocket(socket, FIONREAD, &n) != NO_ERROR)
        return -1;
    return n;
#else
    int n = 0;
    if (ioctl(socket, FIONREAD, &n) < 0)
        return -1;
    return n;
#endif
}

int socket_join_multicast_group(int socket, const struct sockaddr *sa)
{
    if (sa->sa_family == AF_INET6)
    {
        struct sockaddr_in6 *sa6 = (struct sockaddr_in6 *)sa;
        struct ipv6_mreq mreq6 = {0};
        memcpy(&mreq6.ipv6mr_multiaddr, &sa6->sin6_addr,
            sizeof(struct in6_addr));
        return setsockopt(socket, IPPROTO_IPV6, IPV6_JOIN_GROUP,
            (char *)&mreq6, sizeof(mreq6));
    }
    else if (sa->sa_family == AF_INET)
    {
        struct sockaddr_in *sa4 = (struct sockaddr_in *)sa;
        struct ip_mreq mreq = {0};
        mreq.imr_multiaddr.s_addr = sa4->sin_addr.s_addr;
        mreq.imr_interface.s_addr = INADDR_ANY;
        return setsockopt(socket, IPPROTO_IP, IP_ADD_MEMBERSHIP,
            (char *)&mreq, sizeof(mreq));
    }
    return -1;
}

int socket_leave_multicast_group(int socket, const struct sockaddr *sa)
{
    if (sa->sa_family == AF_INET6)
    {
        struct sockaddr_in6 *sa6 = (struct sockaddr_in6 *)sa;
        struct ipv6_mreq mreq6 = {0};
        memcpy(&mreq6.ipv6mr_multiaddr, &sa6->sin6_addr,
            sizeof(struct in6_addr));
        return setsockopt(socket, IPPROTO_IPV6, IPV6_LEAVE_GROUP,
            (char *)&mreq6, sizeof(mreq6));
    }
    else if (sa->sa_family == AF_INET)
    {
        struct sockaddr_in *sa4 = (struct sockaddr_in *)sa;
        struct ip_mreq mreq = {0};
        mreq.imr_multiaddr.s_addr = sa4->sin_addr.s_addr;
        mreq.imr_interface.s_addr = INADDR_ANY;
        return setsockopt(socket, IPPROTO_IP, IP_DROP_MEMBERSHIP,
            (char *)&mreq, sizeof(mreq));
    }
    return -1;
}

int socket_errno(void)
{
#ifdef _WIN32
    int err = WSAGetLastError();
    if (err == 10044) /* WSAESOCKTNOSUPPORT */
    {
        fprintf(stderr,
            "Warning: you might not have TCP/IP \"networking\" turned on\n");
    }
    return err;
#else
    return errno;
#endif
}

int socket_errno_udp(void)
{
#ifdef _WIN32
    int err = socket_errno();
    /* ignore WSAECONNRESET to prevent UDP sockets
       from closing in case of a missing receiver */
    if (err == 10054)
        return 0;
    else
        return err;
#else
    return socket_errno();
#endif
}

void socket_strerror(int err, char *buf, int size)
{
    if (size > 0)
    {
    #ifdef _WIN32
        wchar_t wbuf[1024];
        int count = FormatMessageW(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
            0, err, MAKELANGID (LANG_NEUTRAL, SUBLANG_DEFAULT), wbuf, 1024, NULL);
        if (!count || !WideCharToMultiByte(CP_UTF8, 0, wbuf, count+1, buf, size, 0, 0))
            *buf = '\0';
    #else
        snprintf(buf, size, "%s", strerror(err));
    #endif
    }
}
