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

/*-*- Mode: C; c-basic-offset: 8; indent-tabs-mode: nil -*-*/

/***
  This file is part of systemd.

  Copyright 2013 David Strauss

  systemd is free software; you can redistribute it and/or modify it
  under the terms of the GNU Lesser General Public License as published by
  the Free Software Foundation; either version 2.1 of the License, or
  (at your option) any later version.

  systemd 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
  Lesser General Public License for more details.

  You should have received a copy of the GNU Lesser General Public License
  along with systemd; If not, see <http://www.gnu.org/licenses/>.
 ***/

#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 "fd-util.h"
#include "log.h"
#include "path-util.h"
#include "set.h"
#include "socket-util.h"
#include "string-util.h"
#include "util.h"

#define BUFFER_SIZE (256 * 1024)
#define 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) {
        sd_event_source *es;
        Connection *c;

        assert(context);

        while ((es = set_steal_first(context->listen)))
                sd_event_source_unref(es);

        while ((c = set_first(context->connections)))
                connection_free(c);

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

        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 || errno == EPIPE || errno == ECONNRESET) {
                                *from_source = sd_event_source_unref(*from_source);
                                *from = safe_close(*from);
                        } else if (errno != EAGAIN && errno != 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 || errno == EPIPE || errno == ECONNRESET) {
                                *to_source = sd_event_source_unref(*to_source);
                                *to = safe_close(*to);
                        } else if (errno != EAGAIN && errno != 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;
        socklen_t salen;
        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)-1);
                sa.un.sun_path[sizeof(sa.un.sun_path)-1] = 0;

                salen = offsetof(union sockaddr_union, un.sun_path) + strlen(sa.un.sun_path);

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

        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)-2);
                sa.un.sun_path[sizeof(sa.un.sun_path)-1] = 0;

                salen = offsetof(union sockaddr_union, un.sun_path) + 1 + strlen(sa.un.sun_path + 1);

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

        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) > 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, &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"
               "  -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[] = {
                { "help",       no_argument, NULL, 'h'           },
                { "version",    no_argument, NULL, ARG_VERSION   },
                {}
        };

        int c;

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

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

                switch (c) {

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

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