/*
 * Simple LD_PRELOAD wrapper for connect() which uses tlsproxy as proxy for
 * programs which don't support setting a TLS proxy.
 *
 * The following environment variables can be used:
 *
 * - TLSPROXYHELPER_PORT: port where tlsproxy is running (default 9000)
 *
 * If an error occurs -1 is returned and errno is set (either EAFNOSUPPORT if
 * another protocol besides IPv4/IPv6 was used, or EINVAL if another error
 * occurred). The error is written to stderr with "tlsproxyhelper:" as prefix.
 */

/*
 * Copyright (C) 2011-2014  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/>.
 */


#include <config.h>

/* Necessary for RTLD_NEXT. */
#define _GNU_SOURCE

/* Hack to prevent problems with different declarations of connect(). */
#define connect connect_real

#include <arpa/inet.h>
#include <dlfcn.h>
#include <errno.h>
#include <poll.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#undef connect


#define POLL_TIMEOUT 30000 /* 30 seconds */
#define LOG_PREFIX "tlsproxyhelper: "

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


int connect(int sockfd, const struct sockaddr *addr, socklen_t addrlen);

static int get_tlsproxy_port(void);
static int poll_for(int fd, int read);


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

    char ip[INET6_ADDRSTRLEN];
    char port[5 + 1];

    int tlsproxy_port = 0;

    int non_blocking;
    struct sockaddr *my_addr;
    socklen_t        my_addr_len;
    struct sockaddr_in  my_addr_ipv4;
    struct sockaddr_in6 my_addr_ipv6;

    char buffer[128]; /* response from tlsproxy */
    ssize_t size;

    TLSPROXY_LOAD_FUNCTION(real_connect, "connect");

    tlsproxy_port = get_tlsproxy_port();

    /* Is the socket in non-blocking mode? */
    non_blocking = 0;

    /* IPv4. */
    if (addr->sa_family == AF_INET) {
        struct sockaddr_in *ipv4_addr = (struct sockaddr_in *)addr;

        /* Extract IP and port. */
        snprintf(port, sizeof(port), "%u", ntohs(ipv4_addr->sin_port));
        if (inet_ntop(AF_INET, &ipv4_addr->sin_addr, ip, sizeof(ip)) == NULL) {
            perror(LOG_PREFIX "inet_ntop()");
            errno = EINVAL;
            return -1;
        }

        /* Setup connection parameters to connect to tlsproxy. */
        memset(&my_addr_ipv4, 0, sizeof(my_addr_ipv4));
        my_addr_ipv4.sin_family      = AF_INET;
        my_addr_ipv4.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        my_addr_ipv4.sin_port        = htons((uint16_t)tlsproxy_port);

        my_addr     = (struct sockaddr *)&my_addr_ipv4;
        my_addr_len = sizeof(my_addr_ipv4);

    /* IPv6. */
    } else if (addr->sa_family == AF_INET6) {
        struct sockaddr_in6 *ipv6_addr = (struct sockaddr_in6 *)addr;

        /* Extract IP and port. */
        snprintf(port, sizeof(port), "%u", ntohs(ipv6_addr->sin6_port));
        if (inet_ntop(AF_INET6, &ipv6_addr->sin6_addr, ip, sizeof(ip)) == NULL) {
            perror(LOG_PREFIX "inet_ntop()");
            errno = EINVAL;
            return -1;
        }

        /* Setup connection parameters to connect to tlsproxy. */
        memset(&my_addr_ipv6, 0, sizeof(my_addr_ipv6));
        my_addr_ipv6.sin6_family = AF_INET6;
        my_addr_ipv6.sin6_addr   = in6addr_loopback;
        my_addr_ipv6.sin6_port   = htons((uint16_t)tlsproxy_port);

        my_addr     = (struct sockaddr *)&my_addr_ipv6;
        my_addr_len = sizeof(my_addr_ipv6);

    } else {
        fprintf(stderr, LOG_PREFIX "unknown protocol family: %d\n",
                        addr->sa_family);
        errno = EAFNOSUPPORT;
        return -1;
    }

    /* Simple way to pass through DNS requests. */
    if (!strcmp(port, "53")) {
        return real_connect(sockfd, addr, addrlen);
    }

    /* Perform the real connect. */
    if (real_connect(sockfd, my_addr, my_addr_len) != 0) {
        /* Handle non-blocking sockets. */
        if (errno == EINPROGRESS) {
            non_blocking = 1;
        } else {
            perror(LOG_PREFIX "connect()");
            /* errno is correctly set by real connect(). */
            return -1;
        }
    }

    /* Make sure we can write. */
    if (non_blocking) {
        int optval;
        socklen_t optlen;

        if (poll_for(sockfd, 0 /* write */) != 0) {
            /* poll_for() writes the error message */
            errno = EINVAL;
            return -1;
        }

        /* Check if connect() was successful. */
        optlen = sizeof(optval);
        if (getsockopt(sockfd, SOL_SOCKET, SO_ERROR, &optval, &optlen) != 0) {
            perror(LOG_PREFIX "getsockopt()");
            errno = EINVAL;
            return -1;
        }
        if (optval != 0) {
            fprintf(stderr, LOG_PREFIX "connect() failed (SO_ERROR)\n");
            errno = EINVAL;
            return -1;
        }
    }

    /* Send "CONNECT ip:port HTTP/1.0" to the proxy. */
    snprintf(buffer, sizeof(buffer),
             "CONNECT %s:%s HTTP/1.0\r\n\r\n", ip, port);
    size = write(sockfd, buffer, strlen(buffer));
    if (size != (int)strlen(buffer)) {
        perror(LOG_PREFIX "write()");
        errno = EINVAL;
        return -1;
    }
    fsync(sockfd);

    /* Make sure we can read. */
    if (non_blocking) {
        if (poll_for(sockfd, 1 /* read */) != 0) {
            /* poll_for() writes the error message */
            errno = EINVAL;
            return -1;
        }
    }

    /* Read the response from the proxy and check if it's fine. */
    size = read(sockfd, buffer, sizeof(buffer));
    if (size < 0) {
        perror(LOG_PREFIX "read()");
        errno = EINVAL;
        return -1;
    }
    if (strncmp(buffer, "HTTP/1.0 200 Connection established\r\n", 37)) {
        fprintf(stderr, LOG_PREFIX "invalid proxy response: %s", buffer);
        errno = EINVAL;
        return -1;
    }

    /* connect() returns 0 on success. */
    return 0;
}

static int get_tlsproxy_port(void) {
    const char *string;
    int port = 0;

    /* Extract port from the environment variable. */
    string = getenv("TLSPROXYHELPER_PORT");
    if (string != NULL) {
        port = atoi(string);
    }
    /* No valid port specified, use the default one. */
    if (port <= 0 || port > 0xffff) {
        port = 9000;
    }

    return port;
}

/* Poll for read (mode = 1) or write (mode = 0) possibility on file descriptor
 * sockfd. Returns 0 on success, -1 on failure. */
static int poll_for(int sockfd, int mode) {
    int result;
    struct pollfd fds[1];

    fds[0].fd      = sockfd;
    fds[0].events  = POLLERR | POLLHUP;
    /* Either poll for read or write possibility. */
    if (mode == 1) {
        fds[0].events |= POLLIN | POLLPRI;
    } else if (mode == 0) {
        fds[0].events |= POLLOUT;
    } else {
        abort();
    }
    fds[0].revents = 0;

    result = poll(fds, 1 /* fd count */, POLL_TIMEOUT);
    if (result < 0) {
        perror(LOG_PREFIX "poll()");
        return -1;
    } else if (result == 0) {
        fprintf(stderr, LOG_PREFIX "poll() timeout\n");
        return -1;
    }

    if (fds[0].revents & POLLERR) {
        fprintf(stderr, LOG_PREFIX "poll(): POLLERR\n");
        return -1;
    }
    if (fds[0].revents & POLLHUP) {
        fprintf(stderr, LOG_PREFIX "poll(): POLLHUP\n");
        return -1;
    }

    return 0;
}