[thirdparty/systemd.git] / src / socket-proxy / socket-proxyd.c

/* SPDX-License-Identifier: LGPL-2.1+ */
/***
  This file is part of systemd.

  Copyright 2013 David Strauss
 ***/

#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <netdb.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>

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

#include "alloc-util.h"
#include "fd-util.h"
#include "log.h"
#include "path-util.h"
#include "set.h"
#include "socket-util.h"
#include "string-util.h"
#include "parse-util.h"
#include "util.h"

#define BUFFER_SIZE (256 * 1024)
static unsigned arg_connections_max = 256;

static const char *arg_remote_host = NULL;

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

        Set *listen;
        Set *connections;
} Context;

typedef struct Connection {
        Context *context;

        int server_fd, client_fd;
        int server_to_client_buffer[2]; /* a pipe */
        int client_to_server_buffer[2]; /* a pipe */

        size_t server_to_client_buffer_full, client_to_server_buffer_full;
        size_t server_to_client_buffer_size, client_to_server_buffer_size;

        sd_event_source *server_event_source, *client_event_source;

        sd_resolve_query *resolve_query;
} Connection;

static void connection_free(Connection *c) {
        assert(c);

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

        sd_event_source_unref(c->server_event_source);
        sd_event_source_unref(c->client_event_source);

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

        safe_close_pair(c->server_to_client_buffer);
        safe_close_pair(c->client_to_server_buffer);

        sd_resolve_query_unref(c->resolve_query);

        free(c);
}

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

        set_free_with_destructor(context->listen, sd_event_source_unref);
        set_free_with_destructor(context->connections, connection_free);

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

static int connection_create_pipes(Connection *c, int buffer[2], size_t *sz) {
        int r;

        assert(c);
        assert(buffer);
        assert(sz);

        if (buffer[0] >= 0)
                return 0;

        r = pipe2(buffer, O_CLOEXEC|O_NONBLOCK);
        if (r < 0)
                return log_error_errno(errno, "Failed to allocate pipe buffer: %m");

        (void) fcntl(buffer[0], F_SETPIPE_SZ, BUFFER_SIZE);

        r = fcntl(buffer[0], F_GETPIPE_SZ);
        if (r < 0)
                return log_error_errno(errno, "Failed to get pipe buffer size: %m");

        assert(r > 0);
        *sz = r;

        return 0;
}

static int connection_shovel(
                Connection *c,
                int *from, int buffer[2], int *to,
                size_t *full, size_t *sz,
                sd_event_source **from_source, sd_event_source **to_source) {

        bool shoveled;

        assert(c);
        assert(from);
        assert(buffer);
        assert(buffer[0] >= 0);
        assert(buffer[1] >= 0);
        assert(to);
        assert(full);
        assert(sz);
        assert(from_source);
        assert(to_source);

        do {
                ssize_t z;

                shoveled = false;

                if (*full < *sz && *from >= 0 && *to >= 0) {
                        z = splice(*from, NULL, buffer[1], NULL, *sz - *full, SPLICE_F_MOVE|SPLICE_F_NONBLOCK);
                        if (z > 0) {
                                *full += z;
                                shoveled = true;
                        } else if (z == 0 || IN_SET(errno, EPIPE, ECONNRESET)) {
                                *from_source = sd_event_source_unref(*from_source);
                                *from = safe_close(*from);
                        } else if (!IN_SET(errno, EAGAIN, EINTR))
                                return log_error_errno(errno, "Failed to splice: %m");
                }

                if (*full > 0 && *to >= 0) {
                        z = splice(buffer[0], NULL, *to, NULL, *full, SPLICE_F_MOVE|SPLICE_F_NONBLOCK);
                        if (z > 0) {
                                *full -= z;
                                shoveled = true;
                        } else if (z == 0 || IN_SET(errno, EPIPE, ECONNRESET)) {
                                *to_source = sd_event_source_unref(*to_source);
                                *to = safe_close(*to);
                        } else if (!IN_SET(errno, EAGAIN, EINTR))
                                return log_error_errno(errno, "Failed to splice: %m");
                }
        } while (shoveled);

        return 0;
}

static int connection_enable_event_sources(Connection *c);

static int traffic_cb(sd_event_source *s, int fd, uint32_t revents, void *userdata) {
        Connection *c = userdata;
        int r;

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

        r = connection_shovel(c,
                              &c->server_fd, c->server_to_client_buffer, &c->client_fd,
                              &c->server_to_client_buffer_full, &c->server_to_client_buffer_size,
                              &c->server_event_source, &c->client_event_source);
        if (r < 0)
                goto quit;

        r = connection_shovel(c,
                              &c->client_fd, c->client_to_server_buffer, &c->server_fd,
                              &c->client_to_server_buffer_full, &c->client_to_server_buffer_size,
                              &c->client_event_source, &c->server_event_source);
        if (r < 0)
                goto quit;

        /* EOF on both sides? */
        if (c->server_fd == -1 && c->client_fd == -1)
                goto quit;

        /* Server closed, and all data written to client? */
        if (c->server_fd == -1 && c->server_to_client_buffer_full <= 0)
                goto quit;

        /* Client closed, and all data written to server? */
        if (c->client_fd == -1 && c->client_to_server_buffer_full <= 0)
                goto quit;

        r = connection_enable_event_sources(c);
        if (r < 0)
                goto quit;

        return 1;

quit:
        connection_free(c);
        return 0; /* ignore errors, continue serving */
}

static int connection_enable_event_sources(Connection *c) {
        uint32_t a = 0, b = 0;
        int r;

        assert(c);

        if (c->server_to_client_buffer_full > 0)
                b |= EPOLLOUT;
        if (c->server_to_client_buffer_full < c->server_to_client_buffer_size)
                a |= EPOLLIN;

        if (c->client_to_server_buffer_full > 0)
                a |= EPOLLOUT;
        if (c->client_to_server_buffer_full < c->client_to_server_buffer_size)
                b |= EPOLLIN;

        if (c->server_event_source)
                r = sd_event_source_set_io_events(c->server_event_source, a);
        else if (c->server_fd >= 0)
                r = sd_event_add_io(c->context->event, &c->server_event_source, c->server_fd, a, traffic_cb, c);
        else
                r = 0;

        if (r < 0)
                return log_error_errno(r, "Failed to set up server event source: %m");

        if (c->client_event_source)
                r = sd_event_source_set_io_events(c->client_event_source, b);
        else if (c->client_fd >= 0)
                r = sd_event_add_io(c->context->event, &c->client_event_source, c->client_fd, b, traffic_cb, c);
        else
                r = 0;

        if (r < 0)
                return log_error_errno(r, "Failed to set up client event source: %m");

        return 0;
}

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

        assert(c);

        r = connection_create_pipes(c, c->server_to_client_buffer, &c->server_to_client_buffer_size);
        if (r < 0)
                goto fail;

        r = connection_create_pipes(c, c->client_to_server_buffer, &c->client_to_server_buffer_size);
        if (r < 0)
                goto fail;

        r = connection_enable_event_sources(c);
        if (r < 0)
                goto fail;

        return 0;

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

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

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

        solen = sizeof(error);
        r = getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &solen);
        if (r < 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->client_event_source = sd_event_source_unref(c->client_event_source);

        return connection_complete(c);

fail:
        connection_free(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) {
                log_error_errno(errno, "Failed to get remote socket: %m");
                goto fail;
        }

        r = connect(c->client_fd, sa, salen);
        if (r < 0) {
                if (errno == EINPROGRESS) {
                        r = sd_event_add_io(c->context->event, &c->client_event_source, c->client_fd, EPOLLOUT, connect_cb, c);
                        if (r < 0) {
                                log_error_errno(r, "Failed to add connection socket: %m");
                                goto fail;
                        }

                        r = sd_event_source_set_enabled(c->client_event_source, SD_EVENT_ONESHOT);
                        if (r < 0) {
                                log_error_errno(r, "Failed to enable oneshot event source: %m");
                                goto fail;
                        }
                } else {
                        log_error_errno(errno, "Failed to connect to remote host: %m");
                        goto fail;
                }
        } else {
                r = connection_complete(c);
                if (r < 0)
                        goto fail;
        }

        return 0;

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

static int resolve_cb(sd_resolve_query *q, int ret, const struct addrinfo *ai, void *userdata) {
        Connection *c = userdata;

        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);

        return connection_start(c, ai->ai_addr, ai->ai_addrlen);

fail:
        connection_free(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,
                .ai_flags = AI_ADDRCONFIG
        };

        union sockaddr_union sa = {};
        const char *node, *service;
        int r;

        if (path_is_absolute(arg_remote_host)) {
                sa.un.sun_family = AF_UNIX;
                strncpy(sa.un.sun_path, arg_remote_host, sizeof(sa.un.sun_path));
                return connection_start(c, &sa.sa, SOCKADDR_UN_LEN(sa.un));
        }

        if (arg_remote_host[0] == '@') {
                sa.un.sun_family = AF_UNIX;
                sa.un.sun_path[0] = 0;
                strncpy(sa.un.sun_path+1, arg_remote_host+1, sizeof(sa.un.sun_path)-1);
                return connection_start(c, &sa.sa, SOCKADDR_UN_LEN(sa.un));
        }

        service = strrchr(arg_remote_host, ':');
        if (service) {
                node = strndupa(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 = sd_resolve_getaddrinfo(c->context->resolve, &c->resolve_query, node, service, &hints, resolve_cb, c);
        if (r < 0) {
                log_error_errno(r, "Failed to resolve remote host: %m");
                goto fail;
        }

        return 0;

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

static int add_connection_socket(Context *context, int fd) {
        Connection *c;
        int r;

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

        if (set_size(context->connections) > arg_connections_max) {
                log_warning("Hit connection limit, refusing connection.");
                safe_close(fd);
                return 0;
        }

        r = set_ensure_allocated(&context->connections, NULL);
        if (r < 0) {
                log_oom();
                return 0;
        }

        c = new0(Connection, 1);
        if (!c) {
                log_oom();
                return 0;
        }

        c->context = context;
        c->server_fd = fd;
        c->client_fd = -1;
        c->server_to_client_buffer[0] = c->server_to_client_buffer[1] = -1;
        c->client_to_server_buffer[0] = c->client_to_server_buffer[1] = -1;

        r = set_put(context->connections, c);
        if (r < 0) {
                free(c);
                log_oom();
                return 0;
        }

        return resolve_remote(c);
}

static int accept_cb(sd_event_source *s, int fd, uint32_t revents, void *userdata) {
        _cleanup_free_ char *peer = NULL;
        Context *context = userdata;
        int nfd = -1, r;

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

        nfd = accept4(fd, NULL, NULL, SOCK_NONBLOCK|SOCK_CLOEXEC);
        if (nfd < 0) {
                if (errno != -EAGAIN)
                        log_warning_errno(errno, "Failed to accept() socket: %m");
        } else {
                getpeername_pretty(nfd, true, &peer);
                log_debug("New connection from %s", strna(peer));

                r = add_connection_socket(context, nfd);
                if (r < 0) {
                        log_error_errno(r, "Failed to accept connection, ignoring: %m");
                        safe_close(fd);
                }
        }

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

        return 1;
}

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

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

        r = set_ensure_allocated(&context->listen, NULL);
        if (r < 0) {
                log_oom();
                return r;
        }

        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) {
                log_error("Passed in socket is not a stream socket.");
                return -EINVAL;
        }

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

        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 = set_put(context->listen, source);
        if (r < 0) {
                log_error_errno(r, "Failed to add source to set: %m");
                sd_event_source_unref(source);
                return r;
        }

        /* 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");

        return 0;
}

static void help(void) {
        printf("%1$s [HOST:PORT]\n"
               "%1$s [SOCKET]\n\n"
               "Bidirectionally proxy local sockets to another (possibly remote) socket.\n\n"
               "  -c --connections-max=  Set the maximum number of connections to be accepted\n"
               "  -h --help              Show this help\n"
               "     --version           Show package version\n",
               program_invocation_short_name);
}

static int parse_argv(int argc, char *argv[]) {

        enum {
                ARG_VERSION = 0x100,
                ARG_IGNORE_ENV
        };

        static const struct option options[] = {
                { "connections-max", required_argument, NULL, 'c'           },
                { "help",            no_argument,       NULL, 'h'           },
                { "version",         no_argument,       NULL, ARG_VERSION   },
                {}
        };

        int c, r;

        assert(argc >= 0);
        assert(argv);

        while ((c = getopt_long(argc, argv, "c:h", options, NULL)) >= 0)

                switch (c) {

                case 'h':
                        help();
                        return 0;

                case 'c':
                        r = safe_atou(optarg, &arg_connections_max);
                        if (r < 0) {
                                log_error("Failed to parse --connections-max= argument: %s", optarg);
                                return r;
                        }

                        if (arg_connections_max < 1) {
                                log_error("Connection limit is too low.");
                                return -EINVAL;
                        }

                        break;

                case ARG_VERSION:
                        return version();

                case '?':
                        return -EINVAL;

                default:
                        assert_not_reached("Unhandled option");
                }

        if (optind >= argc) {
                log_error("Not enough parameters.");
                return -EINVAL;
        }

        if (argc != optind+1) {
                log_error("Too many parameters.");
                return -EINVAL;
        }

        arg_remote_host = argv[optind];
        return 1;
}

int main(int argc, char *argv[]) {
        Context context = {};
        int r, n, fd;

        log_parse_environment();
        log_open();

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

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

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

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

        sd_event_set_watchdog(context.event, true);

        n = sd_listen_fds(1);
        if (n < 0) {
                log_error("Failed to receive sockets from parent.");
                r = n;
                goto finish;
        } else if (n == 0) {
                log_error("Didn't get any sockets passed in.");
                r = -EINVAL;
                goto finish;
        }

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

        r = sd_event_loop(context.event);
        if (r < 0) {
                log_error_errno(r, "Failed to run event loop: %m");
                goto finish;
        }

finish:
        context_free(&context);

        return r < 0 ? EXIT_FAILURE : EXIT_SUCCESS;
}
Commit	Line	Data
53e1b683	1	/* SPDX-License-Identifier: LGPL-2.1+ */
912b54ad DS	2	/***
	3	This file is part of systemd.
	4
	5	Copyright 2013 David Strauss
912b54ad DS	6	***/
912b54ad DS	7
912b54ad	8	#include <errno.h>
3f6fd1ba	9	#include <fcntl.h>
912b54ad	10	#include <getopt.h>
3f6fd1ba	11	#include <netdb.h>
912b54ad DS	12	#include <stdio.h>
	13	#include <stdlib.h>
	14	#include <string.h>
912b54ad DS	15	#include <sys/socket.h>
	16	#include <sys/un.h>
	17	#include <unistd.h>
	18
912b54ad DS	19	#include "sd-daemon.h"
912b54ad DS	20	#include "sd-event.h"
fb69d709	21	#include "sd-resolve.h"
3f6fd1ba	22
b5efdb8a	23	#include "alloc-util.h"
3ffd4af2	24	#include "fd-util.h"
96c374d0	25	#include "log.h"
3f6fd1ba LP	26	#include "path-util.h"
3f6fd1ba LP	27	#include "set.h"
912b54ad	28	#include "socket-util.h"
07630cea	29	#include "string-util.h"
dc3b8afb	30	#include "parse-util.h"
912b54ad	31	#include "util.h"
8569a776 LP	32
8569a776 LP	33	#define BUFFER_SIZE (256 * 1024)
dc3b8afb	34	static unsigned arg_connections_max = 256;
912b54ad	35
fb69d709 LP	36	static const char *arg_remote_host = NULL;
fb69d709 LP	37
8569a776	38	typedef struct Context {
fb69d709 LP	39	sd_event *event;
	40	sd_resolve *resolve;
	41
8569a776 LP	42	Set *listen;
	43	Set *connections;
	44	} Context;
912b54ad	45
8569a776	46	typedef struct Connection {
e633ea1c LP	47	Context *context;
e633ea1c LP	48
8569a776 LP	49	int server_fd, client_fd;
	50	int server_to_client_buffer[2]; /* a pipe */
	51	int client_to_server_buffer[2]; /* a pipe */
912b54ad	52
8569a776 LP	53	size_t server_to_client_buffer_full, client_to_server_buffer_full;
	54	size_t server_to_client_buffer_size, client_to_server_buffer_size;
	55
	56	sd_event_source server_event_source, client_event_source;
912b54ad	57
fb69d709 LP	58	sd_resolve_query *resolve_query;
fb69d709 LP	59	} Connection;
912b54ad	60
8569a776 LP	61	static void connection_free(Connection *c) {
8569a776 LP	62	assert(c);
32d3c809	63
e633ea1c LP	64	if (c->context)
	65	set_remove(c->context->connections, c);
	66
8569a776 LP	67	sd_event_source_unref(c->server_event_source);
8569a776 LP	68	sd_event_source_unref(c->client_event_source);
32d3c809	69
03e334a1 LP	70	safe_close(c->server_fd);
03e334a1 LP	71	safe_close(c->client_fd);
32d3c809	72
3d94f76c LP	73	safe_close_pair(c->server_to_client_buffer);
3d94f76c LP	74	safe_close_pair(c->client_to_server_buffer);
32d3c809	75
fb69d709 LP	76	sd_resolve_query_unref(c->resolve_query);
fb69d709 LP	77
8569a776 LP	78	free(c);
8569a776 LP	79	}
32d3c809	80
8569a776	81	static void context_free(Context *context) {
8569a776 LP	82	assert(context);
8569a776 LP	83
224b0e7a ZJS	84	set_free_with_destructor(context->listen, sd_event_source_unref);
224b0e7a ZJS	85	set_free_with_destructor(context->connections, connection_free);
32d3c809	86
fb69d709 LP	87	sd_event_unref(context->event);
fb69d709 LP	88	sd_resolve_unref(context->resolve);
8569a776	89	}
912b54ad	90
8569a776 LP	91	static int connection_create_pipes(Connection c, int buffer[2], size_t sz) {
8569a776 LP	92	int r;
912b54ad	93
8569a776 LP	94	assert(c);
	95	assert(buffer);
	96	assert(sz);
912b54ad	97
8569a776 LP	98	if (buffer[0] >= 0)
8569a776 LP	99	return 0;
912b54ad	100
8569a776	101	r = pipe2(buffer, O_CLOEXEC\|O_NONBLOCK);
4a62c710 MS	102	if (r < 0)
4a62c710 MS	103	return log_error_errno(errno, "Failed to allocate pipe buffer: %m");
912b54ad	104
25dbe4f5	105	(void) fcntl(buffer[0], F_SETPIPE_SZ, BUFFER_SIZE);
8569a776 LP	106
8569a776 LP	107	r = fcntl(buffer[0], F_GETPIPE_SZ);
4a62c710 MS	108	if (r < 0)
4a62c710 MS	109	return log_error_errno(errno, "Failed to get pipe buffer size: %m");
912b54ad	110
8569a776 LP	111	assert(r > 0);
	112	*sz = r;
	113
912b54ad DS	114	return 0;
	115	}
	116
8569a776 LP	117	static int connection_shovel(
	118	Connection *c,
	119	int from, int buffer[2], int to,
	120	size_t full, size_t sz,
	121	sd_event_source from_source, sd_event_source to_source) {
	122
	123	bool shoveled;
	124
	125	assert(c);
	126	assert(from);
	127	assert(buffer);
	128	assert(buffer[0] >= 0);
	129	assert(buffer[1] >= 0);
	130	assert(to);
	131	assert(full);
	132	assert(sz);
	133	assert(from_source);
	134	assert(to_source);
	135
	136	do {
	137	ssize_t z;
	138
	139	shoveled = false;
	140
	141	if (full < sz && from >= 0 && to >= 0) {
	142	z = splice(from, NULL, buffer[1], NULL, sz - *full, SPLICE_F_MOVE\|SPLICE_F_NONBLOCK);
	143	if (z > 0) {
	144	*full += z;
	145	shoveled = true;
4c701096	146	} else if (z == 0 \|\| IN_SET(errno, EPIPE, ECONNRESET)) {
8569a776	147	from_source = sd_event_source_unref(from_source);
03e334a1	148	from = safe_close(from);
ec2ce0c5	149	} else if (!IN_SET(errno, EAGAIN, EINTR))
4a62c710	150	return log_error_errno(errno, "Failed to splice: %m");
912b54ad DS	151	}
912b54ad DS	152
8569a776 LP	153	if (full > 0 && to >= 0) {
	154	z = splice(buffer[0], NULL, to, NULL, full, SPLICE_F_MOVE\|SPLICE_F_NONBLOCK);
	155	if (z > 0) {
	156	*full -= z;
	157	shoveled = true;
4c701096	158	} else if (z == 0 \|\| IN_SET(errno, EPIPE, ECONNRESET)) {
8569a776	159	to_source = sd_event_source_unref(to_source);
03e334a1	160	to = safe_close(to);
ec2ce0c5	161	} else if (!IN_SET(errno, EAGAIN, EINTR))
4a62c710	162	return log_error_errno(errno, "Failed to splice: %m");
8569a776 LP	163	}
8569a776 LP	164	} while (shoveled);
912b54ad	165
8569a776	166	return 0;
912b54ad DS	167	}
912b54ad DS	168
fb69d709	169	static int connection_enable_event_sources(Connection *c);
8569a776 LP	170
	171	static int traffic_cb(sd_event_source s, int fd, uint32_t revents, void userdata) {
	172	Connection *c = userdata;
912b54ad DS	173	int r;
912b54ad DS	174
8569a776 LP	175	assert(s);
	176	assert(fd >= 0);
	177	assert(c);
912b54ad	178
8569a776 LP	179	r = connection_shovel(c,
	180	&c->server_fd, c->server_to_client_buffer, &c->client_fd,
	181	&c->server_to_client_buffer_full, &c->server_to_client_buffer_size,
	182	&c->server_event_source, &c->client_event_source);
	183	if (r < 0)
	184	goto quit;
912b54ad	185
8569a776 LP	186	r = connection_shovel(c,
	187	&c->client_fd, c->client_to_server_buffer, &c->server_fd,
	188	&c->client_to_server_buffer_full, &c->client_to_server_buffer_size,
	189	&c->client_event_source, &c->server_event_source);
	190	if (r < 0)
	191	goto quit;
912b54ad	192
8569a776 LP	193	/* EOF on both sides? */
	194	if (c->server_fd == -1 && c->client_fd == -1)
	195	goto quit;
912b54ad	196
8569a776 LP	197	/* Server closed, and all data written to client? */
	198	if (c->server_fd == -1 && c->server_to_client_buffer_full <= 0)
	199	goto quit;
912b54ad	200
8569a776 LP	201	/* Client closed, and all data written to server? */
	202	if (c->client_fd == -1 && c->client_to_server_buffer_full <= 0)
	203	goto quit;
912b54ad	204
fb69d709	205	r = connection_enable_event_sources(c);
8569a776 LP	206	if (r < 0)
8569a776 LP	207	goto quit;
912b54ad	208
8569a776	209	return 1;
912b54ad	210
8569a776 LP	211	quit:
	212	connection_free(c);
	213	return 0; /* ignore errors, continue serving */
912b54ad DS	214	}
912b54ad DS	215
fb69d709	216	static int connection_enable_event_sources(Connection *c) {
8569a776 LP	217	uint32_t a = 0, b = 0;
8569a776 LP	218	int r;
912b54ad	219
8569a776	220	assert(c);
912b54ad	221
8569a776 LP	222	if (c->server_to_client_buffer_full > 0)
	223	b \|= EPOLLOUT;
	224	if (c->server_to_client_buffer_full < c->server_to_client_buffer_size)
	225	a \|= EPOLLIN;
912b54ad	226
8569a776 LP	227	if (c->client_to_server_buffer_full > 0)
	228	a \|= EPOLLOUT;
	229	if (c->client_to_server_buffer_full < c->client_to_server_buffer_size)
	230	b \|= EPOLLIN;
912b54ad	231
8569a776 LP	232	if (c->server_event_source)
	233	r = sd_event_source_set_io_events(c->server_event_source, a);
	234	else if (c->server_fd >= 0)
fb69d709	235	r = sd_event_add_io(c->context->event, &c->server_event_source, c->server_fd, a, traffic_cb, c);
8569a776 LP	236	else
	237	r = 0;
	238
f647962d MS	239	if (r < 0)
f647962d MS	240	return log_error_errno(r, "Failed to set up server event source: %m");
912b54ad	241
8569a776 LP	242	if (c->client_event_source)
	243	r = sd_event_source_set_io_events(c->client_event_source, b);
	244	else if (c->client_fd >= 0)
fb69d709	245	r = sd_event_add_io(c->context->event, &c->client_event_source, c->client_fd, b, traffic_cb, c);
8569a776 LP	246	else
8569a776 LP	247	r = 0;
912b54ad	248
f647962d MS	249	if (r < 0)
f647962d MS	250	return log_error_errno(r, "Failed to set up client event source: %m");
912b54ad	251
8569a776	252	return 0;
912b54ad DS	253	}
912b54ad DS	254
fb69d709 LP	255	static int connection_complete(Connection *c) {
	256	int r;
	257
	258	assert(c);
	259
	260	r = connection_create_pipes(c, c->server_to_client_buffer, &c->server_to_client_buffer_size);
	261	if (r < 0)
	262	goto fail;
	263
	264	r = connection_create_pipes(c, c->client_to_server_buffer, &c->client_to_server_buffer_size);
	265	if (r < 0)
	266	goto fail;
	267
	268	r = connection_enable_event_sources(c);
	269	if (r < 0)
	270	goto fail;
	271
	272	return 0;
	273
	274	fail:
	275	connection_free(c);
	276	return 0; /* ignore errors, continue serving */
	277	}
	278
8569a776 LP	279	static int connect_cb(sd_event_source s, int fd, uint32_t revents, void userdata) {
	280	Connection *c = userdata;
	281	socklen_t solen;
	282	int error, r;
	283
	284	assert(s);
	285	assert(fd >= 0);
	286	assert(c);
	287
	288	solen = sizeof(error);
	289	r = getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &solen);
	290	if (r < 0) {
56f64d95	291	log_error_errno(errno, "Failed to issue SO_ERROR: %m");
40976028 DS	292	goto fail;
40976028 DS	293	}
912b54ad	294
8569a776	295	if (error != 0) {
279d3c9c	296	log_error_errno(error, "Failed to connect to remote host: %m");
912b54ad DS	297	goto fail;
	298	}
	299
8569a776 LP	300	c->client_event_source = sd_event_source_unref(c->client_event_source);
8569a776 LP	301
fb69d709	302	return connection_complete(c);
912b54ad	303
fb69d709 LP	304	fail:
	305	connection_free(c);
	306	return 0; /* ignore errors, continue serving */
	307	}
912b54ad	308
fb69d709 LP	309	static int connection_start(Connection c, struct sockaddr sa, socklen_t salen) {
	310	int r;
	311
	312	assert(c);
	313	assert(sa);
	314	assert(salen);
	315
	316	c->client_fd = socket(sa->sa_family, SOCK_STREAM\|SOCK_NONBLOCK\|SOCK_CLOEXEC, 0);
	317	if (c->client_fd < 0) {
56f64d95	318	log_error_errno(errno, "Failed to get remote socket: %m");
8569a776	319	goto fail;
fb69d709 LP	320	}
	321
	322	r = connect(c->client_fd, sa, salen);
	323	if (r < 0) {
	324	if (errno == EINPROGRESS) {
	325	r = sd_event_add_io(c->context->event, &c->client_event_source, c->client_fd, EPOLLOUT, connect_cb, c);
	326	if (r < 0) {
da927ba9	327	log_error_errno(r, "Failed to add connection socket: %m");
fb69d709 LP	328	goto fail;
	329	}
	330
	331	r = sd_event_source_set_enabled(c->client_event_source, SD_EVENT_ONESHOT);
	332	if (r < 0) {
da927ba9	333	log_error_errno(r, "Failed to enable oneshot event source: %m");
fb69d709 LP	334	goto fail;
	335	}
	336	} else {
56f64d95	337	log_error_errno(errno, "Failed to connect to remote host: %m");
fb69d709 LP	338	goto fail;
	339	}
	340	} else {
	341	r = connection_complete(c);
	342	if (r < 0)
	343	goto fail;
	344	}
912b54ad	345
8569a776	346	return 0;
912b54ad	347
8569a776 LP	348	fail:
	349	connection_free(c);
	350	return 0; /* ignore errors, continue serving */
	351	}
	352
fb69d709 LP	353	static int resolve_cb(sd_resolve_query q, int ret, const struct addrinfo ai, void *userdata) {
	354	Connection *c = userdata;
	355
	356	assert(q);
	357	assert(c);
	358
	359	if (ret != 0) {
	360	log_error("Failed to resolve host: %s", gai_strerror(ret));
	361	goto fail;
	362	}
	363
	364	c->resolve_query = sd_resolve_query_unref(c->resolve_query);
	365
	366	return connection_start(c, ai->ai_addr, ai->ai_addrlen);
	367
	368	fail:
	369	connection_free(c);
	370	return 0; /* ignore errors, continue serving */
	371	}
	372
	373	static int resolve_remote(Connection *c) {
	374
	375	static const struct addrinfo hints = {
	376	.ai_family = AF_UNSPEC,
	377	.ai_socktype = SOCK_STREAM,
	378	.ai_flags = AI_ADDRCONFIG
	379	};
	380
1ec6af16	381	union sockaddr_union sa = {};
fb69d709	382	const char node, service;
fb69d709 LP	383	int r;
	384
	385	if (path_is_absolute(arg_remote_host)) {
	386	sa.un.sun_family = AF_UNIX;
fc2fffe7 LP	387	strncpy(sa.un.sun_path, arg_remote_host, sizeof(sa.un.sun_path));
fc2fffe7 LP	388	return connection_start(c, &sa.sa, SOCKADDR_UN_LEN(sa.un));
fb69d709 LP	389	}
	390
	391	if (arg_remote_host[0] == '@') {
	392	sa.un.sun_family = AF_UNIX;
	393	sa.un.sun_path[0] = 0;
fc2fffe7 LP	394	strncpy(sa.un.sun_path+1, arg_remote_host+1, sizeof(sa.un.sun_path)-1);
fc2fffe7 LP	395	return connection_start(c, &sa.sa, SOCKADDR_UN_LEN(sa.un));
fb69d709 LP	396	}
	397
	398	service = strrchr(arg_remote_host, ':');
	399	if (service) {
	400	node = strndupa(arg_remote_host, service - arg_remote_host);
313cefa1	401	service++;
fb69d709 LP	402	} else {
	403	node = arg_remote_host;
	404	service = "80";
	405	}
	406
	407	log_debug("Looking up address info for %s:%s", node, service);
	408	r = sd_resolve_getaddrinfo(c->context->resolve, &c->resolve_query, node, service, &hints, resolve_cb, c);
	409	if (r < 0) {
da927ba9	410	log_error_errno(r, "Failed to resolve remote host: %m");
fb69d709 LP	411	goto fail;
	412	}
	413
	414	return 0;
	415
	416	fail:
	417	connection_free(c);
	418	return 0; /* ignore errors, continue serving */
	419	}
	420
	421	static int add_connection_socket(Context *context, int fd) {
8569a776 LP	422	Connection *c;
	423	int r;
	424
	425	assert(context);
8569a776 LP	426	assert(fd >= 0);
8569a776 LP	427
dc3b8afb	428	if (set_size(context->connections) > arg_connections_max) {
8569a776	429	log_warning("Hit connection limit, refusing connection.");
03e334a1	430	safe_close(fd);
8569a776	431	return 0;
912b54ad DS	432	}
912b54ad DS	433
d5099efc	434	r = set_ensure_allocated(&context->connections, NULL);
fb69d709 LP	435	if (r < 0) {
	436	log_oom();
	437	return 0;
	438	}
912b54ad	439
8569a776	440	c = new0(Connection, 1);
fb69d709 LP	441	if (!c) {
	442	log_oom();
	443	return 0;
	444	}
8569a776	445
e633ea1c	446	c->context = context;
8569a776 LP	447	c->server_fd = fd;
	448	c->client_fd = -1;
	449	c->server_to_client_buffer[0] = c->server_to_client_buffer[1] = -1;
	450	c->client_to_server_buffer[0] = c->client_to_server_buffer[1] = -1;
	451
e633ea1c LP	452	r = set_put(context->connections, c);
	453	if (r < 0) {
	454	free(c);
fb69d709 LP	455	log_oom();
fb69d709 LP	456	return 0;
8569a776	457	}
912b54ad	458
fb69d709	459	return resolve_remote(c);
40976028 DS	460	}
	461
	462	static int accept_cb(sd_event_source s, int fd, uint32_t revents, void userdata) {
7b77ed8c	463	_cleanup_free_ char *peer = NULL;
8569a776 LP	464	Context *context = userdata;
8569a776 LP	465	int nfd = -1, r;
40976028	466
8569a776 LP	467	assert(s);
8569a776 LP	468	assert(fd >= 0);
40976028	469	assert(revents & EPOLLIN);
8569a776 LP	470	assert(context);
	471
	472	nfd = accept4(fd, NULL, NULL, SOCK_NONBLOCK\|SOCK_CLOEXEC);
7b77ed8c LP	473	if (nfd < 0) {
7b77ed8c LP	474	if (errno != -EAGAIN)
56f64d95	475	log_warning_errno(errno, "Failed to accept() socket: %m");
7b77ed8c	476	} else {
366b7db4	477	getpeername_pretty(nfd, true, &peer);
8569a776	478	log_debug("New connection from %s", strna(peer));
40976028	479
fb69d709	480	r = add_connection_socket(context, nfd);
f4bd42aa	481	if (r < 0) {
da927ba9	482	log_error_errno(r, "Failed to accept connection, ignoring: %m");
03e334a1	483	safe_close(fd);
f4bd42aa	484	}
7b77ed8c	485	}
8569a776	486
40976028 DS	487	r = sd_event_source_set_enabled(s, SD_EVENT_ONESHOT);
40976028 DS	488	if (r < 0) {
da927ba9	489	log_error_errno(r, "Error while re-enabling listener with ONESHOT: %m");
fb69d709	490	sd_event_exit(context->event, r);
40976028 DS	491	return r;
	492	}
	493
96c374d0	494	return 1;
912b54ad DS	495	}
912b54ad DS	496
fb69d709	497	static int add_listen_socket(Context *context, int fd) {
8569a776 LP	498	sd_event_source *source;
	499	int r;
	500
	501	assert(context);
8569a776	502	assert(fd >= 0);
912b54ad	503
d5099efc	504	r = set_ensure_allocated(&context->listen, NULL);
96c374d0	505	if (r < 0) {
8569a776 LP	506	log_oom();
	507	return r;
	508	}
	509
	510	r = sd_is_socket(fd, 0, SOCK_STREAM, 1);
f647962d MS	511	if (r < 0)
f647962d MS	512	return log_error_errno(r, "Failed to determine socket type: %m");
8569a776 LP	513	if (r == 0) {
	514	log_error("Passed in socket is not a stream socket.");
	515	return -EINVAL;
	516	}
912b54ad	517
8569a776	518	r = fd_nonblock(fd, true);
f647962d MS	519	if (r < 0)
f647962d MS	520	return log_error_errno(r, "Failed to mark file descriptor non-blocking: %m");
912b54ad	521
fb69d709	522	r = sd_event_add_io(context->event, &source, fd, EPOLLIN, accept_cb, context);
f647962d MS	523	if (r < 0)
f647962d MS	524	return log_error_errno(r, "Failed to add event source: %m");
912b54ad	525
8569a776	526	r = set_put(context->listen, source);
96c374d0	527	if (r < 0) {
da927ba9	528	log_error_errno(r, "Failed to add source to set: %m");
8569a776	529	sd_event_source_unref(source);
40976028 DS	530	return r;
	531	}
	532
	533	/* Set the watcher to oneshot in case other processes are also
	534	* watching to accept(). */
8569a776	535	r = sd_event_source_set_enabled(source, SD_EVENT_ONESHOT);
f647962d MS	536	if (r < 0)
f647962d MS	537	return log_error_errno(r, "Failed to enable oneshot mode: %m");
912b54ad	538
8569a776	539	return 0;
912b54ad DS	540	}
912b54ad DS	541
601185b4 ZJS	542	static void help(void) {
	543	printf("%1$s [HOST:PORT]\n"
	544	"%1$s [SOCKET]\n\n"
8569a776	545	"Bidirectionally proxy local sockets to another (possibly remote) socket.\n\n"
23d0fff7	546	" -c --connections-max= Set the maximum number of connections to be accepted\n"
8cf030b3 LP	547	" -h --help Show this help\n"
8cf030b3 LP	548	" --version Show package version\n",
912b54ad	549	program_invocation_short_name);
912b54ad DS	550	}
912b54ad DS	551
8569a776	552	static int parse_argv(int argc, char *argv[]) {
912b54ad DS	553
912b54ad DS	554	enum {
8cf030b3 LP	555	ARG_VERSION = 0x100,
8cf030b3 LP	556	ARG_IGNORE_ENV
912b54ad DS	557	};
	558
	559	static const struct option options[] = {
dc3b8afb DK	560	{ "connections-max", required_argument, NULL, 'c' },
	561	{ "help", no_argument, NULL, 'h' },
	562	{ "version", no_argument, NULL, ARG_VERSION },
eb9da376	563	{}
912b54ad DS	564	};
912b54ad DS	565
dc3b8afb	566	int c, r;
912b54ad DS	567
	568	assert(argc >= 0);
	569	assert(argv);
	570
dc3b8afb	571	while ((c = getopt_long(argc, argv, "c:h", options, NULL)) >= 0)
912b54ad DS	572
	573	switch (c) {
	574
	575	case 'h':
601185b4 ZJS	576	help();
601185b4 ZJS	577	return 0;
912b54ad	578
dc3b8afb DK	579	case 'c':
	580	r = safe_atou(optarg, &arg_connections_max);
	581	if (r < 0) {
	582	log_error("Failed to parse --connections-max= argument: %s", optarg);
	583	return r;
	584	}
	585
	586	if (arg_connections_max < 1) {
	587	log_error("Connection limit is too low.");
	588	return -EINVAL;
	589	}
	590
	591	break;
	592
912b54ad	593	case ARG_VERSION:
3f6fd1ba	594	return version();
912b54ad	595
eb9da376	596	case '?':
912b54ad	597	return -EINVAL;
eb9da376 LP	598
	599	default:
	600	assert_not_reached("Unhandled option");
912b54ad	601	}
912b54ad	602
8569a776 LP	603	if (optind >= argc) {
8569a776 LP	604	log_error("Not enough parameters.");
912b54ad DS	605	return -EINVAL;
	606	}
	607
8569a776 LP	608	if (argc != optind+1) {
8569a776 LP	609	log_error("Too many parameters.");
912b54ad DS	610	return -EINVAL;
	611	}
	612
8569a776	613	arg_remote_host = argv[optind];
912b54ad DS	614	return 1;
	615	}
	616
	617	int main(int argc, char *argv[]) {
8569a776 LP	618	Context context = {};
8569a776 LP	619	int r, n, fd;
912b54ad DS	620
	621	log_parse_environment();
	622	log_open();
	623
8569a776	624	r = parse_argv(argc, argv);
912b54ad DS	625	if (r <= 0)
	626	goto finish;
	627
fb69d709	628	r = sd_event_default(&context.event);
8569a776	629	if (r < 0) {
da927ba9	630	log_error_errno(r, "Failed to allocate event loop: %m");
8569a776 LP	631	goto finish;
8569a776 LP	632	}
912b54ad	633
fb69d709 LP	634	r = sd_resolve_default(&context.resolve);
fb69d709 LP	635	if (r < 0) {
da927ba9	636	log_error_errno(r, "Failed to allocate resolver: %m");
fb69d709 LP	637	goto finish;
	638	}
	639
	640	r = sd_resolve_attach_event(context.resolve, context.event, 0);
	641	if (r < 0) {
da927ba9	642	log_error_errno(r, "Failed to attach resolver: %m");
fb69d709 LP	643	goto finish;
	644	}
	645
	646	sd_event_set_watchdog(context.event, true);
cde93897	647
8cf030b3 LP	648	n = sd_listen_fds(1);
	649	if (n < 0) {
	650	log_error("Failed to receive sockets from parent.");
	651	r = n;
	652	goto finish;
	653	} else if (n == 0) {
	654	log_error("Didn't get any sockets passed in.");
	655	r = -EINVAL;
	656	goto finish;
	657	}
	658
	659	for (fd = SD_LISTEN_FDS_START; fd < SD_LISTEN_FDS_START + n; fd++) {
fb69d709	660	r = add_listen_socket(&context, fd);
8569a776	661	if (r < 0)
96c374d0	662	goto finish;
912b54ad DS	663	}
912b54ad DS	664
fb69d709	665	r = sd_event_loop(context.event);
912b54ad	666	if (r < 0) {
da927ba9	667	log_error_errno(r, "Failed to run event loop: %m");
912b54ad DS	668	goto finish;
	669	}
	670
912b54ad	671	finish:
8569a776 LP	672	context_free(&context);
8569a776 LP	673
912b54ad DS	674	return r < 0 ? EXIT_FAILURE : EXIT_SUCCESS;
912b54ad DS	675	}