/*
 * Simple LD_PRELOAD wrapper to "convert" network sockets to UNIX sockets;
 * works for clients and servers. See README for details.
 *
 * Copyright (C) 2013  Simon Ruderich
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * 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, see <http://www.gnu.org/licenses/>.
 */


/* Necessary for RTLD_NEXT. */
#define _GNU_SOURCE

#include <config.h>

#include <assert.h>
#include <dlfcn.h>
#include <errno.h>
#include <netinet/in.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/un.h>
#include <unistd.h>


/* CONSTANTS */

#define LOG_LEVEL_ERROR   1
#define LOG_LEVEL_WARNING 2
#define LOG_LEVEL_DEBUG   3
#define LOG_LEVEL_MASK    LOG_LEVEL_DEBUG

#define LOG_LEVEL_PERROR  42


/* GLOBAL VARIABLES */

struct list {
    int unix_sockfd;

    int orig_domain;
    int orig_type;

    struct list *next;
};

static struct list socket_list = {
    .unix_sockfd = -1, /* must not match a valid sockfd */
};


/* LOG FUNCTIONS/MACROS */

static int get_log_level(void);

static void log_helper(int action, const char *file, int line, const char *format, va_list ap) {
    int saved_errno = errno;

    static int log_level;
    if (!log_level) {
        log_level = get_log_level();
    }

    int level = action & LOG_LEVEL_MASK;
    if (level > log_level) {
        return;
    }

    const char *prefix;
    if (level == LOG_LEVEL_DEBUG) {
        prefix = "DEBUG";
    } else if (level == LOG_LEVEL_WARNING) {
        prefix = "WARN ";
    } else if (level == LOG_LEVEL_ERROR) {
        prefix = "ERROR";
    } else {
        prefix = "UNKNOWN";
    }

    /* Prevent other threads from interrupting the printf()s. */
    flockfile(stderr);

    fprintf(stderr, "socket2unix [%s] ", prefix);
    fprintf(stderr, "[%s:%3d] ", file, line);
    vfprintf(stderr, format, ap);

    if ((action & ~LOG_LEVEL_MASK) == LOG_LEVEL_PERROR) {
        fprintf(stderr, ": ");

        errno = saved_errno;
        perror("");
    }

    funlockfile(stderr);

    if (level == LOG_LEVEL_ERROR) {
        fprintf(stderr, "Aborting.\n");
        exit(EXIT_FAILURE);
    }
}

static void log_(int level, const char *file, int line, const char *format, ...)
    __attribute__((format(printf, 4, 5)));
static void log_(int level, const char *file, int line, const char *format, ...) {
    va_list ap;

    va_start(ap, format);
    log_helper(level, file, line, format, ap);
    va_end(ap);
}

#define ERROR(...) \
    log_(LOG_LEVEL_ERROR,   __FILE__, __LINE__, __VA_ARGS__)
#define WARN(...) \
    log_(LOG_LEVEL_WARNING, __FILE__, __LINE__, __VA_ARGS__)
#define DBG(...) \
    log_(LOG_LEVEL_DEBUG,   __FILE__, __LINE__, __VA_ARGS__)

#define DIE(...) \
    log_(LOG_LEVEL_ERROR | LOG_LEVEL_PERROR, __FILE__, __LINE__, __VA_ARGS__)


/* LD_PRELOAD */

/* Load the function name using dlsym() if necessary and store it in pointer.
 * Terminate program on failure. */
#define LOAD_FUNCTION(pointer, name) \
    if ((pointer) == NULL) { \
        char *error; \
        dlerror(); /* Clear possibly existing error. */ \
        \
        *(void **) (&(pointer)) = dlsym(RTLD_NEXT, (name)); \
        \
        if ((error = dlerror()) != NULL) { \
            ERROR("%s\n", error); \
        } \
    }


/* OTHER FUNCTIONS */

static struct list *find_sockfd(int sockfd) {
    struct list *e;

    if (sockfd == socket_list.unix_sockfd) {
        return NULL;
    }

    for (e = &socket_list; e != NULL; e = e->next) {
        if (e->unix_sockfd == sockfd) {
            return e;
        }
    }
    return NULL;
}
static struct list *remove_sockfd(int sockfd) {
    struct list *e, *p;

    if (sockfd == socket_list.unix_sockfd) {
        return NULL;
    }

    for (e = &socket_list, p = NULL; e != NULL; p = e, e = e->next) {
        if (e->unix_sockfd == sockfd) {
            p->next = e->next;
            return e;
        }
    }
    return NULL;
}

static const char *get_socket_path(void) {
    const char *path = getenv("SOCKET2UNIX_PATH");
    if (!path) {
        ERROR("SOCKET2UNIX_PATH environment variable not defined\n");
    }
    if (path[0] != '/') {
        ERROR("SOCKET2UNIX_PATH '%s' must be an absolute path\n", path);
    }
    return path;
}
static int get_log_level(void) {
    const char *level = getenv("SOCKET2UNIX_DEBUG");
    if (!level) {
#ifdef DEBUG
        return LOG_LEVEL_DEBUG;
#else
        return LOG_LEVEL_WARNING;
#endif
    }
    int number = atoi(level);
    if (number <= 0 || number > LOG_LEVEL_DEBUG) {
        number = LOG_LEVEL_DEBUG;
    }
    return number;
}

static const char *af_to_name(int af) {
    if (af == AF_UNIX) {
        return "AF_UNIX";
    } else if (af == AF_LOCAL) {
        return "AF_LOCAL";
    } else if (af == AF_INET) {
        return "AF_INET";
    } else if (af == AF_INET6) {
        return "AF_INET6";
    } else if (af == AF_IPX) {
        return "AF_IPX";
#ifdef AF_NETLINK
    } else if (af == AF_NETLINK) {
        return "AF_NETLINK";
#endif
#ifdef AF_X25
    } else if (af == AF_X25) {
        return "AF_X25";
#endif
#ifdef AF_AX25
    } else if (af == AF_AX25) {
        return "AF_AX25";
#endif
#ifdef AF_ATMPVC
    } else if (af == AF_ATMPVC) {
        return "AF_ATMPVC";
#endif
    } else if (af == AF_APPLETALK) {
        return "AF_APPLETALK";
#ifdef AF_PACKET
    } else if (af == AF_PACKET) {
        return "AF_PACKET";
#endif
    } else {
        return "AF_UNKNOWN";
    }
}
static const char *sock_to_name(int sock) {
    if (sock & SOCK_STREAM) {
        return "SOCK_STREAM";
    } else if (sock & SOCK_DGRAM) {
        return "SOCK_DGRAM";
    } else if (sock & SOCK_SEQPACKET) {
        return "SOCK_SEQPACKET";
    } else if (sock & SOCK_RAW) {
        return "SOCK_RAW";
    } else if (sock & SOCK_RDM) {
        return "SOCK_RDM";
#ifdef SOCK_PACKET
    } else if (sock & SOCK_PACKET) {
        return "SOCK_PACKET";
#endif
    } else {
        return "SOCK_UNKNOWN";
    }
}
/* for getsockopt()/setsockopt(). */
static const char *level_to_name(int level) {
    if (level == SOL_SOCKET) {
        return "SOL_SOCKET";
#ifdef SOL_IP
    } else if (level == SOL_IP) {
        return "SOL_IP";
#endif
#ifdef SOL_IPV6
    } else if (level == SOL_IPV6) {
        return "SOL_IPV6";
#endif
    } else if (level == IPPROTO_TCP) {
        return "IPPROTO_TCP";
    } else if (level == IPPROTO_UDP) {
        return "IPPROTO_UDP";
    } else {
        return "SOL_UNKNOWN";
    }
}


static int set_sockaddr_un(struct sockaddr_un *sockaddr,
                           const struct sockaddr *addr, socklen_t addrlen) {
    /* Just in case ... */
    if ((addr->sa_family == AF_INET
                && addrlen < sizeof(struct sockaddr_in))
            || (addr->sa_family == AF_INET6
                && addrlen < sizeof(struct sockaddr_in6))) {
        WARN("invalid addrlen from program\n");
        return -1;
    }

    const char *socket_path = get_socket_path();

    /* The program may open multiple sockets, e.g. IPv4 and IPv6 and on
     * multiple ports. Create unique paths. */
    const char *af;
    int port;
    if (addr->sa_family == AF_INET) {
        af   = "v4";
        port = ntohs(((struct sockaddr_in *)addr)->sin_port);
    } else if (addr->sa_family == AF_INET6) {
        af   = "v6";
        port = ntohs(((struct sockaddr_in6 *)addr)->sin6_port);
    } else {
        af   = "unknown";
        port = 0;
        WARN("unknown sa_family '%s' (%d)\n",
             af_to_name(addr->sa_family), addr->sa_family);
    }

    /* Initialize sockaddr_un. */
    sockaddr->sun_family = AF_UNIX;
    int written = snprintf(sockaddr->sun_path, sizeof(sockaddr->sun_path),
                           "%s-%s-%d", socket_path, af, port);
    /* The maximum length is quite short, check it. */
    if (written >= (int)sizeof(sockaddr->sun_path)) {
        ERROR("path '%s-%s-%d' too long for UNIX socket",
              socket_path, af, port);
    }

    return 0;
}


/* FUNCTIONS OVERWRITTEN BY LD_PRELOAD */

int socket(int domain, int type, int protocol) {
    static int (*real_socket)(int, int, int);
    LOAD_FUNCTION(real_socket, "socket");

    if (domain == AF_UNIX || domain == AF_LOCAL) {
        return real_socket(domain, type, protocol);
    }

    DBG("socket(%s, %s, %d)\n",
        af_to_name(domain), sock_to_name(type), protocol);

    /* We must return the replacement socket in case the program uses select()
     * or similar on it. */

    int unix_sockfd = real_socket(AF_UNIX, type, 0);
    if (unix_sockfd < 0) {
        DIE("bind(): failed to create UNIX socket");
    }

    struct list *entry = malloc(sizeof(*entry));
    if (!entry) {
        DIE("socket(): malloc");
    }
    memset(entry, 0, sizeof(*entry));

    entry->unix_sockfd = unix_sockfd;
    entry->orig_domain = domain;
    entry->orig_type   = type;

    entry->next = socket_list.next;
    socket_list.next = entry;

    return unix_sockfd;
}

int close(int fd) {
    static int (*real_close)(int);
    LOAD_FUNCTION(real_close, "close");

    DBG("close(%d)\n", fd);

    struct list *entry = remove_sockfd(fd);
    if (entry == NULL) {
        DBG("close(%d): sockfd not found\n", fd);
        return real_close(fd);
    }
    assert(fd == entry->unix_sockfd);
    free(entry);

    return real_close(fd);
}

int bind(int sockfd, const struct sockaddr *addr, socklen_t addrlen) {
    static int (*real_bind)(int, const struct sockaddr *, socklen_t);
    LOAD_FUNCTION(real_bind, "bind");

    DBG("bind(%d, ..)\n", sockfd);

    if (addr == NULL || addrlen < sizeof(addr->sa_family)
            || addr->sa_family == AF_UNIX
            || addr->sa_family == AF_LOCAL) {
        return real_bind(sockfd, addr, addrlen);
    }

    struct list *entry = find_sockfd(sockfd);
    if (!entry) {
        DBG("bind(%d, ..): sockfd not found\n", sockfd);
        return real_bind(sockfd, addr, addrlen);
    }
    assert(sockfd == entry->unix_sockfd);
    DBG("bind(%d, ..): %s %s\n",
        sockfd,
        af_to_name(entry->orig_domain), sock_to_name(entry->orig_type));

    struct sockaddr_un sockaddr;
    if (set_sockaddr_un(&sockaddr, addr, addrlen) != 0) {
        ERROR("connect(%d, ..) failed\n", sockfd);
    }

    DBG("bind(%d, ..): using path '%s'\n", sockfd, sockaddr.sun_path);

    int attempts = 0;
    while (attempts < 10) {
        if (real_bind(entry->unix_sockfd, (struct sockaddr *)&sockaddr,
                                          sizeof(sockaddr)) == 0) {
            /* Success. */
            return 0;
        }
        if (errno != EADDRINUSE) {
            DIE("bind(%d, ..): failed to bind to '%s'",
                sockfd, sockaddr.sun_path);
        }

        /* File already exists, unlink it if it's a socket. This has a race
         * condition, but the worst case is that we delete a file created by
         * the user at the path he told us to use. Tough luck .. */

        struct stat buf;
        if (lstat(sockaddr.sun_path, &buf) != 0) {
            /* Looks like a race, better abort. */
            DIE("bind(%d, ..): lstat on UNIX socket '%s' failed",
                sockfd, sockaddr.sun_path);
        }

        if (!S_ISSOCK(buf.st_mode)) {
            ERROR("bind(%d, ..): path '%s' exits and is no socket\n",
                  sockfd, sockaddr.sun_path);
        }

        WARN("bind(%d, ..): unlinking '%s'\n", sockfd, sockaddr.sun_path);
        if (unlink(sockaddr.sun_path) != 0) {
            DIE("bind(%d, ..): unlink '%s' failed",
                sockfd, sockaddr.sun_path);
        }

        attempts++;
    }

    ERROR("bind(%d, ..): failed to create UNIX socket file\n", sockfd);
    return -1; /* never reached */
}

int listen(int sockfd, int backlog) {
    static int (*real_listen)(int, int);
    LOAD_FUNCTION(real_listen, "listen");

    DBG("listen(%d, %d)\n", sockfd, backlog);

    struct list *entry = find_sockfd(sockfd);
    if (!entry) {
        DBG("listen(%d, %d): sockfd not found\n", sockfd, backlog);
        return real_listen(sockfd, backlog);
    }
    assert(sockfd == entry->unix_sockfd);
    DBG("listen(%d, %d): %s %s\n",
        sockfd, backlog,
        af_to_name(entry->orig_domain), sock_to_name(entry->orig_type));

    if (real_listen(entry->unix_sockfd, backlog) != 0) {
        DIE("listen(): failed to listen");
    }

    return 0;
}

int accept(int sockfd, struct sockaddr *addr, socklen_t *addrlen) {
    static int (*real_accept)(int, struct sockaddr *, socklen_t *);
    LOAD_FUNCTION(real_accept, "accept");

    DBG("accept(%d, ..)\n", sockfd);

    struct list *entry = find_sockfd(sockfd);
    if (!entry) {
        DBG("accept(%d, ..): sockfd not found\n", sockfd);
        return real_accept(sockfd, addr, addrlen);
    }
    assert(sockfd == entry->unix_sockfd);
    DBG("accept(%d, ..): %s %s\n",
        sockfd,
        af_to_name(entry->orig_domain), sock_to_name(entry->orig_type));

    struct sockaddr_un sockaddr;
    socklen_t size = sizeof(sockaddr);
    int sock = real_accept(entry->unix_sockfd, (struct sockaddr *)&sockaddr,
                                               &size);
    if (sock < 0) {
        DIE("accept(%d, ..): failed to accept", sockfd);
    }

    if (addr == NULL || addrlen == NULL) {
        return sock;
    }
    DBG("accept(%d, ..): caller requested sockaddr\n", sockfd);

    if (*addrlen < size) {
        WARN("accept(%d, ..): invalid addrlen from program", sockfd);
        errno = EINVAL;
        return -1;
    }

    /* This is not the protocol the program asked for (AF_* vs. AF_UNIX), but
     * it should work most of the time. */
    memcpy(addr, &sockaddr, size);
    *addrlen = size;

    /* TODO: is this enough? */

    return sock;
}

int connect(int sockfd, const struct sockaddr *addr, socklen_t addrlen) {
    static int (*real_connect)(int, const struct sockaddr *, socklen_t);
    LOAD_FUNCTION(real_connect, "connect");

    DBG("connect(%d, ..)\n", sockfd);

    if (addr == NULL || addrlen < sizeof(addr->sa_family)
            || addr->sa_family == AF_UNIX
            || addr->sa_family == AF_LOCAL) {
        return real_connect(sockfd, addr, addrlen);
    }

    struct list *entry = find_sockfd(sockfd);
    if (!entry) {
        DBG("connect(%d, ..): sockfd not found\n", sockfd);
        return real_connect(sockfd, addr, addrlen);
    }
    assert(sockfd == entry->unix_sockfd);
    DBG("connect(%d, ..): %s %s\n",
        sockfd,
        af_to_name(entry->orig_domain), sock_to_name(entry->orig_type));

    struct sockaddr_un sockaddr;
    if (set_sockaddr_un(&sockaddr, addr, addrlen) != 0) {
        ERROR("connect(%d, ..) failed\n", sockfd);
    }

    DBG("connect(%d, ..): using path '%s'\n", sockfd, sockaddr.sun_path);

    if (real_connect(entry->unix_sockfd, (struct sockaddr *)&sockaddr,
                                         sizeof(sockaddr)) != 0) {
        DIE("connect(%d, ..): failed to connect", sockfd);
    }

    return 0;
}


int getsockname(int sockfd, struct sockaddr *addr, socklen_t *addrlen) {
    static int (*real_getsockname)(int, struct sockaddr *, socklen_t *);
    LOAD_FUNCTION(real_getsockname, "getsockname");

    DBG("getsockname(%d, ..)\n", sockfd);

    return real_getsockname(sockfd, addr, addrlen);
}

int getpeername(int sockfd, struct sockaddr *addr, socklen_t *addrlen) {
    static int (*real_getpeername)(int, struct sockaddr *, socklen_t *);
    LOAD_FUNCTION(real_getpeername, "getpeername");

    DBG("getpeername(%d, ..)\n", sockfd);

    return real_getpeername(sockfd, addr, addrlen);
}

int getsockopt(int sockfd, int level, int optname, void *optval, socklen_t *optlen) {
    static int (*real_getsockopt)(int, int, int, void *, socklen_t *);
    LOAD_FUNCTION(real_getsockopt, "getsockopt");

    DBG("getsockopt(%d, %d %s, %d, ..)\n",
        sockfd, level, level_to_name(level), optname);

    return real_getsockopt(sockfd, level, optname, optval, optlen);
}
int setsockopt(int sockfd, int level, int optname, const void *optval, socklen_t optlen) {
    static int (*real_setsockopt)(int, int, int, const void *, socklen_t);
    LOAD_FUNCTION(real_setsockopt, "setsockopt");

    DBG("setsockopt(%d, %d %s, %d, ..)\n",
        sockfd, level, level_to_name(level), optname);

    return real_setsockopt(sockfd, level, optname, optval, optlen);
}