/* SPDX-License-Identifier: LGPL-2.1-or-later */

#include <netdb.h>
#include <stdio.h>
#include <unistd.h>

#include "sd-daemon.h"
#include "sd-event.h"
#include "sd-resolve.h"

#include "alloc-util.h"
#include "build.h"
#include "daemon-util.h"
#include "errno-util.h"
#include "event-util.h"
#include "fd-util.h"
#include "format-table.h"
#include "in-addr-util.h"
#include "io-util.h"
#include "log.h"
#include "main-func.h"
#include "options.h"
#include "parse-util.h"
#include "pretty-print.h"
#include "resolve-private.h"
#include "set.h"
#include "socket-forward.h"
#include "socket-util.h"
#include "string-table.h"
#include "string-util.h"
#include "strv.h"
#include "time-util.h"

static unsigned arg_connections_max = 256;
static const char *arg_remote_host = NULL;
static usec_t arg_exit_idle_time = USEC_INFINITY;

typedef enum ProxyProtocol {
        PROXY_NONE,
        PROXY_V1,
        _PROXY_PROTOCOL_MAX,
        _PROXY_PROTOCOL_INVALID = -EINVAL,
} ProxyProtocol;

static const char* const proxy_protocol_table[_PROXY_PROTOCOL_MAX] = {
        [PROXY_V1] = "v1",
};

static ProxyProtocol arg_proxy_protocol = PROXY_NONE;

DEFINE_PRIVATE_STRING_TABLE_LOOKUP_FROM_STRING(proxy_protocol, ProxyProtocol);

typedef struct Context {
        sd_event *event;
        sd_resolve *resolve;
        sd_event_source *idle_time;

        Set *listen;
        Set *connections;
} Context;

typedef struct Connection {
        Context *context;

        int server_fd, client_fd;

        sd_event_source *connect_event_source;

        SocketForward *forward;

        sd_resolve_query *resolve_query;
} Connection;

static Connection* connection_free(Connection *c) {
        if (!c)
                return NULL;

        if (c->context)
                set_remove(c->context->connections, c);

        sd_event_source_unref(c->connect_event_source);
        socket_forward_free(c->forward);

        safe_close(c->server_fd);
        safe_close(c->client_fd);

        sd_resolve_query_unref(c->resolve_query);

        return mfree(c);
}

DEFINE_TRIVIAL_CLEANUP_FUNC(Connection*, connection_free);

DEFINE_PRIVATE_HASH_OPS_WITH_VALUE_DESTRUCTOR(
                connection_hash_ops,
                void, trivial_hash_func, trivial_compare_func,
                Connection, connection_free);

static void context_done(Context *context) {
        assert(context);

        set_free(context->listen);
        set_free(context->connections);

        sd_event_unref(context->event);
        sd_resolve_unref(context->resolve);
        sd_event_source_unref(context->idle_time);
}

static int idle_time_cb(sd_event_source *s, uint64_t usec, void *userdata) {
        Context *c = userdata;
        int r;

        if (!set_isempty(c->connections)) {
                log_warning("Idle timer fired even though there are connections, ignoring");
                return 0;
        }

        r = sd_event_exit(c->event, 0);
        if (r < 0) {
                log_warning_errno(r, "Error while stopping event loop, ignoring: %m");
                return 0;
        }
        return 0;
}

static void context_reset_timer(Context *context) {
        int r;

        assert(context);

        if (arg_exit_idle_time < USEC_INFINITY && set_isempty(context->connections)) {
                r = event_reset_time_relative(
                                context->event, &context->idle_time, CLOCK_MONOTONIC,
                                arg_exit_idle_time, 0, idle_time_cb, context,
                                SD_EVENT_PRIORITY_NORMAL, "idle-timer", /* force_reset= */ true);
                if (r < 0)
                        log_warning_errno(r, "Failed to reset idle timer, ignoring: %m");
        }
}

static void connection_release(Connection *c) {
        Context *context = ASSERT_PTR(ASSERT_PTR(c)->context);

        connection_free(c);
        context_reset_timer(context);
}

static int connection_forward_done(SocketForward *sf, int error, void *userdata) {
        Connection *c = ASSERT_PTR(userdata);

        if (error < 0)
                log_error_errno(error, "Forwarding failed: %m");

        connection_release(c);
        return 0; /* ignore errors, continue serving */
}

static int send_proxy_protocol_v1(Connection *c) {
        _cleanup_free_ char *header = NULL;
        int r, header_len;
        union sockaddr_union local_sa, remote_sa;
        socklen_t sa_len;

        assert(c);

        r = sd_is_socket(c->server_fd, AF_UNSPEC, SOCK_STREAM, /* listening= */ 0);
        if (r < 0) {
                log_warning_errno(r, "Failed to issue SO_TYPE, reporting fallback proxy data 'UNKNOWN': %m");
                goto unknown;
        }
        if (r == 0) {
                log_warning("Only TCP is supported by the PROXY protocol, reporting fallback proxy data 'UNKNOWN'.");
                goto unknown;
        }

        sa_len = sizeof(local_sa);
        if (getsockname(c->server_fd, &local_sa.sa, &sa_len) < 0) {
                log_warning_errno(errno, "Failed to get local address (getsockname), reporting fallback proxy data 'UNKNOWN': %m");
                goto unknown;
        }

        sa_len = sizeof(remote_sa);
        if (getpeername(c->server_fd, &remote_sa.sa, &sa_len) < 0) {
                log_warning_errno(errno, "Failed to get remote address (getpeername), reporting fallback proxy data 'UNKNOWN': %m");
                goto unknown;
        }

        const char *proto = NULL;
        switch (remote_sa.sa.sa_family) {

        case AF_INET:
                proto = "TCP4";
                break;

        case AF_INET6:
                proto = "TCP6";
                break;

        default:
                log_warning("Only TCP over IPv4 and IPv6 are supported, reporting fallback proxy data 'UNKNOWN'.");
                goto unknown;
        }

        const union in_addr_union *remote_addr = sockaddr_in_addr(&remote_sa.sa);
        const union in_addr_union *local_addr = sockaddr_in_addr(&local_sa.sa);

        unsigned remote_port, local_port;
        assert_se(sockaddr_port(&remote_sa.sa, &remote_port) >= 0);
        assert_se(sockaddr_port(&local_sa.sa, &local_port) >= 0);

        header_len = asprintf(
                         &header, "PROXY %s %s %s %u %u\r\n",
                         proto,
                         IN_ADDR_TO_STRING(remote_sa.sa.sa_family, remote_addr),
                         IN_ADDR_TO_STRING(local_sa.sa.sa_family, local_addr),
                         remote_port,
                         local_port);
        if (header_len < 0)
                return log_oom();

        r = loop_write_full(c->client_fd, header, header_len, 10 * USEC_PER_SEC);
        if (r < 0)
                return log_error_errno(r, "Failed to write to backend host: %m");

        /* success */
        return 0;

unknown:
        /* ignore previous errors - server can decide to deny UNKNOWN connections */
        r = loop_write_full(c->client_fd, "PROXY UNKNOWN\r\n", SIZE_MAX, 10 * USEC_PER_SEC);
        if (r < 0)
                return log_error_errno(r, "Failed to write to backend host: %m");

        return 0;
}

static int connection_complete(Connection *c) {
        int r;

        assert(c);

        if (arg_proxy_protocol == PROXY_V1) {
                r = send_proxy_protocol_v1(c);
                if (r < 0)
                        return r;
        }

        r = socket_forward_new(
                        c->context->event,
                        TAKE_FD(c->server_fd),
                        TAKE_FD(c->client_fd),
                        connection_forward_done, c,
                        &c->forward);
        if (r < 0)
                return log_error_errno(r, "Failed to set up forwarding: %m");

        return 0;
}

static int connect_cb(sd_event_source *s, int fd, uint32_t revents, void *userdata) {
        Connection *c = ASSERT_PTR(userdata);
        socklen_t solen;
        int error;

        assert(s);
        assert(fd >= 0);

        solen = sizeof(error);
        if (getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &solen) < 0) {
                log_error_errno(errno, "Failed to issue SO_ERROR: %m");
                goto fail;
        }

        if (error != 0) {
                log_error_errno(error, "Failed to connect to remote host: %m");
                goto fail;
        }

        c->connect_event_source = sd_event_source_unref(c->connect_event_source);

        if (connection_complete(c) < 0)
                goto fail;

        return 0;

fail:
        connection_release(c);
        return 0; /* ignore errors, continue serving */
}

static int connection_start(Connection *c, struct sockaddr *sa, socklen_t salen) {
        int r;

        assert(c);
        assert(sa);
        assert(salen);

        c->client_fd = socket(sa->sa_family, SOCK_STREAM|SOCK_NONBLOCK|SOCK_CLOEXEC, 0);
        if (c->client_fd < 0)
                return log_error_errno(errno, "Failed to get remote socket: %m");

        r = connect(c->client_fd, sa, salen);
        if (r < 0) {
                if (errno != EINPROGRESS)
                        return log_error_errno(errno, "Failed to connect to remote host: %m");

                r = sd_event_add_io(c->context->event, &c->connect_event_source, c->client_fd, EPOLLOUT, connect_cb, c);
                if (r < 0)
                        return log_error_errno(r, "Failed to add connection socket: %m");

                r = sd_event_source_set_enabled(c->connect_event_source, SD_EVENT_ONESHOT);
                if (r < 0)
                        return log_error_errno(r, "Failed to enable oneshot event source: %m");

                return 0;
        }

        return connection_complete(c);
}

static int resolve_handler(sd_resolve_query *q, int ret, const struct addrinfo *ai, Connection *c) {
        assert(q);
        assert(c);

        if (ret != 0) {
                log_error("Failed to resolve host: %s", gai_strerror(ret));
                goto fail;
        }

        c->resolve_query = sd_resolve_query_unref(c->resolve_query);

        if (connection_start(c, ai->ai_addr, ai->ai_addrlen) < 0)
                goto fail;

        return 0;

fail:
        connection_release(c);
        return 0; /* ignore errors, continue serving */
}

static int resolve_remote(Connection *c) {

        static const struct addrinfo hints = {
                .ai_family = AF_UNSPEC,
                .ai_socktype = SOCK_STREAM,
        };

        const char *node, *service;
        int r;

        if (IN_SET(arg_remote_host[0], '/', '@')) {
                union sockaddr_union sa;
                int sa_len;

                r = sockaddr_un_set_path(&sa.un, arg_remote_host);
                if (r < 0)
                        return log_error_errno(r, "Specified address doesn't fit in an AF_UNIX address, refusing: %m");
                sa_len = r;

                return connection_start(c, &sa.sa, sa_len);
        }

        service = strrchr(arg_remote_host, ':');
        if (service) {
                node = strndupa_safe(arg_remote_host,
                                     service - arg_remote_host);
                service++;
        } else {
                node = arg_remote_host;
                service = "80";
        }

        log_debug("Looking up address info for %s:%s", node, service);
        r = resolve_getaddrinfo(c->context->resolve, &c->resolve_query, node, service, &hints, resolve_handler, NULL, c);
        if (r < 0)
                return log_error_errno(r, "Failed to resolve remote host: %m");

        return 0;
}

static int context_add_connection(Context *context, int fd) {
        int r;

        assert(context);

        _cleanup_close_ int nfd = accept4(fd, NULL, NULL, SOCK_NONBLOCK|SOCK_CLOEXEC);
        if (nfd < 0) {
                if (!ERRNO_IS_ACCEPT_AGAIN(errno))
                        log_warning_errno(errno, "Failed to accept() socket, ignoring: %m");

                return -errno;
        }

        if (DEBUG_LOGGING) {
                _cleanup_free_ char *peer = NULL;
                (void) getpeername_pretty(nfd, true, &peer);
                log_debug("New connection from %s", strna(peer));
        }

        if (set_size(context->connections) > arg_connections_max)
                return log_warning_errno(SYNTHETIC_ERRNO(EBUSY), "Hit connection limit, refusing connection.");

        r = sd_event_source_set_enabled(context->idle_time, SD_EVENT_OFF);
        if (r < 0)
                log_warning_errno(r, "Unable to disable idle timer, continuing: %m");

        _cleanup_(connection_freep) Connection *c = new(Connection, 1);
        if (!c)
                return log_oom();

        *c = (Connection) {
                .server_fd = TAKE_FD(nfd),
                .client_fd = -EBADF,
        };

        r = set_ensure_put(&context->connections, &connection_hash_ops, c);
        if (r < 0)
                return log_oom();

        c->context = context;

        r = resolve_remote(c);
        if (r < 0)
                return r;

        TAKE_PTR(c);
        return 0;
}

static int accept_cb(sd_event_source *s, int fd, uint32_t revents, void *userdata) {
        Context *context = ASSERT_PTR(userdata);
        int r;

        assert(s);
        assert(fd >= 0);
        assert(revents & EPOLLIN);

        if (context_add_connection(context, fd) < 0)
                context_reset_timer(context);

        r = sd_event_source_set_enabled(s, SD_EVENT_ONESHOT);
        if (r < 0)
                return log_error_errno(r, "Error while re-enabling listener with ONESHOT: %m");

        return 1;
}

static int add_listen_socket(Context *context, int fd) {
        int r;

        assert(context);
        assert(fd >= 0);

        r = sd_is_socket(fd, 0, SOCK_STREAM, 1);
        if (r < 0)
                return log_error_errno(r, "Failed to determine socket type: %m");
        if (r == 0)
                return log_error_errno(SYNTHETIC_ERRNO(EINVAL),
                                       "Passed in socket is not a stream socket.");

        r = fd_nonblock(fd, true);
        if (r < 0)
                return log_error_errno(r, "Failed to mark file descriptor non-blocking: %m");

        _cleanup_(sd_event_source_unrefp) sd_event_source *source = NULL;
        r = sd_event_add_io(context->event, &source, fd, EPOLLIN, accept_cb, context);
        if (r < 0)
                return log_error_errno(r, "Failed to add event source: %m");

        r = sd_event_source_set_exit_on_failure(source, true);
        if (r < 0)
                return log_error_errno(r, "Failed to enable exit-on-failure logic: %m");

        /* Set the watcher to oneshot in case other processes are also
         * watching to accept(). */
        r = sd_event_source_set_enabled(source, SD_EVENT_ONESHOT);
        if (r < 0)
                return log_error_errno(r, "Failed to enable oneshot mode: %m");

        r = set_ensure_consume(&context->listen, &event_source_hash_ops, TAKE_PTR(source));
        if (r < 0)
                return log_error_errno(r, "Failed to add source to set: %m");

        return 0;
}

static int help(void) {
        _cleanup_free_ char *link = NULL, *time_link = NULL;
        _cleanup_(table_unrefp) Table *options = NULL;
        int r;

        r = terminal_urlify_man("systemd-socket-proxyd", "8", &link);
        if (r < 0)
                return log_oom();
        r = terminal_urlify_man("systemd.time", "7", &time_link);
        if (r < 0)
                return log_oom();

        r = option_parser_get_help_table(&options);
        if (r < 0)
                return r;

        printf("%1$s [HOST:PORT]\n"
               "%1$s [SOCKET]\n"
               "\n%2$sBidirectionally proxy local sockets to another (possibly remote) socket.%3$s\n\n",
               program_invocation_short_name,
               ansi_highlight(),
               ansi_normal());

        r = table_print_or_warn(options);
        if (r < 0)
                return r;

        printf("\nSee %s for --exit-idle-time= time span format.\n"
               "See the %s for details.\n",
               time_link, link);
        return 0;
}

static int parse_argv(int argc, char *argv[]) {
        assert(argc >= 0);
        assert(argv);

        OptionParser opts = { argc, argv };
        int r;

        FOREACH_OPTION_OR_RETURN(c, &opts)
                switch (c) {

                OPTION_COMMON_HELP:
                        return help();

                OPTION_COMMON_VERSION:
                        return version();

                OPTION('c', "connections-max", "NUMBER",
                       "Set the maximum number of connections to be accepted"):
                        r = safe_atou(opts.arg, &arg_connections_max);
                        if (r < 0)
                                return log_error_errno(r, "Failed to parse --connections-max= argument: %s", opts.arg);

                        if (arg_connections_max < 1)
                                return log_error_errno(SYNTHETIC_ERRNO(EINVAL),
                                                       "Connection limit is too low.");

                        break;

                OPTION_LONG("exit-idle-time", "TIME",
                            "Exit when without a connection for this duration"):
                        r = parse_sec(opts.arg, &arg_exit_idle_time);
                        if (r < 0)
                                return log_error_errno(r, "Failed to parse --exit-idle-time= argument: %s", opts.arg);
                        break;

                OPTION_LONG("proxy-protocol", "v1",
                            "Enable PROXY protocol: v1"):
                        arg_proxy_protocol = proxy_protocol_from_string(opts.arg);
                        if (arg_proxy_protocol < 0)
                                return log_error_errno(arg_proxy_protocol, "Failed to parse --proxy-protocol= argument: %s", opts.arg);
                        break;
                }

        char **args = option_parser_get_args(&opts);
        size_t n = strv_length(args);
        if (n < 1)
                return log_error_errno(SYNTHETIC_ERRNO(EINVAL), "Not enough parameters.");
        if (n > 1)
                return log_error_errno(SYNTHETIC_ERRNO(EINVAL), "Too many parameters.");

        arg_remote_host = args[0];
        return 1;
}

static int run(int argc, char *argv[]) {
        _cleanup_(context_done) Context context = {};
        _unused_ _cleanup_(notify_on_cleanup) const char *notify_stop = NULL;
        int r, n, fd;

        log_setup();

        r = parse_argv(argc, argv);
        if (r <= 0)
                return r;

        r = sd_event_default(&context.event);
        if (r < 0)
                return log_error_errno(r, "Failed to allocate event loop: %m");

        r = sd_resolve_default(&context.resolve);
        if (r < 0)
                return log_error_errno(r, "Failed to allocate resolver: %m");

        r = sd_resolve_attach_event(context.resolve, context.event, 0);
        if (r < 0)
                return log_error_errno(r, "Failed to attach resolver: %m");

        sd_event_set_watchdog(context.event, true);

        r = sd_listen_fds(1);
        if (r < 0)
                return log_error_errno(r, "Failed to receive sockets from parent.");
        if (r == 0)
                return log_error_errno(SYNTHETIC_ERRNO(EINVAL), "Didn't get any sockets passed in.");

        n = r;

        for (fd = SD_LISTEN_FDS_START; fd < SD_LISTEN_FDS_START + n; fd++) {
                r = add_listen_socket(&context, fd);
                if (r < 0)
                        return r;
        }

        notify_stop = notify_start(NOTIFY_READY_MESSAGE, NOTIFY_STOPPING_MESSAGE);
        r = sd_event_loop(context.event);
        if (r < 0)
                return log_error_errno(r, "Failed to run event loop: %m");

        return 0;
}

DEFINE_MAIN_FUNCTION(run);
